adtsenc: Check frame size.

Inspired by work from: Michael Niedermayer <michaelni@gmx.at>.

Signed-off-by: Alex Converse <alex.converse@gmail.com>

2 files changed, 21 insertions, 4 deletions

diff --git a/libavformat/adtsenc.c b/libavformat/adtsenc.c
index ce002fe..55fece5 100644
--- a/libavformat/adtsenc.c
+++ b/libavformat/adtsenc.c
@@ -27,6 +27,8 @@
 #include "avformat.h"
 #include "adts.h"
 
+#define ADTS_MAX_FRAME_BYTES ((1 << 13) - 1)
+
 int ff_adts_decode_extradata(AVFormatContext *s, ADTSContext *adts, uint8_t *buf, int size)
 {
     GetBitContext gb;
@@ -93,6 +95,13 @@ int ff_adts_write_frame_header(ADTSContext *ctx,
 {
     PutBitContext pb;
 
+    unsigned full_frame_size = (unsigned)ADTS_HEADER_SIZE + size + pce_size;
+    if (full_frame_size > ADTS_MAX_FRAME_BYTES) {
+        av_log(NULL, AV_LOG_ERROR, "ADTS frame size too large: %u (max %d)\n",
+               full_frame_size, ADTS_MAX_FRAME_BYTES);
+        return AVERROR_INVALIDDATA;
+    }
+
     init_put_bits(&pb, buf, ADTS_HEADER_SIZE);
 
     /* adts_fixed_header */
@@ -110,7 +119,7 @@ int ff_adts_write_frame_header(ADTSContext *ctx,
     /* adts_variable_header */
     put_bits(&pb, 1, 0);        /* copyright_identification_bit */
     put_bits(&pb, 1, 0);        /* copyright_identification_start */
-    put_bits(&pb, 13, ADTS_HEADER_SIZE + size + pce_size); /* aac_frame_length */
+    put_bits(&pb, 13, full_frame_size); /* aac_frame_length */
     put_bits(&pb, 11, 0x7ff);   /* adts_buffer_fullness */
     put_bits(&pb, 2, 0);        /* number_of_raw_data_blocks_in_frame */
 
@@ -128,7 +137,10 @@ static int adts_write_packet(AVFormatContext *s, AVPacket *pkt)
     if (!pkt->size)
         return 0;
     if (adts->write_adts) {
-        ff_adts_write_frame_header(adts, buf, pkt->size, adts->pce_size);
+        int err = ff_adts_write_frame_header(adts, buf, pkt->size,
+                                             adts->pce_size);
+        if (err < 0)
+            return err;
         avio_write(pb, buf, ADTS_HEADER_SIZE);
         if (adts->pce_size) {
             avio_write(pb, adts->pce_data, adts->pce_size);
diff --git a/libavformat/mpegtsenc.c b/libavformat/mpegtsenc.c
index 9f45f13..f28b719 100644
--- a/libavformat/mpegtsenc.c
+++ b/libavformat/mpegtsenc.c
@@ -971,7 +971,7 @@ static int mpegts_write_packet(AVFormatContext *s, AVPacket *pkt)
             return -1;
         if ((AV_RB16(pkt->data) & 0xfff0) != 0xfff0) {
             ADTSContext *adts = ts_st->adts;
-            int new_size;
+            int new_size, err;
             if (!adts) {
                 av_log(s, AV_LOG_ERROR, "aac bitstream not in adts format "
                        "and extradata missing\n");
@@ -983,7 +983,12 @@ static int mpegts_write_packet(AVFormatContext *s, AVPacket *pkt)
             data = av_malloc(new_size);
             if (!data)
                 return AVERROR(ENOMEM);
-            ff_adts_write_frame_header(adts, data, pkt->size, adts->pce_size);
+            err = ff_adts_write_frame_header(adts, data, pkt->size,
+                                             adts->pce_size);
+            if (err < 0) {
+                av_free(data);
+                return err;
+            }
             if (adts->pce_size) {
                 memcpy(data+ADTS_HEADER_SIZE, adts->pce_data, adts->pce_size);
                 adts->pce_size = 0;