snapshot of https://github.com/pion/webrtc.git (main, 2026-09-15; blobs > 131072B omitted)

pkg/media/h265writer/h265writer.go

@@ -0,0 +1,159 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+// Package h265writer implements H265/HEVC media container writer
+package h265writer
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"os"
+
+	"github.com/pion/rtp"
+	"github.com/pion/rtp/codecs"
+	"github.com/pion/webrtc/v4/pkg/media/h265reader"
+)
+
+const (
+	typeAP = 48 // Aggregation Packet
+	typeFU = 49 // Fragmentation Unit
+)
+
+// H265Writer is used to take H.265/HEVC RTP packets defined in RFC 7798, parse them and
+// write the data to an io.Writer.
+type H265Writer struct {
+	writer       io.Writer
+	hasKeyFrame  bool
+	cachedPacket *codecs.H265Depacketizer
+}
+
+// New builds a new H265 writer.
+func New(filename string) (*H265Writer, error) {
+	f, err := os.Create(filename) //nolint:gosec
+	if err != nil {
+		return nil, err
+	}
+
+	return NewWith(f), nil
+}
+
+// NewWith initializes a new H265 writer with an io.Writer output.
+func NewWith(w io.Writer) *H265Writer {
+	return &H265Writer{
+		writer: w,
+	}
+}
+
+// WriteRTP adds a new packet and writes the appropriate headers for it.
+func (h *H265Writer) WriteRTP(packet *rtp.Packet) error {
+	if len(packet.Payload) == 0 {
+		return nil
+	}
+
+	if !h.hasKeyFrame {
+		if h.hasKeyFrame = isKeyFrame(packet.Payload); !h.hasKeyFrame {
+			// key frame not defined yet. discarding packet
+			return nil
+		}
+	}
+
+	if h.cachedPacket == nil {
+		h.cachedPacket = &codecs.H265Depacketizer{}
+	}
+
+	data, err := h.cachedPacket.Unmarshal(packet.Payload)
+	if err != nil || len(data) == 0 {
+		return err
+	}
+
+	_, err = h.writer.Write(data)
+
+	return err
+}
+
+// Close closes the underlying writer.
+func (h *H265Writer) Close() error {
+	h.cachedPacket = nil
+	if h.writer != nil {
+		if closer, ok := h.writer.(io.Closer); ok {
+			return closer.Close()
+		}
+	}
+
+	return nil
+}
+
+func isKeyFrame(data []byte) bool {
+	if len(data) < 2 {
+		return false
+	}
+
+	// Get NAL unit type from first byte (bits 6-1)
+	naluType := (data[0] & 0x7E) >> 1
+	if isKeyFrameNalu(h265reader.NalUnitType(naluType)) {
+		return true
+	}
+
+	// Check for parameter sets or IDR frames
+	switch naluType {
+	case typeAP:
+		// For aggregation packets, check if any contained NAL is a key frame
+		return checkAggregationPacketForKeyFrame(data)
+	case typeFU:
+		// For fragmentation units, check the NAL type in the FU header
+		if len(data) < 3 {
+			return false
+		}
+		fuNaluType := h265reader.NalUnitType((data[2] & 0x7E) >> 1)
+
+		return isKeyFrameNalu(fuNaluType)
+	}
+
+	return false
+}
+
+func checkAggregationPacketForKeyFrame(data []byte) bool {
+	// Skip the payload header (2 bytes for H.265)
+	offset := 2
+
+	for offset < len(data) {
+		if offset+2 > len(data) {
+			break
+		}
+
+		// Read NAL unit size (2 bytes in network byte order)
+		var naluSize uint16
+		buf := bytes.NewReader(data[offset : offset+2])
+		if err := binary.Read(buf, binary.BigEndian, &naluSize); err != nil {
+			break
+		}
+		offset += 2
+
+		if offset+int(naluSize) > len(data) {
+			break
+		}
+
+		if naluSize > 0 {
+			// Check NAL unit type
+			naluType := h265reader.NalUnitType((data[offset] & 0x7E) >> 1)
+			if isKeyFrameNalu(naluType) {
+				return true
+			}
+		}
+
+		offset += int(naluSize)
+	}
+
+	return false
+}
+
+func isKeyFrameNalu(naluType h265reader.NalUnitType) bool {
+	switch naluType {
+	case h265reader.NalUnitTypeVps, h265reader.NalUnitTypeSps, h265reader.NalUnitTypePps,
+		h265reader.NalUnitTypeIdrWRadl, h265reader.NalUnitTypeIdrNLp:
+		return true
+	default:
+		return false
+	}
+}

