fix formats.H265.PTSEqualsDTS()

bluenviron · Aug 8, 2023 · fc98ca5 · fc98ca5
1 parent ea96245
commit fc98ca5
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 11 deletions.
diff --git a/pkg/formats/h264.go b/pkg/formats/h264.go
@@ -41,17 +41,15 @@ func rtpH264ContainsIDR(pkt *rtp.Packet) bool {
 				return false
 			}
 
-			nalu := payload[:size]
-			payload = payload[size:]
+			var nalu []byte
+			nalu, payload = payload[:size], payload[size:]
 
 			typ = h264.NALUType(nalu[0] & 0x1F)
 			if typ == h264.NALUTypeIDR {
 				return true
 			}
 		}
 
-		return false
-
 	case 28: // FU-A
 		if len(pkt.Payload) < 2 {
 			return false
@@ -63,11 +61,12 @@ func rtpH264ContainsIDR(pkt *rtp.Packet) bool {
 		}
 
 		typ := h264.NALUType(pkt.Payload[1] & 0x1F)
-		return (typ == h264.NALUTypeIDR)
-
-	default:
-		return false
+		if typ == h264.NALUTypeIDR {
+			return true
+		}
 	}
+
+	return false
 }
 
 // H264 is a RTP format for the H264 codec, defined in MPEG-4 part 10.

diff --git a/pkg/formats/h265.go b/pkg/formats/h265.go
@@ -6,11 +6,67 @@ import (
 	"strconv"
 	"sync"
 
+	"github.com/bluenviron/mediacommon/pkg/codecs/h265"
 	"github.com/pion/rtp"
 
 	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtph265"
 )
 
+// check whether a RTP/H265 packet is random access, without decoding the packet.
+func rtpH265IsRandomAccess(pkt *rtp.Packet) bool {
+	if len(pkt.Payload) == 0 {
+		return false
+	}
+
+	typ := h265.NALUType((pkt.Payload[0] >> 1) & 0b111111)
+
+	switch typ {
+	case h265.NALUType_IDR_W_RADL, h265.NALUType_IDR_N_LP, h265.NALUType_CRA_NUT:
+		return true
+
+	case h265.NALUType_AggregationUnit:
+		payload := pkt.Payload[2:]
+
+		for len(payload) > 0 {
+			if len(payload) < 2 {
+				return false
+			}
+
+			size := uint16(payload[0])<<8 | uint16(payload[1])
+			payload = payload[2:]
+
+			if size == 0 || int(size) > len(payload) {
+				return false
+			}
+
+			var nalu []byte
+			nalu, payload = payload[:size], payload[size:]
+
+			typ = h265.NALUType((nalu[0] >> 1) & 0b111111)
+			if typ == h265.NALUType_IDR_W_RADL || typ == h265.NALUType_IDR_N_LP || typ == h265.NALUType_CRA_NUT {
+				return true
+			}
+		}
+
+	case h265.NALUType_FragmentationUnit:
+		if len(pkt.Payload) < 3 {
+			return false
+		}
+
+		start := pkt.Payload[1] >> 7
+		if start != 1 {
+			return false
+		}
+
+		typ := h265.NALUType(pkt.Payload[2] & 0b111111)
+		if typ == h265.NALUType_IDR_W_RADL || typ == h265.NALUType_IDR_N_LP || typ == h265.NALUType_CRA_NUT {
+			return true
+		}
+	}
+
+	return false
+}
+
 // H265 is a RTP format for the H265 codec.
 // Specification: https://datatracker.ietf.org/doc/html/rfc7798
 type H265 struct {
@@ -111,8 +167,8 @@ func (f *H265) FMTP() map[string]string {
 }
 
 // PTSEqualsDTS implements Format.
-func (f *H265) PTSEqualsDTS(*rtp.Packet) bool {
-	return true
+func (f *H265) PTSEqualsDTS(pkt *rtp.Packet) bool {
+	return rtpH265IsRandomAccess(pkt)
 }
 
 // CreateDecoder creates a decoder able to decode the content of the format.

diff --git a/pkg/formats/h265_test.go b/pkg/formats/h265_test.go
@@ -3,6 +3,7 @@ package formats
 import (
 	"testing"
 
+	"github.com/bluenviron/mediacommon/pkg/codecs/h265"
 	"github.com/pion/rtp"
 	"github.com/stretchr/testify/require"
 )
@@ -16,7 +17,6 @@ func TestH265Attributes(t *testing.T) {
 	}
 	require.Equal(t, "H265", format.Codec())
 	require.Equal(t, 90000, format.ClockRate())
-	require.Equal(t, true, format.PTSEqualsDTS(&rtp.Packet{}))
 
 	vps, sps, pps := format.SafeParams()
 	require.Equal(t, []byte{0x01, 0x02}, vps)
@@ -31,6 +31,22 @@ func TestH265Attributes(t *testing.T) {
 	require.Equal(t, []byte{0x0B, 0x0C}, pps)
 }
 
+func TestH265PTSEqualsDTS(t *testing.T) {
+	format := &H265{
+		PayloadTyp: 96,
+		VPS:        []byte{0x01, 0x02},
+		SPS:        []byte{0x03, 0x04},
+		PPS:        []byte{0x05, 0x06},
+	}
+
+	require.Equal(t, true, format.PTSEqualsDTS(&rtp.Packet{
+		Payload: []byte{byte(h265.NALUType_CRA_NUT) << 1},
+	}))
+	require.Equal(t, false, format.PTSEqualsDTS(&rtp.Packet{
+		Payload: []byte{byte(h265.NALUType_TRAIL_N) << 1},
+	}))
+}
+
 func TestH265DecEncoder(t *testing.T) {
 	format := &H265{}