Fix AV1 IVF writer using RTP depacketizer

1. Replaced the OBU reader with the AV1 depacketizer to actually handle RTP input. 2. Now explicitly constructing OBUTemporalDelimiter to ensure proper separation of frames.
pion · Feb 24, 2025 · 46c0b5e · 46c0b5e
1 parent 44062a7
commit 46c0b5e
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 25 deletions.
diff --git a/pkg/media/ivfwriter/ivfwriter.go b/pkg/media/ivfwriter/ivfwriter.go
@@ -12,7 +12,7 @@ import (
 
 	"github.com/pion/rtp"
 	"github.com/pion/rtp/codecs"
-	"github.com/pion/rtp/codecs/av1/frame"
+	"github.com/pion/rtp/codecs/av1/obu"
 )
 
 var (
@@ -44,7 +44,7 @@ type (
 		currentFrame []byte
 
 		// AV1
-		av1Frame frame.AV1
+		av1Depacketizer *codecs.AV1Depacketizer
 	}
 )
 
@@ -97,7 +97,6 @@ func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
 	if writer.codec == codecUnset {
 		writer.codec = codecVP8
 	}
-
 	if err := writer.writeHeader(); err != nil {
 		return nil, err
 	}
@@ -228,21 +227,38 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error { //nolint:cyclop, gocogn
 		}
 		i.currentFrame = nil
 	case codecAV1:
-		av1Packet := &codecs.AV1Packet{}
-		if _, err := av1Packet.Unmarshal(packet.Payload); err != nil {
-			return err
+		if i.av1Depacketizer == nil {
+			i.av1Depacketizer = &codecs.AV1Depacketizer{}
 		}
 
-		obus, err := i.av1Frame.ReadFrames(av1Packet)
+		payload, err := i.av1Depacketizer.Unmarshal(packet.Payload)
 		if err != nil {
 			return err
 		}
 
-		for j := range obus {
-			if err := i.writeFrame(obus[j], relativeTstampMs); err != nil {
-				return err
-			}
+		if !i.seenKeyFrame && !i.av1Depacketizer.N {
+			return nil
+		}
+
+		i.seenKeyFrame = true
+
+		i.currentFrame = append(i.currentFrame, payload...)
+
+		if !packet.Marker {
+			return nil
+		}
+
+		delimiter := obu.Header{
+			Type:         obu.OBUTemporalDelimiter,
+			HasSizeField: true,
 		}
+		frame := append(delimiter.Marshal(), 0)
+		frame = append(frame, i.currentFrame...)
+
+		if err := i.writeFrame(frame, relativeTstampMs); err != nil {
+			return err
+		}
+		i.currentFrame = nil
 	default:
 		return errCodecUnset
 	}

diff --git a/pkg/media/ivfwriter/ivfwriter_test.go b/pkg/media/ivfwriter/ivfwriter_test.go
@@ -263,15 +263,22 @@ func TestIVFWriter_AV1(t *testing.T) {
 		writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
 		assert.NoError(t, err)
 
-		assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x00, 0x01, 0xFF}}))
+		assert.NoError(
+			t,
+			writer.WriteRTP(
+				&rtp.Packet{
+					Header: rtp.Header{Marker: true},
+					// N = 1, Length = 1, OBU_TYPE = 4
+					Payload: []byte{0x08, 0x01, 0x20},
+				}),
+		)
+
 		assert.NoError(t, writer.Close())
 		assert.Equal(t, buffer.Bytes(), []byte{
-			0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20,
-			0x0, 0x41, 0x56, 0x30, 0x31, 0x80, 0x2,
-			0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0,
-			0x0, 0x0, 0x84, 0x3, 0x0, 0x0, 0x0, 0x0,
-			0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
-			0x0, 0x0, 0x0, 0x0, 0x0, 0xff,
+			0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31,
+			0x80, 0x2, 0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84,
+			0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x4, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+			0x0, 0x0, 0x0, 0x0, 0x0, 0x12, 0x0, 0x22, 0x0,
 		})
 	})
 
@@ -281,8 +288,13 @@ func TestIVFWriter_AV1(t *testing.T) {
 		writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
 		assert.NoError(t, err)
 
-		for _, p := range [][]byte{{0x40, 0x02, 0x00, 0x01}, {0xc0, 0x02, 0x02, 0x03}, {0xc0, 0x02, 0x04, 0x04}} {
-			assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: p}))
+		for _, p := range [][]byte{
+			{0x48, 0x02, 0x00, 0x01}, // Y=true
+			{0xc0, 0x02, 0x02, 0x03}, // Z=true, Y=true
+			{0xc0, 0x02, 0x04, 0x04}, // Z=true, Y=true
+			{0x80, 0x01, 0x05},       // Z=true, Y=false (But we still don't set Marker to true)
+		} {
+			assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: p, Header: rtp.Header{Marker: false}}))
 			assert.Equal(t, buffer.Bytes(), []byte{
 				0x44, 0x4b, 0x49, 0x46, 0x0,
 				0x0, 0x20, 0x0, 0x41, 0x56, 0x30,
@@ -292,15 +304,52 @@ func TestIVFWriter_AV1(t *testing.T) {
 				0x0, 0x0,
 			})
 		}
-		assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x80, 0x01, 0x05}}))
+		assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0x01, 0x20}, Header: rtp.Header{Marker: true}}))
 		assert.Equal(t, buffer.Bytes(), []byte{
-			0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31, 0x80,
-			0x2, 0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84, 0x3, 0x0, 0x0,
-			0x0, 0x0, 0x0, 0x0, 0x7, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
-			0x0, 0x1, 0x2, 0x3, 0x4, 0x4, 0x5,
+			0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31, 0x80, 0x2, 0xe0, 0x1, 0x1e,
+			0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84, 0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0xc, 0x0, 0x0, 0x0,
+			0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x12, 0x0, 0x2, 0x6, 0x1, 0x2, 0x3, 0x4, 0x4, 0x5, 0x22, 0x0,
 		})
 		assert.NoError(t, writer.Close())
 	})
+
+	t.Run("Invalid OBU", func(t *testing.T) {
+		buffer := &bytes.Buffer{}
+
+		writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
+		assert.NoError(t, err)
+
+		assert.Error(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0x02, 0xff}}))
+		assert.Error(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x08, 0x01, 0xff}}))
+	})
+
+	t.Run("Skips middle sequence start", func(t *testing.T) {
+		buffer := &bytes.Buffer{}
+
+		writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
+		assert.NoError(t, err)
+
+		assert.NoError(t, writer.WriteRTP(&rtp.Packet{Header: rtp.Header{Marker: true}, Payload: []byte{0x00, 0x01, 0x20}}))
+
+		assert.NoError(
+			t,
+			writer.WriteRTP(
+				&rtp.Packet{
+					Header: rtp.Header{Marker: true},
+					// N = 1, Length = 1, OBU_TYPE = 4
+					Payload: []byte{0x08, 0x01, 0x20},
+				},
+			),
+		)
+
+		assert.NoError(t, writer.Close())
+		assert.Equal(t, buffer.Bytes(), []byte{
+			0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31,
+			0x80, 0x2, 0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84,
+			0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x4, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+			0x0, 0x0, 0x0, 0x0, 0x0, 0x12, 0x0, 0x22, 0x0,
+		})
+	})
 }
 
 func TestIVFWriter_VP9(t *testing.T) {