From 59c0ef520feeda0f083e62d85d6aab268b9e0ff7 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Thu, 5 Jan 2023 13:10:28 +0100
Subject: [PATCH] fix crash when H264 or H265 parameters are nil (#1155)

in H264 and H264, parameters (VPS, SPS and PPS) are automatically
prepended to random access units. When parameters were nil, they were
prepended anyway, leading to a crash. Parameters are not prepended only
if they are valid.
---
 internal/formatprocessor/h264.go      | 15 +++++++++++----
 internal/formatprocessor/h264_test.go | 23 ++++++++++++++++++-----
 internal/formatprocessor/h265.go      | 20 +++++++++++++++-----
 internal/formatprocessor/h265_test.go | 23 ++++++++++++++++++++++-
 4 files changed, 66 insertions(+), 15 deletions(-)

diff --git a/internal/formatprocessor/h264.go b/internal/formatprocessor/h264.go
index a673c311..3c628836 100644
--- a/internal/formatprocessor/h264.go
+++ b/internal/formatprocessor/h264.go
@@ -138,6 +138,8 @@ func (t *formatProcessorH264) updateTrackParametersFromNALUs(nalus [][]byte) {
 }
 
 func (t *formatProcessorH264) remuxAccessUnit(nalus [][]byte) [][]byte {
+	var sps []byte
+	var pps []byte
 	addParameters := false
 	n := 0
 
@@ -154,7 +156,12 @@ func (t *formatProcessorH264) remuxAccessUnit(nalus [][]byte) [][]byte {
 		case h264.NALUTypeIDR: // prepend parameters if there's at least an IDR
 			if !addParameters {
 				addParameters = true
-				n += 2
+				sps = t.format.SafeSPS()
+				pps = t.format.SafePPS()
+
+				if sps != nil && pps != nil {
+					n += 2
+				}
 			}
 		}
 		n++
@@ -167,9 +174,9 @@ func (t *formatProcessorH264) remuxAccessUnit(nalus [][]byte) [][]byte {
 	filteredNALUs := make([][]byte, n)
 	i := 0
 
-	if addParameters {
-		filteredNALUs[0] = t.format.SafeSPS()
-		filteredNALUs[1] = t.format.SafePPS()
+	if addParameters && sps != nil && pps != nil {
+		filteredNALUs[0] = sps
+		filteredNALUs[1] = pps
 		i = 2
 	}
 
diff --git a/internal/formatprocessor/h264_test.go b/internal/formatprocessor/h264_test.go
index 88d58983..fa1f3b61 100644
--- a/internal/formatprocessor/h264_test.go
+++ b/internal/formatprocessor/h264_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/aler9/gortsplib/v2/pkg/codecs/h264"
 	"github.com/aler9/gortsplib/v2/pkg/format"
 	"github.com/pion/rtp"
 	"github.com/stretchr/testify/require"
@@ -20,13 +21,14 @@ func TestH264DynamicParams(t *testing.T) {
 
 	enc := forma.CreateEncoder()
 
-	pkts, err := enc.Encode([][]byte{{7, 1, 2, 3}}, 0) // SPS
+	pkts, err := enc.Encode([][]byte{{byte(h264.NALUTypeIDR)}}, 0)
 	require.NoError(t, err)
-	p.Process(&DataH264{RTPPackets: []*rtp.Packet{pkts[0]}}, false)
+	data := &DataH264{RTPPackets: []*rtp.Packet{pkts[0]}}
+	p.Process(data, true)
 
-	pkts, err = enc.Encode([][]byte{{8}}, 0) // PPS
-	require.NoError(t, err)
-	p.Process(&DataH264{RTPPackets: []*rtp.Packet{pkts[0]}}, false)
+	require.Equal(t, [][]byte{
+		{byte(h264.NALUTypeIDR)},
+	}, data.AU)
 
 	pkts, err = enc.Encode([][]byte{{7, 4, 5, 6}}, 0) // SPS
 	require.NoError(t, err)
@@ -38,6 +40,17 @@ func TestH264DynamicParams(t *testing.T) {
 
 	require.Equal(t, []byte{7, 4, 5, 6}, forma.SPS)
 	require.Equal(t, []byte{8, 1}, forma.PPS)
+
+	pkts, err = enc.Encode([][]byte{{byte(h264.NALUTypeIDR)}}, 0)
+	require.NoError(t, err)
+	data = &DataH264{RTPPackets: []*rtp.Packet{pkts[0]}}
+	p.Process(data, true)
+
+	require.Equal(t, [][]byte{
+		{0x07, 4, 5, 6},
+		{0x08, 1},
+		{byte(h264.NALUTypeIDR)},
+	}, data.AU)
 }
 
 func TestH264OversizedPackets(t *testing.T) {
diff --git a/internal/formatprocessor/h265.go b/internal/formatprocessor/h265.go
index 0d158f6d..e634f68f 100644
--- a/internal/formatprocessor/h265.go
+++ b/internal/formatprocessor/h265.go
@@ -154,6 +154,9 @@ func (t *formatProcessorH265) updateTrackParametersFromNALUs(nalus [][]byte) {
 }
 
 func (t *formatProcessorH265) remuxAccessUnit(nalus [][]byte) [][]byte {
+	var vps []byte
+	var sps []byte
+	var pps []byte
 	addParameters := false
 	n := 0
 
@@ -171,7 +174,14 @@ func (t *formatProcessorH265) remuxAccessUnit(nalus [][]byte) [][]byte {
 		case h265.NALUType_IDR_W_RADL, h265.NALUType_IDR_N_LP, h265.NALUType_CRA_NUT:
 			if !addParameters {
 				addParameters = true
-				n += 3
+
+				vps = t.format.SafeVPS()
+				sps = t.format.SafeSPS()
+				pps = t.format.SafePPS()
+
+				if vps != nil && sps != nil && pps != nil {
+					n += 3
+				}
 			}
 		}
 		n++
@@ -184,10 +194,10 @@ func (t *formatProcessorH265) remuxAccessUnit(nalus [][]byte) [][]byte {
 	filteredNALUs := make([][]byte, n)
 	i := 0
 
-	if addParameters {
-		filteredNALUs[0] = t.format.SafeVPS()
-		filteredNALUs[1] = t.format.SafeSPS()
-		filteredNALUs[2] = t.format.SafePPS()
+	if addParameters && vps != nil && sps != nil && pps != nil {
+		filteredNALUs[0] = vps
+		filteredNALUs[1] = sps
+		filteredNALUs[2] = pps
 		i = 3
 	}
 
diff --git a/internal/formatprocessor/h265_test.go b/internal/formatprocessor/h265_test.go
index 5172a358..e3c70c08 100644
--- a/internal/formatprocessor/h265_test.go
+++ b/internal/formatprocessor/h265_test.go
@@ -20,7 +20,16 @@ func TestH265DynamicParams(t *testing.T) {
 
 	enc := forma.CreateEncoder()
 
-	pkts, err := enc.Encode([][]byte{{byte(h265.NALUType_VPS_NUT) << 1, 1, 2, 3}}, 0)
+	pkts, err := enc.Encode([][]byte{{byte(h265.NALUType_CRA_NUT) << 1, 0}}, 0)
+	require.NoError(t, err)
+	data := &DataH265{RTPPackets: []*rtp.Packet{pkts[0]}}
+	p.Process(data, true)
+
+	require.Equal(t, [][]byte{
+		{byte(h265.NALUType_CRA_NUT) << 1, 0},
+	}, data.AU)
+
+	pkts, err = enc.Encode([][]byte{{byte(h265.NALUType_VPS_NUT) << 1, 1, 2, 3}}, 0)
 	require.NoError(t, err)
 	p.Process(&DataH265{RTPPackets: []*rtp.Packet{pkts[0]}}, false)
 
@@ -35,6 +44,18 @@ func TestH265DynamicParams(t *testing.T) {
 	require.Equal(t, []byte{byte(h265.NALUType_VPS_NUT) << 1, 1, 2, 3}, forma.VPS)
 	require.Equal(t, []byte{byte(h265.NALUType_SPS_NUT) << 1, 4, 5, 6}, forma.SPS)
 	require.Equal(t, []byte{byte(h265.NALUType_PPS_NUT) << 1, 7, 8, 9}, forma.PPS)
+
+	pkts, err = enc.Encode([][]byte{{byte(h265.NALUType_CRA_NUT) << 1, 0}}, 0)
+	require.NoError(t, err)
+	data = &DataH265{RTPPackets: []*rtp.Packet{pkts[0]}}
+	p.Process(data, true)
+
+	require.Equal(t, [][]byte{
+		{byte(h265.NALUType_VPS_NUT) << 1, 1, 2, 3},
+		{byte(h265.NALUType_SPS_NUT) << 1, 4, 5, 6},
+		{byte(h265.NALUType_PPS_NUT) << 1, 7, 8, 9},
+		{byte(h265.NALUType_CRA_NUT) << 1, 0},
+	}, data.AU)
 }
 
 func TestH265OversizedPackets(t *testing.T) {