Make GetRTCPPackets and GetRTPHeader nil-safe

attributes.go

@@ -33,7 +33,9 @@ func (a Attributes) Set(key any, val any) {
 }
 
 // GetRTPHeader gets the RTP header if present. If it is not present, it will be
-// unmarshalled from the raw byte slice and stored in the attributes.
+// unmarshalled from the raw byte slice and stored in the attributes. If the
+// provided attributes is nil, GetRTPHeader will always unmarshal but will not
+// cache the result.
 func (a Attributes) GetRTPHeader(raw []byte) (*rtp.Header, error) {
 	if val, ok := a[rtpHeaderKey]; ok {
 		if header, ok := val.(*rtp.Header); ok {
@@ -46,14 +48,17 @@ func (a Attributes) GetRTPHeader(raw []byte) (*rtp.Header, error) {
 	if _, err := header.Unmarshal(raw); err != nil {
 		return nil, err
 	}
-	a[rtpHeaderKey] = header
+	if a != nil {
+		a[rtpHeaderKey] = header
+	}
 
 	return header, nil
 }
 
 // GetRTCPPackets gets the RTCP packets if present. If the packet slice is not
 // present, it will be unmarshalled from the raw byte slice and stored in the
-// attributes.
+// attributes. If the provided attributes is nil, GetRTCPPackets will always
+// unmarshal but will not cache the result.
 func (a Attributes) GetRTCPPackets(raw []byte) ([]rtcp.Packet, error) {
 	if val, ok := a[rtcpPacketsKey]; ok {
 		if packets, ok := val.([]rtcp.Packet); ok {
@@ -66,7 +71,9 @@ func (a Attributes) GetRTCPPackets(raw []byte) ([]rtcp.Packet, error) {
 	if err != nil {
 		return nil, err
 	}
-	a[rtcpPacketsKey] = pkts
+	if a != nil {
+		a[rtcpPacketsKey] = pkts
+	}
 
 	return pkts, nil
 }

attributes_test.go

@@ -39,45 +39,67 @@ func TestAttributesGetRTPHeader(t *testing.T) {
 	})
 
 	t.Run("NotPresent", func(t *testing.T) {
-		attributes := Attributes{}
-		hdr := &rtp.Header{
-			Version:          0,
-			Padding:          false,
-			Extension:        false,
-			Marker:           false,
-			PayloadType:      0,
-			SequenceNumber:   0,
-			Timestamp:        0,
-			SSRC:             0,
-			ExtensionProfile: 0,
-			Extensions:       nil,
+		for msg, nilAttributes := range map[string]bool{"NilAttributes": true, "EmptyAttributes": false} {
+			t.Run(msg, func(t *testing.T) {
+				var attributes Attributes
+				if !nilAttributes {
+					attributes = Attributes{}
+				}
+				hdr := &rtp.Header{
+					Version:          0,
+					Padding:          false,
+					Extension:        false,
+					Marker:           false,
+					PayloadType:      0,
+					SequenceNumber:   0,
+					Timestamp:        0,
+					SSRC:             0,
+					ExtensionProfile: 0,
+					Extensions:       nil,
+				}
+				buf, err := hdr.Marshal()
+				assert.NoError(t, err)
+				header, err := attributes.GetRTPHeader(buf)
+				assert.NoError(t, err)
+				assert.Equal(t, hdr, header)
+
+				// If attributes were provided, make sure GetRTPHeader returns cached results.
+				if !nilAttributes {
+					amortizedAllocs := testing.AllocsPerRun(100, func() {
+						_, _ = attributes.GetRTPHeader(buf)
+					})
+					assert.Zero(t, amortizedAllocs)
+				}
+			})
 		}
-		buf, err := hdr.Marshal()
-		assert.NoError(t, err)
-		header, err := attributes.GetRTPHeader(buf)
-		assert.NoError(t, err)
-		assert.Equal(t, hdr, header)
 	})
 
 	t.Run("NotPresentFromFullRTPPacket", func(t *testing.T) {
-		attributes := Attributes{}
-		pkt := &rtp.Packet{Header: rtp.Header{
-			Version:          0,
-			Padding:          false,
-			Extension:        false,
-			Marker:           false,
-			PayloadType:      0,
-			SequenceNumber:   0,
-			Timestamp:        0,
-			SSRC:             0,
-			ExtensionProfile: 0,
-			Extensions:       nil,
-		}, Payload: make([]byte, 1000)}
-		buf, err := pkt.Marshal()
-		assert.NoError(t, err)
-		header, err := attributes.GetRTPHeader(buf)
-		assert.NoError(t, err)
-		assert.Equal(t, &pkt.Header, header)
+		for msg, nilAttributes := range map[string]bool{"NilAttributes": true, "EmptyAttributes": false} {
+			t.Run(msg, func(t *testing.T) {
+				var attributes Attributes
+				if !nilAttributes {
+					attributes = Attributes{}
+				}
+				pkt := &rtp.Packet{Header: rtp.Header{
+					Version:          0,
+					Padding:          false,
+					Extension:        false,
+					Marker:           false,
+					PayloadType:      0,
+					SequenceNumber:   0,
+					Timestamp:        0,
+					SSRC:             0,
+					ExtensionProfile: 0,
+					Extensions:       nil,
+				}, Payload: make([]byte, 1000)}
+				buf, err := pkt.Marshal()
+				assert.NoError(t, err)
+				header, err := attributes.GetRTPHeader(buf)
+				assert.NoError(t, err)
+				assert.Equal(t, &pkt.Header, header)
+			})
+		}
 	})
 }
 
@@ -110,18 +132,33 @@ func TestAttributesGetRTCPPackets(t *testing.T) {
 	})
 
 	t.Run("NotPresent", func(t *testing.T) {
-		attributes := Attributes{}
-		sr := &rtcp.SenderReport{
-			SSRC:        0,
-			NTPTime:     0,
-			RTPTime:     0,
-			PacketCount: 0,
-			OctetCount:  0,
+		for msg, nilAttributes := range map[string]bool{"NilAttributes": true, "EmptyAttributes": false} {
+			t.Run(msg, func(t *testing.T) {
+				var attributes Attributes
+				if !nilAttributes {
+					attributes = Attributes{}
+				}
+				sr := &rtcp.SenderReport{
+					SSRC:        0,
+					NTPTime:     0,
+					RTPTime:     0,
+					PacketCount: 0,
+					OctetCount:  0,
+				}
+				buf, err := sr.Marshal()
+				assert.NoError(t, err)
+				packets, err := attributes.GetRTCPPackets(buf)
+				assert.NoError(t, err)
+				assert.Equal(t, []rtcp.Packet{sr}, packets)
+
+				// If attributes were provided, make sure GetRTCPPackets returns cached results.
+				if !nilAttributes {
+					amortizedAllocs := testing.AllocsPerRun(100, func() {
+						_, _ = attributes.GetRTCPPackets(buf)
+					})
+					assert.Zero(t, amortizedAllocs)
+				}
+			})
 		}
-		buf, err := sr.Marshal()
-		assert.NoError(t, err)
-		packets, err := attributes.GetRTCPPackets(buf)
-		assert.NoError(t, err)
-		assert.Equal(t, []rtcp.Packet{sr}, packets)
 	})
 }