From 18f83daa9a2162b0288eac53e58007b7c894fa10 Mon Sep 17 00:00:00 2001 From: Kevin Caffrey Date: Fri, 24 Jul 2026 15:15:04 -0400 Subject: [PATCH] Make GetRTCPPackets and GetRTPHeader nil-safe Allow both `GetRTCPPackets()` and `GetRTPHeader()` to be called on nil attributes without panicking. Instead, they will unmarshal but not cache the results when called on a nil receiver. --- attributes.go | 15 ++++-- attributes_test.go | 131 +++++++++++++++++++++++++++++---------------- 2 files changed, 95 insertions(+), 51 deletions(-) diff --git a/attributes.go b/attributes.go index 1eb4f109..0aa6f9de 100644 --- a/attributes.go +++ b/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 } diff --git a/attributes_test.go b/attributes_test.go index dd0c107c..5fc01a7b 100644 --- a/attributes_test.go +++ b/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) }) }