Update AV1 Payloader to work with libwebrtc

Don't use leb128 with long fragments. Just set the W bit.

Cache the SequenceHeader and deliver it with the subsequent frame.

codecs/av1_packet.go

@@ -20,52 +20,79 @@ const (
 	nMask     = byte(0b00001000)
 	nBitshift = 3
 
+	obuFrameTypeMask     = byte(0b01111000)
+	obuFrameTypeBitshift = 3
+
+	obuFameTypeSequenceHeader = 1
+
 	av1PayloaderHeadersize = 1
+
+	leb128Size = 1
 )
 
 // AV1Payloader payloads AV1 packets
-type AV1Payloader struct{}
+type AV1Payloader struct {
+	sequenceHeader []byte
+}
 
 // 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) {
-	maxFragmentSize := int(mtu) - av1PayloaderHeadersize - 2
-	payloadDataRemaining := len(payload)
 	payloadDataIndex := 0
+	payloadDataRemaining := len(payload)
 
-	// Make sure the fragment/payload size is correct
-	if min(maxFragmentSize, payloadDataRemaining) <= 0 {
+	// Payload Data and MTU is non-zero
+	if mtu <= 0 || payloadDataRemaining <= 0 {
 		return payloads
 	}
 
+	// Cache Sequence Header and packetize with next payload
+	frameType := (payload[0] & obuFrameTypeMask) >> obuFrameTypeBitshift
+	if frameType == obuFameTypeSequenceHeader {
+		p.sequenceHeader = payload
+		return
+	}
+
 	for payloadDataRemaining > 0 {
-		currentFragmentSize := min(maxFragmentSize, payloadDataRemaining)
-		leb128Size := 1
-		if currentFragmentSize >= 127 {
-			leb128Size = 2
+		obuCount := byte(1)
+		metadataSize := av1PayloaderHeadersize
+		if len(p.sequenceHeader) != 0 {
+			obuCount++
+			metadataSize += leb128Size + len(p.sequenceHeader)
 		}
 
-		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)
-		}
+		out := make([]byte, min(int(mtu), payloadDataRemaining+metadataSize))
+		outOffset := av1PayloaderHeadersize
+		out[0] = obuCount << wBitshift
 
-		copy(out[av1PayloaderHeadersize+leb128Size:], payload[payloadDataIndex:payloadDataIndex+currentFragmentSize])
-		payloads = append(payloads, out)
+		if obuCount == 2 {
+			// This Payload contain the start of a Coded Video Sequence
+			out[0] ^= nMask
+
+			out[1] = byte(obu.EncodeLEB128(uint(len(p.sequenceHeader))))
+			copy(out[2:], p.sequenceHeader)
 
-		payloadDataRemaining -= currentFragmentSize
-		payloadDataIndex += currentFragmentSize
+			outOffset += leb128Size + len(p.sequenceHeader)
+
+			p.sequenceHeader = nil
+		}
 
-		if len(payloads) > 1 {
+		outBufferRemaining := len(out) - outOffset
+		copy(out[outOffset:], payload[payloadDataIndex:payloadDataIndex+outBufferRemaining])
+		payloadDataRemaining -= outBufferRemaining
+		payloadDataIndex += outBufferRemaining
+
+		// Does this Fragment contain an OBU that started in a previous payload
+		if len(payloads) > 0 {
 			out[0] ^= zMask
 		}
+
+		// This OBU will be continued in next Payload
 		if payloadDataRemaining != 0 {
 			out[0] ^= yMask
 		}
+
+		payloads = append(payloads, out)
 	}
 
 	return payloads

codecs/av1_packet_test.go

@@ -4,6 +4,7 @@
 package codecs
 
 import (
+	"bytes"
 	"errors"
 	"reflect"
 	"testing"
@@ -12,28 +13,72 @@ import (
 )
 
 func TestAV1_Marshal(t *testing.T) {
-	const mtu = 5
+	p := &AV1Payloader{}
 
-	for _, test := range []struct {
-		input  []byte
-		output [][]byte
-	}{
-		{[]byte{0x01}, [][]byte{{0x00, 0x01, 0x01}}},
-		{[]byte{0x00, 0x01, 0x02, 0x03, 0x04, 0x04, 0x05}, [][]byte{{0x40, 0x02, 0x00, 0x01}, {0xc0, 0x02, 0x02, 0x03}, {0xc0, 0x02, 0x04, 0x04}, {0x80, 0x01, 0x05}}},
-	} {
-		test := test
+	t.Run("Unfragmented OBU", func(t *testing.T) {
+		OBU := []byte{0x00, 0x01, 0x2, 0x3, 0x4, 0x5}
+		payloads := p.Payload(100, OBU)
 
-		p := &AV1Payloader{}
-		if payloads := p.Payload(mtu, test.input); !reflect.DeepEqual(payloads, test.output) {
-			t.Fatalf("Expected(%02x) did not equal actual(%02x)", test.output, payloads)
+		if len(payloads) != 1 || len(payloads[0]) != 7 {
+			t.Fatal("Expected one unfragmented Payload")
 		}
-	}
 
-	p := &AV1Payloader{}
-	zeroMtuPayload := p.Payload(0, []byte{0x0A, 0x0B, 0x0C})
-	if zeroMtuPayload != nil {
-		t.Fatal("Unexpected output from zero MTU AV1 Payloader")
-	}
+		if payloads[0][0] != 0x10 {
+			t.Fatal("Only W bit should be set")
+		}
+
+		if !bytes.Equal(OBU, payloads[0][1:]) {
+			t.Fatal("OBU modified during packetization")
+		}
+	})
+
+	t.Run("Fragmented OBU", func(t *testing.T) {
+		OBU := []byte{0x00, 0x01, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8}
+		payloads := p.Payload(4, OBU)
+
+		if len(payloads) != 3 || len(payloads[0]) != 4 || len(payloads[1]) != 4 || len(payloads[2]) != 4 {
+			t.Fatal("Expected three fragmented Payload")
+		}
+
+		if payloads[0][0] != 0x10|yMask {
+			t.Fatal("W and Y bit should be set")
+		}
+
+		if payloads[1][0] != 0x10|yMask|zMask {
+			t.Fatal("W, Y and Z bit should be set")
+		}
+
+		if payloads[2][0] != 0x10|zMask {
+			t.Fatal("W and Z bit should be set")
+		}
+
+		if !bytes.Equal(OBU[0:3], payloads[0][1:]) || !bytes.Equal(OBU[3:6], payloads[1][1:]) || !bytes.Equal(OBU[6:9], payloads[2][1:]) {
+			t.Fatal("OBU modified during packetization")
+		}
+	})
+
+	t.Run("Sequence Header Caching", func(t *testing.T) {
+		sequenceHeaderFrame := []byte{0xb, 0xA, 0xB, 0xC}
+		normalFrame := []byte{0x0, 0x1, 0x2, 0x3}
+
+		payloads := p.Payload(100, sequenceHeaderFrame)
+		if len(payloads) != 0 {
+			t.Fatal("Sequence Header was not properly cached")
+		}
+
+		payloads = p.Payload(100, normalFrame)
+		if len(payloads) != 1 {
+			t.Fatal("Expected one payload")
+		}
+
+		if payloads[0][0] != 0x20|nMask {
+			t.Fatal("W and N bit should be set")
+		}
+
+		if !bytes.Equal(sequenceHeaderFrame, payloads[0][2:6]) || !bytes.Equal(normalFrame, payloads[0][6:10]) {
+			t.Fatal("OBU modified during packetization")
+		}
+	})
 }
 
 func TestAV1_Unmarshal_Error(t *testing.T) {