From 30c7f6a15b5fef5283520fcc0867eef08018bfdb Mon Sep 17 00:00:00 2001
From: "Alfred E. Heggestad" <114750+alfredh@users.noreply.github.com>
Date: Wed, 10 Apr 2024 14:38:23 +0200
Subject: [PATCH] h264: add STAP-A decode with long startcodes

---
 include/re_h264.h |  1 +
 src/h264/nal.c    | 53 ++++++++++++++++++++++-------
 test/h264.c       | 87 ++++++++++++++++++++++++++++++++++-------------
 3 files changed, 106 insertions(+), 35 deletions(-)

diff --git a/include/re_h264.h b/include/re_h264.h
index 1268e1172..e5674ea8e 100644
--- a/include/re_h264.h
+++ b/include/re_h264.h
@@ -104,3 +104,4 @@ bool h264_is_keyframe(int type);
 int  h264_stap_encode(struct mbuf *mb, const uint8_t *frame,
 		      size_t frame_sz);
 int  h264_stap_decode_annexb(struct mbuf *mb_frame, struct mbuf *mb_pkt);
+int  h264_stap_decode_annexb_long(struct mbuf *mb_frame, struct mbuf *mb_pkt);
diff --git a/src/h264/nal.c b/src/h264/nal.c
index 180cedd33..1ac333a09 100644
--- a/src/h264/nal.c
+++ b/src/h264/nal.c
@@ -365,17 +365,8 @@ int h264_stap_encode(struct mbuf *mb, const uint8_t *frame,
 }
 
 
-/**
- * Decode STAP-A payload and convert to Annex-B NAL units
- *
- * @param mb_frame Target mbuffer for frame with multiple NAL units
- * @param mb_pkt   Input packet with STAP-A payload
- *
- * @return 0 if success, otherwise errorcode
- *
- * NOTE: The NAL header must be decoded outside
- */
-int h264_stap_decode_annexb(struct mbuf *mb_frame, struct mbuf *mb_pkt)
+static int stap_decode_annexb_int(struct mbuf *mb_frame, struct mbuf *mb_pkt,
+				  bool long_startcode)
 {
 	if (!mb_frame || !mb_pkt)
 		return EINVAL;
@@ -398,8 +389,13 @@ int h264_stap_decode_annexb(struct mbuf *mb_frame, struct mbuf *mb_pkt)
 #endif
 
 		static const uint8_t sc3[] = {0, 0, 1};
+		static const uint8_t sc4[] = {0, 0, 0, 1};
 
-		int err = mbuf_write_mem(mb_frame, sc3, sizeof(sc3));
+		int err;
+		if (long_startcode)
+			err = mbuf_write_mem(mb_frame, sc4, sizeof(sc4));
+		else
+			err = mbuf_write_mem(mb_frame, sc3, sizeof(sc3));
 		err |= mbuf_write_mem(mb_frame, mbuf_buf(mb_pkt), len);
 		if (err)
 			return err;
@@ -409,3 +405,36 @@ int h264_stap_decode_annexb(struct mbuf *mb_frame, struct mbuf *mb_pkt)
 
 	return 0;
 }
+
+
+/**
+ * Decode STAP-A payload and convert to Annex-B NAL units
+ *
+ * @param mb_frame Target mbuffer for frame with multiple NAL units
+ * @param mb_pkt   Input packet with STAP-A payload
+ *
+ * @return 0 if success, otherwise errorcode
+ *
+ * NOTE: The NAL header must be decoded outside
+ */
+int h264_stap_decode_annexb(struct mbuf *mb_frame, struct mbuf *mb_pkt)
+{
+	return stap_decode_annexb_int(mb_frame, mb_pkt, false);
+}
+
+
+/**
+ * Decode STAP-A payload and convert to Annex-B NAL units
+ * with long startcode (0x00 0x00 0x00 0x01).
+ *
+ * @param mb_frame Target mbuffer for frame with multiple NAL units
+ * @param mb_pkt   Input packet with STAP-A payload
+ *
+ * @return 0 if success, otherwise errorcode
+ *
+ * NOTE: The NAL header must be decoded outside
+ */
+int h264_stap_decode_annexb_long(struct mbuf *mb_frame, struct mbuf *mb_pkt)
+{
+	return stap_decode_annexb_int(mb_frame, mb_pkt, true);
+}
diff --git a/test/h264.c b/test/h264.c
index d10738937..186724b10 100644
--- a/test/h264.c
+++ b/test/h264.c
@@ -69,26 +69,9 @@ static void dump_rtp(const uint8_t *p, size_t size)
 #endif
 
 
