github.com/pion/webrtc/v4@v4.0.1/pkg/media/h264reader/h264reader_test.go (about)

     1  // SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
     2  // SPDX-License-Identifier: MIT
     3  
     4  package h264reader
     5  
     6  import (
     7  	"bytes"
     8  	"io"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func CreateReader(h264 []byte, require *require.Assertions) *H264Reader {
    15  	reader, err := NewReader(bytes.NewReader(h264))
    16  
    17  	require.Nil(err)
    18  	require.NotNil(reader)
    19  
    20  	return reader
    21  }
    22  
    23  func TestDataDoesNotStartWithH264Header(t *testing.T) {
    24  	require := require.New(t)
    25  
    26  	testFunction := func(input []byte, expectedErr error) {
    27  		reader := CreateReader(input, require)
    28  		nal, err := reader.NextNAL()
    29  		require.ErrorIs(err, expectedErr)
    30  		require.Nil(nal)
    31  	}
    32  
    33  	h264Bytes1 := []byte{2}
    34  	testFunction(h264Bytes1, io.EOF)
    35  
    36  	h264Bytes2 := []byte{0, 2}
    37  	testFunction(h264Bytes2, io.EOF)
    38  
    39  	h264Bytes3 := []byte{0, 0, 2}
    40  	testFunction(h264Bytes3, io.EOF)
    41  
    42  	h264Bytes4 := []byte{0, 0, 2, 0}
    43  	testFunction(h264Bytes4, errDataIsNotH264Stream)
    44  
    45  	h264Bytes5 := []byte{0, 0, 0, 2}
    46  	testFunction(h264Bytes5, errDataIsNotH264Stream)
    47  }
    48  
    49  func TestParseHeader(t *testing.T) {
    50  	require := require.New(t)
    51  	h264Bytes := []byte{0x0, 0x0, 0x1, 0xAB}
    52  
    53  	reader := CreateReader(h264Bytes, require)
    54  
    55  	nal, err := reader.NextNAL()
    56  	require.Nil(err)
    57  
    58  	require.Equal(1, len(nal.Data))
    59  	require.True(nal.ForbiddenZeroBit)
    60  	require.Equal(uint32(0), nal.PictureOrderCount)
    61  	require.Equal(uint8(1), nal.RefIdc)
    62  	require.Equal(NalUnitTypeEndOfStream, nal.UnitType)
    63  }
    64  
    65  func TestEOF(t *testing.T) {
    66  	require := require.New(t)
    67  
    68  	testFunction := func(input []byte) {
    69  		reader := CreateReader(input, require)
    70  
    71  		nal, err := reader.NextNAL()
    72  		require.Equal(io.EOF, err)
    73  		require.Nil(nal)
    74  	}
    75  
    76  	h264Bytes1 := []byte{0, 0, 0, 1}
    77  	testFunction(h264Bytes1)
    78  
    79  	h264Bytes2 := []byte{0, 0, 1}
    80  	testFunction(h264Bytes2)
    81  
    82  	h264Bytes3 := []byte{}
    83  	testFunction(h264Bytes3)
    84  }
    85  
    86  func TestSkipSEI(t *testing.T) {
    87  	require := require.New(t)
    88  	h264Bytes := []byte{
    89  		0x0, 0x0, 0x0, 0x1, 0xAA,
    90  		0x0, 0x0, 0x0, 0x1, 0x6, // SEI
    91  		0x0, 0x0, 0x0, 0x1, 0xAB,
    92  	}
    93  
    94  	reader := CreateReader(h264Bytes, require)
    95  
    96  	nal, err := reader.NextNAL()
    97  	require.Nil(err)
    98  	require.Equal(byte(0xAA), nal.Data[0])
    99  
   100  	nal, err = reader.NextNAL()
   101  	require.Nil(err)
   102  	require.Equal(byte(0xAB), nal.Data[0])
   103  }
   104  
   105  func TestIssue1734_NextNal(t *testing.T) {
   106  	tt := [...][]byte{
   107  		[]byte("\x00\x00\x010\x00\x00\x01\x00\x00\x01"),
   108  		[]byte("\x00\x00\x00\x01\x00\x00\x01"),
   109  	}
   110  
   111  	for _, cur := range tt {
   112  		r, err := NewReader(bytes.NewReader(cur))
   113  		require.NoError(t, err)
   114  
   115  		// Just make sure it doesn't crash
   116  		for {
   117  			nal, err := r.NextNAL()
   118  
   119  			if err != nil || nal == nil {
   120  				break
   121  			}
   122  		}
   123  	}
   124  }
   125  
   126  func TestTrailing01AfterStartCode(t *testing.T) {
   127  	r, err := NewReader(bytes.NewReader([]byte{
   128  		0x0, 0x0, 0x0, 0x1, 0x01,
   129  		0x0, 0x0, 0x0, 0x1, 0x01,
   130  	}))
   131  	require.NoError(t, err)
   132  
   133  	for i := 0; i <= 1; i++ {
   134  		nal, err := r.NextNAL()
   135  		require.NoError(t, err)
   136  		require.NotNil(t, nal)
   137  	}
   138  }