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

pkg/media/h264writer/h264writer.go

@@ -0,0 +1,108 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+// Package h264writer implements H264 media container writer
+package h264writer
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"os"
+
+	"github.com/pion/rtp"
+	"github.com/pion/rtp/codecs"
+)
+
+type (
+	// H264Writer is used to take RTP packets, parse them and
+	// write the data to an io.Writer.
+	// Currently it only supports non-interleaved mode
+	// Therefore, only 1-23, 24 (STAP-A), 28 (FU-A) NAL types are allowed.
+	// https://tools.ietf.org/html/rfc6184#section-5.2
+	H264Writer struct {
+		writer       io.Writer
+		hasKeyFrame  bool
+		cachedPacket *codecs.H264Packet
+	}
+)
+
+// New builds a new H264 writer.
+func New(filename string) (*H264Writer, error) {
+	f, err := os.Create(filename) //nolint:gosec
+	if err != nil {
+		return nil, err
+	}
+
+	return NewWith(f), nil
+}
+
+// NewWith initializes a new H264 writer with an io.Writer output.
+func NewWith(w io.Writer) *H264Writer {
+	return &H264Writer{
+		writer: w,
+	}
+}
+
+// WriteRTP adds a new packet and writes the appropriate headers for it.
+func (h *H264Writer) 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.H264Packet{}
+	}
+
+	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 *H264Writer) 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 {
+	const (
+		typeSTAPA       = 24
+		typeSPS         = 7
+		naluTypeBitmask = 0x1F
+	)
+
+	var word uint32
+
+	payload := bytes.NewReader(data)
+	if err := binary.Read(payload, binary.BigEndian, &word); err != nil {
+		return false
+	}
+
+	naluType := (word >> 24) & naluTypeBitmask
+	if naluType == typeSTAPA && word&naluTypeBitmask == typeSPS {
+		return true
+	} else if naluType == typeSPS {
+		return true
+	}
+
+	return false
+}