diff mbox

[FFmpeg-devel,2/7] avcodec/cbs: Check for overflow when assembling fragments

Message ID 20191209222604.28920-2-andreas.rheinhardt@gmail.com
State New
Headers show

Commit Message

Andreas Rheinhardt Dec. 9, 2019, 10:25 p.m. UTC
Up until now, the various functions for assembling fragments did not
check for overflow in the computation of the size of the output fragment
or whether the size they allocate actually fits into an int (as it must
given that the AVBuffer API is based on size fields of type int).

In order to optimize this, the overflow check is partially done
generically: The size of all the units is just added while checking for
it to be <= INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE. The result is
accessible to the function doing the actual assembling.

Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@gmail.com>
---
In case of H.26x, the new allocation might be slightly larger than the
old one (e.g. two units of odd size would be independently rounded down
when multiplying by 3 followed by dividing by two). But that hardly
matters.

 libavcodec/cbs.c       | 10 ++++++++++
 libavcodec/cbs_av1.c   |  7 +------
 libavcodec/cbs_h2645.c | 12 +++++++-----
 libavcodec/cbs_jpeg.c  |  9 ++++++---
 libavcodec/cbs_mpeg2.c |  9 +++++----
 libavcodec/cbs_vp9.c   |  6 ++++--
 6 files changed, 33 insertions(+), 20 deletions(-)
diff mbox

Patch

diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index 805049404b..e70e11fb99 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -343,6 +343,7 @@  static int cbs_write_unit_data(CodedBitstreamContext *ctx,
 int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
                                CodedBitstreamFragment *frag)
 {
+    size_t fragment_size;
     int err, i;
 
     for (i = 0; i < frag->nb_units; i++) {
@@ -366,6 +367,15 @@  int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
     av_buffer_unref(&frag->data_ref);
     frag->data = NULL;
 
+    fragment_size = 0;
+    for (i = 0; i < frag->nb_units; i++) {
+        if (frag->units[i].data_size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+                                               - fragment_size)
+            return AVERROR(ERANGE);
+        fragment_size += frag->units[i].data_size;
+    }
+    frag->data_size = fragment_size;
+
     err = ctx->codec->assemble_fragment(ctx, frag);
     if (err < 0) {
         av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to assemble fragment.\n");
diff --git a/libavcodec/cbs_av1.c b/libavcodec/cbs_av1.c
index bbe4461130..1c51839f16 100644
--- a/libavcodec/cbs_av1.c
+++ b/libavcodec/cbs_av1.c
@@ -1223,13 +1223,9 @@  static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
 static int cbs_av1_assemble_fragment(CodedBitstreamContext *ctx,
                                      CodedBitstreamFragment *frag)
 {
-    size_t size, pos;
+    size_t size = frag->data_size, pos;
     int i;
 
-    size = 0;
-    for (i = 0; i < frag->nb_units; i++)
-        size += frag->units[i].data_size;
-
     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
     if (!frag->data_ref)
         return AVERROR(ENOMEM);
@@ -1243,7 +1239,6 @@  static int cbs_av1_assemble_fragment(CodedBitstreamContext *ctx,
         pos += frag->units[i].data_size;
     }
     av_assert0(pos == size);
-    frag->data_size = size;
 
     return 0;
 }
diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index 5f71d80584..4b17780836 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -1392,11 +1392,13 @@  static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
         av_assert0(frag->units[i].data);
     }
 
-    max_size = 0;
-    for (i = 0; i < frag->nb_units; i++) {
-        // Start code + content with worst-case emulation prevention.
-        max_size += 4 + frag->units[i].data_size * 3 / 2;
-    }
+    // Content with worst-case emulation prevention.
+    max_size = frag->data_size + frag->data_size / 2;
+    if (frag->nb_units > (int)(INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+                                       - max_size + 3) / 4)
+        return AVERROR(ERANGE);
+    // Start codes
+    max_size += 4 * frag->nb_units;
 
     data = av_realloc(NULL, max_size + AV_INPUT_BUFFER_PADDING_SIZE);
     if (!data)
diff --git a/libavcodec/cbs_jpeg.c b/libavcodec/cbs_jpeg.c
index 2bb6c8d18c..6c131d2e2e 100644
--- a/libavcodec/cbs_jpeg.c
+++ b/libavcodec/cbs_jpeg.c
@@ -391,13 +391,12 @@  static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
 {
     const CodedBitstreamUnit *unit;
     uint8_t *data;
-    size_t size, dp, sp;
+    size_t size = frag->data_size, dp, sp;
     int i;
 
-    size = 4; // SOI + EOI.
+    size += 4; // SOI + EOI.
     for (i = 0; i < frag->nb_units; i++) {
         unit = &frag->units[i];
-        size += 2 + unit->data_size;
         if (unit->type == JPEG_MARKER_SOS) {
             for (sp = 0; sp < unit->data_size; sp++) {
                 if (unit->data[sp] == 0xff)
@@ -405,6 +404,10 @@  static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
             }
         }
     }
+    if (frag->nb_units > (int)(INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+                                       - size + 1) / 2)
+        return AVERROR(ERANGE);
+    size += 2 * frag->nb_units;
 
     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
     if (!frag->data_ref)
diff --git a/libavcodec/cbs_mpeg2.c b/libavcodec/cbs_mpeg2.c
index 255f033734..e42c602476 100644
--- a/libavcodec/cbs_mpeg2.c
+++ b/libavcodec/cbs_mpeg2.c
@@ -384,12 +384,13 @@  static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
                                        CodedBitstreamFragment *frag)
 {
     uint8_t *data;
-    size_t size, dp;
+    size_t dp, size = frag->data_size;
     int i;
 
-    size = 0;
-    for (i = 0; i < frag->nb_units; i++)
-        size += 3 + frag->units[i].data_size;
+    if (frag->nb_units > (INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+                                  - size + 2) / 3)
+        return AVERROR(ERANGE);
+    size += 3 * frag->nb_units;
 
     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
     if (!frag->data_ref)
diff --git a/libavcodec/cbs_vp9.c b/libavcodec/cbs_vp9.c
index f6cfaa3b36..ea7747182e 100644
--- a/libavcodec/cbs_vp9.c
+++ b/libavcodec/cbs_vp9.c
@@ -595,9 +595,11 @@  static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
         sfi.bytes_per_framesize_minus_1  = size_len - 1;
         sfi.frames_in_superframe_minus_1 = frag->nb_units - 1;
 
-        size = 2;
+        size = 2 + size_len * frag->nb_units + frag->data_size;
+        if (size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE)
+            return AVERROR(ERANGE);
+
         for (i = 0; i < frag->nb_units; i++) {
-            size += size_len + frag->units[i].data_size;
             sfi.frame_sizes[i] = frag->units[i].data_size;
         }