Fix H265 support for MSE in Safari

pkg/h265/helper.go

@@ -2,19 +2,57 @@ package h265
 
 import (
 	"encoding/base64"
+	"encoding/binary"
 	"github.com/AlexxIT/go2rtc/pkg/streamer"
 )
 
 const (
-	NALUnitTypeIFrame = 19
+	NALUTypePFrame  = 1
+	NALUTypeIFrame  = 19
+	NALUTypeIFrame2 = 20
+	NALUTypeIFrame3 = 21
+	NALUTypeVPS     = 32
+	NALUTypeSPS     = 33
+	NALUTypePPS     = 34
+	NALUTypeFU      = 49
 )
 
-func NALUnitType(b []byte) byte {
+func NALUType(b []byte) byte {
-	return b[4] >> 1
+	return (b[4] >> 1) & 0x3F
 }
 
 func IsKeyframe(b []byte) bool {
-	return NALUnitType(b) == NALUnitTypeIFrame
+	for {
+		switch NALUType(b) {
+		case NALUTypePFrame:
+			return false
+		case NALUTypeIFrame, NALUTypeIFrame2, NALUTypeIFrame3:
+			return true
+		}
+
+		size := int(binary.BigEndian.Uint32(b)) + 4
+		if size < len(b) {
+			b = b[size:]
+			continue
+		} else {
+			return false
+		}
+	}
+}
+
+func Types(data []byte) []byte {
+	var types []byte
+	for {
+		types = append(types, NALUType(data))
+
+		size := 4 + int(binary.BigEndian.Uint32(data))
+		if size < len(data) {
+			data = data[size:]
+		} else {
+			break
+		}
+	}
+	return types
 }
 
 func GetParameterSet(fmtp string) (vps, sps, pps []byte) {

pkg/h265/rtp.go

@@ -1,6 +1,7 @@
 package h265
 
 import (
+	"encoding/binary"
 	"github.com/AlexxIT/go2rtc/pkg/h264"
 	"github.com/AlexxIT/go2rtc/pkg/streamer"
 	"github.com/deepch/vdk/codec/h265parser"
@@ -8,77 +9,59 @@ import (
 )
 
 func RTPDepay(track *streamer.Track) streamer.WrapperFunc {
-	vps, sps, pps := GetParameterSet(track.Codec.FmtpLine)
+	//vps, sps, pps := GetParameterSet(track.Codec.FmtpLine)
+	//ps := h264.EncodeAVC(vps, sps, pps)
 
-	var buffer []byte
+	buf := make([]byte, 0, 512*1024) // 512K
+	var nuStart int
 
 	return func(push streamer.WriterFunc) streamer.WriterFunc {
 		return func(packet *rtp.Packet) error {
-			nut := (packet.Payload[0] >> 1) & 0x3f
+			data := packet.Payload
-			//fmt.Printf(
+			nuType := (data[0] >> 1) & 0x3F
-			//	"[RTP] codec: %s, nalu: %2d, size: %6d, ts: %10d, pt: %2d, ssrc: %d, seq: %d\n",
+			//log.Printf("[RTP] codec: %s, nalu: %2d, size: %6d, ts: %10d, pt: %2d, ssrc: %d, seq: %d, %v", track.Codec.Name, nuType, len(packet.Payload), packet.Timestamp, packet.PayloadType, packet.SSRC, packet.SequenceNumber, packet.Marker)
-			//	track.Codec.Name, nut, len(packet.Payload), packet.Timestamp,
-			//	packet.PayloadType, packet.SSRC, packet.SequenceNumber,
-			//)
 
-			switch nut {
+			if nuType == NALUTypeFU {
-			case h265parser.NAL_UNIT_UNSPECIFIED_49:
-				data := packet.Payload
 				switch data[2] >> 6 {
 				case 2: // begin
-					buffer = []byte{
+					nuType = data[2] & 0x3F
-						(data[0] & 0x81) | (data[2] & 0x3f << 1), data[1],
+
-					}
+					// push PS data before keyframe
-					buffer = append(buffer, data[3:]...)
+					//if len(buf) == 0 && nuType >= 19 && nuType <= 21 {
+					//	buf = append(buf, ps...)
+					//}
+
+					nuStart = len(buf)
+					buf = append(buf, 0, 0, 0, 0) // NAL unit size
+					buf = append(buf, (data[0]&0x81)|(nuType<<1), data[1])
+					buf = append(buf, data[3:]...)
 					return nil
 				case 0: // continue
-					buffer = append(buffer, data[3:]...)
+					buf = append(buf, data[3:]...)
 					return nil
 				case 1: // end
-					packet.Payload = append(buffer, data[3:]...)
+					buf = append(buf, data[3:]...)
+					binary.BigEndian.PutUint32(buf[nuStart:], uint32(len(buf)-nuStart-4))
 				}
-			case h265parser.NAL_UNIT_VPS:
+			} else {
-				vps = packet.Payload
+				nuStart = len(buf)
-				return nil
+				buf = append(buf, 0, 0, 0, 0) // NAL unit size
-			case h265parser.NAL_UNIT_SPS:
+				buf = append(buf, data...)
-				sps = packet.Payload
+				binary.BigEndian.PutUint32(buf[nuStart:], uint32(len(data)))
-				return nil
-			case h265parser.NAL_UNIT_PPS:
-				pps = packet.Payload
-				return nil
-			default:
-				//panic("not implemented")
 			}
 
-			var clone rtp.Packet
+			// collect all NAL Units for Access Unit
-
+			if !packet.Marker {
-			nut = (packet.Payload[0] >> 1) & 0x3f
+				return nil
-			if nut >= h265parser.NAL_UNIT_CODED_SLICE_BLA_W_LP && nut <= h265parser.NAL_UNIT_CODED_SLICE_CRA {
-				clone = *packet
-				clone.Version = h264.RTPPacketVersionAVC
-				clone.Payload = h264.EncodeAVC(vps)
-				if err := push(&clone); err != nil {
-					return err
-				}
-
-				clone = *packet
-				clone.Version = h264.RTPPacketVersionAVC
-				clone.Payload = h264.EncodeAVC(sps)
-				if err := push(&clone); err != nil {
-					return err
-				}
-
-				clone = *packet
-				clone.Version = h264.RTPPacketVersionAVC
-				clone.Payload = h264.EncodeAVC(pps)
-				if err := push(&clone); err != nil {
-					return err
-				}
 			}
 
-			clone = *packet
+			//log.Printf("[HEVC] %v, len: %d", Types(buf), len(buf))
+
+			clone := *packet
 			clone.Version = h264.RTPPacketVersionAVC
-			clone.Payload = h264.EncodeAVC(packet.Payload)
+			clone.Payload = buf
+
+			buf = buf[:0]
 
 			return push(&clone)
 		}
@@ -106,13 +89,13 @@ func SafariPay(mtu uint16) streamer.WrapperFunc {
 
 			var start byte
 
-			nut := (data[4] >> 1) & 0b111111
+			nuType := (data[4] >> 1) & 0b111111
 			//fmt.Printf("[H265] nut: %2d, size: %6d, data: %16x\n", nut, len(data), data[4:20])
 			switch {
-			case nut >= h265parser.NAL_UNIT_VPS && nut <= h265parser.NAL_UNIT_PPS:
+			case nuType >= h265parser.NAL_UNIT_VPS && nuType <= h265parser.NAL_UNIT_PPS:
 				buffer = append(buffer, data...)
 				return nil
-			case nut >= h265parser.NAL_UNIT_CODED_SLICE_BLA_W_LP && nut <= h265parser.NAL_UNIT_CODED_SLICE_CRA:
+			case nuType >= h265parser.NAL_UNIT_CODED_SLICE_BLA_W_LP && nuType <= h265parser.NAL_UNIT_CODED_SLICE_CRA:
 				buffer = append([]byte{3}, buffer...)
 				data = append(buffer, data...)
 				start = 1