cbs: Always check for bitstream end before reading
[libav.git] / libavcodec / cbs.c
1 /*
2 * This file is part of Libav.
3 *
4 * Libav is free software; you can redistribute it and/or
5 * modify it under the terms of the GNU Lesser General Public
6 * License as published by the Free Software Foundation; either
7 * version 2.1 of the License, or (at your option) any later version.
8 *
9 * Libav is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
12 * Lesser General Public License for more details.
13 *
14 * You should have received a copy of the GNU Lesser General Public
15 * License along with Libav; if not, write to the Free Software
16 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
17 */
18
19 #include <string.h>
20
21 #include "config.h"
22
23 #include "libavutil/avassert.h"
24 #include "libavutil/common.h"
25
26 #include "cbs.h"
27 #include "cbs_internal.h"
28
29
30 static const CodedBitstreamType *cbs_type_table[] = {
31 #if CONFIG_CBS_H264
32 &ff_cbs_type_h264,
33 #endif
34 #if CONFIG_CBS_H265
35 &ff_cbs_type_h265,
36 #endif
37 #if CONFIG_CBS_MPEG2
38 &ff_cbs_type_mpeg2,
39 #endif
40 };
41
42 int ff_cbs_init(CodedBitstreamContext *ctx,
43 enum AVCodecID codec_id, void *log_ctx)
44 {
45 const CodedBitstreamType *type;
46 int i;
47
48 type = NULL;
49 for (i = 0; i < FF_ARRAY_ELEMS(cbs_type_table); i++) {
50 if (cbs_type_table[i]->codec_id == codec_id) {
51 type = cbs_type_table[i];
52 break;
53 }
54 }
55 if (!type)
56 return AVERROR(EINVAL);
57
58 ctx->log_ctx = log_ctx;
59 ctx->codec = type;
60
61 ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
62 if (!ctx->priv_data)
63 return AVERROR(ENOMEM);
64
65 ctx->decompose_unit_types = NULL;
66
67 ctx->trace_enable = 0;
68 ctx->trace_level = AV_LOG_TRACE;
69
70 return 0;
71 }
72
73 void ff_cbs_close(CodedBitstreamContext *ctx)
74 {
75 if (ctx->codec && ctx->codec->close)
76 ctx->codec->close(ctx);
77
78 av_freep(&ctx->priv_data);
79 }
80
81 static void cbs_unit_uninit(CodedBitstreamContext *ctx,
82 CodedBitstreamUnit *unit)
83 {
84 if (ctx->codec->free_unit && unit->content && !unit->content_external)
85 ctx->codec->free_unit(unit);
86
87 av_freep(&unit->data);
88 unit->data_size = 0;
89 unit->data_bit_padding = 0;
90 }
91
92 void ff_cbs_fragment_uninit(CodedBitstreamContext *ctx,
93 CodedBitstreamFragment *frag)
94 {
95 int i;
96
97 for (i = 0; i < frag->nb_units; i++)
98 cbs_unit_uninit(ctx, &frag->units[i]);
99 av_freep(&frag->units);
100 frag->nb_units = 0;
101
102 av_freep(&frag->data);
103 frag->data_size = 0;
104 frag->data_bit_padding = 0;
105 }
106
107 static int cbs_read_fragment_content(CodedBitstreamContext *ctx,
108 CodedBitstreamFragment *frag)
109 {
110 int err, i, j;
111
112 for (i = 0; i < frag->nb_units; i++) {
113 if (ctx->decompose_unit_types) {
114 for (j = 0; j < ctx->nb_decompose_unit_types; j++) {
115 if (ctx->decompose_unit_types[j] == frag->units[i].type)
116 break;
117 }
118 if (j >= ctx->nb_decompose_unit_types)
119 continue;
120 }
121
122 err = ctx->codec->read_unit(ctx, &frag->units[i]);
123 if (err == AVERROR(ENOSYS)) {
124 av_log(ctx->log_ctx, AV_LOG_WARNING,
125 "Decomposition unimplemented for unit %d "
126 "(type %d).\n", i, frag->units[i].type);
127 } else if (err < 0) {
128 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to read unit %d "
129 "(type %d).\n", i, frag->units[i].type);
130 return err;
131 }
132 }
133
134 return 0;
135 }
136
137 int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
138 CodedBitstreamFragment *frag,
139 const AVCodecParameters *par)
140 {
141 int err;
142
143 memset(frag, 0, sizeof(*frag));
144
145 frag->data = par->extradata;
146 frag->data_size = par->extradata_size;
147
148 err = ctx->codec->split_fragment(ctx, frag, 1);
149 if (err < 0)
150 return err;
151
152 frag->data = NULL;
153 frag->data_size = 0;
154
155 return cbs_read_fragment_content(ctx, frag);
156 }
157
158 int ff_cbs_read_packet(CodedBitstreamContext *ctx,
159 CodedBitstreamFragment *frag,
160 const AVPacket *pkt)
161 {
162 int err;
163
164 memset(frag, 0, sizeof(*frag));
165
166 frag->data = pkt->data;
167 frag->data_size = pkt->size;
168
169 err = ctx->codec->split_fragment(ctx, frag, 0);
170 if (err < 0)
171 return err;
172
173 frag->data = NULL;
174 frag->data_size = 0;
175
176 return cbs_read_fragment_content(ctx, frag);
177 }
178
179 int ff_cbs_read(CodedBitstreamContext *ctx,
180 CodedBitstreamFragment *frag,
181 const uint8_t *data, size_t size)
182 {
183 int err;
184
185 memset(frag, 0, sizeof(*frag));
186
187 // (We won't write to this during split.)
188 frag->data = (uint8_t*)data;
189 frag->data_size = size;
190
191 err = ctx->codec->split_fragment(ctx, frag, 0);
192 if (err < 0)
193 return err;
194
195 frag->data = NULL;
196 frag->data_size = 0;
197
198 return cbs_read_fragment_content(ctx, frag);
199 }
200
201
202 int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
203 CodedBitstreamFragment *frag)
204 {
205 int err, i;
206
207 for (i = 0; i < frag->nb_units; i++) {
208 if (!frag->units[i].content)
209 continue;
210
211 err = ctx->codec->write_unit(ctx, &frag->units[i]);
212 if (err < 0) {
213 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to write unit %d "
214 "(type %d).\n", i, frag->units[i].type);
215 return err;
216 }
217 }
218
219 err = ctx->codec->assemble_fragment(ctx, frag);
220 if (err < 0) {
221 av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to assemble fragment.\n");
222 return err;
223 }
224
225 return 0;
226 }
227
228 int ff_cbs_write_extradata(CodedBitstreamContext *ctx,
229 AVCodecParameters *par,
230 CodedBitstreamFragment *frag)
231 {
232 int err;
233
234 err = ff_cbs_write_fragment_data(ctx, frag);
235 if (err < 0)
236 return err;
237
238 av_freep(&par->extradata);
239
240 par->extradata = av_malloc(frag->data_size +
241 AV_INPUT_BUFFER_PADDING_SIZE);
242 if (!par->extradata)
243 return AVERROR(ENOMEM);
244
245 memcpy(par->extradata, frag->data, frag->data_size);
246 memset(par->extradata + frag->data_size, 0,
247 AV_INPUT_BUFFER_PADDING_SIZE);
248 par->extradata_size = frag->data_size;
249
250 return 0;
251 }
252
253 int ff_cbs_write_packet(CodedBitstreamContext *ctx,
254 AVPacket *pkt,
255 CodedBitstreamFragment *frag)
256 {
257 int err;
258
259 err = ff_cbs_write_fragment_data(ctx, frag);
260 if (err < 0)
261 return err;
262
263 av_new_packet(pkt, frag->data_size);
264 if (err < 0)
265 return err;
266
267 memcpy(pkt->data, frag->data, frag->data_size);
268 pkt->size = frag->data_size;
269
270 return 0;
271 }
272
273
274 void ff_cbs_trace_header(CodedBitstreamContext *ctx,
275 const char *name)
276 {
277 if (!ctx->trace_enable)
278 return;
279
280 av_log(ctx->log_ctx, ctx->trace_level, "%s\n", name);
281 }
282
283 void ff_cbs_trace_syntax_element(CodedBitstreamContext *ctx, int position,
284 const char *name, const char *bits,
285 int64_t value)
286 {
287 size_t name_len, bits_len;
288 int pad;
289
290 if (!ctx->trace_enable)
291 return;
292
293 av_assert0(value >= INT_MIN && value <= UINT32_MAX);
294
295 name_len = strlen(name);
296 bits_len = strlen(bits);
297
298 if (name_len + bits_len > 60)
299 pad = bits_len + 2;
300 else
301 pad = 61 - name_len;
302
303 av_log(ctx->log_ctx, ctx->trace_level, "%-10d %s%*s = %"PRId64"\n",
304 position, name, pad, bits, value);
305 }
306
307 int ff_cbs_read_unsigned(CodedBitstreamContext *ctx, BitstreamContext *bc,
308 int width, const char *name, uint32_t *write_to,
309 uint32_t range_min, uint32_t range_max)
310 {
311 uint32_t value;
312 int position;
313
314 av_assert0(width <= 32);
315
316 if (bitstream_bits_left(bc) < width) {
317 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid value at "
318 "%s: bitstream ended.\n", name);
319 return AVERROR_INVALIDDATA;
320 }
321
322 if (ctx->trace_enable)
323 position = bitstream_tell(bc);
324
325 value = bitstream_read(bc, width);
326
327 if (ctx->trace_enable) {
328 char bits[33];
329 int i;
330 for (i = 0; i < width; i++)
331 bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
332 bits[i] = 0;
333
334 ff_cbs_trace_syntax_element(ctx, position, name, bits, value);
335 }
336
337 if (value < range_min || value > range_max) {
338 av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
339 "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
340 name, value, range_min, range_max);
341 return AVERROR_INVALIDDATA;
342 }
343
344 *write_to = value;
345 return 0;
346 }
347
348 int ff_cbs_write_unsigned(CodedBitstreamContext *ctx, PutBitContext *pbc,
349 int width, const char *name, uint32_t value,
350 uint32_t range_min, uint32_t range_max)
351 {
352 av_assert0(width <= 32);
353
354 if (value < range_min || value > range_max) {
355 av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
356 "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
357 name, value, range_min, range_max);
358 return AVERROR_INVALIDDATA;
359 }
360
361 if (put_bits_left(pbc) < width)
362 return AVERROR(ENOSPC);
363
364 if (ctx->trace_enable) {
365 char bits[33];
366 int i;
367 for (i = 0; i < width; i++)
368 bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
369 bits[i] = 0;
370
371 ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc), name, bits, value);
372 }
373
374 if (width < 32)
375 put_bits(pbc, width, value);
376 else
377 put_bits32(pbc, value);
378
379 return 0;
380 }
381
382
383 static int cbs_insert_unit(CodedBitstreamContext *ctx,
384 CodedBitstreamFragment *frag,
385 int position)
386 {
387 CodedBitstreamUnit *units;
388
389 units = av_malloc_array(frag->nb_units + 1, sizeof(*units));
390 if (!units)
391 return AVERROR(ENOMEM);
392
393 if (position > 0)
394 memcpy(units, frag->units, position * sizeof(*units));
395 if (position < frag->nb_units)
396 memcpy(units + position + 1, frag->units + position,
397 (frag->nb_units - position) * sizeof(*units));
398
399 memset(units + position, 0, sizeof(*units));
400
401 av_freep(&frag->units);
402 frag->units = units;
403 ++frag->nb_units;
404
405 return 0;
406 }
407
408 int ff_cbs_insert_unit_content(CodedBitstreamContext *ctx,
409 CodedBitstreamFragment *frag,
410 int position, uint32_t type,
411 void *content)
412 {
413 int err;
414
415 if (position == -1)
416 position = frag->nb_units;
417 av_assert0(position >= 0 && position <= frag->nb_units);
418
419 err = cbs_insert_unit(ctx, frag, position);
420 if (err < 0)
421 return err;
422
423 frag->units[position].type = type;
424 frag->units[position].content = content;
425 frag->units[position].content_external = 1;
426
427 return 0;
428 }
429
430 int ff_cbs_insert_unit_data(CodedBitstreamContext *ctx,
431 CodedBitstreamFragment *frag,
432 int position, uint32_t type,
433 uint8_t *data, size_t data_size)
434 {
435 int err;
436
437 if (position == -1)
438 position = frag->nb_units;
439 av_assert0(position >= 0 && position <= frag->nb_units);
440
441 err = cbs_insert_unit(ctx, frag, position);
442 if (err < 0)
443 return err;
444
445 frag->units[position].type = type;
446 frag->units[position].data = data;
447 frag->units[position].data_size = data_size;
448
449 return 0;
450 }
451
452 int ff_cbs_delete_unit(CodedBitstreamContext *ctx,
453 CodedBitstreamFragment *frag,
454 int position)
455 {
456 if (position < 0 || position >= frag->nb_units)
457 return AVERROR(EINVAL);
458
459 cbs_unit_uninit(ctx, &frag->units[position]);
460
461 --frag->nb_units;
462
463 if (frag->nb_units == 0) {
464 av_freep(&frag->units);
465
466 } else {
467 memmove(frag->units + position,
468 frag->units + position + 1,
469 (frag->nb_units - position) * sizeof(*frag->units));
470
471 // Don't bother reallocating the unit array.
472 }
473
474 return 0;
475 }