Create av1_sample_buffer_support.go

codecs/av1/frame/av1_sample_buffer_support.go

@@ -0,0 +1,97 @@
+package frame
+
+import (
+	"github.com/pion/rtp/codecs"
+	"github.com/pion/rtp/codecs/av1/obu"
+)
+
+type AV1PacketSampleBufferSupport struct {
+	popFrame bool
+	avFrame  *AV1
+}
+
+func (d *AV1PacketSampleBufferSupport) IsPartitionTail(marker bool, _ []byte) bool {
+	d.popFrame = true
+	return marker
+}
+
+// IsPartitionHead checks whether if this is a head of the AV1 partition
+func (d *AV1PacketSampleBufferSupport) IsPartitionHead(payload []byte) bool {
+	d.popFrame = true
+	if len(payload) == 0 {
+		return false
+	}
+	return (payload[0] & byte(0b10000000)) == 0
+}
+
+func (d *AV1PacketSampleBufferSupport) Unmarshal(payload []byte) ([]byte, error) {
+
+	if d.popFrame {
+		d.avFrame = &AV1{}
+		d.popFrame = false // start frame assembling
+	}
+
+	packet := codecs.AV1Packet{}
+	_, err := packet.Unmarshal(payload)
+
+	if err != nil {
+		return nil, err
+	}
+
+	OBUs, _ := d.avFrame.ReadFrames(&packet)
+
+	if len(OBUs) == 0 {
+		return nil, nil
+	}
+
+	var payloadSize uint = 0
+
+	for i := range OBUs {
+		obuLength := uint(len(OBUs[i]))
+		if obuLength == 0 {
+			continue
+		}
+		payloadSize += obuLength
+		payloadSize += obu.SizeLeb128(obu.EncodeLEB128(obuLength - 1))
+	}
+
+	result := make([]byte, payloadSize)
+
+	offset := 0
+	for i := range OBUs {
+		obuLength := len(OBUs[i])
+
+		if obuLength == 0 {
+			continue
+		}
+
+		lenMinus := obuLength - 1
+
+		result[offset] = OBUs[i][0] ^ 2 // mark size header exists
+		offset++
+		payloadSize := obu.EncodeLEB128(uint(lenMinus))
+
+		switch obu.SizeLeb128(payloadSize) {
+		case 4:
+			result[offset] = byte(payloadSize >> 24)
+			offset++
+			fallthrough
+		case 3:
+			result[offset] = byte(payloadSize >> 16)
+			offset++
+			fallthrough
+		case 2:
+			result[offset] = byte(payloadSize >> 8)
+			offset++
+			fallthrough
+		case 1:
+			result[offset] = byte(payloadSize)
+			offset++
+		}
+
+		copy(result[offset:], OBUs[i][1:])
+		offset += lenMinus
+	}
+
+	return result, nil
+}

codecs/av1/frame/av1_sample_buffer_support_test.go

