diff --git a/test/vp9_boolcoder_test.cc b/test/vp9_boolcoder_test.cc index e002a9cd26d..aeff0d7a50e 100644 --- a/test/vp9_boolcoder_test.cc +++ b/test/vp9_boolcoder_test.cc @@ -65,7 +65,7 @@ TEST(VP9, TestBitIO) { vpx_write(&bw, bit, static_cast(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); +} diff --git a/vp9/encoder/vp9_bitstream.c b/vp9/encoder/vp9_bitstream.c index 704fcd6509d..dddf1bc4683 100644 --- a/vp9/encoder/vp9_bitstream.c +++ b/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,7 +1411,8 @@ 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; @@ -1415,14 +1420,14 @@ void vp9_pack_bitstream(VP9_COMP *cpi, uint8_t *dest, size_t dest_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); diff --git a/vpx_dsp/bitwriter.c b/vpx_dsp/bitwriter.c index 37859d3f110..d3ef9bd89af 100644 --- a/vpx_dsp/bitwriter.c +++ b/vpx_dsp/bitwriter.c @@ -9,6 +9,7 @@ */ #include +#include #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; } diff --git a/vpx_dsp/bitwriter.h b/vpx_dsp/bitwriter.h index f6c00ab91b4..daff331daf6 100644 --- a/vpx_dsp/bitwriter.h +++ b/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;