[FFmpeg-devel] [PATCH] RTP: get the buffer information of the RTP extension header.

[Wujian(Chin)]

Signed-off-by: wujian_nanjing <wujian2 at huawei.com>
---
 libavcodec/avpacket.c | 44 ++++++++++++++++++++++++++++++++++++++++++++
 libavcodec/packet.h   | 17 +++++++++++++++++
 libavformat/demux.c   |  7 +++++++
 libavformat/rtpdec.c  | 37 ++++++++++++++++++++++++++++---------
 4 files changed, 96 insertions(+), 9 deletions(-)

diff --git a/libavcodec/avpacket.c b/libavcodec/avpacket.c
index bcbc416..45c131a 100644
--- a/libavcodec/avpacket.c
+++ b/libavcodec/avpacket.c
@@ -46,6 +46,8 @@ void av_init_packet(AVPacket *pkt)
     pkt->opaque               = NULL;
     pkt->opaque_ref           = NULL;
     pkt->time_base            = av_make_q(0, 1);
+    pkt->ext                  = NULL;
+    pkt->extlen               = 0;
 }
 #endif
 
@@ -109,6 +111,46 @@ int av_new_packet(AVPacket *pkt, int size)
     return 0;
 }
 
+static void av_packet_free_ext(AVPacket *pkt)
+{
+    if (!pkt)
+        return;
+
+    if (pkt->ext) {
+        av_free(pkt->ext);
+        pkt->ext = NULL;
+    }
+
+    pkt->extlen = 0;
+}
+
+int av_set_packet_ext(AVPacket *pkt, uint8_t *buf, int len)
+{
+    if (!buf || (len <= 0))
+        return 0;
+
+    if (pkt->ext == buf) {
+        pkt->extlen = len;
+        return 0;
+    }
+
+    pkt->ext = av_malloc(len);
+    pkt->extlen = len;
+    (void)memcpy(pkt->ext, buf, len);
+
+    return 0;
+}
+
+static int av_copy_packet_ext(AVPacket *pkt, const AVPacket *src)
+{
+    pkt->ext = NULL;
+    pkt->extlen = 0;
+
+    av_set_packet_ext(pkt, src->ext, src->extlen);
+
+    return 0;
+}
+
 void av_shrink_packet(AVPacket *pkt, int size)
 {
     if (pkt->size <= size)
@@ -422,6 +464,7 @@ int av_packet_copy_props(AVPacket *dst, const AVPacket *src)
 void av_packet_unref(AVPacket *pkt)
 {
     av_packet_free_side_data(pkt);
+    av_packet_free_ext(pkt);
     av_buffer_unref(&pkt->opaque_ref);
     av_buffer_unref(&pkt->buf);
     get_packet_defaults(pkt);
@@ -456,6 +499,7 @@ int av_packet_ref(AVPacket *dst, const AVPacket *src)
     }
 
     dst->size = src->size;
+    av_copy_packet_ext(dst, src);
 
     return 0;
 fail:
diff --git a/libavcodec/packet.h b/libavcodec/packet.h
index 404d520..137d1fa 100644
--- a/libavcodec/packet.h
+++ b/libavcodec/packet.h
@@ -374,6 +374,13 @@ typedef struct AVPacket {
     uint8_t *data;
     int   size;
     int   stream_index;
+
+    /**
+     * Extension header data and size
+     */
+    uint8_t *ext;
+    uint32_t extlen;
+
     /**
      * A combination of AV_PKT_FLAG values
      */
@@ -523,6 +530,16 @@ void av_init_packet(AVPacket *pkt);
 int av_new_packet(AVPacket *pkt, int size);
 
 /**
+ * Allocate the ext data from buf with len lenght.
+ * Make sure pkt->ext is null, otherwise cause memory leak.
+ *
+ * @param pkt packet
+ * @param buf buffer with ext data
+ * @param len buffer size
+ */
+int av_set_packet_ext(AVPacket *pkt, uint8_t *buf, int len);
+
+/**
  * Reduce packet size, correctly zeroing padding
  *
  * @param pkt packet
diff --git a/libavformat/demux.c b/libavformat/demux.c
index 1620716..2a8b201 100644
--- a/libavformat/demux.c
+++ b/libavformat/demux.c
@@ -1175,6 +1175,13 @@ static int parse_packet(AVFormatContext *s, AVPacket *pkt,
             pkt->side_data_elems    = 0;
         }
 
+        if (pkt->extlen > 0) {
+            out_pkt->extlen = pkt->extlen;
+            out_pkt->ext = pkt->ext;
+            pkt->extlen = 0;
+            pkt->ext = NULL;
+        }
+
         /* set the duration */
         out_pkt->duration = (sti->parser->flags & PARSER_FLAG_COMPLETE_FRAMES) ? pkt->duration : 0;
         if (st->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
diff --git a/libavformat/rtpdec.c b/libavformat/rtpdec.c
index fa7544c..2a539e0 100644
--- a/libavformat/rtpdec.c
+++ b/libavformat/rtpdec.c
@@ -699,13 +699,14 @@ static int rtp_parse_packet_internal(RTPDemuxContext *s, AVPacket *pkt,
 {
     unsigned int ssrc;
     int payload_type, seq, flags = 0;
-    int ext, csrc;
+    int has_ext, csrc, extlen;
+    uint8_t *extbuf;
     AVStream *st;
     uint32_t timestamp;
     int rv = 0;
 
     csrc         = buf[0] & 0x0f;
-    ext          = buf[0] & 0x10;
+    has_ext      = buf[0] & 0x10;
     payload_type = buf[1] & 0x7f;
     if (buf[1] & 0x80)
         flags |= RTP_FLAG_MARKER;
@@ -744,18 +745,25 @@ static int rtp_parse_packet_internal(RTPDemuxContext *s, AVPacket *pkt,
         return AVERROR_INVALIDDATA;
 
     /* RFC 3550 Section 5.3.1 RTP Header Extension handling */
-    if (ext) {
+    if (has_ext) {
         if (len < 4)
             return -1;
         /* calculate the header extension length (stored as number
          * of 32-bit words) */
-        ext = (AV_RB16(buf + 2) + 1) << 2;
+        extlen = (AV_RB16(buf + 2) + 1) << 2;
 
-        if (len < ext)
+        if (len < extlen)
             return -1;
-        // skip past RTP header extension
-        len -= ext;
-        buf += ext;
+
+         // Save RTP header extension data
+        if (extlen > 0) {
+            extbuf = av_malloc(extlen);
+            if (extbuf)
+                (void)memcpy(extbuf, buf, extlen);
+        }
+
+        len -= extlen;
+        buf += extlen;
     }
 
     if (s->handler && s->handler->parse_packet) {
@@ -763,14 +771,25 @@ static int rtp_parse_packet_internal(RTPDemuxContext *s, AVPacket *pkt,
                                       s->st, pkt, &timestamp, buf, len, seq,
                                       flags);
     } else if (st) {
-        if ((rv = av_new_packet(pkt, len)) < 0)
+        if ((rv = av_new_packet(pkt, len)) < 0) {
+            if ((extlen > 0) && extbuf)
+                av_free(extbuf);
             return rv;
+        }
         memcpy(pkt->data, buf, len);
         pkt->stream_index = st->index;
     } else {
+        if ((extlen > 0) && extbuf)
+            av_free(extbuf);
         return AVERROR(EINVAL);
     }
 
+    // attach ext data
+    if ((extlen > 0) && extbuf) {
+        av_set_packet_ext(pkt, extbuf, extlen);
+        av_free(extbuf);
+    }
+
     // now perform timestamp things....
     finalize_packet(s, pkt, timestamp);
 
-- 
2.7.4