Update codec parameters from first keyframe

pkg/h264/avcc.go

@@ -15,7 +15,22 @@ func RepairAVCC(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 	sps, pps := GetParameterSet(codec.FmtpLine)
 	ps := JoinNALU(sps, pps)
 
+	fmtpLineUpdated := false
+
 	return func(packet *rtp.Packet) {
+		// Update FmtpLine from first keyframe with parameter sets
+		// This fixes MSE aspect ratio issues when RTSP cameras don't send SPS/PPS in DESCRIBE
+		if !fmtpLineUpdated && ContainsParameterSets(packet.Payload) {
+			newFmtpLine := GetFmtpLine(packet.Payload)
+			if newFmtpLine != "" {
+				codec.FmtpLine = newFmtpLine
+				// Re-extract SPS/PPS with updated FmtpLine
+				sps, pps = GetParameterSet(codec.FmtpLine)
+				ps = JoinNALU(sps, pps)
+			}
+			fmtpLineUpdated = true
+		}
+
 		// this can happen for FLV from FFmpeg
 		if NALUType(packet.Payload) == NALUTypeSEI {
 			size := int(binary.BigEndian.Uint32(packet.Payload)) + 4
@@ -70,7 +85,7 @@ func SplitNALU(avcc []byte) [][]byte {
 
 func NALUTypes(avcc []byte) []byte {
 	var types []byte
-	for {
+	for len(avcc) >= 5 { // minimum: 4 bytes length + 1 byte NAL header
 		types = append(types, NALUType(avcc))
 
 		size := 4 + int(binary.BigEndian.Uint32(avcc))

pkg/h264/h264.go

@@ -125,15 +125,19 @@ func GetParameterSet(fmtp string) (sps, pps []byte) {
 func GetFmtpLine(avc []byte) string {
 	s := "packetization-mode=1"
 
-	for {
+	for len(avc) >= 5 { // minimum: 4 bytes length + 1 byte NAL header
 		size := 4 + int(binary.BigEndian.Uint32(avc))
 
 		switch NALUType(avc) {
 		case NALUTypeSPS:
-			s += ";profile-level-id=" + hex.EncodeToString(avc[5:8])
-			s += ";sprop-parameter-sets=" + base64.StdEncoding.EncodeToString(avc[4:size])
+			if len(avc) >= 8 && size <= len(avc) { // need at least profile-level-id
+				s += ";profile-level-id=" + hex.EncodeToString(avc[5:8])
+				s += ";sprop-parameter-sets=" + base64.StdEncoding.EncodeToString(avc[4:size])
+			}
 		case NALUTypePPS:
-			s += "," + base64.StdEncoding.EncodeToString(avc[4:size])
+			if size <= len(avc) {
+				s += "," + base64.StdEncoding.EncodeToString(avc[4:size])
+			}
 		}
 
 		if size < len(avc) {
@@ -142,4 +146,22 @@ func GetFmtpLine(avc []byte) string {
 			return s
 		}
 	}
+	return s
+}
+
+// ContainsParameterSets checks if payload contains both SPS and PPS
+func ContainsParameterSets(payload []byte) bool {
+	types := NALUTypes(payload)
+	hasSPS, hasPPS := false, false
+
+	for _, nalType := range types {
+		switch nalType {
+		case NALUTypeSPS:
+			hasSPS = true
+		case NALUTypePPS:
+			hasPPS = true
+		}
+	}
+
+	return hasSPS && hasPPS
 }

pkg/h264/rtp.go

@@ -21,6 +21,8 @@ func RTPDepay(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 
 	buf := make([]byte, 0, 512*1024) // 512K
 
+	fmtpLineUpdated := false
+
 	return func(packet *rtp.Packet) {
 		//log.Printf("[RTP] codec: %s, nalu: %2d, size: %6d, ts: %10d, pt: %2d, ssrc: %d, seq: %d, %v", codec.Name, packet.Payload[0]&0x1F, len(packet.Payload), packet.Timestamp, packet.PayloadType, packet.SSRC, packet.SequenceNumber, packet.Marker)
 
@@ -97,6 +99,19 @@ func RTPDepay(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 
 		//log.Printf("[AVC] %v, len: %d, ts: %10d, seq: %d", NALUTypes(payload), len(payload), packet.Timestamp, packet.SequenceNumber)
 
+		// Update FmtpLine from first keyframe with parameter sets
+		// This fixes MSE aspect ratio issues when RTSP cameras don't send SPS/PPS in DESCRIBE
+		if !fmtpLineUpdated && ContainsParameterSets(payload) {
+			newFmtpLine := GetFmtpLine(payload)
+			if newFmtpLine != "" {
+				codec.FmtpLine = newFmtpLine
+				// Re-extract SPS/PPS with updated FmtpLine
+				sps, pps = GetParameterSet(codec.FmtpLine)
+				ps = JoinNALU(sps, pps)
+			}
+			fmtpLineUpdated = true
+		}
+
 		clone := *packet
 		clone.Version = RTPPacketVersionAVC
 		clone.Payload = payload

pkg/h265/avcc.go

@@ -15,7 +15,22 @@ func RepairAVCC(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 	vds, sps, pps := GetParameterSet(codec.FmtpLine)
 	ps := h264.JoinNALU(vds, sps, pps)
 
+	fmtpLineUpdated := false
+
 	return func(packet *rtp.Packet) {
+		// Update FmtpLine from first keyframe with parameter sets
+		// This fixes MSE aspect ratio issues when RTSP cameras don't send VPS/SPS/PPS in DESCRIBE
+		if !fmtpLineUpdated && ContainsParameterSets(packet.Payload) {
+			newFmtpLine := GetFmtpLine(packet.Payload)
+			if newFmtpLine != "" {
+				codec.FmtpLine = newFmtpLine
+				// Re-extract VPS/SPS/PPS with updated FmtpLine
+				vds, sps, pps = GetParameterSet(codec.FmtpLine)
+				ps = h264.JoinNALU(vds, sps, pps)
+			}
+			fmtpLineUpdated = true
+		}
+
 		switch NALUType(packet.Payload) {
 		case NALUTypeIFrame, NALUTypeIFrame2, NALUTypeIFrame3:
 			clone := *packet

pkg/h265/helper.go

@@ -45,7 +45,7 @@ func IsKeyframe(b []byte) bool {
 
 func Types(data []byte) []byte {
 	var types []byte
-	for {
+	for len(data) >= 5 { // minimum: 4 bytes length + 1 byte NAL header
 		types = append(types, NALUType(data))
 
 		size := 4 + int(binary.BigEndian.Uint32(data))
@@ -74,3 +74,60 @@ func GetParameterSet(fmtp string) (vps, sps, pps []byte) {
 
 	return
 }
+
+// ContainsParameterSets checks if payload contains VPS, SPS and PPS
+func ContainsParameterSets(payload []byte) bool {
+	types := Types(payload)
+	hasVPS, hasSPS, hasPPS := false, false, false
+
+	for _, nalType := range types {
+		switch nalType {
+		case NALUTypeVPS:
+			hasVPS = true
+		case NALUTypeSPS:
+			hasSPS = true
+		case NALUTypePPS:
+			hasPPS = true
+		}
+	}
+
+	return hasVPS && hasSPS && hasPPS
+}
+
+// GetFmtpLine extracts VPS, SPS, PPS from AVCC payload and returns FmtpLine
+func GetFmtpLine(avcc []byte) string {
+	var vps, sps, pps []byte
+
+	for len(avcc) >= 5 { // minimum: 4 bytes length + 1 byte NAL header
+		size := 4 + int(binary.BigEndian.Uint32(avcc))
+
+		if size > len(avcc) {
+			break
+		}
+
+		switch NALUType(avcc) {
+		case NALUTypeVPS:
+			vps = avcc[4:size]
+		case NALUTypeSPS:
+			sps = avcc[4:size]
+		case NALUTypePPS:
+			pps = avcc[4:size]
+		}
+
+		if size < len(avcc) {
+			avcc = avcc[size:]
+		} else {
+			break
+		}
+	}
+
+	if len(vps) == 0 || len(sps) == 0 || len(pps) == 0 {
+		return ""
+	}
+
+	fmtp := "sprop-vps=" + base64.StdEncoding.EncodeToString(vps)
+	fmtp += ";sprop-sps=" + base64.StdEncoding.EncodeToString(sps)
+	fmtp += ";sprop-pps=" + base64.StdEncoding.EncodeToString(pps)
+
+	return fmtp
+}

pkg/h265/rtp.go

@@ -16,6 +16,8 @@ func RTPDepay(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 	var nuStart int
 	var seqNum uint16
 
+	fmtpLineUpdated := false
+
 	return func(packet *rtp.Packet) {
 		data := packet.Payload
 		if len(data) < 3 {
@@ -104,6 +106,19 @@ func RTPDepay(codec *core.Codec, handler core.HandlerFunc) core.HandlerFunc {
 
 		//log.Printf("[HEVC] %v, len: %d", Types(buf), len(buf))
 
+		// Update FmtpLine from first keyframe with parameter sets
+		// This fixes MSE aspect ratio issues when RTSP cameras don't send VPS/SPS/PPS in DESCRIBE
+		if !fmtpLineUpdated && ContainsParameterSets(buf) {
+			newFmtpLine := GetFmtpLine(buf)
+			if newFmtpLine != "" {
+				codec.FmtpLine = newFmtpLine
+				// Re-extract VPS/SPS/PPS with updated FmtpLine
+				vps, sps, pps = GetParameterSet(codec.FmtpLine)
+				ps = h264.JoinNALU(vps, sps, pps)
+			}
+			fmtpLineUpdated = true
+		}
+
 		clone := *packet
 		clone.Version = h264.RTPPacketVersionAVC
 		clone.Payload = buf

pkg/probe/consumer.go

@@ -5,6 +5,8 @@ import (
 	"strings"
 
 	"github.com/AlexxIT/go2rtc/pkg/core"
+	"github.com/AlexxIT/go2rtc/pkg/h264"
+	"github.com/AlexxIT/go2rtc/pkg/h265"
 )
 
 type Probe struct {
@@ -40,9 +42,29 @@ func Create(name string, query url.Values) *Probe {
 
 func (p *Probe) AddTrack(media *core.Media, codec *core.Codec, track *core.Receiver) error {
 	sender := core.NewSender(media, track.Codec)
-	sender.Handler = func(pkt *core.Packet) {
+
+	handler := func(pkt *core.Packet) {
 		p.Send += len(pkt.Payload)
 	}
+
+	// Apply format handlers to update FmtpLine from first keyframe
+	// This fixes MSE aspect ratio issues when RTSP cameras don't send SPS/PPS in DESCRIBE
+	switch track.Codec.Name {
+	case core.CodecH264:
+		if track.Codec.IsRTP() {
+			handler = h264.RTPDepay(track.Codec, handler)
+		} else {
+			handler = h264.RepairAVCC(track.Codec, handler)
+		}
+	case core.CodecH265:
+		if track.Codec.IsRTP() {
+			handler = h265.RTPDepay(track.Codec, handler)
+		} else {
+			handler = h265.RepairAVCC(track.Codec, handler)
+		}
+	}
+
+	sender.Handler = handler
 	sender.HandleRTP(track)
 	p.Senders = append(p.Senders, sender)
 	return nil