From 73830e29be960d248436ee595dfcbb9b4a88dbe5 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sun, 24 Jul 2022 01:07:25 +0200
Subject: [PATCH] rtph264: support streams that encode NALUs into Annex-B
 (https://github.com/aler9/rtsp-simple-server/issues/1029)

---
 pkg/rtph264/decoder.go      | 68 +++++++++++++++++++++++++++++++++++--
 pkg/rtph264/rtph264_test.go | 29 ++++++++++++++++
 2 files changed, 95 insertions(+), 2 deletions(-)

diff --git a/pkg/rtph264/decoder.go b/pkg/rtph264/decoder.go
index 14122655..5f3a56df 100644
--- a/pkg/rtph264/decoder.go
+++ b/pkg/rtph264/decoder.go
@@ -1,6 +1,7 @@
 package rtph264
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"time"
@@ -28,6 +29,8 @@ type Decoder struct {
 	fragmentedMode      bool
 	fragmentedParts     [][]byte
 	fragmentedSize      int
+	firstNALUParsed     bool
+	annexBMode          bool
 
 	// for DecodeUntilMarker()
 	naluBuffer [][]byte
@@ -46,6 +49,7 @@ func (d *Decoder) Decode(pkt *rtp.Packet) ([][]byte, time.Duration, error) {
 		}
 
 		typ := naluType(pkt.Payload[0] & 0x1F)
+
 		switch typ {
 		case naluTypeSTAPA:
 			var nalus [][]byte
@@ -77,6 +81,13 @@ func (d *Decoder) Decode(pkt *rtp.Packet) ([][]byte, time.Duration, error) {
 			}
 
 			d.firstPacketReceived = true
+
+			var err error
+			nalus, err = d.finalize(nalus)
+			if err != nil {
+				return nil, 0, err
+			}
+
 			return nalus, d.timeDecoder.Decode(pkt.Timestamp), nil
 
 		case naluTypeFUA: // first packet of a fragmented NALU
@@ -112,8 +123,17 @@ func (d *Decoder) Decode(pkt *rtp.Packet) ([][]byte, time.Duration, error) {
 			return nil, 0, fmt.Errorf("packet type not supported (%v)", typ)
 		}
 
+		nalus := [][]byte{pkt.Payload}
+
 		d.firstPacketReceived = true
-		return [][]byte{pkt.Payload}, d.timeDecoder.Decode(pkt.Timestamp), nil
+
+		var err error
+		nalus, err = d.finalize(nalus)
+		if err != nil {
+			return nil, 0, err
+		}
+
+		return nalus, d.timeDecoder.Decode(pkt.Timestamp), nil
 	}
 
 	// we are decoding a fragmented NALU
@@ -157,10 +177,18 @@ func (d *Decoder) Decode(pkt *rtp.Packet) ([][]byte, time.Duration, error) {
 	for _, p := range d.fragmentedParts {
 		n += copy(ret[n:], p)
 	}
+	nalus := [][]byte{ret}
 
 	d.fragmentedParts = d.fragmentedParts[:0]
 	d.fragmentedMode = false
-	return [][]byte{ret}, d.timeDecoder.Decode(pkt.Timestamp), nil
+
+	var err error
+	nalus, err = d.finalize(nalus)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	return nalus, d.timeDecoder.Decode(pkt.Timestamp), nil
 }
 
 // DecodeUntilMarker decodes NALUs from a RTP/H264 packet and puts them in a buffer.
@@ -188,3 +216,39 @@ func (d *Decoder) DecodeUntilMarker(pkt *rtp.Packet) ([][]byte, time.Duration, e
 
 	return ret, pts, nil
 }
+
+func (d *Decoder) finalize(nalus [][]byte) ([][]byte, error) {
+	// some cameras / servers wrap NALUs into Annex-B
+	if !d.firstNALUParsed {
+		d.firstNALUParsed = true
+
+		if len(nalus) == 1 {
+			nalu := nalus[0]
+
+			i := bytes.Index(nalu, []byte{0x00, 0x00, 0x00, 0x01})
+			if i >= 0 {
+				d.annexBMode = true
+
+				if !bytes.HasPrefix(nalu, []byte{0x00, 0x00, 0x00, 0x01}) {
+					nalu = append([]byte{0x00, 0x00, 0x00, 0x01}, nalu...)
+				}
+
+				return h264.AnnexBUnmarshal(nalu)
+			}
+		}
+	} else if d.annexBMode {
+		if len(nalus) != 1 {
+			return nil, fmt.Errorf("multiple NALUs in Annex-B mode are not supported")
+		}
+
+		nalu := nalus[0]
+
+		if !bytes.HasPrefix(nalu, []byte{0x00, 0x00, 0x00, 0x01}) {
+			nalu = append([]byte{0x00, 0x00, 0x00, 0x01}, nalu...)
+		}
+
+		return h264.AnnexBUnmarshal(nalu)
+	}
+
+	return nalus, nil
+}
diff --git a/pkg/rtph264/rtph264_test.go b/pkg/rtph264/rtph264_test.go
index 561c5a76..633626df 100644
--- a/pkg/rtph264/rtph264_test.go
+++ b/pkg/rtph264/rtph264_test.go
@@ -418,6 +418,35 @@ func TestDecodeSTAPAWithPadding(t *testing.T) {
 	}, nalus)
 }
 
+func TestDecodeAnnexB(t *testing.T) {
+	d := &Decoder{}
+	d.Init()
+
+	for i := 0; i < 2; i++ {
+		nalus, _, err := d.Decode(&rtp.Packet{
+			Header: rtp.Header{
+				Version:        2,
+				Marker:         true,
+				PayloadType:    96,
+				SequenceNumber: 17647,
+				Timestamp:      2289531307,
+				SSRC:           0x9dbb7812,
+			},
+			Payload: mergeBytes(
+				[]byte{0x00, 0x00, 0x00, 0x01},
+				[]byte{0x01, 0x02, 0x03, 0x04},
+				[]byte{0x00, 0x00, 0x00, 0x01},
+				[]byte{0x01, 0x02, 0x03, 0x04},
+			),
+		})
+		require.NoError(t, err)
+		require.Equal(t, [][]byte{
+			{0x01, 0x02, 0x03, 0x04},
+			{0x01, 0x02, 0x03, 0x04},
+		}, nalus)
+	}
+}
+
 func TestDecodeErrors(t *testing.T) {
 	for _, ca := range []struct {
 		name string