From 4cbf2d3ebdd8d5fe3ced2e224bdaf8c323fb419a Mon Sep 17 00:00:00 2001 From: cnderrauber Date: Tue, 3 Nov 2026 16:17:33 +0800 Subject: [PATCH] Shrink NALU size when stripping h264/h265 aggregation packet trailer (#5125) Fix #5082 --- pkg/sfu/downtrack.go | 8 +- pkg/sfu/packettrailer/packet_trailer.go | 50 ++++++- pkg/sfu/packettrailer/packet_trailer_test.go | 145 ++++++++++++++++++- 3 files changed, 199 insertions(+), 4 deletions(-) diff --git a/pkg/sfu/downtrack.go b/pkg/sfu/downtrack.go index f41567009..9ac8d6baa 100644 --- a/pkg/sfu/downtrack.go +++ b/pkg/sfu/downtrack.go @@ -1078,7 +1078,7 @@ func (d *DownTrack) WriteRTP(extPkt *buffer.ExtPacket, layer int32) int32 { trailerStripped := 0 if d.params.StripPacketTrailer { - if trailerStripped = packettrailer.StripTrailer(payload, tp.marker || tp.isEndOfLayerFrame); trailerStripped > 0 { + if trailerStripped = packettrailer.StripTrailer(payload, d.Mime(), tp.marker || tp.isEndOfLayerFrame); trailerStripped > 0 { payload = payload[:len(payload)-trailerStripped] } } @@ -2349,6 +2349,12 @@ func (d *DownTrack) retransmitPacket(epm *extPacketMeta, sourcePkt []byte, isPro "trailerStripped", epm.trailerStripped, "payloadSize", len(payload)-rtxOffset, ) + } else if !packettrailer.ShrinkLastAggregatedNALU(payload[rtxOffset:], d.Mime(), int(epm.trailerStripped)) { + d.params.Logger.Warnw( + "could not shrink aggregated NAL unit for packet trailer", nil, + "trailerStripped", epm.trailerStripped, + "payloadSize", len(payload)-rtxOffset, + ) } else { payload = payload[:len(payload)-int(epm.trailerStripped)] } diff --git a/pkg/sfu/packettrailer/packet_trailer.go b/pkg/sfu/packettrailer/packet_trailer.go index d86a4cddc..1aada46a7 100644 --- a/pkg/sfu/packettrailer/packet_trailer.go +++ b/pkg/sfu/packettrailer/packet_trailer.go @@ -14,12 +14,21 @@ package packettrailer +import ( + "encoding/binary" + + "github.com/livekit/protocol/codecs/mime" +) + var Magic = [4]byte{'L', 'K', 'T', 'S'} const ( xorByte = 0xFF envelopeSize = 5 // 1B trailer_len + 4B magic + + h264STAPA = 0x18 + h265AP = 0x30 ) // StripTrailer returns the number of bytes to strip from the end of an RTP @@ -28,7 +37,11 @@ const ( // immediately before it. isEndOfFrame must be set only for packets ending an // encoded frame, i. e. where a trailer could have been appended. Returns 0 if // absent or ineligible. -func StripTrailer(payload []byte, isEndOfFrame bool) int { +// +// The trailer is appended to the encoded frame before packetization, so in an +// aggregation packet it is counted in the size of the last NAL unit. That size +// is reduced in place, see ShrinkLastAggregatedNALU. +func StripTrailer(payload []byte, mimeType mime.MimeType, isEndOfFrame bool) int { if !isEndOfFrame || len(payload) < envelopeSize { return 0 } @@ -44,5 +57,40 @@ func StripTrailer(payload []byte, isEndOfFrame bool) int { return 0 } + if !ShrinkLastAggregatedNALU(payload, mimeType, trailerLen) { + return 0 + } + return trailerLen } + +// ShrinkLastAggregatedNALU reduces, in place, the size field of the last NAL +// unit in an H.264 STAP-A or H.265 aggregation packet by n bytes, so that the +// packet stays well formed once n bytes are cut from its end. Other payloads +// are left untouched. Returns false if the payload is an aggregation packet +// whose last NAL unit does not end the payload or is not longer than n. +func ShrinkLastAggregatedNALU(payload []byte, mimeType mime.MimeType, n int) bool { + var offset int + switch { + case mimeType == mime.MimeTypeH264 && len(payload) >= 1 && payload[0]&0x1F == h264STAPA: + offset = 1 + case mimeType == mime.MimeTypeH265 && len(payload) >= 2 && (payload[0]>>1)&0x3F == h265AP: + offset = 2 + default: + return true + } + + for offset+2 <= len(payload) { + size := int(binary.BigEndian.Uint16(payload[offset:])) + end := offset + 2 + size + if end == len(payload) { + if size <= n { + return false + } + binary.BigEndian.PutUint16(payload[offset:], uint16(size-n)) + return true + } + offset = end + } + return false +} diff --git a/pkg/sfu/packettrailer/packet_trailer_test.go b/pkg/sfu/packettrailer/packet_trailer_test.go index 653b71fa2..e8da7a409 100644 --- a/pkg/sfu/packettrailer/packet_trailer_test.go +++ b/pkg/sfu/packettrailer/packet_trailer_test.go @@ -15,8 +15,13 @@ package packettrailer import ( + "bytes" "encoding/binary" "testing" + + "github.com/pion/rtp/codecs" + + "github.com/livekit/protocol/codecs/mime" ) const ( @@ -78,7 +83,7 @@ func makeTimestampOnlyTrailer(timestampUs int64) []byte { } func TestStripTrailer(t *testing.T) { - fullTrailerSize := 21 // (1+1+8) + (1+1+4) + 5 + fullTrailerSize := 21 // (1+1+8) + (1+1+4) + 5 tsOnlyTrailerSize := 15 // (1+1+8) + 5 tests := []struct { @@ -173,10 +178,146 @@ func TestStripTrailer(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got := StripTrailer(tt.payload, tt.marker) + got := StripTrailer(tt.payload, mime.MimeTypeVP8, tt.marker) if got != tt.wantStrip { t.Errorf("StripTrailer() = %d, want %d", got, tt.wantStrip) } }) } } + +// makeAggregate builds an aggregation packet with the given payload header and +// NAL units, appending trailer to the last NAL unit and counting it in its size. +func makeAggregate(header []byte, nalus [][]byte, trailer []byte) []byte { + p := append([]byte{}, header...) + for i, nalu := range nalus { + size := len(nalu) + if i == len(nalus)-1 { + size += len(trailer) + } + p = binary.BigEndian.AppendUint16(p, uint16(size)) + p = append(p, nalu...) + } + return append(p, trailer...) +} + +func TestStripTrailerAggregationPacket(t *testing.T) { + h264NALUs := [][]byte{ + {0x67, 0x42, 0xC0, 0x1F}, // SPS + {0x68, 0xCE, 0x3C, 0x80}, // PPS + {0x65, 0x88, 0x84, 0x00, 0x33}, // IDR slice + } + h265NALUs := [][]byte{ + {0x40, 0x01, 0x0C, 0x01}, // VPS + {0x42, 0x01, 0x01, 0x01}, // SPS + {0x44, 0x01, 0xC1, 0x72}, // PPS + {0x26, 0x01, 0xAF, 0x08, 0x40}, // IDR slice + } + trailer := makeTrailer(1700000000000000, 42) + + unmarshalH264 := func(p []byte) ([]byte, error) { + return (&codecs.H264Packet{}).Unmarshal(p) + } + unmarshalH265 := func(p []byte) ([]byte, error) { + d := &codecs.H265Depacketizer{} + return d.Unmarshal(p) + } + + tests := []struct { + name string + mimeType mime.MimeType + payload []byte + want []byte + unmarshal func([]byte) ([]byte, error) + }{ + { + name: "H.264 STAP-A", + mimeType: mime.MimeTypeH264, + payload: makeAggregate([]byte{0x78}, h264NALUs, trailer), + want: makeAggregate([]byte{0x78}, h264NALUs, nil), + unmarshal: unmarshalH264, + }, + { + name: "H.264 STAP-A with a single NAL unit", + mimeType: mime.MimeTypeH264, + payload: makeAggregate([]byte{0x78}, h264NALUs[2:], trailer), + want: makeAggregate([]byte{0x78}, h264NALUs[2:], nil), + unmarshal: unmarshalH264, + }, + { + name: "H.265 aggregation packet", + mimeType: mime.MimeTypeH265, + payload: makeAggregate([]byte{0x60, 0x01}, h265NALUs, trailer), + want: makeAggregate([]byte{0x60, 0x01}, h265NALUs, nil), + unmarshal: unmarshalH265, + }, + { + name: "H.264 single NAL unit", + mimeType: mime.MimeTypeH264, + payload: append(append([]byte{}, h264NALUs[2]...), trailer...), + want: h264NALUs[2], + unmarshal: unmarshalH264, + }, + { + name: "H.264 FU-A end fragment", + mimeType: mime.MimeTypeH264, + payload: append([]byte{0x7C, 0x45, 0x01, 0x02, 0x03}, trailer...), + want: []byte{0x7C, 0x45, 0x01, 0x02, 0x03}, + unmarshal: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + n := StripTrailer(tt.payload, tt.mimeType, true) + if n != len(trailer) { + t.Fatalf("StripTrailer() = %d, want %d", n, len(trailer)) + } + got := tt.payload[:len(tt.payload)-n] + if !bytes.Equal(got, tt.want) { + t.Fatalf("stripped payload = %x, want %x", got, tt.want) + } + if tt.unmarshal != nil { + if _, err := tt.unmarshal(got); err != nil { + t.Fatalf("depacketizer rejected stripped payload: %v", err) + } + } + }) + } +} + +func TestStripTrailerMalformedAggregationPacket(t *testing.T) { + trailer := makeTrailer(1700000000000000, 42) + nalus := [][]byte{{0x67, 0x42, 0xC0, 0x1F}, {0x65, 0x88, 0x84}} + + tests := []struct { + name string + payload []byte + }{ + { + // trailer not counted in the last NAL unit, e.g. split off by a start code inside it + name: "last NAL unit does not end the payload", + payload: append(makeAggregate([]byte{0x78}, nalus, nil), trailer...), + }, + { + name: "trailer is the whole last NAL unit", + payload: makeAggregate([]byte{0x78}, [][]byte{nalus[0], {}}, trailer), + }, + { + name: "truncated size field", + payload: append([]byte{0x78, 0x00}, trailer...), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + orig := append([]byte{}, tt.payload...) + if n := StripTrailer(tt.payload, mime.MimeTypeH264, true); n != 0 { + t.Fatalf("StripTrailer() = %d, want 0", n) + } + if !bytes.Equal(tt.payload, orig) { + t.Fatalf("payload modified: %x, want %x", tt.payload, orig) + } + }) + } +}