From 178e759e147eb10509db13ff11370463bf39caaa Mon Sep 17 00:00:00 2001
From: Robert Edmonds <edmonds@users.noreply.github.com>
Date: Tue, 19 Mar 2024 01:22:22 -0400
Subject: [PATCH] H264RtpDepacketizer: De-packetize access units rather than
 individual NALUs

This commit updates the `H264RtpDepacketizer` to accumulate the NALUs for a particular RTP timestamp into a single output message, rather than returning each NALU as an individual message. This helps decoders which may want to see the non-VCL SPS/PPS/etc. NALUs in the same access unit as a VCL NALU rather than as standalone messages.

Each NALU in the access unit buffer is prepended with an H.264 Annex B start code 0x00, 0x00, 0x01.
---
 src/h264rtpdepacketizer.cpp | 83 +++++++++++++++++--------------------
 1 file changed, 38 insertions(+), 45 deletions(-)

diff --git a/src/h264rtpdepacketizer.cpp b/src/h264rtpdepacketizer.cpp
index 024a46b3b..a624916d0 100644
--- a/src/h264rtpdepacketizer.cpp
+++ b/src/h264rtpdepacketizer.cpp
@@ -10,23 +10,12 @@
 
 #include "h264rtpdepacketizer.hpp"
 #include "nalunit.hpp"
-#include "track.hpp"
 
-#include "impl/logcounter.hpp"
-
-#include <cmath>
-#include <utility>
-
-#ifdef _WIN32
-#include <winsock2.h>
-#else
-#include <arpa/inet.h>
-#endif
+#include "impl/internals.hpp"
 
 namespace rtc {
 
-const unsigned long stapaHeaderSize = 1;
-const auto fuaHeaderSize = 2;
+const binary naluStartCode = {byte{0}, byte{0}, byte{1}};
 
 const uint8_t naluTypeSTAPA = 24;
 const uint8_t naluTypeFUA = 28;
@@ -34,54 +23,53 @@ const uint8_t naluTypeFUA = 28;
 message_vector H264RtpDepacketizer::buildFrames(message_vector::iterator begin,
                                                 message_vector::iterator end, uint32_t timestamp) {
 	message_vector out = {};
-	auto fua_buffer = std::vector<std::byte>{};
+	auto accessUnit = binary{};
 	auto frameInfo = std::make_shared<FrameInfo>(timestamp);
+	auto nFrags = 0;
 
-	for (auto it = begin; it != end; it++) {
+	for (auto it = begin; it != end; ++it) {
 		auto pkt = it->get();
 		auto pktParsed = reinterpret_cast<const rtc::RtpHeader *>(pkt->data());
-		auto headerSize =
-		    sizeof(rtc::RtpHeader) + pktParsed->csrcCount() + pktParsed->getExtensionHeaderSize();
-		auto nalUnitHeader = NalUnitHeader{std::to_integer<uint8_t>(pkt->at(headerSize))};
-
-		if (fua_buffer.size() != 0 || nalUnitHeader.unitType() == naluTypeFUA) {
-			if (fua_buffer.size() == 0) {
-				fua_buffer.push_back(std::byte(0));
-			}
-
-			auto nalUnitFragmentHeader =
-			    NalUnitFragmentHeader{std::to_integer<uint8_t>(pkt->at(headerSize + 1))};
+		auto rtpHeaderSize = pktParsed->getSize() + pktParsed->getExtensionHeaderSize();
+		auto nalUnitHeader = NalUnitHeader{std::to_integer<uint8_t>(pkt->at(rtpHeaderSize))};
 
-			std::copy(pkt->begin() + headerSize + fuaHeaderSize, pkt->end(),
-			          std::back_inserter(fua_buffer));
+		if (nalUnitHeader.unitType() == naluTypeFUA) {
+			auto nalUnitFragmentHeader = NalUnitFragmentHeader{
+			    std::to_integer<uint8_t>(pkt->at(rtpHeaderSize + sizeof(NalUnitHeader)))};
 
-			if (nalUnitFragmentHeader.isEnd()) {
-				fua_buffer.at(0) =
-				    std::byte(nalUnitHeader.idc() | nalUnitFragmentHeader.unitType());
+			if (nFrags++ == 0) {
+				std::copy(naluStartCode.begin(), naluStartCode.end(),
+				          std::back_inserter(accessUnit));
 
-				out.push_back(
-				    make_message(std::move(fua_buffer), Message::Binary, 0, nullptr, frameInfo));
-				fua_buffer.clear();
+				accessUnit.emplace_back(
+				    byte(nalUnitHeader.idc() | nalUnitFragmentHeader.unitType()));
 			}
+
+			std::copy(pkt->begin() + rtpHeaderSize + sizeof(NalUnitHeader) +
+			              sizeof(NalUnitFragmentHeader),
+			          pkt->end(), std::back_inserter(accessUnit));
 		} else if (nalUnitHeader.unitType() > 0 && nalUnitHeader.unitType() < 24) {
-			out.push_back(make_message(pkt->begin() + headerSize, pkt->end(), Message::Binary, 0,
-			                           nullptr, frameInfo));
+			std::copy(naluStartCode.begin(), naluStartCode.end(), std::back_inserter(accessUnit));
+			std::copy(pkt->begin() + rtpHeaderSize, pkt->end(), std::back_inserter(accessUnit));
 		} else if (nalUnitHeader.unitType() == naluTypeSTAPA) {
-			auto currOffset = stapaHeaderSize + headerSize;
+			auto currOffset = rtpHeaderSize + sizeof(NalUnitHeader);
 
-			while (currOffset < pkt->size()) {
-				auto naluSize =
-				    uint16_t(pkt->at(currOffset)) << 8 | uint8_t(pkt->at(currOffset + 1));
+			while (currOffset + sizeof(uint16_t) < pkt->size()) {
+				auto naluSize = std::to_integer<uint16_t>(pkt->at(currOffset)) << 8 |
+				                std::to_integer<uint16_t>(pkt->at(currOffset + 1));
 
-				currOffset += 2;
+				currOffset += sizeof(uint16_t);
 
 				if (pkt->size() < currOffset + naluSize) {
-					throw std::runtime_error("STAP-A declared size is larger then buffer");
+					throw std::runtime_error("H264 STAP-A declared size is larger than buffer");
 				}
 
-				out.push_back(make_message(pkt->begin() + currOffset,
-				                           pkt->begin() + currOffset + naluSize, Message::Binary, 0,
-				                           nullptr, frameInfo));
+				std::copy(naluStartCode.begin(), naluStartCode.end(),
+				          std::back_inserter(accessUnit));
+
+				std::copy(pkt->begin() + currOffset, pkt->begin() + currOffset + naluSize,
+				          std::back_inserter(accessUnit));
+
 				currOffset += naluSize;
 			}
 		} else {
@@ -89,6 +77,11 @@ message_vector H264RtpDepacketizer::buildFrames(message_vector::iterator begin,
 		}
 	}
 
+	if (!accessUnit.empty()) {
+		out.emplace_back(make_message(accessUnit.begin(), accessUnit.end(), Message::Binary, 0,
+		                              nullptr, frameInfo));
+	}
+
 	return out;
 }