@@ -526,6 +526,7 @@ static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
PutBitContext *pbc)
{
VP9RawFrame *frame = unit->content;
+ size_t data_size, header_size;
int err;
err = cbs_vp9_write_frame(ctx, pbc, frame);
@@ -535,16 +536,28 @@ static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
// Frame must be byte-aligned.
av_assert0(put_bits_count(pbc) % 8 == 0);
+ flush_put_bits(pbc);
+ data_size = header_size = put_bits_count(pbc) / 8;
+ unit->data_bit_padding = 0;
+
if (frame->data) {
- if (frame->data_size > put_bits_left(pbc) / 8)
- return AVERROR(ENOSPC);
+ if (frame->data_size > INT_MAX - AV_INPUT_BUFFER_PADDING_SIZE
+ - header_size)
+ return AVERROR(ENOMEM);
- flush_put_bits(pbc);
- memcpy(put_bits_ptr(pbc), frame->data, frame->data_size);
- skip_put_bytes(pbc, frame->data_size);
+ data_size += frame->data_size;
}
- return ff_cbs_default_write_unit_data(ctx, unit, pbc);
+ err = ff_cbs_alloc_unit_data(ctx, unit, data_size);
+ if (err < 0)
+ return err;
+
+ memcpy(unit->data, pbc->buf, header_size);
+
+ if (frame->data)
+ memcpy(unit->data + header_size, frame->data, frame->data_size);
+
+ return 0;
}
static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
Writing a unit (always a frame) in cbs_vp9 used an intermediate buffer to write the frame header followed by the frame data that was copied into said buffer. Afterwards, the final buffer for the frame was allocated and everything copied into this buffer. But it is trivial to compute the needed size of the final buffer after having written the header, so one can allocate the final buffer immediately and copy the frame data directly into it. Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@gmail.com> --- libavcodec/cbs_vp9.c | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-)