Perform bounds checks in vpx_writer

In the vpx_writer struct, change the buffer_end field to the size field.
Change vpx_stop_encode() to return true on success, false on failure
(output buffer full).

In write_compressed_header(), remove the assertion
assert(header_bc.pos <= 0xffff). The caller (vp9_pack_bitstream()) will
check that condition.

In vp9_pack_bitstream(), the variable "first_part_size" is renamed
"compressed_hdr_size".

Bug: webm:1844
Change-Id: I4ed6ab905a707ad44d875e53036d5a42523a65d0
(cherry picked from commit 73703c1)

test/vp9_boolcoder_test.cc

@@ -65,7 +65,7 @@ TEST(VP9, TestBitIO) {
           vpx_write(&bw, bit, static_cast<int>(probas[i]));
         }
 
-        vpx_stop_encode(&bw);
+        GTEST_ASSERT_EQ(vpx_stop_encode(&bw), 0);
         // vpx_reader_fill() may read into uninitialized data that
         // isn't used meaningfully, but may trigger an MSan warning.
         memset(bw_buffer + bw.pos, 0, sizeof(BD_VALUE) - 1);
@@ -90,3 +90,24 @@ TEST(VP9, TestBitIO) {
     }
   }
 }
+
+TEST(VP9, TestBitIOBufferSize0) {
+  vpx_writer bw;
+  uint8_t bw_buffer[1];
+  vpx_start_encode(&bw, bw_buffer, 0);
+  GTEST_ASSERT_EQ(vpx_stop_encode(&bw), -1);
+}
+
+TEST(VP9, TestBitIOBufferSize1) {
+  vpx_writer bw;
+  uint8_t bw_buffer[1];
+  vpx_start_encode(&bw, bw_buffer, sizeof(bw_buffer));
+  GTEST_ASSERT_EQ(vpx_stop_encode(&bw), -1);
+}
+
+TEST(VP9, TestBitIOBufferSize2) {
+  vpx_writer bw;
+  uint8_t bw_buffer[2];
+  vpx_start_encode(&bw, bw_buffer, sizeof(bw_buffer));
+  GTEST_ASSERT_EQ(vpx_stop_encode(&bw), 0);
+}

vp9/encoder/vp9_bitstream.c

@@ -948,8 +948,7 @@ static int encode_tile_worker(void *arg1, void *arg2) {
   write_modes(cpi, xd, &cpi->tile_data[data->tile_idx].tile_info,
               &data->bit_writer, tile_row, data->tile_idx,
               &data->max_mv_magnitude, data->interp_filter_selected);
-  vpx_stop_encode(&data->bit_writer);
-  return 1;
+  return vpx_stop_encode(&data->bit_writer) == 0;
 }
 
 void vp9_bitstream_encode_tiles_buffer_dealloc(VP9_COMP *const cpi) {
@@ -1123,7 +1122,10 @@ static size_t encode_tiles(VP9_COMP *cpi, uint8_t *data_ptr, size_t data_size) {
                   tile_row, tile_col, &cpi->max_mv_magnitude,
                   cpi->interp_filter_selected);
 
-      vpx_stop_encode(&residual_bc);
+      if (vpx_stop_encode(&residual_bc)) {
+        vpx_internal_error(&cm->error, VPX_CODEC_ERROR,
+                           "encode_tiles: output buffer full");
+      }
       if (tile_col < tile_cols - 1 || tile_row < tile_rows - 1) {
         // size of this tile
         mem_put_be32(data_ptr + total_size, residual_bc.pos);
@@ -1377,8 +1379,10 @@ static size_t write_compressed_header(VP9_COMP *cpi, uint8_t *data,
                         &counts->mv);
   }
 
-  vpx_stop_encode(&header_bc);
-  assert(header_bc.pos <= 0xffff);
+  if (vpx_stop_encode(&header_bc)) {
+    vpx_internal_error(&cm->error, VPX_CODEC_ERROR,
+                       "write_compressed_header: output buffer full");
+  }
 
   return header_bc.pos;
 }
