Skip to content

Commit

Permalink
Add Unlacer (#90)
Browse files Browse the repository at this point in the history
  • Loading branch information
at-wat authored Dec 16, 2019
1 parent de19aaf commit 947b61d
Show file tree
Hide file tree
Showing 7 changed files with 458 additions and 21 deletions.
36 changes: 27 additions & 9 deletions block.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ package ebml
import (
"errors"
"io"
"io/ioutil"
)

var (
Expand Down Expand Up @@ -76,17 +75,20 @@ type Lace struct {
}

// UnmarshalBlock unmarshals EBML Block structure.
func UnmarshalBlock(r io.Reader) (*Block, error) {
func UnmarshalBlock(r io.Reader, n uint64) (*Block, error) {
var b Block
var err error
if b.TrackNumber, _, err = readVInt(r); err != nil {
var nRead int
if b.TrackNumber, nRead, err = readVInt(r); err != nil {
return nil, err
}
n -= uint64(nRead)
if v, err := readInt(r, 2); err == nil {
b.Timecode = int16(v.(int64))
} else {
return nil, err
}
n -= 2

var bs [1]byte
switch _, err := io.ReadFull(r, bs[:]); err {
Expand All @@ -96,6 +98,8 @@ func UnmarshalBlock(r io.Reader) (*Block, error) {
default:
return nil, err
}
n--

if bs[0]&blockFlagMaskKeyframe != 0 {
b.Keyframe = true
}
Expand All @@ -107,16 +111,30 @@ func UnmarshalBlock(r io.Reader) (*Block, error) {
}
b.Lacing = LacingMode((bs[0] & blockFlagMaskLacing) >> 1)

if b.Lacing != LacingNo {
return nil, errLaceUnimplemented
var ul Unlacer
switch b.Lacing {
case LacingNo:
ul, err = NewNoUnlacer(r, n)
case LacingXiph:
ul, err = NewXiphUnlacer(r, n)
case LacingEBML:
ul, err = NewEBMLUnlacer(r, n)
case LacingFixed:
ul, err = NewFixedUnlacer(r, n)
}

b.Data = [][]byte{{}}
b.Data[0], err = ioutil.ReadAll(r)
if err != nil {
return nil, err
}
return &b, nil
for {
frame, err := ul.Read()
if err == io.EOF {
return &b, nil
}
if err != nil {
return nil, err
}
b.Data = append(b.Data, frame)
}
}

// MarshalBlock marshals EBML Block structure.
Expand Down
57 changes: 46 additions & 11 deletions block_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,10 +38,31 @@ func TestUnmarshalBlock(t *testing.T) {
[]byte{0x21, 0x23, 0x45, 0x00, 0x02, 0x00},
Block{0x012345, 0x0002, false, false, LacingNo, false, nil, [][]byte{{}}},
},
"FixedLace": {
[]byte{0x82, 0x01, 0x23, 0x02, 0x02, 0x0A, 0x0B, 0x0C},
Block{
0x02, 0x0123, false, false, LacingFixed, false, nil,
[][]byte{{0x0A}, {0x0B}, {0x0C}},
},
},
"XiphLace": {
[]byte{0x82, 0x01, 0x23, 0x04, 0x02, 0x01, 0x02, 0x0A, 0x0B, 0x1B, 0x0C},
Block{
0x02, 0x0123, false, false, LacingXiph, false, nil,
[][]byte{{0x0A}, {0x0B, 0x1B}, {0x0C}},
},
},
"EBMLLace": {
[]byte{0x82, 0x01, 0x23, 0x06, 0x02, 0x81, 0x82, 0x0A, 0x0B, 0x1B, 0x0C},
Block{
0x02, 0x0123, false, false, LacingEBML, false, nil,
[][]byte{{0x0A}, {0x0B, 0x1B}, {0x0C}},
},
},
}
for n, c := range testCases {
t.Run(n, func(t *testing.T) {
block, err := UnmarshalBlock(bytes.NewBuffer(c.input))
block, err := UnmarshalBlock(bytes.NewBuffer(c.input), uint64(len(c.input)))
if err != nil {
t.Fatalf("Failed to unmarshal block: %v", err)
}
Expand All @@ -53,18 +74,32 @@ func TestUnmarshalBlock(t *testing.T) {
}

func TestUnmarshalBlock_Error(t *testing.T) {
input := []byte{0x21, 0x23, 0x45, 0x00, 0x02, 0x00}

t.Run("EOF",
func(t *testing.T) {
for l := 0; l < len(input); l++ {
if _, err := UnmarshalBlock(bytes.NewBuffer(input[:l])); err != io.ErrUnexpectedEOF {
t.Errorf("UnmarshalBlock should return %v against short data (%d bytes), but got %v",
io.ErrUnexpectedEOF, l, err)
}
t.Run("EOF", func(t *testing.T) {
input := []byte{0x21, 0x23, 0x45, 0x00, 0x02, 0x00}
for l := 0; l < len(input); l++ {
if _, err := UnmarshalBlock(bytes.NewBuffer(input[:l]), uint64(len(input))); err != io.ErrUnexpectedEOF {
t.Errorf("UnmarshalBlock should return %v against short data (%d bytes), but got %v",
io.ErrUnexpectedEOF, l, err)
}
}
})
testCases := map[string]struct {
input []byte
err error
}{
"UndivisibleFixedLace": {
[]byte{0x82, 0x00, 0x00, 0x02, 0x02, 0x00, 0x00},
errFixedLaceUndivisible,
},
)
}
for n, c := range testCases {
t.Run(n, func(t *testing.T) {
_, err := UnmarshalBlock(bytes.NewBuffer(c.input), uint64(len(c.input)))
if err != c.err {
t.Errorf("Unexpected error, expected: %v, got: %v", c.err, err)
}
})
}
}

func TestMarshalBlock(t *testing.T) {
Expand Down
17 changes: 17 additions & 0 deletions testutils_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package ebml

import (
"bytes"
"io"
)

type limitedDummyWriter struct {
Expand All @@ -16,3 +17,19 @@ func (s *limitedDummyWriter) Write(b []byte) (int, error) {
}
return len(b), nil
}

type delayedBrokenReader struct {
b []byte
n int
limit int
}

func (s *delayedBrokenReader) Read(b []byte) (int, error) {
p := s.n
s.n += len(b)
if s.n > s.limit {
return len(b) - (s.n - s.limit), io.ErrClosedPipe
}
copy(b, s.b[p:p+len(b)])
return len(b), nil
}
158 changes: 158 additions & 0 deletions unlacer.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
// Copyright 2019 The ebml-go authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ebml

import (
"errors"
"io"
)

var (
errFixedLaceUndivisible = errors.New("undivisible fixed lace")
)

// Unlacer is the interface to read laced frames in Block.
type Unlacer interface {
Read() ([]byte, error)
}

type unlacer struct {
r io.Reader
i int
size []int
}

func (u *unlacer) Read() ([]byte, error) {
if u.i >= len(u.size) {
return nil, io.EOF
}
n := u.size[u.i]
u.i++

b := make([]byte, n)
_, err := io.ReadFull(u.r, b)
return b, err
}

// NewNoUnlacer creates pass-through Unlacer for not laced data.
func NewNoUnlacer(r io.Reader, n uint64) (Unlacer, error) {
return &unlacer{r: r, size: []int{int(n)}}, nil
}

// NewXiphUnlacer creates Unlacer for Xiph laced data.
func NewXiphUnlacer(r io.Reader, n uint64) (Unlacer, error) {
var nFrame int
var b [1]byte
switch _, err := io.ReadFull(r, b[:]); err {
case nil:
nFrame = int(b[0]) + 1
case io.EOF:
return nil, io.ErrUnexpectedEOF
default:
return nil, err
}
n--

ul := &unlacer{
r: r,
size: make([]int, nFrame),
}
for i := 0; i < nFrame-1; i++ {
for {
var b [1]byte
switch _, err := io.ReadFull(ul.r, b[:]); err {
case nil:
case io.EOF:
return nil, io.ErrUnexpectedEOF
default:
return nil, err
}
n--
ul.size[i] += int(b[0])
if b[0] != 0xFF {
ul.size[nFrame-1] -= ul.size[i]
break
}
}
}
ul.size[nFrame-1] += int(n)
if ul.size[nFrame-1] <= 0 {
return nil, io.ErrUnexpectedEOF
}

return ul, nil
}

// NewFixedUnlacer creates Unlacer for Fixed laced data.
func NewFixedUnlacer(r io.Reader, n uint64) (Unlacer, error) {
var nFrame int
var b [1]byte
switch _, err := io.ReadFull(r, b[:]); err {
case nil:
nFrame = int(b[0]) + 1
case io.EOF:
return nil, io.ErrUnexpectedEOF
default:
return nil, err
}

ul := &unlacer{
r: r,
size: make([]int, nFrame),
}
ul.size[0] = (int(n) - 1) / nFrame
for i := 1; i < nFrame; i++ {
ul.size[i] = ul.size[0]
}
if ul.size[0]*nFrame+1 != int(n) {
return nil, errFixedLaceUndivisible
}
return ul, nil
}

// NewEBMLUnlacer creates Unlacer for EBML laced data.
func NewEBMLUnlacer(r io.Reader, n uint64) (Unlacer, error) {
var nFrame int
var b [1]byte
switch _, err := io.ReadFull(r, b[:]); err {
case nil:
nFrame = int(b[0]) + 1
case io.EOF:
return nil, io.ErrUnexpectedEOF
default:
return nil, err
}
n--

ul := &unlacer{
r: r,
size: make([]int, nFrame),
}
for i := 0; i < nFrame-1; i++ {
n64, nRead, err := readVInt(ul.r)
if err != nil {
return nil, err
}
n -= uint64(nRead)
ul.size[i] = int(n64)
ul.size[nFrame-1] -= int(n64)
}
ul.size[nFrame-1] += int(n)
if ul.size[nFrame-1] <= 0 {
return nil, io.ErrUnexpectedEOF
}

return ul, nil
}
Loading

0 comments on commit 947b61d

Please sign in to comment.