-static int test_h264_stap_a_encode(void)
+static int test_h264_stap_a_encode_base(const uint8_t *frame, size_t len,
+				   bool long_startcode)
 {
-	static const uint8_t frame[] = {
-
-		/* AUD */
-		0x00, 0x00, 0x01,
-		0x09, 0x10,
-
-		/* SPS */
-		0x00, 0x00, 0x01,
-		0x67, 0x42, 0xc0, 0x1f, 0x8c, 0x8d, 0x40,
-
-		/* PPS */
-		0x00, 0x00, 0x01,
-		0x68, 0xce, 0x3c, 0x80,
-
-		/* IDR_SLICE */
-		0x00, 0x00, 0x01,
-		0x65, 0xb8, 0x00, 0x04, 0x00, 0x00, 0x05, 0x39,
-	};
 	enum { MAX_NRI = 3 };
 	struct mbuf *mb_pkt   = mbuf_alloc(256);
 	struct mbuf *mb_frame = mbuf_alloc(256);
@@ -100,7 +83,7 @@ static int test_h264_stap_a_encode(void)
 		goto out;
 	}
 
-	err = h264_stap_encode(mb_pkt, frame, sizeof(frame));
+	err = h264_stap_encode(mb_pkt, frame, len);
 	if (err)
 		goto out;
 
@@ -112,10 +95,16 @@ static int test_h264_stap_a_encode(void)
 	ASSERT_EQ(MAX_NRI,          hdr.nri);              /* NOTE: max NRI */
 	ASSERT_EQ(H264_NALU_STAP_A, hdr.type);
 
-	err = h264_stap_decode_annexb(mb_frame, mb_pkt);
-	ASSERT_EQ(0, err);
+	if (long_startcode) {
+		err = h264_stap_decode_annexb_long(mb_frame, mb_pkt);
+		ASSERT_EQ(0, err);
+	}
+	else {
+		err = h264_stap_decode_annexb(mb_frame, mb_pkt);
+		ASSERT_EQ(0, err);
+	}
 
-	TEST_MEMCMP(frame, sizeof(frame), mb_frame->buf, mb_frame->end);
+	TEST_MEMCMP(frame, len, mb_frame->buf, mb_frame->end);
 
  out:
 	mem_deref(mb_frame);
@@ -125,6 +114,58 @@ static int test_h264_stap_a_encode(void)
 }
 
 
+static int test_h264_stap_a_encode(void)
+{
+	static const uint8_t frame[] = {
+
+		/* AUD */
+		0x00, 0x00, 0x01,
+		0x09, 0x10,
+
+		/* SPS */
+		0x00, 0x00, 0x01,
+		0x67, 0x42, 0xc0, 0x1f, 0x8c, 0x8d, 0x40,
+
+		/* PPS */
+		0x00, 0x00, 0x01,
+		0x68, 0xce, 0x3c, 0x80,
+
+		/* IDR_SLICE */
+		0x00, 0x00, 0x01,
+		0x65, 0xb8, 0x00, 0x04, 0x00, 0x00, 0x05, 0x39,
+	};
+	static const uint8_t frame_long[] = {
+
+		/* AUD */
+		0x00, 0x00, 0x00, 0x01,
+		0x09, 0x10,
+
+		/* SPS */
+		0x00, 0x00, 0x00, 0x01,
+		0x67, 0x42, 0xc0, 0x1f, 0x8c, 0x8d, 0x40,
+
+		/* PPS */
+		0x00, 0x00, 0x00, 0x01,
+		0x68, 0xce, 0x3c, 0x80,
+
+		/* IDR_SLICE */
+		0x00, 0x00, 0x00, 0x01,
+		0x65, 0xb8, 0x00, 0x04, 0x00, 0x00, 0x05, 0x39,
+	};
+	int err;
+
+	err = test_h264_stap_a_encode_base(frame, sizeof(frame), false);
+	TEST_ERR(err);
+
+	err = test_h264_stap_a_encode_base(frame_long, sizeof(frame_long),
+					   true);
+	TEST_ERR(err);
+
+ out:
+	return err;
+}
+
+
 static int test_h264_stap_a_decode(void)
 {
 	static const uint8_t pkt[] = {