Add multi OBU support for AV1Payloader

codecs/av1_packet.go

@@ -23,9 +23,38 @@ const (
 // AV1Payloader payloads AV1 packets
 type AV1Payloader struct{}
 
+func packetObuNoSize(obu obu.OBU) []byte {
+	// https://aomediacodec.github.io/av1-rtp-spec/#45-payload-structure
+	// To minimize overhead, the obu_has_size_field flag SHOULD be set to zero in all OBUs.
+	data := obu.Header.Marshal()
+	data[0] &= 0b1111_1101
+	data = append(data, obu.Data...)
+	return data
+}
+
+func (p *AV1Payloader) Payload(mtu uint16, payload []byte) (payloads [][]byte) {
+	// Split frame data into obus
+	obus, err := obu.SplitOBU(payload)
+	if err != nil {
+		return
+	}
+	for _, part := range obus {
+		// Packetization Rules (https://aomediacodec.github.io/av1-rtp-spec/#5-packetization-rules)
+		// Each RTP packet MUST NOT contain OBUs that belong to different temporal units.
+		// The temporal delimiter OBU, if present, SHOULD be removed when transmitting, and MUST be ignored by receivers.
+		// Tile list OBUs are not supported. They SHOULD be removed when transmitted, and MUST be ignored by receivers.
+		if part.Header.ObuType == obu.OBU_TEMPORAL_DELIMITER || part.Header.ObuType == obu.OBU_TILE_LIST {
+			continue
+		}
+		chunk := packetObuNoSize(part)
+		payloads = append(payloads, p.payload(mtu, chunk)...)
+	}
+	return
+}
+
 // Payload fragments a AV1 packet across one or more byte arrays
 // See AV1Packet for description of AV1 Payload Header
-func (p *AV1Payloader) Payload(mtu uint16, payload []byte) (payloads [][]byte) {
+func (p *AV1Payloader) payload(mtu uint16, payload []byte) (payloads [][]byte) {
 	maxFragmentSize := int(mtu) - av1PayloaderHeadersize - 2
 	payloadDataRemaining := len(payload)
 	payloadDataIndex := 0
@@ -37,21 +66,9 @@ func (p *AV1Payloader) Payload(mtu uint16, payload []byte) (payloads [][]byte) {
 
 	for payloadDataRemaining > 0 {
 		currentFragmentSize := min(maxFragmentSize, payloadDataRemaining)
-		leb128Size := 1
-		if currentFragmentSize >= 127 {
-			leb128Size = 2
-		}
-
-		out := make([]byte, av1PayloaderHeadersize+leb128Size+currentFragmentSize)
-		leb128Value := obu.EncodeLEB128(uint(currentFragmentSize))
-		if leb128Size == 1 {
-			out[1] = byte(leb128Value)
-		} else {
-			out[1] = byte(leb128Value >> 8)
-			out[2] = byte(leb128Value)
-		}
-
-		copy(out[av1PayloaderHeadersize+leb128Size:], payload[payloadDataIndex:payloadDataIndex+currentFragmentSize])
+		// Current implement always handle one obu a time, so we can just set W=1 and omit leb128Size
+		out := make([]byte, av1PayloaderHeadersize+currentFragmentSize)
+		copy(out[av1PayloaderHeadersize:], payload[payloadDataIndex:payloadDataIndex+currentFragmentSize])
 		payloads = append(payloads, out)
 
 		payloadDataRemaining -= currentFragmentSize
@@ -63,6 +80,8 @@ func (p *AV1Payloader) Payload(mtu uint16, payload []byte) (payloads [][]byte) {
 		if payloadDataRemaining != 0 {
 			out[0] ^= yMask
 		}
+		// set W = 1
+		out[0] |= 1 << wBitshift
 	}
 
 	return payloads

codecs/av1_packet_test.go

@@ -2,7 +2,6 @@ package codecs
 
 import (
 	"errors"
-	"fmt"
 	"reflect"
 	"testing"
 
@@ -49,7 +48,7 @@ func TestAV1_Unmarshal_Error(t *testing.T) {
 		av1Pkt := &AV1Packet{}
 
 		if _, err := av1Pkt.Unmarshal(test.input); !errors.Is(err, test.expectedError) {
-			t.Fatal(fmt.Sprintf("Expected error(%s) but got (%s)", test.expectedError, err))
+			t.Fatalf("Expected error(%s) but got (%s)", test.expectedError, err)
 		}
 	}
 }

pkg/obu/leb128.go

@@ -11,23 +11,33 @@ const (
 // ErrFailedToReadLEB128 indicates that a buffer ended before a LEB128 value could be successfully read
 var ErrFailedToReadLEB128 = errors.New("payload ended before LEB128 was finished")
 
-// EncodeLEB128 encodes a uint as LEB128
-func EncodeLEB128(in uint) (out uint) {
+// AppendUleb128 appends v to b using unsigned LEB128 encoding.
+func AppendUleb128(b []byte, v uint) []byte {
+	// If it's less than or equal to 7-bit
+	if v < 0x80 {
+		return append(b, byte(v))
+	}
+
 	for {
-		// Copy seven bits from in and discard
-		// what we have copied from in
-		out |= (in & sevenLsbBitmask)
-		in >>= 7
-
-		// If we have more bits to encode set MSB
-		// otherwise we are done
-		if in != 0 {
-			out |= msbBitmask
-			out <<= 8
-		} else {
-			return out
+		c := uint8(v & 0x7f)
+		v >>= 7
+
+		if v != 0 {
+			c |= 0x80
+		}
+
+		b = append(b, c)
+
+		if c&0x80 == 0 {
+			break
 		}
 	}
+
+	return b
+}
+
+func EncodeLEB128(in uint) []byte {
+	return AppendUleb128(make([]byte, 0), in)
 }
 
 func decodeLEB128(in uint) (out uint) {

pkg/obu/leb128_test.go

@@ -1,27 +1,29 @@
 package obu
 
 import (
+	"encoding/hex"
 	"errors"
 	"testing"
 )
 
 func TestLEB128(t *testing.T) {
 	for _, test := range []struct {
 		Value   uint
-		Encoded uint
+		Encoded string
 	}{
-		{0, 0},
-		{5, 5},
-		{999999, 0xBF843D},
+		{0, "00"},
+		{5, "05"},
+		{999999, "bf843d"},
 	} {
 		test := test
 
 		encoded := EncodeLEB128(test.Value)
-		if encoded != test.Encoded {
-			t.Fatalf("Actual(%d) did not equal expected(%d)", encoded, test.Encoded)
+		encodedHex := hex.EncodeToString(encoded)
+		if encodedHex != test.Encoded {
+			t.Fatalf("Actual(%s) did not equal expected(%s)", encodedHex, test.Encoded)
 		}
 
-		decoded := decodeLEB128(encoded)
+		decoded, _, _ := ReadLeb128(encoded)
 		if decoded != test.Value {
 			t.Fatalf("Actual(%d) did not equal expected(%d)", decoded, test.Value)
 		}

pkg/obu/split.go

@@ -0,0 +1,152 @@
+package obu
+
+import (
+	"errors"
+	"io"
+)
+
+type OBUType uint8
+
+type OBUHeader struct {
+	ObuType         OBUType // 1-4
+	ExtensionFlag   bool    // 5
+	HasSizeField    bool    // 6
+	ExtensionHeader byte    // 8-16, if ExtensionFlag=1
+}
+
+type OBUReader struct {
+	buffer []byte
+	idx    uint
+	size   uint
+}
+
+var (
+	errInvalidObuData = errors.New("invalid obu data")
+	errForbidenBit    = errors.New("forbidenBit=1 in OBU Header")
+)
+
+const (
+	forbiddenBitMask  = uint8(0b10000000)
+	typeMask          = uint8(0b01111000)
+	typeShift         = 3
+	extensionFlagMask = uint8(0b00000100)
+	hasSizeFlagMask   = uint8(0b00000010)
+	reserved1BitMask  = uint8(0b00000001)
+)
+
+const (
+	OBU_SEQUENCE_HEADER        OBUType = 1
+	OBU_TEMPORAL_DELIMITER     OBUType = 2
+	OBU_FRAME_HEADER           OBUType = 3
+	OBU_TILE_GROUP             OBUType = 4
+	OBU_METADATA               OBUType = 5
+	OBU_FRAME                  OBUType = 6
+	OBU_REDUNDANT_FRAME_HEADER OBUType = 7
+	OBU_TILE_LIST              OBUType = 8
+	OBU_PADDING                OBUType = 15
+	// Others are Reserved
+)
+
+type OBU struct {
+	Header *OBUHeader
+	Data   []byte
+}
+
+func (h *OBUHeader) Marshal() []byte {
+	// header size
+	size := 1
+	if h.ExtensionFlag {
+		size = 2
+	}
+	data := make([]byte, size)
+	// Type
+	data[0] |= byte(h.ObuType << typeShift)
+	if h.HasSizeField {
+		data[0] |= hasSizeFlagMask
+	}
+	if h.ExtensionFlag {
+		data[0] |= extensionFlagMask
+		data[1] = h.ExtensionHeader
+	}
+	return data
+}
+
+func (or *OBUReader) ReadLeb128() (uint, error) {
+	val, nread, err := ReadLeb128(or.buffer[or.idx:])
+	or.idx += nread
+	return val, err
+}
+
+func (or *OBUReader) ReadHeader() (header OBUHeader, err error) {
+	num := or.buffer[or.idx]
+	or.idx += 1
+	// Check ForbidenBit
+	if num&0x80 != 0 {
+		err = errForbidenBit
+		return
+	}
+	header.ObuType = OBUType((num & typeMask) >> typeShift)
+	header.ExtensionFlag = (num & extensionFlagMask) != 0
+	header.HasSizeField = (num & hasSizeFlagMask) != 0
+
+	if header.ExtensionFlag {
+		num = or.buffer[or.idx]
+		or.idx += 1
+		header.ExtensionHeader = num
+	}
+	return
+}
+
+// read next obu
+func (or *OBUReader) ParseNext() (*OBU, error) {
+	if or.idx == or.size {
+		return nil, io.EOF
+	} else if or.idx > or.size {
+		return nil, errInvalidObuData
+	}
+	var obuData OBU
+	header, err := or.ReadHeader()
+	if err != nil {
+		return nil, err
+	}
+	obuData.Header = &header
+	if header.HasSizeField {
+		size, err := or.ReadLeb128()
+		if err != nil {
+			return nil, err
+		}
+		obuData.Data = or.buffer[or.idx : or.idx+size]
+		or.idx += size
+	} else {
+		obuData.Data = or.buffer[or.idx:]
+		or.idx = or.size
+	}
+	return &obuData, nil
+}
+
+func (obu *OBU) Marshal() []byte {
+	// https://aomediacodec.github.io/av1-rtp-spec/#45-payload-structure
+	// To minimize overhead, the obu_has_size_field flag SHOULD be set to zero in all OBUs.
+	data := obu.Header.Marshal()
+	if obu.Header.HasSizeField {
+		AppendUleb128(data, uint(len(obu.Data)))
+	}
+	data = append(data, obu.Data...)
+	return data
+}
+
+// Extract obus from frame data
+func SplitOBU(payload []byte) (obus []OBU, err error) {
+	reader := OBUReader{buffer: payload, size: uint(len(payload))}
+	for {
+		obu, err := reader.ParseNext()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+		obus = append(obus, *obu)
+	}
+	return obus, nil
+}