Implement STAP-A support for H.264

include/formats/h264.hh

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <deque>
+
 #include "frame.hh"
 #include "queue.hh"
 #include "formats/h26x.hh"
@@ -8,6 +10,17 @@ namespace uvg_rtp {
 
     namespace formats {
 
+        enum H264_NAL_TYPES {
+            H264_PKT_AGGR = 24,
+            H264_PKT_FRAG = 28
+        };
+
+        struct h264_aggregation_packet {
+            uint8_t fu_indicator[uvg_rtp::frame::HEADER_SIZE_H264_FU];
+            uvg_rtp::buf_vec nalus;  /* discrete NAL units */
+            uvg_rtp::buf_vec aggr_pkt; /* crafted aggregation packet */
+        };
+
         struct h264_headers {
             uint8_t fu_indicator[uvg_rtp::frame::HEADER_SIZE_H264_FU];
 
@@ -43,11 +56,11 @@ namespace uvg_rtp {
         } h264_info_t;
 
         typedef struct {
+            std::deque<uvg_rtp::frame::rtp_frame *> queued;
             std::unordered_map<uint32_t, h264_info_t> frames;
             std::unordered_set<uint32_t> dropped;
         } h264_frame_info_t;
 
-
         class h264 : public h26x {
             public:
                 h264(uvg_rtp::socket *socket, uvg_rtp::rtp *rtp, int flags);
@@ -72,6 +85,15 @@ namespace uvg_rtp {
                  * Return RTP_GENERIC_ERROR if the packet was corrupted in some way */
                 static rtp_error_t packet_handler(void *arg, int flags, frame::rtp_frame **frame);
 
+                /* If the packet handler must return more than one frame, it can install a frame getter
+                 * that is called by the auxiliary handler caller if packet_handler() returns RTP_MULTIPLE_PKTS_READY
+                 *
+                 * "arg" is the same that is passed to packet_handler
+                 *
+                 * Return RTP_PKT_READY if "frame" contains a frame that can be returned to user
+                 * Return RTP_NOT_FOUND if there are no more frames */
+                static rtp_error_t frame_getter(void *arg, frame::rtp_frame **frame);
+
                 /* Return pointer to the internal frame info structure which is relayed to packet handler */
                 h264_frame_info_t *get_h264_frame_info();
 
@@ -79,7 +101,14 @@ namespace uvg_rtp {
                 rtp_error_t push_nal_unit(uint8_t *data, size_t data_len, bool more);
 
             private:
+                /* Construct an aggregation packet from data in "aggr_pkt_info_" */
+                rtp_error_t make_aggregation_pkt();
+
+                /* Clear aggregation buffers */
+                void clear_aggregation_info();
+
                 h264_frame_info_t finfo_;
+                h264_aggregation_packet aggr_pkt_info_;
         };
     };
 };

include/formats/h265.hh

@@ -20,7 +20,7 @@ namespace uvg_rtp {
 
         struct h265_aggregation_packet {
             uint8_t nal_header[uvg_rtp::frame::HEADER_SIZE_H265_NAL];
-            uvg_rtp::buf_vec buffers;  /* discrete NAL units */
+            uvg_rtp::buf_vec nalus;  /* discrete NAL units */
             uvg_rtp::buf_vec aggr_pkt; /* crafted aggregation packet */
         };
 

src/formats/h264.cc

@@ -14,6 +14,63 @@
 
 #include "formats/h264.hh"
 
+void uvg_rtp::formats::h264::clear_aggregation_info()
+{
+    aggr_pkt_info_.nalus.clear();
+    aggr_pkt_info_.aggr_pkt.clear();
+}
+
+rtp_error_t uvg_rtp::formats::h264::make_aggregation_pkt()
+{
+    rtp_error_t ret;
+    uint8_t nri = 0;
+
+    if (aggr_pkt_info_.nalus.empty())
+        return RTP_INVALID_VALUE;
+
+    /* Only one buffer in the vector -> no need to create an aggregation packet */
+    if (aggr_pkt_info_.nalus.size() == 1) {
+        if ((ret = fqueue_->enqueue_message(aggr_pkt_info_.nalus)) != RTP_OK) {
+            LOG_ERROR("Failed to enqueue Single NAL Unit packet!");
+            return ret;
+        }
+
+        return fqueue_->flush_queue();
+    }
+
+    /* find maximum NRI from given NALUs,
+     * it is going to be the NRI value of theSTAP-A header */
+    for (auto& nalu : aggr_pkt_info_.nalus) {
+        if (((nalu.second[0] >> 5) & 0x3) > nri)
+            nri = (nalu.second[0] >> 5) & 0x3;
+    }
+
+    /* create header for the packet and craft the aggregation packet
+     * according to the format defined in RFC 6184 */
+    aggr_pkt_info_.fu_indicator[0] = (0 << 7) | ((nri & 0x3) << 5) | H264_PKT_AGGR;
+
+    aggr_pkt_info_.aggr_pkt.push_back(
+        std::make_pair(
+            uvg_rtp::frame::HEADER_SIZE_H264_FU,
+            aggr_pkt_info_.fu_indicator
+        )
+    );
+
+    for (auto& nalu: aggr_pkt_info_.nalus) {
+        auto pkt_size = nalu.first;
+        nalu.first    = htons(nalu.first);
+
+        aggr_pkt_info_.aggr_pkt.push_back(std::make_pair(sizeof(uint16_t), (uint8_t *)&nalu.first));
+        aggr_pkt_info_.aggr_pkt.push_back(std::make_pair(pkt_size, nalu.second));
+    }
+
+    if ((ret = fqueue_->enqueue_message(aggr_pkt_info_.aggr_pkt)) != RTP_OK) {
+        LOG_ERROR("Failed to enqueue NALUs of an aggregation packet!");
+    }
+
+    return ret;
+}
+
 rtp_error_t uvg_rtp::formats::h264::push_nal_unit(uint8_t *data, size_t data_len, bool more)
 {
     if (data_len <= 3)
@@ -26,14 +83,29 @@ rtp_error_t uvg_rtp::formats::h264::push_nal_unit(uint8_t *data, size_t data_len
     size_t payload_size = rtp_ctx_->get_payload_size();
 
     if (data_len - 3 <= payload_size) {
-        if ((ret = fqueue_->enqueue_message(data, data_len)) != RTP_OK) {
-            LOG_ERROR("enqeueu failed for small packet");
-            return ret;
-        }
-
-        if (more)
+        /* If there is more data coming in (possibly another small packet)
+         * create entry to "aggr_pkt_info_" to construct an aggregation packet */
+        if (more) {
+            aggr_pkt_info_.nalus.push_back(std::make_pair(data_len, data));
             return RTP_NOT_READY;
-        return fqueue_->flush_queue();
+        } else {
+            if (aggr_pkt_info_.nalus.empty()) {
+                if ((ret = fqueue_->enqueue_message(data, data_len)) != RTP_OK) {
+                    LOG_ERROR("Failed to enqueue Single NAL Unit packet!");
+                    return ret;
+                }
+                return fqueue_->flush_queue();
+            } else {
+                (void)make_aggregation_pkt();
+                ret = fqueue_->flush_queue();
+                clear_aggregation_info();
+                return ret;
+            }
+        }
+    } else {
+        /* If smaller NALUs were queued before this NALU,
+         * send them in an aggregation packet before proceeding with fragmentation */
+        (void)make_aggregation_pkt();
     }
 
     /* The payload is larger than MTU (1500 bytes) so we must split it into smaller RTP frames
@@ -47,7 +119,7 @@ rtp_error_t uvg_rtp::formats::h264::push_nal_unit(uint8_t *data, size_t data_len
     auto buffers = fqueue_->get_buffer_vector();
     auto headers = (uvg_rtp::formats::h264_headers *)fqueue_->get_media_headers();
 
-    headers->fu_indicator[0] = (data[0] & 0xe0) | 28;
+    headers->fu_indicator[0] = (data[0] & 0xe0) | H264_PKT_FRAG;
 
     headers->fu_headers[0] = (uint8_t)((1 << 7) | nal_type);
     headers->fu_headers[1] = nal_type;
@@ -107,3 +179,16 @@ uvg_rtp::formats::h264_frame_info_t *uvg_rtp::formats::h264::get_h264_frame_info
 {
     return &finfo_;
 }
+
+rtp_error_t uvg_rtp::formats::h264::frame_getter(void *arg, uvg_rtp::frame::rtp_frame **frame)
+{
+    auto finfo = (uvg_rtp::formats::h264_frame_info_t *)arg;
+
+    if (finfo->queued.size()) {
+        *frame = finfo->queued.front();
+        finfo->queued.pop_front();
+        return RTP_PKT_READY;
+    }
+
+    return RTP_NOT_FOUND;
+}

src/formats/h264_pkt_handler.cc

@@ -36,10 +36,10 @@ static int __get_frag(uvg_rtp::frame::rtp_frame *frame)
     bool first_frag = frame->payload[1] & 0x80;
     bool last_frag  = frame->payload[1] & 0x40;
 
-    if ((frame->payload[0] & 0x1f) == 24)
+    if ((frame->payload[0] & 0x1f) == uvg_rtp::formats::H264_PKT_AGGR)
         return FT_STAP_A;
 
-    if ((frame->payload[0] & 0x1f) != 28)
+    if ((frame->payload[0] & 0x1f) != uvg_rtp::formats::H264_PKT_FRAG)
         return FT_NOT_FRAG;
 
     if (first_frag && last_frag)
@@ -54,13 +54,38 @@ static int __get_frag(uvg_rtp::frame::rtp_frame *frame)
     return FT_MIDDLE;
 }
 
-/* TODO: This requires additional support from packet dispatcher.
- * Auxiliary handlers must be able to return more than one packet
- * or auxiliary handlers must provide additional hooking function
- * for the pkt dispatcher so it can query all received packets */
-static rtp_error_t __handle_stap_a(uvg_rtp::frame::rtp_frame **frame)
+static rtp_error_t __handle_stap_a(uvg_rtp::formats::h264_frame_info_t *finfo, uvg_rtp::frame::rtp_frame **out)
 {
-    return RTP_PKT_READY;
+    uvg_rtp::buf_vec nalus;
+
+    size_t size  = 0;
+    auto  *frame = *out;
+
+    for (size_t i = uvg_rtp::frame::HEADER_SIZE_H264_FU; i < frame->payload_len; ) {
+        nalus.push_back(
+            std::make_pair(
+                ntohs(*(uint16_t *)&frame->payload[i]),
+                &frame->payload[i] + sizeof(uint16_t)
+            )
+        );
+
+        size += ntohs(*(uint16_t *)&frame->payload[i]);
+        i    += ntohs(*(uint16_t *)&frame->payload[i]) + sizeof(uint16_t);
+    }
+
+    for (size_t i = 0; i < nalus.size(); ++i) {
+        auto retframe = uvg_rtp::frame::alloc_rtp_frame(nalus[i].first);
+
+        std::memcpy(
+            retframe->payload,
+            nalus[i].second,
+            nalus[i].first
+        );
+
+        finfo->queued.push_back(retframe);
+    }
+
+    return RTP_MULTIPLE_PKTS_READY;
 }
 
 static inline uint8_t __get_nal(uvg_rtp::frame::rtp_frame *frame)
@@ -124,7 +149,7 @@ rtp_error_t uvg_rtp::formats::h264::packet_handler(void *arg, int flags, uvg_rtp
     uint8_t nal_type = __get_nal(frame);
 
     if (frag_type == FT_STAP_A)
-        return __handle_stap_a(out);
+        return __handle_stap_a(finfo, out);
 
     if (frag_type == FT_NOT_FRAG)
         return RTP_PKT_READY;

src/formats/h265.cc

@@ -16,7 +16,7 @@
 
 void uvg_rtp::formats::h265::clear_aggregation_info()
 {
-    aggr_pkt_info_.buffers.clear();
+    aggr_pkt_info_.nalus.clear();
     aggr_pkt_info_.aggr_pkt.clear();
 }
 
@@ -24,12 +24,12 @@ rtp_error_t uvg_rtp::formats::h265::make_aggregation_pkt()
 {
     rtp_error_t ret;
 
-    if (aggr_pkt_info_.buffers.empty())
+    if (aggr_pkt_info_.nalus.empty())
         return RTP_INVALID_VALUE;
 
     /* Only one buffer in the vector -> no need to create an aggregation packet */
-    if (aggr_pkt_info_.buffers.size() == 1) {
-        if ((ret = fqueue_->enqueue_message(aggr_pkt_info_.buffers)) != RTP_OK) {
+    if (aggr_pkt_info_.nalus.size() == 1) {
+        if ((ret = fqueue_->enqueue_message(aggr_pkt_info_.nalus)) != RTP_OK) {
             LOG_ERROR("Failed to enqueue Single NAL Unit packet!");
             return ret;
         }
@@ -49,21 +49,21 @@ rtp_error_t uvg_rtp::formats::h265::make_aggregation_pkt()
         )
     );
 
-    for (size_t i = 0; i < aggr_pkt_info_.buffers.size(); ++i) {
-        auto pkt_size                   = aggr_pkt_info_.buffers[i].first;
-        aggr_pkt_info_.buffers[i].first = htons(aggr_pkt_info_.buffers[i].first);
+    for (size_t i = 0; i < aggr_pkt_info_.nalus.size(); ++i) {
+        auto pkt_size                   = aggr_pkt_info_.nalus[i].first;
+        aggr_pkt_info_.nalus[i].first = htons(aggr_pkt_info_.nalus[i].first);
 
         aggr_pkt_info_.aggr_pkt.push_back(
             std::make_pair(
                 sizeof(uint16_t),
-                (uint8_t *)&aggr_pkt_info_.buffers[i].first
+                (uint8_t *)&aggr_pkt_info_.nalus[i].first
             )
         );
 
         aggr_pkt_info_.aggr_pkt.push_back(
             std::make_pair(
                 pkt_size,
-                aggr_pkt_info_.buffers[i].second
+                aggr_pkt_info_.nalus[i].second
             )
         );
     }
@@ -91,10 +91,10 @@ rtp_error_t uvg_rtp::formats::h265::push_nal_unit(uint8_t *data, size_t data_len
         /* If there is more data coming in (possibly another small packet)
          * create entry to "aggr_pkt_info_" to construct an aggregation packet */
         if (more) {
-            aggr_pkt_info_.buffers.push_back(std::make_pair(data_len, data));
+            aggr_pkt_info_.nalus.push_back(std::make_pair(data_len, data));
             return RTP_NOT_READY;
         } else {
-            if (aggr_pkt_info_.buffers.empty()) {
+            if (aggr_pkt_info_.nalus.empty()) {
                 if ((ret = fqueue_->enqueue_message(data, data_len)) != RTP_OK) {
                     LOG_ERROR("Failed to enqueue Single NAL Unit packet!");
                     return ret;

src/formats/h265_pkt_handler.cc

@@ -90,7 +90,7 @@ static rtp_error_t __handle_ap(uvg_rtp::formats::h265_frame_info_t *finfo, uvg_r
     size_t size  = 0;
     auto  *frame = *out;
 
-    for (size_t i = 2; i < frame->payload_len; ) {
+    for (size_t i = uvg_rtp::frame::HEADER_SIZE_H265_NAL; i < frame->payload_len; ) {
         nalus.push_back(
             std::make_pair(
                 ntohs(*(uint16_t *)&frame->payload[i]),

src/media_stream.cc

@@ -150,9 +150,9 @@ rtp_error_t uvg_rtp::media_stream::init()
             media_ = new uvg_rtp::formats::h264(socket_, rtp_, ctx_config_.flags);
             pkt_dispatcher_->install_aux_handler(
                 rtp_handler_key_,
-                dynamic_cast<uvg_rtp::formats::h265 *>(media_)->get_h265_frame_info(),
+                dynamic_cast<uvg_rtp::formats::h264 *>(media_)->get_h264_frame_info(),
                 dynamic_cast<uvg_rtp::formats::h264 *>(media_)->packet_handler,
-                nullptr
+                dynamic_cast<uvg_rtp::formats::h264 *>(media_)->frame_getter
             );
             break;
 
@@ -269,9 +269,9 @@ rtp_error_t uvg_rtp::media_stream::init(uvg_rtp::zrtp *zrtp)
             media_ = new uvg_rtp::formats::h264(socket_, rtp_, ctx_config_.flags);
             pkt_dispatcher_->install_aux_handler(
                 rtp_handler_key_,
-                dynamic_cast<uvg_rtp::formats::h265 *>(media_)->get_h265_frame_info(),
+                dynamic_cast<uvg_rtp::formats::h264 *>(media_)->get_h264_frame_info(),
                 dynamic_cast<uvg_rtp::formats::h264 *>(media_)->packet_handler,
-                nullptr
+                dynamic_cast<uvg_rtp::formats::h264 *>(media_)->frame_getter
             );
             break;
 
@@ -384,9 +384,9 @@ rtp_error_t uvg_rtp::media_stream::add_srtp_ctx(uint8_t *key, uint8_t *salt)
             media_ = new uvg_rtp::formats::h264(socket_, rtp_, ctx_config_.flags);
             pkt_dispatcher_->install_aux_handler(
                 rtp_handler_key_,
-                dynamic_cast<uvg_rtp::formats::h265 *>(media_)->get_h265_frame_info(),
+                dynamic_cast<uvg_rtp::formats::h264 *>(media_)->get_h264_frame_info(),
                 dynamic_cast<uvg_rtp::formats::h264 *>(media_)->packet_handler,
-                nullptr
+                dynamic_cast<uvg_rtp::formats::h264 *>(media_)->frame_getter
             );
             break;