From c91910351fc7b744a247dc403465238ee431c7c2 Mon Sep 17 00:00:00 2001
From: Heikki Tampio <heikki.tampio@tuni.fi>
Date: Tue, 20 Feb 2024 13:00:48 +0200
Subject: [PATCH] formats: Develop solution to H26x aggregation when there are
 multiple non-consecutive small NAL units

---
 src/formats/h26x.cc     | 10 +++++++---
 src/formats/h26x.hh     |  1 +
 test/test_4_formats.cpp |  7 ++++---
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/formats/h26x.cc b/src/formats/h26x.cc
index e6a77c24..16228262 100644
--- a/src/formats/h26x.cc
+++ b/src/formats/h26x.cc
@@ -356,6 +356,7 @@ rtp_error_t uvgrtp::formats::h26x::push_media_frame(sockaddr_in& addr, sockaddr_
         nal.prefix_len = 0;
         nal.size = data_len;
         nal.aggregate = false;
+        nal.was_aggregated = false;
 
         nals.push_back(nal);
     }
@@ -381,6 +382,7 @@ rtp_error_t uvgrtp::formats::h26x::push_media_frame(sockaddr_in& addr, sockaddr_
         {
             if (nal.aggregate)
             {
+                nal.was_aggregated = true;
                 if ((ret = add_aggregate_packet(data + nal.offset, nal.size)) != RTP_OK)
                 {
                     clear_aggregation_info();
@@ -388,6 +390,9 @@ rtp_error_t uvgrtp::formats::h26x::push_media_frame(sockaddr_in& addr, sockaddr_
                     return ret;
                 }
             }
+            else {
+                break;
+            }
         }
 
         (void)finalize_aggregation_pkt();
@@ -398,8 +403,7 @@ rtp_error_t uvgrtp::formats::h26x::push_media_frame(sockaddr_in& addr, sockaddr_
 
     for (auto& nal : nals) // non-aggregatable NAL units
     {
-        //UVG_LOG_DEBUG("NAL size %u", nal.size);
-        if (do_not_aggr || !nal.aggregate || !should_aggregate)
+        if (do_not_aggr || !nal.was_aggregated || !should_aggregate)
         {
             if ((ret = fqueue_->init_transaction(data + nal.offset, true)) != RTP_OK) {
                 UVG_LOG_ERROR("Invalid frame queue or failed to initialize transaction!");
@@ -816,7 +820,7 @@ rtp_error_t uvgrtp::formats::h26x::packet_handler(void* args, int rce_flags, uin
         if (e && continuous) {                   
             size_t nal_size = 0; // Find size of the complete reconstructed NAL unit       
             for (auto p : reconstructed_fragments.at(start).seqs) {
-                nal_size += fragments_[p]->payload_len;
+                nal_size += (fragments_[p]->payload_len - sizeof_fu_headers);
                 au.fragments_info[p].reconstructed = true;
             }
             /* Work in progress feature: here we discard inter frames if their references were not received correctly */
diff --git a/src/formats/h26x.hh b/src/formats/h26x.hh
index c9141325..86ddb435 100644
--- a/src/formats/h26x.hh
+++ b/src/formats/h26x.hh
@@ -74,6 +74,7 @@ namespace uvgrtp {
             size_t prefix_len = 0;
             size_t size = 0;
             bool aggregate = false;
+            bool was_aggregated = false;
         };
 
         class h26x : public media {
diff --git a/test/test_4_formats.cpp b/test/test_4_formats.cpp
index 17ebe5af..2ab4d91c 100644
--- a/test/test_4_formats.cpp
+++ b/test/test_4_formats.cpp
@@ -708,7 +708,7 @@ TEST(FormatTests, h265_aggregation)
     uvgrtp::media_stream* receiver = nullptr;
 
     aggr_received = 0;
-    int expected = 3;
+    int expected = 5;
 
     if (sess)
     {
@@ -721,10 +721,10 @@ TEST(FormatTests, h265_aggregation)
     rtp_format_t format = RTP_FORMAT_H265;
     int test_runs = 1;
 
-    std::vector<size_t> test_sizes = { 100, 200, 300 };
+    std::vector<size_t> test_sizes = { 100, 200, 1700, 300, 400};
 
     size_t total_size = 0;
-    std::unique_ptr<uint8_t[]> test_frame = std::unique_ptr<uint8_t[]>(new uint8_t[700]);
+    std::unique_ptr<uint8_t[]> test_frame = std::unique_ptr<uint8_t[]>(new uint8_t[2700]);
 
     for (auto& size : test_sizes)
     {
@@ -841,6 +841,7 @@ TEST(FormatTests, h265_disable_aggr)
 
 inline void aggr_receive_hook(void* arg, uvgrtp::frame::rtp_frame* frame)
 {
+    std::cout << "Rec frame size " << frame->payload_len << std::endl;
     aggr_received++;
     (void)uvgrtp::frame::dealloc_frame(frame);
 }
\ No newline at end of file