From 0ab12600c51cf2edc410f4ef9d527897ee0fed60 Mon Sep 17 00:00:00 2001
From: Juliusz Chroboczek <jch@irif.fr>
Date: Mon, 15 Apr 2024 18:38:21 +0200
Subject: [PATCH] Make parsing of the OBU list in AV1Packet optional

In Pion, the (*XXXPacket).Unmarshal operation is supposed to be
fast and perform no allocation, as it is sometimes used to check
just a single flag in the packet header.  The AV1Packet
implementation breaks this property by parsing the whole list
of OBUs at Unmarshal time, even though it is likely to be unneeded.

The proper solution would be to remove the parsing of OBUs from
Unmarshal, for consistency with the other packet formats.
Unfortunately, this would break backwards compatibility.  We
therefore introduce a field DontParseOBUs which controls the
parsing.  Unfortunately, this means that the inefficient path
is the default, which smells really badly, but there's not much
we can do about it at this stage.
---
 codecs/av1_packet.go | 43 +++++++++++++++++++++++++++++++++++--------
 1 file changed, 35 insertions(+), 8 deletions(-)

diff --git a/codecs/av1_packet.go b/codecs/av1_packet.go
index 3a78b90..622c8b1 100644
--- a/codecs/av1_packet.go
+++ b/codecs/av1_packet.go
@@ -131,11 +131,22 @@ type AV1Packet struct {
 	// N: MUST be set to 1 if the packet is the first packet of a coded video sequence, and MUST be set to 0 otherwise.
 	N bool
 
+	// If true, don't populate the OBUElements field.
+	dontParseBody bool
+
 	// Each AV1 RTP Packet is a collection of OBU Elements. Each OBU Element may be a full OBU, or just a fragment of one.
 	// AV1Frame provides the tools to construct a collection of OBUs from a collection of OBU Elements
+	//
+	// Deprecated: use the ParseOBUElements method instead.
 	OBUElements [][]byte
 }
 
+// SetDontParseBody specifies whether to skip populating the OBUElements field.
+// Applications should set this to true and use the ParseBody function instead.
+func (p *AV1Packet) SetDontParseBody(dont bool) {
+	p.dontParseBody = true
+}
+
 // Unmarshal parses the passed byte slice and stores the result in the AV1Packet this method is called upon
 func (p *AV1Packet) Unmarshal(payload []byte) ([]byte, error) {
 	if payload == nil {
@@ -153,13 +164,28 @@ func (p *AV1Packet) Unmarshal(payload []byte) ([]byte, error) {
 		return nil, errIsKeyframeAndFragment
 	}
 
-	currentIndex := uint(1)
-	p.OBUElements = [][]byte{}
+	if !p.dontParseBody {
+		elts, err := p.ParseBody(payload[1:])
+		if err != nil {
+			return nil, err
+		}
+		p.OBUElements = elts
+	}
+
+	return payload[1:], nil
+}
+
+// ParseOBUElements returns the list of OBUs in p
+func (p *AV1Packet) ParseBody(payload []byte) ([][]byte, error) {
+	if p.OBUElements != nil {
+		return p.OBUElements, nil
+	}
+
+	elts := [][]byte{}
 
-	var (
-		obuElementLength, bytesRead uint
-		err                         error
-	)
+	obuElementLength := uint(0)
+	bytesRead := uint(0)
+	currentIndex := uint(0)
 	for i := 1; ; i++ {
 		if currentIndex == uint(len(payload)) {
 			break
@@ -170,6 +196,7 @@ func (p *AV1Packet) Unmarshal(payload []byte) ([]byte, error) {
 			bytesRead = 0
 			obuElementLength = uint(len(payload)) - currentIndex
 		} else {
+			var err error
 			obuElementLength, bytesRead, err = obu.ReadLeb128(payload[currentIndex:])
 			if err != nil {
 				return nil, err
@@ -180,9 +207,9 @@ func (p *AV1Packet) Unmarshal(payload []byte) ([]byte, error) {
 		if uint(len(payload)) < currentIndex+obuElementLength {
 			return nil, errShortPacket
 		}
-		p.OBUElements = append(p.OBUElements, payload[currentIndex:currentIndex+obuElementLength])
+		elts = append(elts, payload[currentIndex:currentIndex+obuElementLength])
 		currentIndex += obuElementLength
 	}
 
-	return payload[1:], nil
+	return elts, nil
 }