rtpdec_hevc: Share the implementation of fragmented packets with h264

Signed-off-by: Martin Storsjö <martin@martin.st>

3 files changed, 29 insertions, 35 deletions

diff --git a/libavformat/rtpdec_formats.h b/libavformat/rtpdec_formats.h
index 4dcc1986ea..43b4b5b91b 100644
--- a/libavformat/rtpdec_formats.h
+++ b/libavformat/rtpdec_formats.h
@@ -42,6 +42,9 @@ int ff_h264_handle_aggregated_packet(AVFormatContext *ctx, AVPacket *pkt,
                                      const uint8_t *buf, int len,
                                      int start_skip, int *nal_counters,
                                      int nal_mask);
+int ff_h264_handle_frag_packet(AVPacket *pkt, const uint8_t *buf, int len,
+                               int start_bit, const uint8_t *nal_header,
+                               int nal_header_len);
 void ff_h264_parse_framesize(AVCodecContext *codec, const char *p);
 
 extern RTPDynamicProtocolHandler ff_ac3_dynamic_handler;
diff --git a/libavformat/rtpdec_h264.c b/libavformat/rtpdec_h264.c
index ba166b87b6..e624cc962e 100644
--- a/libavformat/rtpdec_h264.c
+++ b/libavformat/rtpdec_h264.c
@@ -257,12 +257,32 @@ int ff_h264_handle_aggregated_packet(AVFormatContext *ctx, AVPacket *pkt,
     return 0;
 }
 
+int ff_h264_handle_frag_packet(AVPacket *pkt, const uint8_t *buf, int len,
+                               int start_bit, const uint8_t *nal_header,
+                               int nal_header_len)
+{
+    int ret;
+    int tot_len = len;
+    int pos = 0;
+    if (start_bit)
+        tot_len += sizeof(start_sequence) + nal_header_len;
+    if ((ret = av_new_packet(pkt, tot_len)) < 0)
+        return ret;
+    if (start_bit) {
+        memcpy(pkt->data + pos, start_sequence, sizeof(start_sequence));
+        pos += sizeof(start_sequence);
+        memcpy(pkt->data + pos, nal_header, nal_header_len);
+        pos += nal_header_len;
+    }
+    memcpy(pkt->data + pos, buf, len);
+    return 0;
+}
+
 static int h264_handle_packet_fu_a(AVFormatContext *ctx, AVPacket *pkt,
                                    const uint8_t *buf, int len,
                                    int *nal_counters, int nal_mask)
 {
     uint8_t fu_indicator, fu_header, start_bit, nal_type, nal;
-    int ret;
 
     if (len < 3) {
         av_log(ctx, AV_LOG_ERROR, "Too short data for FU-A H264 RTP packet\n");
@@ -279,22 +299,9 @@ static int h264_handle_packet_fu_a(AVFormatContext *ctx, AVPacket *pkt,
     buf += 2;
     len -= 2;
 
-    if (start_bit) {
-        if (nal_counters)
-            nal_counters[nal_type & nal_mask]++;
-        /* copy in the start sequence, and the reconstructed nal */
-        if ((ret = av_new_packet(pkt, sizeof(start_sequence) + sizeof(nal) + len)) < 0)
-            return ret;
-        memcpy(pkt->data, start_sequence, sizeof(start_sequence));
-        pkt->data[sizeof(start_sequence)] = nal;
-        memcpy(pkt->data + sizeof(start_sequence) + sizeof(nal), buf, len);
-    } else {
-        if ((ret = av_new_packet(pkt, len)) < 0)
-            return ret;
-        memcpy(pkt->data, buf, len);
-    }
-
-    return 0;
+    if (start_bit && nal_counters)
+        nal_counters[nal_type & nal_mask]++;
+    return ff_h264_handle_frag_packet(pkt, buf, len, start_bit, &nal, 1);
 }
 
 // return 0 on packet, no more left, 1 on packet, 1 on partial packet
diff --git a/libavformat/rtpdec_hevc.c b/libavformat/rtpdec_hevc.c
index e04725be5d..44cf1d218f 100644
--- a/libavformat/rtpdec_hevc.c
+++ b/libavformat/rtpdec_hevc.c
@@ -331,24 +331,8 @@ static int hevc_handle_packet(AVFormatContext *ctx, PayloadContext *rtp_hevc_ctx
         new_nal_header[0] = (rtp_pl[0] & 0x81) | (fu_type << 1);
         new_nal_header[1] = rtp_pl[1];
 
-        /* start fragment vs. subsequent fragments */
-        if (first_fragment) {
-            /* create A/V packet which is big enough */
-            if ((res = av_new_packet(pkt, sizeof(start_sequence) + sizeof(new_nal_header) + len)) < 0)
-                return res;
-            /* A/V packet: copy start sequence */
-            memcpy(pkt->data, start_sequence, sizeof(start_sequence));
-            /* A/V packet: copy new NAL header */
-            memcpy(pkt->data + sizeof(start_sequence), new_nal_header, sizeof(new_nal_header));
-            /* A/V packet: copy NAL unit data */
-            memcpy(pkt->data + sizeof(start_sequence) + sizeof(new_nal_header), buf, len);
-        } else {
-            /* create A/V packet */
-            if ((res = av_new_packet(pkt, len)) < 0)
-                return res;
-            /* A/V packet: copy NAL unit data */
-            memcpy(pkt->data, buf, len);
-        }
+        res = ff_h264_handle_frag_packet(pkt, buf, len, first_fragment,
+                                         new_nal_header, sizeof(new_nal_header));
 
         break;
     /* PACI packet */