@@ -1388,7 +1392,7 @@ void vp9_pack_bitstream(VP9_COMP *cpi, uint8_t *dest, size_t dest_size,
   VP9_COMMON *const cm = &cpi->common;
   uint8_t *data = dest;
   size_t data_size = dest_size;
-  size_t first_part_size, uncompressed_hdr_size;
+  size_t uncompressed_hdr_size, compressed_hdr_size;
   struct vpx_write_bit_buffer wb = { data, 0 };
   struct vpx_write_bit_buffer saved_wb;
 
@@ -1407,22 +1411,23 @@ void vp9_pack_bitstream(VP9_COMP *cpi, uint8_t *dest, size_t dest_size,
   }
 
   saved_wb = wb;
-  vpx_wb_write_literal(&wb, 0, 16);  // don't know in advance first part. size
+  // don't know in advance compressed header size
+  vpx_wb_write_literal(&wb, 0, 16);
 
   uncompressed_hdr_size = vpx_wb_bytes_written(&wb);
   data += uncompressed_hdr_size;
   data_size -= uncompressed_hdr_size;
 
   vpx_clear_system_state();
 
-  first_part_size = write_compressed_header(cpi, data, data_size);
-  data += first_part_size;
-  data_size -= first_part_size;
-  if (first_part_size > UINT16_MAX) {
+  compressed_hdr_size = write_compressed_header(cpi, data, data_size);
+  data += compressed_hdr_size;
+  data_size -= compressed_hdr_size;
+  if (compressed_hdr_size > UINT16_MAX) {
     vpx_internal_error(&cm->error, VPX_CODEC_ERROR,
-                       "first_part_size > 16 bits");
+                       "compressed_hdr_size > 16 bits");
   }
-  vpx_wb_write_literal(&saved_wb, (int)first_part_size, 16);
+  vpx_wb_write_literal(&saved_wb, (int)compressed_hdr_size, 16);
 
   data += encode_tiles(cpi, data, data_size);
 

vpx_dsp/bitwriter.c

@@ -9,6 +9,7 @@
  */
 
 #include <assert.h>
+#include <limits.h>
 
 #include "./bitwriter.h"
 
@@ -20,13 +21,16 @@ void vpx_start_encode(vpx_writer *br, uint8_t *source, size_t size) {
   br->lowvalue = 0;
   br->range = 255;
   br->count = -24;
-  br->buffer = source;
-  br->buffer_end = source + size;
+  br->error = 0;
   br->pos = 0;
+  // Make sure it is safe to cast br->pos to int in vpx_write().
+  if (size > INT_MAX) size = INT_MAX;
+  br->size = (unsigned int)size;
+  br->buffer = source;
   vpx_write_bit(br, 0);
 }
 
-void vpx_stop_encode(vpx_writer *br) {
+int vpx_stop_encode(vpx_writer *br) {
   int i;
 
 #if CONFIG_BITSTREAM_DEBUG
@@ -35,9 +39,17 @@ void vpx_stop_encode(vpx_writer *br) {
   for (i = 0; i < 32; i++) vpx_write_bit(br, 0);
 
   // Ensure there's no ambigous collision with any index marker bytes
-  if ((br->buffer[br->pos - 1] & 0xe0) == 0xc0) br->buffer[br->pos++] = 0;
+  if (!br->error && (br->buffer[br->pos - 1] & 0xe0) == 0xc0) {
+    if (br->pos < br->size) {
+      br->buffer[br->pos++] = 0;
+    } else {
+      br->error = 1;
+    }
+  }
 
 #if CONFIG_BITSTREAM_DEBUG
   bitstream_queue_set_skip_write(0);
 #endif
+
+  return br->error ? -1 : 0;
 }

vpx_dsp/bitwriter.h

@@ -29,13 +29,19 @@ typedef struct vpx_writer {
   unsigned int lowvalue;
   unsigned int range;
   int count;
+  // Whether there has been an error.
+  int error;
+  // We maintain the invariant that pos <= size, i.e., we never write beyond
+  // the end of the buffer. If pos would be incremented to be greater than
+  // size, leave pos unchanged and set error to 1.
   unsigned int pos;
+  unsigned int size;
   uint8_t *buffer;
-  const uint8_t *buffer_end;
 } vpx_writer;
 
 void vpx_start_encode(vpx_writer *br, uint8_t *source, size_t size);
-void vpx_stop_encode(vpx_writer *br);
+// Returns 0 on success and returns -1 in case of error.
+int vpx_stop_encode(vpx_writer *br);
 
 static INLINE VPX_NO_UNSIGNED_SHIFT_CHECK void vpx_write(vpx_writer *br,
                                                          int bit,
@@ -78,18 +84,25 @@ static INLINE VPX_NO_UNSIGNED_SHIFT_CHECK void vpx_write(vpx_writer *br,
   if (count >= 0) {
     int offset = shift - count;
 
-    if ((lowvalue << (offset - 1)) & 0x80000000) {
-      int x = br->pos - 1;
+    if (!br->error) {
+      if ((lowvalue << (offset - 1)) & 0x80000000) {
+        int x = (int)br->pos - 1;
 
-      while (x >= 0 && br->buffer[x] == 0xff) {
-        br->buffer[x] = 0;
-        x--;
+        while (x >= 0 && br->buffer[x] == 0xff) {
+          br->buffer[x] = 0;
+          x--;
+        }
+
+        // TODO(wtc): How to prove x >= 0?
+        br->buffer[x] += 1;
       }
 
-      br->buffer[x] += 1;
+      if (br->pos < br->size) {
+        br->buffer[br->pos++] = (lowvalue >> (24 - offset)) & 0xff;
+      } else {
+        br->error = 1;
+      }
     }
-
-    br->buffer[br->pos++] = (lowvalue >> (24 - offset)) & 0xff;
     lowvalue <<= offset;
     shift = count;
     lowvalue &= 0xffffff;