Merge pull request #42 from GyeongHoKim/fix/rpi

fix(rtsp): populate H.264 SDP parameter sets from first IDR

Author: GyeongHoKim

internal/rtsp/livesource.go

@@ -30,6 +30,12 @@ type AccessUnit struct {
 // stuck clients do not balloon memory.
 const liveBufferSize = 8
 
+// H.264 NAL unit types used by parameter-set extraction (ITU-T H.264 §7.3.1).
+const (
+	h264NALTypeSPS = 7
+	h264NALTypePPS = 8
+)
+
 // LiveSource is a Source backed by an external H.264 producer. The producer
 // (e.g. internal/rpicamera) pushes AccessUnit values onto Push; the source's
 // Run goroutine encodes them into RTP packets and writes them to the
@@ -55,8 +61,11 @@ type LiveSource struct {
 }
 
 // NewLiveSource builds a LiveSource for the codec described by probe. SPS
-// and PPS may be empty on the probe — clients learn them from in-band NAL
-// units once the stream starts. logger may be nil.
+// and PPS may be empty on the probe — they are extracted from the first
+// IDR access unit and written into the gortsplib H264 format before the
+// ready signal releases DESCRIBE, so the SDP carries sprop-parameter-sets
+// and clients are not stuck on "non-existing PPS referenced". logger may
+// be nil.
 func NewLiveSource(probe *ProbeResult, logger *slog.Logger) *LiveSource {
 	if logger == nil {
 		logger = obs.Discard()
@@ -140,6 +149,7 @@ func (l *LiveSource) Run(ctx context.Context) error {
 				return nil
 			}
 			if hasH264IDR(au.NALs) {
+				l.fillH264ParameterSets(au.NALs)
 				l.readyOnce.Do(func() { close(l.ready) })
 			}
 			select {
@@ -176,6 +186,44 @@ func (l *LiveSource) writeAU(enc *rtph264.Encoder, au AccessUnit) error {
 	return nil
 }
 
+// fillH264ParameterSets pulls SPS (NAL type 7) and PPS (NAL type 8) out of
+// the first IDR access unit and writes them into the attached gortsplib
+// format.H264 so DESCRIBE emits sprop-parameter-sets in the SDP.
+//
+// Without this, clients that subscribe between IDRs never see SPS/PPS in
+// either the SDP (empty) or in-band (already passed) and decode loops
+// indefinitely on "non-existing PPS referenced".
+func (l *LiveSource) fillH264ParameterSets(nals [][]byte) {
+	if l.media == nil || len(l.media.Formats) == 0 {
+		return
+	}
+	h, ok := l.media.Formats[0].(*format.H264)
+	if !ok {
+		return
+	}
+	if len(h.SPS) > 0 && len(h.PPS) > 0 {
+		return
+	}
+	var sps, pps []byte
+	for _, nal := range nals {
+		if len(nal) == 0 {
+			continue
+		}
+		switch nal[0] & 0x1F {
+		case h264NALTypeSPS:
+			sps = nal
+		case h264NALTypePPS:
+			pps = nal
+		}
+	}
+	if len(h.SPS) == 0 && len(sps) > 0 {
+		h.SPS = sps
+	}
+	if len(h.PPS) == 0 && len(pps) > 0 {
+		h.PPS = pps
+	}
+}
+
 // hasH264IDR scans access-unit NAL units for an IDR slice (NAL type 5).
 func hasH264IDR(au [][]byte) bool {
 	for _, nal := range au {

internal/rtsp/livesource_test.go

@@ -1,11 +1,15 @@
 package rtsp
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/bluenviron/gortsplib/v5/pkg/description"
+	"github.com/bluenviron/gortsplib/v5/pkg/format"
 )
 
 func TestHasH264IDR(t *testing.T) {
@@ -31,6 +35,75 @@ func TestHasH264IDR(t *testing.T) {
 	}
 }
 
+func TestLiveSourceFillsH264ParameterSetsFromFirstIDR(t *testing.T) {
+	t.Parallel()
+	sps := []byte{0x67, 0x42, 0xc0, 0x1e}
+	pps := []byte{0x68, 0xce, 0x3c, 0x80}
+	idr := []byte{0x65, 0xb8, 0x00, 0x01}
+
+	cases := []struct {
+		name    string
+		probe   *ProbeResult
+		nals    [][]byte
+		wantSPS []byte
+		wantPPS []byte
+	}{
+		{
+			name:    "writes SPS+PPS when probe is empty",
+			probe:   &ProbeResult{Codec: CodecH264},
+			nals:    [][]byte{sps, pps, idr},
+			wantSPS: sps,
+			wantPPS: pps,
+		},
+		{
+			name:    "missing PPS leaves PPS untouched",
+			probe:   &ProbeResult{Codec: CodecH264},
+			nals:    [][]byte{sps, idr},
+			wantSPS: sps,
+			wantPPS: nil,
+		},
+		{
+			name:    "preserves probe-supplied SPS+PPS",
+			probe:   &ProbeResult{Codec: CodecH264, SPS: []byte{0xAA}, PPS: []byte{0xBB}},
+			nals:    [][]byte{sps, pps, idr},
+			wantSPS: []byte{0xAA},
+			wantPPS: []byte{0xBB},
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			ls := NewLiveSource(tc.probe, nil)
+			h := &format.H264{
+				PayloadTyp:        96,
+				PacketizationMode: 1,
+				SPS:               tc.probe.SPS,
+				PPS:               tc.probe.PPS,
+			}
+			media := &description.Media{
+				Type:    description.MediaTypeVideo,
+				Formats: []format.Format{h},
+			}
+			ls.AttachStream(nil, media)
+
+			ls.fillH264ParameterSets(tc.nals)
+
+			if !bytes.Equal(h.SPS, tc.wantSPS) {
+				t.Fatalf("SPS=%x want %x", h.SPS, tc.wantSPS)
+			}
+			if !bytes.Equal(h.PPS, tc.wantPPS) {
+				t.Fatalf("PPS=%x want %x", h.PPS, tc.wantPPS)
+			}
+		})
+	}
+}
+
+func TestLiveSourceFillH264ParameterSetsNoMediaIsNoOp(t *testing.T) {
+	t.Parallel()
+	ls := NewLiveSource(&ProbeResult{Codec: CodecH264}, nil)
+	ls.fillH264ParameterSets([][]byte{{0x67, 0x42}, {0x68, 0xce}})
+}
+
 func TestLiveSourceReadyBeforeIDRBlocks(t *testing.T) {
 	t.Parallel()
 	probe := &ProbeResult{Codec: CodecH264, Width: 1920, Height: 1080, FPS: 30}