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

pkg/media/ivfwriter/ivfwriter.go

@@ -0,0 +1,406 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+// Package ivfwriter implements IVF media container writer
+package ivfwriter
+
+import (
+	"encoding/binary"
+	"errors"
+	"io"
+	"os"
+
+	"github.com/pion/rtp"
+	"github.com/pion/rtp/codecs"
+	"github.com/pion/rtp/codecs/av1/obu"
+)
+
+var (
+	errFileNotOpened        = errors.New("file not opened")
+	errInvalidNilPacket     = errors.New("invalid nil packet")
+	errCodecUnset           = errors.New("codec is unset")
+	errCodecAlreadySet      = errors.New("codec is already set")
+	errNoSuchCodec          = errors.New("no codec for this MimeType")
+	errInvalidMediaTimebase = errors.New("invalid media timebase")
+)
+
+type (
+	codec int
+
+	// IVFWriter is used to take RTP packets and write them to an IVF on disk.
+	IVFWriter struct {
+		ioWriter     io.Writer
+		count        uint64
+		seenKeyFrame bool
+
+		codec codec
+
+		timebaseDenominator uint32
+		timebaseNumerator   uint32
+		firstFrameTimestamp uint32
+		clockRate           uint64
+		videoWidth          uint16
+		videoHeight         uint16
+
+		directPTS bool
+
+		// VP8, VP9
+		currentFrame []byte
+
+		// AV1
+		av1Depacketizer *codecs.AV1Depacketizer
+	}
+)
+
+const (
+	codecUnset codec = iota
+	codecVP8
+	codecVP9
+	codecAV1
+
+	mimeTypeVP8 = "video/VP8"
+	mimeTypeVP9 = "video/VP9"
+	mimeTypeAV1 = "video/AV1"
+)
+
+// New builds a new IVF writer.
+func New(fileName string, opts ...Option) (*IVFWriter, error) {
+	file, err := os.Create(fileName) //nolint:gosec
+	if err != nil {
+		return nil, err
+	}
+	writer, err := NewWith(file, opts...)
+	if err != nil {
+		return nil, err
+	}
+	writer.ioWriter = file
+
+	return writer, nil
+}
+
+// NewWith initialize a new IVF writer with an io.Writer output.
+func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
+	if out == nil {
+		return nil, errFileNotOpened
+	}
+
+	writer := &IVFWriter{
+		ioWriter:            out,
+		seenKeyFrame:        false,
+		timebaseDenominator: 30,
+		timebaseNumerator:   1,
+		clockRate:           90000,
+		videoWidth:          640,
+		videoHeight:         480,
+	}
+
+	for _, o := range opts {
+		if err := o(writer); err != nil {
+			return nil, err
+		}
+	}
+
+	if writer.codec == codecUnset {
+		writer.codec = codecVP8
+	}
+	if err := writer.writeHeader(); err != nil {
+		return nil, err
+	}
+
+	if writer.timebaseDenominator == 0 {
+		return nil, errInvalidMediaTimebase
+	}
+
+	return writer, nil
+}
+
+func (i *IVFWriter) writeHeader() error {
+	header := make([]byte, 32)
+	copy(header[0:], "DKIF")                      // DKIF
+	binary.LittleEndian.PutUint16(header[4:], 0)  // Version
+	binary.LittleEndian.PutUint16(header[6:], 32) // Header size
+
+	// FOURCC
+	switch i.codec {
+	case codecVP8:
+		copy(header[8:], "VP80")
+	case codecVP9:
+		copy(header[8:], "VP90")
+	case codecAV1:
+		copy(header[8:], "AV01")
+	default:
+		return errCodecUnset
+	}
+
+	binary.LittleEndian.PutUint16(header[12:], i.videoWidth)          // Width in pixels
+	binary.LittleEndian.PutUint16(header[14:], i.videoHeight)         // Height in pixels
+	binary.LittleEndian.PutUint32(header[16:], i.timebaseDenominator) // Framerate denominator
+	binary.LittleEndian.PutUint32(header[20:], i.timebaseNumerator)   // Framerate numerator
+	binary.LittleEndian.PutUint32(header[24:], 900)                   // Frame count, will be updated on first Close() call
+	binary.LittleEndian.PutUint32(header[28:], 0)                     // Unused
+
+	_, err := i.ioWriter.Write(header)
+
+	return err
+}
+
+func (i *IVFWriter) timestampToPts(timestamp uint64) uint64 {
+	return timestamp * uint64(i.timebaseNumerator) / uint64(i.timebaseDenominator)
+}
+
+func (i *IVFWriter) writeFrame(frame []byte, timestamp uint64) error {
+	frameHeader := make([]byte, 12)
+	//nolint:gosec // G115
+	binary.LittleEndian.PutUint32(frameHeader[0:], uint32(len(frame))) // Frame length
+
+	var pts uint64
+	if i.directPTS {
+		// Direct PTS mode: use timestamp directly as PTS.
+		pts = timestamp
+	} else {
+		// Existing behavior: convert using timebase.
+		pts = i.timestampToPts(timestamp)
+	}
+	binary.LittleEndian.PutUint64(frameHeader[4:], pts) // PTS
+
+	i.count++
+
+	if _, err := i.ioWriter.Write(frameHeader); err != nil {
+		return err
+	}
+	_, err := i.ioWriter.Write(frame)
+
+	return err
+}
+
+// WriteRTP adds a new packet and writes the appropriate headers for it.
+func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error {
+	if i.ioWriter == nil {
+		return errFileNotOpened
+	} else if len(packet.Payload) == 0 {
+		return nil
+	}
+
+	if i.count == 0 {
+		i.firstFrameTimestamp = packet.Timestamp
+	}
+
+	var timestamp uint64
+	if i.directPTS {
+		// Direct PTS mode: use RTP timestamp directly (no millisecond conversion).
+		timestamp = uint64(packet.Timestamp - i.firstFrameTimestamp)
+	} else {
+		// Existing behavior: convert to milliseconds first.
+		timestamp = 1000 * uint64(packet.Timestamp-i.firstFrameTimestamp) / i.clockRate
+	}
+
+	switch i.codec {
+	case codecVP8:
+		return i.writeVP8(packet, timestamp)
+	case codecVP9:
+		return i.writeVP9(packet, timestamp)
+	case codecAV1:
+		return i.writeAV1(packet, timestamp)
+	default:
+		return errCodecUnset
+	}
+}
+
+func (i *IVFWriter) writeVP8(packet *rtp.Packet, timestamp uint64) error {
+	vp8Packet := codecs.VP8Packet{}
+	if _, err := vp8Packet.Unmarshal(packet.Payload); err != nil {
+		return err
+	}
+
+	// a payload that carries only the VP8 descriptor has nothing to write.
+	if len(vp8Packet.Payload) == 0 {
+		return nil
+	}
+
+	isKeyFrame := (vp8Packet.Payload[0] & 0x01) == 0
+	switch {
+	case !i.seenKeyFrame && !isKeyFrame:
+		return nil
+	case i.currentFrame == nil && vp8Packet.S != 1:
+		return nil
+	}
+
+	i.seenKeyFrame = true
+	i.currentFrame = append(i.currentFrame, vp8Packet.Payload[0:]...)
+
+	if !packet.Marker {
+		return nil
+	} else if len(i.currentFrame) == 0 {
+		return nil
+	}
+
+	if err := i.writeFrame(i.currentFrame, timestamp); err != nil {
+		return err
+	}
+	i.currentFrame = nil
+
+	return nil
+}
+
+func (i *IVFWriter) writeVP9(packet *rtp.Packet, timestamp uint64) error {
+	vp9Packet := codecs.VP9Packet{}
+	if _, err := vp9Packet.Unmarshal(packet.Payload); err != nil {
+		return err
+	}
+
+	switch {
+	case !i.seenKeyFrame && vp9Packet.P:
+		return nil
+	case i.currentFrame == nil && !vp9Packet.B:
+		return nil
+	}
+
+	i.seenKeyFrame = true
+	i.currentFrame = append(i.currentFrame, vp9Packet.Payload[0:]...)
+
+	if !packet.Marker {
+		return nil
+	} else if len(i.currentFrame) == 0 {
+		return nil
+	}
+
+	// the timestamp must be sequential. webrtc mandates a clock rate of 90000
+	// and we've assumed 30fps in the header.
+	if err := i.writeFrame(i.currentFrame, timestamp); err != nil {
+		return err
+	}
+	i.currentFrame = nil
+
+	return nil
+}
+
+func (i *IVFWriter) writeAV1(packet *rtp.Packet, timestamp uint64) error {
+	if i.av1Depacketizer == nil {
+		i.av1Depacketizer = &codecs.AV1Depacketizer{}
+	}
+
+	payload, err := i.av1Depacketizer.Unmarshal(packet.Payload)
+	if err != nil {
+		return err
+	}
+
+	if !i.seenKeyFrame {
+		isKeyFrame := i.av1Depacketizer.N || (len(payload) > 0 && obu.Type((payload[0]&0x78)>>3) == obu.OBUSequenceHeader)
+		if !isKeyFrame {
+			return nil
+		}
+
+		i.seenKeyFrame = true
+	}
+
+	i.currentFrame = append(i.currentFrame, payload...)
+	if !packet.Marker {
+		return nil
+	}
+
+	delimiter := obu.Header{
+		Type:         obu.OBUTemporalDelimiter,
+		HasSizeField: true,
+	}
+	frame := append(delimiter.Marshal(), 0)
+	frame = append(frame, i.currentFrame...)
+
+	if err := i.writeFrame(frame, timestamp); err != nil {
+		return err
+	}
+	i.currentFrame = nil
+
+	return nil
+}
+
+// Close stops the recording.
+func (i *IVFWriter) Close() error {
+	if i.ioWriter == nil {
+		// Returns no error as it may be convenient to call
+		// Close() multiple times
+		return nil
+	}
+
+	defer func() {
+		i.ioWriter = nil
+	}()
+
+	if ws, ok := i.ioWriter.(io.WriteSeeker); ok {
+		// Update the framecount
+		if _, err := ws.Seek(24, 0); err != nil {
+			return err
+		}
+		buff := make([]byte, 4)
+		binary.LittleEndian.PutUint32(buff, uint32(i.count)) //nolint:gosec // G115
+		if _, err := ws.Write(buff); err != nil {
+			return err
+		}
+	}
+
+	if closer, ok := i.ioWriter.(io.Closer); ok {
+		return closer.Close()
+	}
+
+	return nil
+}
+
+// An Option configures a SampleBuilder.
+type Option func(i *IVFWriter) error
+
+// WithCodec configures if IVFWriter is writing AV1 or VP8 packets to disk.
+func WithCodec(mimeType string) Option {
+	return func(i *IVFWriter) error {
+		if i.codec != codecUnset {
+			return errCodecAlreadySet
+		}
+
+		switch mimeType {
+		case mimeTypeVP8:
+			i.codec = codecVP8
+		case mimeTypeVP9:
+			i.codec = codecVP9
+		case mimeTypeAV1:
+			i.codec = codecAV1
+		default:
+			return errNoSuchCodec
+		}
+
+		return nil
+	}
+}
+
+func WithWidthAndHeight(width, height uint16) Option {
+	return func(i *IVFWriter) error {
+		i.videoWidth = width
+		i.videoHeight = height
+
+		return nil
+	}
+}
+
+func WithFrameRate(numerator, denominator uint32) Option {
+	return func(i *IVFWriter) error {
+		i.timebaseNumerator = numerator
+		i.timebaseDenominator = denominator
+
+		return nil
+	}
+}
+
+// WithDirectPTS enables direct use of RTP timestamps as PTS values
+// without millisecond conversion.
+//
+// When this option is used, RTP timestamps are written directly as PTS values,
+// preserving full timestamp precision. Use WithFrameRate to set the appropriate
+// timebase (e.g., WithFrameRate(1, 90000) for standard 90kHz RTP clock).
+//
+// Example usage for standard RTP video (90kHz clock rate):
+//
+//	NewWith(file, WithFrameRate(1, 90000), WithDirectPTS())
+func WithDirectPTS() Option {
+	return func(i *IVFWriter) error {
+		i.directPTS = true
+
+		return nil
+	}
+}