Cache SPS/PPS/VPS in H265Payloader

codecs/h265_packet.go

@@ -40,6 +40,12 @@ const (
 	h265NaluPACIPacketType         = 50
 	h265AggregatedPacketMaxSize    = ^uint16(0)
 	h265AggregatedPacketLengthSize = 2
+
+	h265NaluVpsType    = 32
+	h265NaluSpsType    = 33
+	h265NaluPpsType    = 34
+	h265NaluAudType    = 35
+	h265NaluFillerType = 38
 )
 
 // H265NALUHeader is a H265 NAL Unit Header.
@@ -1723,12 +1729,16 @@ func (p *H265Packet) Packet() isH265Packet {
 // H265Payloader payloads H265 packets.
 type H265Payloader struct {
 	// Deprecated: Has no effect.
-	AddDONL         bool
+	AddDONL bool
+	// SkipAggregation disables H265 aggregation packets.
 	SkipAggregation bool
+	vpsNalu         *h265SingleNALUnitPacket
+	spsNalu         *h265SingleNALUnitPacket
+	ppsNalu         *h265SingleNALUnitPacket
 }
 
 // Payload fragments a H265 packet across one or more byte arrays.
-func (p *H265Payloader) Payload(mtu uint16, payload []byte) [][]byte { // nolint:cyclop
+func (p *H265Payloader) Payload(mtu uint16, payload []byte) [][]byte { // nolint:cyclop,gocognit
 	// SampleBuilder reuses the payload buffer so this is required
 	tmp := make([]byte, len(payload))
 	copy(tmp, payload)
@@ -1775,33 +1785,84 @@ func (p *H265Payloader) Payload(mtu uint16, payload []byte) [][]byte { // nolint
 			nalu[2:],
 		}
 
-		if len(nalu) > int(mtu) { // nolint: nestif
-			flushBuffer()
-			fragments, err := newH265FragmentationPackets(mtu, &packet)
-			if err != nil {
-				return
-			}
-			for _, fragment := range fragments {
-				payloads = append(payloads, fragment.serialize(make([]byte, 0)))
-			}
-		} else {
-			if p.SkipAggregation {
+		if p.SkipAggregation {
+			p.vpsNalu = nil
+			p.spsNalu = nil
+			p.ppsNalu = nil
+			if len(nalu) > int(mtu) {
+				flushBuffer()
+				fragments, err := newH265FragmentationPackets(mtu, &packet)
+				if err != nil {
+					return
+				}
+				for _, fragment := range fragments {
+					payloads = append(payloads, fragment.serialize(make([]byte, 0)))
+				}
+			} else {
 				payloads = append(payloads, nalu)
-
-				return
 			}
+
+			return
+		}
+
+		switch header.Type() {
+		case h265NaluVpsType:
+			p.vpsNalu = &packet
+
+			return
+		case h265NaluSpsType:
+			p.spsNalu = &packet
+
+			return
+		case h265NaluPpsType:
+			p.ppsNalu = &packet
+
+			return
+		case h265NaluAudType, h265NaluFillerType:
+			return
+		}
+
+		pendingNalus := make([]h265SingleNALUnitPacket, 0, 4)
+		if p.vpsNalu != nil {
+			pendingNalus = append(pendingNalus, *p.vpsNalu)
+			p.vpsNalu = nil
+		}
+		if p.spsNalu != nil {
+			pendingNalus = append(pendingNalus, *p.spsNalu)
+			p.spsNalu = nil
+		}
+		if p.ppsNalu != nil {
+			pendingNalus = append(pendingNalus, *p.ppsNalu)
+			p.ppsNalu = nil
+		}
+
+		if len(nalu) <= int(mtu) {
+			pendingNalus = append(pendingNalus, packet)
+		}
+		for _, pending := range pendingNalus {
 			if len(naluBuffer) == 0 {
-				if canAggregateH265(mtu, &packet) {
-					naluBuffer = append(naluBuffer, packet)
+				if canAggregateH265(mtu, &pending) {
+					naluBuffer = append(naluBuffer, pending)
 				} else {
-					payloads = append(payloads, nalu)
+					payloads = append(payloads, pending.serialize(make([]byte, 0, pending.wireSize())))
 				}
 			} else {
 				// can't fit any more packets, just send what we have and make current first in buffer
-				if shouldAggregateH265Now(mtu, naluBuffer, packet) {
+				if shouldAggregateH265Now(mtu, naluBuffer, pending) {
 					flushBuffer()
 				}
-				naluBuffer = append(naluBuffer, packet)
+				naluBuffer = append(naluBuffer, pending)
+			}
+		}
+
+		if len(nalu) > int(mtu) { // nolint: nestif
+			flushBuffer()
+			fragments, err := newH265FragmentationPackets(mtu, &packet)
+			if err != nil {
+				return
+			}
+			for _, fragment := range fragments {
+				payloads = append(payloads, fragment.serialize(make([]byte, 0)))
 			}
 		}
 	})

codecs/h265_packet_test.go

@@ -21,6 +21,13 @@ func createTestH265Header(pType, layerID, tid uint8, f bool) H265NALUHeader {
 	return H265NALUHeader(uint16(tid) | (uint16(pType) << 9) | (uint16(layerID) << 3) | fVal | zVal)
 }
 
+func createTestH265NALU(pType uint8, payload []byte) []byte {
+	nalu := make([]byte, 0, h265NaluHeaderSize+len(payload))
+	nalu = binary.BigEndian.AppendUint16(nalu, uint16(createTestH265Header(pType, 0, 1, false)))
+
+	return append(nalu, payload...)
+}
+
 func TestH265_NALU_Header(t *testing.T) {
 	tt := [...]struct {
 		RawHeader []byte
@@ -679,6 +686,76 @@ func TestH265Packetizer_Aggregated(t *testing.T) {
 	assert.Equal(t, payload, aggregated[18:], "Expected payloads to match")
 }
 
+func TestH265Payloader_Payload_VPS_SPS_PPS_handling(t *testing.T) {
+	packetizer := H265Payloader{}
+
+	vps := createTestH265NALU(h265NaluVpsType, []byte{0x00, 0x01})
+	sps := createTestH265NALU(h265NaluSpsType, []byte{0x02, 0x03})
+	pps := createTestH265NALU(h265NaluPpsType, []byte{0x04, 0x05})
+	vcl := createTestH265NALU(19, []byte{0x06, 0x07})
+
+	assert.Empty(t, packetizer.Payload(1500, vps), "VPS should be cached")
+	assert.Empty(t, packetizer.Payload(1500, sps), "SPS should be cached")
+	assert.Empty(t, packetizer.Payload(1500, pps), "PPS should be cached")
+
+	payloads := packetizer.Payload(1500, vcl)
+	assert.Len(t, payloads, 1)
+
+	aggregation, err := parseH265AggregationPacket(payloads[0], false)
+	assert.NoError(t, err)
+
+	packets, err := splitH265AggregationPacket(*aggregation)
+	assert.NoError(t, err)
+	assert.Len(t, packets, 4)
+	assert.Equal(t, uint8(h265NaluVpsType), packets[0].payloadHeader.Type())
+	assert.Equal(t, uint8(h265NaluSpsType), packets[1].payloadHeader.Type())
+	assert.Equal(t, uint8(h265NaluPpsType), packets[2].payloadHeader.Type())
+	assert.Equal(t, uint8(19), packets[3].payloadHeader.Type())
+	assert.Equal(t, []byte{0x06, 0x07}, packets[3].payload)
+}
+
+func TestH265Payloader_Payload_VPS_SPS_PPS_before_fragmented(t *testing.T) {
+	packetizer := H265Payloader{}
+
+	assert.Empty(t, packetizer.Payload(20, createTestH265NALU(h265NaluVpsType, []byte{0x00, 0x01})))
+	assert.Empty(t, packetizer.Payload(20, createTestH265NALU(h265NaluSpsType, []byte{0x02, 0x03})))
+	assert.Empty(t, packetizer.Payload(20, createTestH265NALU(h265NaluPpsType, []byte{0x04, 0x05})))
+
+	payload := make([]byte, 50)
+	for i := range payload {
+		payload[i] = uint8(i) //nolint:gosec // G115 test data
+	}
+	payloads := packetizer.Payload(20, createTestH265NALU(19, payload))
+	assert.Greater(t, len(payloads), 1)
+
+	aggregation, err := parseH265AggregationPacket(payloads[0], false)
+	assert.NoError(t, err)
+
+	packets, err := splitH265AggregationPacket(*aggregation)
+	assert.NoError(t, err)
+	assert.Len(t, packets, 3)
+	assert.Equal(t, uint8(h265NaluVpsType), packets[0].payloadHeader.Type())
+	assert.Equal(t, uint8(h265NaluSpsType), packets[1].payloadHeader.Type())
+	assert.Equal(t, uint8(h265NaluPpsType), packets[2].payloadHeader.Type())
+
+	fragmentHeader := H265NALUHeader(binary.BigEndian.Uint16(payloads[1][0:2]))
+	assert.True(t, fragmentHeader.IsFragmentationUnit())
+}
+
+func TestH265Payloader_Payload_SkipAggregation(t *testing.T) {
+	packetizer := H265Payloader{SkipAggregation: true}
+
+	vps := createTestH265NALU(h265NaluVpsType, []byte{0x00, 0x01})
+	sps := createTestH265NALU(h265NaluSpsType, []byte{0x02, 0x03})
+	pps := createTestH265NALU(h265NaluPpsType, []byte{0x04, 0x05})
+	vcl := createTestH265NALU(19, []byte{0x06, 0x07})
+
+	assert.Equal(t, [][]byte{vps}, packetizer.Payload(1500, vps))
+	assert.Equal(t, [][]byte{sps}, packetizer.Payload(1500, sps))
+	assert.Equal(t, [][]byte{pps}, packetizer.Payload(1500, pps))
+	assert.Equal(t, [][]byte{vcl}, packetizer.Payload(1500, vcl))
+}
+
 func TestH265Packetizer_Fragmented(t *testing.T) {
 	initSequence := []byte{0x00, 0x00, 0x00, 0x01, 0x00}