pkg/media/h265writer/h265writer_test.go

@@ -0,0 +1,157 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package h265writer
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/pion/rtp"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestH265Writer_WriteRTP(t *testing.T) {
+	buf := &bytes.Buffer{}
+	writer := NewWith(buf)
+	defer func() {
+		assert.NoError(t, writer.Close())
+	}()
+
+	// Test with empty payload
+	packet := &rtp.Packet{Payload: []byte{}}
+	err := writer.WriteRTP(packet)
+	assert.NoError(t, err)
+
+	// Test with VPS packet (key frame)
+	vpsPayload := []byte{0x40, 0x01, 0x0C, 0x01, 0xFF, 0xFF, 0x01, 0x60}
+	packet = &rtp.Packet{Payload: vpsPayload}
+
+	err = writer.WriteRTP(packet)
+	assert.NoError(t, err)
+
+	// Check that the buffer contains the expected start code + VPS data
+	expectedContent := append([]byte{0x00, 0x00, 0x00, 0x01}, vpsPayload...)
+	assert.Equal(t, expectedContent, buf.Bytes(), "Buffer should contain start code followed by VPS payload")
+}
+
+func TestIsKeyFrame(t *testing.T) {
+	tests := []struct {
+		name     string
+		data     []byte
+		expected bool
+	}{
+		{
+			name:     "VPS NAL unit",
+			data:     []byte{0x40, 0x01, 0x0C, 0x01}, // VPS (type 32)
+			expected: true,
+		},
+		{
+			name:     "SPS NAL unit",
+			data:     []byte{0x42, 0x01, 0x01, 0x01}, // SPS (type 33)
+			expected: true,
+		},
+		{
+			name:     "PPS NAL unit",
+			data:     []byte{0x44, 0x01, 0xC1, 0x73}, // PPS (type 34)
+			expected: true,
+		},
+		{
+			name:     "IDR_W_RADL NAL unit",
+			data:     []byte{0x26, 0x01, 0xAF, 0x06}, // IDR_W_RADL (type 19)
+			expected: true,
+		},
+		{
+			name:     "IDR_N_LP NAL unit",
+			data:     []byte{0x28, 0x01, 0xAF, 0x06}, // IDR_N_LP (type 20)
+			expected: true,
+		},
+		{
+			name:     "TRAIL_R NAL unit",
+			data:     []byte{0x02, 0x01, 0xAF, 0x06}, // TRAIL_R (type 1)
+			expected: false,
+		},
+		{
+			name:     "Empty data",
+			data:     []byte{},
+			expected: false,
+		},
+		{
+			name:     "Single byte",
+			data:     []byte{0x40},
+			expected: false,
+		},
+		{
+			name:     "Fragmentation Unit with VPS",
+			data:     []byte{0x62, 0x01, 0x40}, // FU with VPS NAL type
+			expected: true,
+		},
+		{
+			name:     "Fragmentation Unit with TRAIL_R",
+			data:     []byte{0x62, 0x01, 0x02}, // FU with TRAIL_R NAL type
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := isKeyFrame(tt.data)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestCheckAggregationPacketForKeyFrame(t *testing.T) {
+	tests := []struct {
+		name     string
+		data     []byte
+		expected bool
+	}{
+		{
+			name: "AP with VPS",
+			data: []byte{
+				0x60, 0x01, // AP header
+				0x00, 0x04, // NALU size (4 bytes)
+				0x40, 0x01, 0x0C, 0x01, // VPS NAL unit
+			},
+			expected: true,
+		},
+		{
+			name: "AP with TRAIL_R",
+			data: []byte{
+				0x60, 0x01, // AP header
+				0x00, 0x04, // NALU size (4 bytes)
+				0x02, 0x01, 0xAF, 0x06, // TRAIL_R NAL unit
+			},
+			expected: false,
+		},
+		{
+			name: "AP with multiple NALUs including SPS",
+			data: []byte{
+				0x60, 0x01, // AP header
+				0x00, 0x04, // First NALU size
+				0x02, 0x01, 0xAF, 0x06, // TRAIL_R NAL unit
+				0x00, 0x04, // Second NALU size
+				0x42, 0x01, 0x01, 0x01, // SPS NAL unit
+			},
+			expected: true,
+		},
+		{
+			name:     "Malformed AP - insufficient data",
+			data:     []byte{0x60, 0x01, 0x00}, // AP header + incomplete size
+			expected: false,
+		},
+		{
+			name:     "Empty AP",
+			data:     []byte{0x60, 0x01}, // AP header only
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := checkAggregationPacketForKeyFrame(tt.data)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}