av1: refactor Bitstream; fix marshaling OBUs with size field (bluenvi…

…ron/mediamtx#3886)
bluenviron · Jan 19, 2025 · b996d0e · b996d0e
1 parent a9da823
commit b996d0e
Show file tree

Hide file tree

Showing 6 changed files with 80 additions and 47 deletions.
diff --git a/pkg/codecs/av1/bitstream.go b/pkg/codecs/av1/bitstream.go
@@ -4,64 +4,69 @@ import (
 	"fmt"
 )
 
-func obuRemoveSize(h *OBUHeader, sizeN int, ob []byte) []byte {
-	newOBU := make([]byte, len(ob)-sizeN)
-	newOBU[0] = (byte(h.Type) << 3)
-	copy(newOBU[1:], ob[1+sizeN:])
-	return newOBU
-}
-
 // BitstreamUnmarshal extracts a temporal unit from a bitstream.
 // Optionally, it also removes the size field from OBUs.
+//
+// Deprecated: replacted by Bitstream.Unmarshal. The removeSizeField has no effect anymore.
+func BitstreamUnmarshal(buf []byte, _ bool) ([][]byte, error) {
+	var b Bitstream
+	err := b.Unmarshal(buf)
+	return b, err
+}
+
+// BitstreamMarshal encodes a temporal unit into a bitstream.
+//
+// Deprecated: replacted by Bitstream.Marshal
+func BitstreamMarshal(tu [][]byte) ([]byte, error) {
+	return Bitstream(tu).Marshal()
+}
+
+// Bitstream is an AV1 bitstream.
 // Specification: https://aomediacodec.github.io/av1-spec/#low-overhead-bitstream-format
-func BitstreamUnmarshal(bs []byte, removeSizeField bool) ([][]byte, error) {
-	var ret [][]byte
+type Bitstream [][]byte
 
+// Unmarshal decodes a Bitstream.
+func (bs *Bitstream) Unmarshal(buf []byte) error {
 	for {
 		var h OBUHeader
-		err := h.Unmarshal(bs)
+		err := h.Unmarshal(buf)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		if !h.HasSize {
-			return nil, fmt.Errorf("OBU size not present")
+			return fmt.Errorf("OBU size not present")
 		}
 
 		var size LEB128
-		n, err := size.Unmarshal(bs[1:])
+		n, err := size.Unmarshal(buf[1:])
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		obuLen := 1 + n + int(size)
-		if len(bs) < obuLen {
-			return nil, fmt.Errorf("not enough bytes")
+		if len(buf) < obuLen {
+			return fmt.Errorf("not enough bytes")
 		}
 
-		obu := bs[:obuLen]
-
-		if removeSizeField {
-			obu = obuRemoveSize(&h, n, obu)
-		}
+		var obu []byte
+		obu, buf = buf[:obuLen], buf[obuLen:]
 
-		ret = append(ret, obu)
-		bs = bs[obuLen:]
+		*bs = append(*bs, obu)
 
-		if len(bs) == 0 {
+		if len(buf) == 0 {
 			break
 		}
 	}
 
-	return ret, nil
+	return nil
 }
 
-// BitstreamMarshal encodes a temporal unit into a bitstream.
-// Specification: https://aomediacodec.github.io/av1-spec/#low-overhead-bitstream-format
-func BitstreamMarshal(tu [][]byte) ([]byte, error) {
+// Marshal encodes a Bitstream.
+func (bs Bitstream) Marshal() ([]byte, error) {
 	n := 0
 
-	for _, obu := range tu {
+	for _, obu := range bs {
 		n += len(obu)
 
 		var h OBUHeader
@@ -79,7 +84,7 @@ func BitstreamMarshal(tu [][]byte) ([]byte, error) {
 	buf := make([]byte, n)
 	n = 0
 
-	for _, obu := range tu {
+	for _, obu := range bs {
 		var h OBUHeader
 		h.Unmarshal(obu) //nolint:errcheck
 
@@ -89,6 +94,8 @@ func BitstreamMarshal(tu [][]byte) ([]byte, error) {
 			size := len(obu) - 1
 			n += LEB128(uint32(size)).MarshalTo(buf[n:])
 			n += copy(buf[n:], obu[1:])
+		} else {
+			n += copy(buf[n:], obu)
 		}
 	}
 

diff --git a/pkg/codecs/av1/bitstream_test.go b/pkg/codecs/av1/bitstream_test.go
@@ -9,7 +9,7 @@ import (
 var casesBitstream = []struct {
 	name string
 	enc  []byte
-	dec  [][]byte
+	dec  Bitstream
 }{
 	{
 		"standard",
@@ -21,12 +21,12 @@ var casesBitstream = []struct {
 		},
 		[][]byte{
 			{
-				0x08, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf, 0xc3,
-				0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41,
+				0x0a, 0x0e, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf,
+				0xc3, 0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41,
 			},
 			{
-				0x08, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf, 0xc3,
-				0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41,
+				0x0a, 0x0e, 0x00, 0x00, 0x00, 0x4a, 0xab, 0xbf,
+				0xc3, 0x77, 0x6b, 0xe4, 0x40, 0x40, 0x40, 0x41,
 			},
 		},
 	},
@@ -35,7 +35,8 @@ var casesBitstream = []struct {
 func TestBitstreamUnmarshal(t *testing.T) {
 	for _, ca := range casesBitstream {
 		t.Run(ca.name, func(t *testing.T) {
-			dec, err := BitstreamUnmarshal(ca.enc, true)
+			var dec Bitstream
+			err := dec.Unmarshal(ca.enc)
 			require.NoError(t, err)
 			require.Equal(t, ca.dec, dec)
 		})
@@ -45,7 +46,29 @@ func TestBitstreamUnmarshal(t *testing.T) {
 func TestBitstreamMarshal(t *testing.T) {
 	for _, ca := range casesBitstream {
 		t.Run(ca.name, func(t *testing.T) {
-			enc, err := BitstreamMarshal(ca.dec)
+			enc, err := ca.dec.Marshal()
+			require.NoError(t, err)
+			require.Equal(t, ca.enc, enc)
+		})
+	}
+}
+
+func TestBitstreamMarshalWithoutSize(t *testing.T) {
+	for _, ca := range casesBitstream {
+		t.Run(ca.name, func(t *testing.T) {
+			var tu Bitstream
+			for _, obu := range ca.dec {
+				var size LEB128
+				n, err := size.Unmarshal(obu[1:])
+				require.NoError(t, err)
+
+				newObu := make([]byte, len(obu)-n)
+				newObu[0] = obu[0] & 0b01111000
+				copy(newObu[1:], obu[1+n:])
+				tu = append(tu, newObu)
+			}
+
+			enc, err := tu.Marshal()
 			require.NoError(t, err)
 			require.Equal(t, ca.enc, enc)
 		})
@@ -58,9 +81,10 @@ func FuzzBitstreamUnmarshal(f *testing.F) {
 	}
 
 	f.Fuzz(func(_ *testing.T, b []byte) {
-		tu, err := BitstreamUnmarshal(b, true)
+		var tu Bitstream
+		err := tu.Unmarshal(b)
 		if err == nil {
-			BitstreamMarshal(tu) //nolint:errcheck
+			tu.Marshal() //nolint:errcheck
 		}
 	})
 }
diff --git a/pkg/formats/fmp4/init.go b/pkg/formats/fmp4/init.go
@@ -26,8 +26,9 @@ const (
 	streamTypeAudioStream  = 0x05
 )
 
-func av1FindSequenceHeader(bs []byte) ([]byte, error) {
-	tu, err := av1.BitstreamUnmarshal(bs, true)
+func av1FindSequenceHeader(buf []byte) ([]byte, error) {
+	var tu av1.Bitstream
+	err := tu.Unmarshal(buf)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/formats/fmp4/init_track.go b/pkg/formats/fmp4/init_track.go
@@ -333,8 +333,8 @@ func (it *InitTrack) marshal(w *mp4Writer) error {
 			return err
 		}
 
-		var bs []byte
-		bs, err = av1.BitstreamMarshal([][]byte{codec.SequenceHeader})
+		var enc []byte
+		enc, err = av1.Bitstream([][]byte{codec.SequenceHeader}).Marshal()
 		if err != nil {
 			return err
 		}
@@ -351,7 +351,7 @@ func (it *InitTrack) marshal(w *mp4Writer) error {
 			ChromaSubsamplingX:   boolToUint8(av1SequenceHeader.ColorConfig.SubsamplingX),
 			ChromaSubsamplingY:   boolToUint8(av1SequenceHeader.ColorConfig.SubsamplingY),
 			ChromaSamplePosition: uint8(av1SequenceHeader.ColorConfig.ChromaSamplePosition),
-			ConfigOBUs:           bs,
+			ConfigOBUs:           enc,
 		})
 		if err != nil {
 			return err

diff --git a/pkg/formats/fmp4/part_sample.go b/pkg/formats/fmp4/part_sample.go
@@ -15,7 +15,7 @@ type PartSample struct {
 
 // NewPartSampleAV1 creates a sample with AV1 data.
 func NewPartSampleAV1(sequenceHeaderPresent bool, tu [][]byte) (*PartSample, error) {
-	bs, err := av1.BitstreamMarshal(tu)
+	bs, err := av1.Bitstream(tu).Marshal()
 	if err != nil {
 		return nil, err
 	}
@@ -42,7 +42,8 @@ func NewPartSampleH26x(ptsOffset int32, randomAccessPresent bool, au [][]byte) (
 
 // GetAV1 gets AV1 data from the sample.
 func (ps PartSample) GetAV1() ([][]byte, error) {
-	tu, err := av1.BitstreamUnmarshal(ps.Payload, true)
+	var tu av1.Bitstream
+	err := tu.Unmarshal(ps.Payload)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/formats/pmp4/track.go b/pkg/formats/pmp4/track.go
@@ -355,7 +355,7 @@ func (t *Track) marshal(w *mp4Writer) (*headerTrackMarshalResult, error) {
 		}
 
 		var bs []byte
-		bs, err = av1.BitstreamMarshal([][]byte{codec.SequenceHeader})
+		bs, err = av1.Bitstream([][]byte{codec.SequenceHeader}).Marshal()
 		if err != nil {
 			return nil, err
 		}