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

2026-09-15 20:06:07 +05:00
commit 8acb1a3a29
394 changed files with 78746 additions and 0 deletions
--- a/pkg/media/h264reader/h264reader.go
+++ b/pkg/media/h264reader/h264reader.go
@@ -0,0 +1,239 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+// Package h264reader implements a H264 Annex-B Reader
+package h264reader
+
+import (
+	"bytes"
+	"errors"
+	"io"
+)
+
+// H264Reader reads data from stream and constructs h264 nal units.
+type H264Reader struct {
+	stream                      io.Reader
+	nalBuffer                   []byte
+	countOfConsecutiveZeroBytes int
+	nalPrefixParsed             bool
+	readBuffer                  []byte
+	tmpReadBuf                  []byte
+	includeSEI                  bool
+}
+
+var (
+	errNilReader           = errors.New("stream is nil")
+	errDataIsNotH264Stream = errors.New("data is not a H264 bitstream")
+)
+
+// NewReader creates new H264Reader.
+func NewReader(in io.Reader) (*H264Reader, error) {
+	if in == nil {
+		return nil, errNilReader
+	}
+
+	reader := &H264Reader{
+		stream:          in,
+		nalBuffer:       make([]byte, 0),
+		nalPrefixParsed: false,
+		readBuffer:      make([]byte, 0),
+		tmpReadBuf:      make([]byte, 4096),
+		includeSEI:      false,
+	}
+
+	return reader, nil
+}
+
+// Option configures the behavior of H264Reader.
+type Option func(*H264Reader) error
+
+// NewReaderWithOptions creates new H264Reader with options.
+// The default behavior is to skip SEI NAL units.
+func NewReaderWithOptions(in io.Reader, options ...Option) (*H264Reader, error) {
+	reader, err := NewReader(in)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, option := range options {
+		if err := option(reader); err != nil {
+			return nil, err
+		}
+	}
+
+	return reader, nil
+}
+
+// WithIncludeSEI controls whether SEI (Supplemental Enhancement Information) NAL units are returned.
+// Default is false (SEI is skipped).
+func WithIncludeSEI(include bool) Option {
+	return func(r *H264Reader) error {
+		r.includeSEI = include
+
+		return nil
+	}
+}
+
+// NAL H.264 Network Abstraction Layer.
+type NAL struct {
+	PictureOrderCount uint32
+
+	// NAL header
+	ForbiddenZeroBit bool
+	RefIdc           uint8
+	UnitType         NalUnitType
+
+	Data []byte // header byte + rbsp
+}
+
+func (reader *H264Reader) read(numToRead int) (data []byte, e error) {
+	for len(reader.readBuffer) < numToRead {
+		n, err := reader.stream.Read(reader.tmpReadBuf)
+		if err != nil {
+			return nil, err
+		}
+		if n == 0 {
+			break
+		}
+		reader.readBuffer = append(reader.readBuffer, reader.tmpReadBuf[0:n]...)
+	}
+
+	numShouldRead := min(numToRead, len(reader.readBuffer))
+	data = reader.readBuffer[0:numShouldRead]
+	reader.readBuffer = reader.readBuffer[numShouldRead:]
+
+	return data, nil
+}
+
+func (reader *H264Reader) bitStreamStartsWithH264Prefix() (prefixLength int, e error) {
+	nalPrefix3Bytes := []byte{0, 0, 1}
+	nalPrefix4Bytes := []byte{0, 0, 0, 1}
+
+	prefixBuffer, e := reader.read(4)
+	if e != nil {
+		return prefixLength, e
+	}
+
+	n := len(prefixBuffer)
+
+	if n == 0 {
+		return 0, io.EOF
+	}
+
+	if n < 3 {
+		return 0, errDataIsNotH264Stream
+	}
+
+	nalPrefix3BytesFound := bytes.Equal(nalPrefix3Bytes, prefixBuffer[:3])
+	if n == 3 {
+		if nalPrefix3BytesFound {
+			return 0, io.EOF
+		}
+
+		return 0, errDataIsNotH264Stream
+	}
+
+	// n == 4
+	if nalPrefix3BytesFound {
+		reader.nalBuffer = append(reader.nalBuffer, prefixBuffer[3])
+
+		return 3, nil
+	}
+
+	nalPrefix4BytesFound := bytes.Equal(nalPrefix4Bytes, prefixBuffer)
+	if nalPrefix4BytesFound {
+		return 4, nil
+	}
+
+	return 0, errDataIsNotH264Stream
+}
+
+// NextNAL reads from stream and returns then next NAL,
+// and an error if there is incomplete frame data.
+// Returns all nil values when no more NALs are available.
+func (reader *H264Reader) NextNAL() (*NAL, error) {
+	if !reader.nalPrefixParsed {
+		_, err := reader.bitStreamStartsWithH264Prefix()
+		if err != nil {
+			return nil, err
+		}
+
+		reader.nalPrefixParsed = true
+	}
+
+	for {
+		buffer, err := reader.read(1)
+		if err != nil {
+			break
+		}
+
+		n := len(buffer)
+
+		if n != 1 {
+			break
+		}
+		readByte := buffer[0]
+		nalFound := reader.processByte(readByte)
+		if nalFound {
+			nal := newNal(reader.nalBuffer)
+			nal.parseHeader()
+			if !reader.includeSEI && nal.UnitType == NalUnitTypeSEI {
+				reader.nalBuffer = nil
+
+				continue
+			}
+
+			break
+		}
+
+		reader.nalBuffer = append(reader.nalBuffer, readByte)
+	}
+
+	if len(reader.nalBuffer) == 0 {
+		return nil, io.EOF
+	}
+
+	nal := newNal(reader.nalBuffer)
+	reader.nalBuffer = nil
+	nal.parseHeader()
+
+	return nal, nil
+}
+
+func (reader *H264Reader) processByte(readByte byte) (nalFound bool) {
+	nalFound = false
+
+	switch readByte {
+	case 0:
+		reader.countOfConsecutiveZeroBytes++
+	case 1:
+		if reader.countOfConsecutiveZeroBytes >= 2 {
+			countOfConsecutiveZeroBytesInPrefix := 2
+			if reader.countOfConsecutiveZeroBytes > 2 {
+				countOfConsecutiveZeroBytesInPrefix = 3
+			}
+
+			if nalUnitLength := len(reader.nalBuffer) - countOfConsecutiveZeroBytesInPrefix; nalUnitLength > 0 {
+				reader.nalBuffer = reader.nalBuffer[0:nalUnitLength]
+				nalFound = true
+			}
+		}
+
+		reader.countOfConsecutiveZeroBytes = 0
+	default:
+		reader.countOfConsecutiveZeroBytes = 0
+	}
+
+	return nalFound
+}
+
+func newNal(data []byte) *NAL {
+	return &NAL{PictureOrderCount: 0, ForbiddenZeroBit: false, RefIdc: 0, UnitType: NalUnitTypeUnspecified, Data: data}
+}
+
+func (h *NAL) parseHeader() {
+	firstByte := h.Data[0]
+	h.ForbiddenZeroBit = (((firstByte & 0x80) >> 7) == 1) // 0x80 = 0b10000000
+	h.RefIdc = (firstByte & 0x60) >> 5                    // 0x60 = 0b01100000
+	h.UnitType = NalUnitType((firstByte & 0x1F) >> 0)     // 0x1F = 0b00011111
+}
--- a/pkg/media/h264reader/h264reader_test.go
+++ b/pkg/media/h264reader/h264reader_test.go
@@ -0,0 +1,164 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package h264reader
+
+import (
+	"bytes"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func CreateReader(h264 []byte, require *require.Assertions) *H264Reader {
+	reader, err := NewReader(bytes.NewReader(h264))
+
+	require.Nil(err)
+	require.NotNil(reader)
+
+	return reader
+}
+
+func TestDataDoesNotStartWithH264Header(t *testing.T) {
+	require := require.New(t)
+
+	testFunction := func(input []byte, expectedErr error) {
+		reader := CreateReader(input, require)
+		nal, err := reader.NextNAL()
+		require.ErrorIs(err, expectedErr)
+		require.Nil(nal)
+	}
+
+	h264Bytes1 := []byte{2}
+	testFunction(h264Bytes1, io.EOF)
+
+	h264Bytes2 := []byte{0, 2}
+	testFunction(h264Bytes2, io.EOF)
+
+	h264Bytes3 := []byte{0, 0, 2}
+	testFunction(h264Bytes3, io.EOF)
+
+	h264Bytes4 := []byte{0, 0, 2, 0}
+	testFunction(h264Bytes4, errDataIsNotH264Stream)
+
+	h264Bytes5 := []byte{0, 0, 0, 2}
+	testFunction(h264Bytes5, errDataIsNotH264Stream)
+}
+
+func TestParseHeader(t *testing.T) {
+	require := require.New(t)
+	h264Bytes := []byte{0x0, 0x0, 0x1, 0xAB}
+
+	reader := CreateReader(h264Bytes, require)
+
+	nal, err := reader.NextNAL()
+	require.Nil(err)
+
+	require.Equal(1, len(nal.Data))
+	require.True(nal.ForbiddenZeroBit)
+	require.Equal(uint32(0), nal.PictureOrderCount)
+	require.Equal(uint8(1), nal.RefIdc)
+	require.Equal(NalUnitTypeEndOfStream, nal.UnitType)
+}
+
+func TestEOF(t *testing.T) {
+	require := require.New(t)
+
+	testFunction := func(input []byte) {
+		reader := CreateReader(input, require)
+
+		nal, err := reader.NextNAL()
+		require.Equal(io.EOF, err)
+		require.Nil(nal)
+	}
+
+	h264Bytes1 := []byte{0, 0, 0, 1}
+	testFunction(h264Bytes1)
+
+	h264Bytes2 := []byte{0, 0, 1}
+	testFunction(h264Bytes2)
+
+	h264Bytes3 := []byte{}
+	testFunction(h264Bytes3)
+}
+
+func TestSkipSEI(t *testing.T) {
+	require := require.New(t)
+	h264Bytes := []byte{
+		0x0, 0x0, 0x0, 0x1, 0xAA,
+		0x0, 0x0, 0x0, 0x1, 0x6, // SEI
+		0x0, 0x0, 0x0, 0x1, 0xAB,
+	}
+
+	reader := CreateReader(h264Bytes, require)
+
+	nal, err := reader.NextNAL()
+	require.Nil(err)
+	require.Equal(byte(0xAA), nal.Data[0])
+
+	nal, err = reader.NextNAL()
+	require.Nil(err)
+	require.Equal(byte(0xAB), nal.Data[0])
+}
+
+func TestIncludeSEI(t *testing.T) {
+	require := require.New(t)
+	h264Bytes := []byte{
+		0x0, 0x0, 0x0, 0x1, 0xAA,
+		0x0, 0x0, 0x0, 0x1, 0x6, // SEI
+		0x0, 0x0, 0x0, 0x1, 0xAB,
+	}
+
+	reader, err := NewReaderWithOptions(bytes.NewReader(h264Bytes), WithIncludeSEI(true))
+	require.NoError(err)
+	require.NotNil(reader)
+
+	nal, err := reader.NextNAL()
+	require.NoError(err)
+	require.Equal(byte(0xAA), nal.Data[0])
+
+	nal, err = reader.NextNAL()
+	require.NoError(err)
+	require.Equal(NalUnitTypeSEI, nal.UnitType)
+	require.Equal(byte(0x6), nal.Data[0])
+
+	nal, err = reader.NextNAL()
+	require.NoError(err)
+	require.Equal(byte(0xAB), nal.Data[0])
+}
+
+func TestIssue1734_NextNal(t *testing.T) {
+	tt := [...][]byte{
+		[]byte("\x00\x00\x010\x00\x00\x01\x00\x00\x01"),
+		[]byte("\x00\x00\x00\x01\x00\x00\x01"),
+	}
+
+	for _, cur := range tt {
+		r, err := NewReader(bytes.NewReader(cur))
+		require.NoError(t, err)
+
+		// Just make sure it doesn't crash
+		for {
+			nal, err := r.NextNAL()
+
+			if err != nil || nal == nil {
+				break
+			}
+		}
+	}
+}
+
+func TestTrailing01AfterStartCode(t *testing.T) {
+	reader, err := NewReader(bytes.NewReader([]byte{
+		0x0, 0x0, 0x0, 0x1, 0x01,
+		0x0, 0x0, 0x0, 0x1, 0x01,
+	}))
+	require.NoError(t, err)
+
+	for i := 0; i <= 1; i++ {
+		nal, err := reader.NextNAL()
+		require.NoError(t, err)
+		require.NotNil(t, nal)
+	}
+}
--- a/pkg/media/h264reader/nalunittype.go
+++ b/pkg/media/h264reader/nalunittype.go
@@ -0,0 +1,72 @@
+// SPDX-FileCopyrightText: 2026 The Pion community <https://pion.ly>
+// SPDX-License-Identifier: MIT
+
+package h264reader
+
+import "strconv"
+
+// NalUnitType is the type of a NAL.
+type NalUnitType uint8
+
+// Enums for NalUnitTypes.
+const (
+	NalUnitTypeUnspecified              NalUnitType = 0  // Unspecified
+	NalUnitTypeCodedSliceNonIdr         NalUnitType = 1  // Coded slice of a non-IDR picture
+	NalUnitTypeCodedSliceDataPartitionA NalUnitType = 2  // Coded slice data partition A
+	NalUnitTypeCodedSliceDataPartitionB NalUnitType = 3  // Coded slice data partition B
+	NalUnitTypeCodedSliceDataPartitionC NalUnitType = 4  // Coded slice data partition C
+	NalUnitTypeCodedSliceIdr            NalUnitType = 5  // Coded slice of an IDR picture
+	NalUnitTypeSEI                      NalUnitType = 6  // Supplemental enhancement information (SEI)
+	NalUnitTypeSPS                      NalUnitType = 7  // Sequence parameter set
+	NalUnitTypePPS                      NalUnitType = 8  // Picture parameter set
+	NalUnitTypeAUD                      NalUnitType = 9  // Access unit delimiter
+	NalUnitTypeEndOfSequence            NalUnitType = 10 // End of sequence
+	NalUnitTypeEndOfStream              NalUnitType = 11 // End of stream
+	NalUnitTypeFiller                   NalUnitType = 12 // Filler data
+	NalUnitTypeSpsExt                   NalUnitType = 13 // Sequence parameter set extension
+	NalUnitTypeCodedSliceAux            NalUnitType = 19 // Coded slice of an auxiliary coded picture without partitioning
+	// 14..18                                            // Reserved.
+	// 20..23                                            // Reserved.
+	// 24..31                                            // Unspecified.
+)
+
+func (n *NalUnitType) String() string { //nolint:cyclop
+	var str string
+	switch *n {
+	case NalUnitTypeUnspecified:
+		str = "Unspecified"
+	case NalUnitTypeCodedSliceNonIdr:
+		str = "CodedSliceNonIdr"
+	case NalUnitTypeCodedSliceDataPartitionA:
+		str = "CodedSliceDataPartitionA"
+	case NalUnitTypeCodedSliceDataPartitionB:
+		str = "CodedSliceDataPartitionB"
+	case NalUnitTypeCodedSliceDataPartitionC:
+		str = "CodedSliceDataPartitionC"
+	case NalUnitTypeCodedSliceIdr:
+		str = "CodedSliceIdr"
+	case NalUnitTypeSEI:
+		str = "SEI"
+	case NalUnitTypeSPS:
+		str = "SPS"
+	case NalUnitTypePPS:
+		str = "PPS"
+	case NalUnitTypeAUD:
+		str = "AUD"
+	case NalUnitTypeEndOfSequence:
+		str = "EndOfSequence"
+	case NalUnitTypeEndOfStream:
+		str = "EndOfStream"
+	case NalUnitTypeFiller:
+		str = "Filler"
+	case NalUnitTypeSpsExt:
+		str = "SpsExt"
+	case NalUnitTypeCodedSliceAux:
+		str = "NalUnitTypeCodedSliceAux"
+	default:
+		str = "Unknown"
+	}
+	str = str + "(" + strconv.FormatInt(int64(*n), 10) + ")"
+
+	return str
+}