@@ -0,0 +1,149 @@
+package frame
+
+import (
+	"bytes"
+	"github.com/pion/rtp"
+	"github.com/pion/rtp/codecs/av1/obu"
+	"github.com/pion/webrtc/v3/pkg/media/samplebuilder"
+	"testing"
+	"time"
+)
+
+func buildAv1Payload(data byte, padding int) []byte {
+	dataSize := 0
+
+	if data > 0 {
+		dataSize = 1
+	}
+
+	payloadSize := obu.EncodeLEB128(uint(1 + padding + dataSize))
+	result := make([]byte, 3+obu.SizeLeb128(payloadSize))
+
+	result[0] = 0 // AV1 RTP header
+
+	offset := 1
+
+	switch obu.SizeLeb128(payloadSize) {
+	case 4:
+		result[offset] = byte(payloadSize >> 24)
+		offset++
+		fallthrough
+	case 3:
+		result[offset] = byte(payloadSize >> 16)
+		offset++
+		fallthrough
+	case 2:
+		result[offset] = byte(payloadSize >> 8)
+		offset++
+		fallthrough
+	case 1:
+		result[offset] = byte(payloadSize)
+		offset++
+	}
+
+	result[offset] = 0 // OBU HEADER
+
+	if dataSize > 0 {
+		offset++
+		result[offset] = data
+	}
+	return append(result, make([]byte, padding)...)
+}
+func buildAv1Packages(seqNo *uint16, timestamp *uint32, padding int) []*rtp.Packet {
+	s := *seqNo
+	t := *timestamp
+
+	*timestamp += 1800
+	*seqNo += 5
+
+	return []*rtp.Packet{
+		{Header: rtp.Header{SequenceNumber: s, Timestamp: t}, Payload: buildAv1Payload(1, 0)},
+		{Header: rtp.Header{SequenceNumber: s + 1, Timestamp: t}, Payload: buildAv1Payload(2, 0)},
+		{Header: rtp.Header{SequenceNumber: s + 2, Timestamp: t}, Payload: buildAv1Payload(3, 0)},
+		{Header: rtp.Header{SequenceNumber: s + 3, Timestamp: t}, Payload: buildAv1Payload(4, 0)},
+		{Header: rtp.Header{SequenceNumber: s + 4, Timestamp: t, Marker: true}, Payload: buildAv1Payload(5, padding)},
+	}
+}
+func TestAV1SampleBufferSupport(t *testing.T) {
+
+	assembledAv1Frame := []byte{2, 1, 1, 2, 1, 2, 2, 1, 3, 2, 1, 4, 2, 1, 5}
+	t.Run("AV1 Sample Buffer returning OBU stream", func(t *testing.T) {
+		videoStreamBuilder := samplebuilder.New(100, &AV1PacketSampleBufferSupport{}, 90000,
+			samplebuilder.WithMaxTimeDelay(time.Millisecond*100))
+		var seqNo uint16 = 0
+		var timestamp uint32 = 0
+
+		for i := 0; i < 4; i++ {
+			for _, pkt := range buildAv1Packages(&seqNo, &timestamp, 0) {
+				sample := videoStreamBuilder.Pop()
+				if nil != sample {
+					if !bytes.Equal(sample.Data, assembledAv1Frame) {
+						t.Fatal("issue in unmarshalling")
+					}
+				}
+				videoStreamBuilder.Push(pkt)
+			}
+		}
+
+		for i := 1; i < 16400; i++ { // check OBU len up to 3 bytes
+			for _, pkt := range buildAv1Packages(&seqNo, &timestamp, i) {
+				sample := videoStreamBuilder.Pop()
+				if nil != sample {
+
+					if !bytes.Equal(sample.Data[0:12], assembledAv1Frame[0:12]) {
+						t.Fatal("issue in unmarshalling")
+					}
+
+					if sample.Data[len(sample.Data)-i] != 5 {
+						t.Fatal("issue in unmarshalling")
+					}
+
+					if i > 0 {
+						padding := make([]byte, i-1)
+						if !bytes.Equal(sample.Data[len(sample.Data)-i+1:], padding) {
+							t.Fatal("issue in unmarshalling")
+						}
+					}
+				}
+				videoStreamBuilder.Push(pkt)
+			}
+		}
+	})
+}
+
+func buildHeaderOnlyAv1Packets(seqNo *uint16, timestamp *uint32, padding int) []*rtp.Packet {
+	s := *seqNo
+	t := *timestamp
+
+	*timestamp += 1800
+	*seqNo += 3
+
+	return []*rtp.Packet{
+		// Two header-only OBUs
+		{Header: rtp.Header{SequenceNumber: s, Timestamp: t}, Payload: buildAv1Payload(0, 0)},
+		{Header: rtp.Header{SequenceNumber: s + 1, Timestamp: t}, Payload: buildAv1Payload(0, 0)},
+		{Header: rtp.Header{SequenceNumber: s + 2, Timestamp: t, Marker: true}, Payload: buildAv1Payload(5, padding)},
+	}
+}
+
+func TestAV1SampleBufferSupport_OBUWIthoutPayload(t *testing.T) {
+	assembledAv1Frame := []byte{2, 0, 2, 0, 2, 1, 5}
+	t.Run("AV1 Sample Buffer with header only OBU elements", func(t *testing.T) {
+		videoStreamBuilder := samplebuilder.New(100, &AV1PacketSampleBufferSupport{}, 90000,
+			samplebuilder.WithMaxTimeDelay(time.Millisecond*100))
+		var seqNo uint16 = 0
+		var timestamp uint32 = 0
+
+		for i := 0; i < 4; i++ {
+			for _, pkt := range buildHeaderOnlyAv1Packets(&seqNo, &timestamp, 0) {
+				sample := videoStreamBuilder.Pop()
+				if nil != sample {
+					if !bytes.Equal(sample.Data, assembledAv1Frame) {
+						t.Fatal("issue in unmarshalling")
+					}
+				}
+				videoStreamBuilder.Push(pkt)
+			}
+		}
+	})
+}

codecs/av1/obu/leb128.go

@@ -68,6 +68,17 @@
 	return 0, 0, ErrFailedToReadLEB128
 }
 
+func SizeLeb128(leb128 uint) uint {
+	if (leb128 >> 24) > 0 {
+		return 4
+	} else if (leb128 >> 16) > 0 {
+		return 3
+	} else if (leb128 >> 8) > 0 {
+		return 2
+	}
+	return 1
+}
+
 // WriteToLeb128 writes a uint to a LEB128 encoded byte slice.
 func WriteToLeb128(in uint) []byte {
 	b := make([]byte, 10)

go.mod

@@ -2,4 +2,7 @@ module github.com/pion/rtp
 
 go 1.19
 
-require github.com/pion/randutil v0.1.0
+require (
+	github.com/pion/randutil v0.1.0
+	github.com/pion/webrtc/v3 v3.2.21
+)