• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * This file is part of FFmpeg.
3  *
4  * FFmpeg 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  * FFmpeg 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 FFmpeg; 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/buffer.h"
25 #include "libavutil/common.h"
26 #include "libavutil/opt.h"
27 
28 #include "cbs.h"
29 #include "cbs_internal.h"
30 
31 
32 static const CodedBitstreamType *const cbs_type_table[] = {
33 #if CONFIG_CBS_AV1
34     &ff_cbs_type_av1,
35 #endif
36 #if CONFIG_CBS_H264
37     &ff_cbs_type_h264,
38 #endif
39 #if CONFIG_CBS_H265
40     &ff_cbs_type_h265,
41 #endif
42 #if CONFIG_CBS_JPEG
43     &ff_cbs_type_jpeg,
44 #endif
45 #if CONFIG_CBS_MPEG2
46     &ff_cbs_type_mpeg2,
47 #endif
48 #if CONFIG_CBS_VP9
49     &ff_cbs_type_vp9,
50 #endif
51 };
52 
53 const enum AVCodecID ff_cbs_all_codec_ids[] = {
54 #if CONFIG_CBS_AV1
55     AV_CODEC_ID_AV1,
56 #endif
57 #if CONFIG_CBS_H264
58     AV_CODEC_ID_H264,
59 #endif
60 #if CONFIG_CBS_H265
61     AV_CODEC_ID_H265,
62 #endif
63 #if CONFIG_CBS_JPEG
64     AV_CODEC_ID_MJPEG,
65 #endif
66 #if CONFIG_CBS_MPEG2
67     AV_CODEC_ID_MPEG2VIDEO,
68 #endif
69 #if CONFIG_CBS_VP9
70     AV_CODEC_ID_VP9,
71 #endif
72     AV_CODEC_ID_NONE
73 };
74 
ff_cbs_init(CodedBitstreamContext ** ctx_ptr,enum AVCodecID codec_id,void * log_ctx)75 int ff_cbs_init(CodedBitstreamContext **ctx_ptr,
76                 enum AVCodecID codec_id, void *log_ctx)
77 {
78     CodedBitstreamContext *ctx;
79     const CodedBitstreamType *type;
80     int i;
81 
82     type = NULL;
83     for (i = 0; i < FF_ARRAY_ELEMS(cbs_type_table); i++) {
84         if (cbs_type_table[i]->codec_id == codec_id) {
85             type = cbs_type_table[i];
86             break;
87         }
88     }
89     if (!type)
90         return AVERROR(EINVAL);
91 
92     ctx = av_mallocz(sizeof(*ctx));
93     if (!ctx)
94         return AVERROR(ENOMEM);
95 
96     ctx->log_ctx = log_ctx;
97     ctx->codec   = type; /* Must be before any error */
98 
99     if (type->priv_data_size) {
100         ctx->priv_data = av_mallocz(ctx->codec->priv_data_size);
101         if (!ctx->priv_data) {
102             av_freep(&ctx);
103             return AVERROR(ENOMEM);
104         }
105         if (type->priv_class) {
106             *(const AVClass **)ctx->priv_data = type->priv_class;
107             av_opt_set_defaults(ctx->priv_data);
108         }
109     }
110 
111     ctx->decompose_unit_types = NULL;
112 
113     ctx->trace_enable = 0;
114     ctx->trace_level  = AV_LOG_TRACE;
115 
116     *ctx_ptr = ctx;
117     return 0;
118 }
119 
ff_cbs_flush(CodedBitstreamContext * ctx)120 void ff_cbs_flush(CodedBitstreamContext *ctx)
121 {
122     if (ctx->codec->flush)
123         ctx->codec->flush(ctx);
124 }
125 
ff_cbs_close(CodedBitstreamContext ** ctx_ptr)126 void ff_cbs_close(CodedBitstreamContext **ctx_ptr)
127 {
128     CodedBitstreamContext *ctx = *ctx_ptr;
129 
130     if (!ctx)
131         return;
132 
133     if (ctx->codec->close)
134         ctx->codec->close(ctx);
135 
136     av_freep(&ctx->write_buffer);
137 
138     if (ctx->codec->priv_class && ctx->priv_data)
139         av_opt_free(ctx->priv_data);
140 
141     av_freep(&ctx->priv_data);
142     av_freep(ctx_ptr);
143 }
144 
cbs_unit_uninit(CodedBitstreamUnit * unit)145 static void cbs_unit_uninit(CodedBitstreamUnit *unit)
146 {
147     av_buffer_unref(&unit->content_ref);
148     unit->content = NULL;
149 
150     av_buffer_unref(&unit->data_ref);
151     unit->data             = NULL;
152     unit->data_size        = 0;
153     unit->data_bit_padding = 0;
154 }
155 
ff_cbs_fragment_reset(CodedBitstreamFragment * frag)156 void ff_cbs_fragment_reset(CodedBitstreamFragment *frag)
157 {
158     int i;
159 
160     for (i = 0; i < frag->nb_units; i++)
161         cbs_unit_uninit(&frag->units[i]);
162     frag->nb_units = 0;
163 
164     av_buffer_unref(&frag->data_ref);
165     frag->data             = NULL;
166     frag->data_size        = 0;
167     frag->data_bit_padding = 0;
168 }
169 
ff_cbs_fragment_free(CodedBitstreamFragment * frag)170 void ff_cbs_fragment_free(CodedBitstreamFragment *frag)
171 {
172     ff_cbs_fragment_reset(frag);
173 
174     av_freep(&frag->units);
175     frag->nb_units_allocated = 0;
176 }
177 
cbs_read_fragment_content(CodedBitstreamContext * ctx,CodedBitstreamFragment * frag)178 static int cbs_read_fragment_content(CodedBitstreamContext *ctx,
179                                      CodedBitstreamFragment *frag)
180 {
181     int err, i, j;
182 
183     for (i = 0; i < frag->nb_units; i++) {
184         CodedBitstreamUnit *unit = &frag->units[i];
185 
186         if (ctx->decompose_unit_types) {
187             for (j = 0; j < ctx->nb_decompose_unit_types; j++) {
188                 if (ctx->decompose_unit_types[j] == unit->type)
189                     break;
190             }
191             if (j >= ctx->nb_decompose_unit_types)
192                 continue;
193         }
194 
195         av_buffer_unref(&unit->content_ref);
196         unit->content = NULL;
197 
198         av_assert0(unit->data && unit->data_ref);
199 
200         err = ctx->codec->read_unit(ctx, unit);
201         if (err == AVERROR(ENOSYS)) {
202             av_log(ctx->log_ctx, AV_LOG_VERBOSE,
203                    "Decomposition unimplemented for unit %d "
204                    "(type %"PRIu32").\n", i, unit->type);
205         } else if (err == AVERROR(EAGAIN)) {
206             av_log(ctx->log_ctx, AV_LOG_VERBOSE,
207                    "Skipping decomposition of unit %d "
208                    "(type %"PRIu32").\n", i, unit->type);
209             av_buffer_unref(&unit->content_ref);
210             unit->content = NULL;
211         } else if (err < 0) {
212             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to read unit %d "
213                    "(type %"PRIu32").\n", i, unit->type);
214             return err;
215         }
216     }
217 
218     return 0;
219 }
220 
cbs_fill_fragment_data(CodedBitstreamFragment * frag,const uint8_t * data,size_t size)221 static int cbs_fill_fragment_data(CodedBitstreamFragment *frag,
222                                   const uint8_t *data, size_t size)
223 {
224     av_assert0(!frag->data && !frag->data_ref);
225 
226     frag->data_ref =
227         av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
228     if (!frag->data_ref)
229         return AVERROR(ENOMEM);
230 
231     frag->data      = frag->data_ref->data;
232     frag->data_size = size;
233 
234     memcpy(frag->data, data, size);
235     memset(frag->data + size, 0,
236            AV_INPUT_BUFFER_PADDING_SIZE);
237 
238     return 0;
239 }
240 
cbs_read_data(CodedBitstreamContext * ctx,CodedBitstreamFragment * frag,AVBufferRef * buf,const uint8_t * data,size_t size,int header)241 static int cbs_read_data(CodedBitstreamContext *ctx,
242                          CodedBitstreamFragment *frag,
243                          AVBufferRef *buf,
244                          const uint8_t *data, size_t size,
245                          int header)
246 {
247     int err;
248 
249     if (buf) {
250         frag->data_ref = av_buffer_ref(buf);
251         if (!frag->data_ref)
252             return AVERROR(ENOMEM);
253 
254         frag->data      = (uint8_t *)data;
255         frag->data_size = size;
256 
257     } else {
258         err = cbs_fill_fragment_data(frag, data, size);
259         if (err < 0)
260             return err;
261     }
262 
263     err = ctx->codec->split_fragment(ctx, frag, header);
264     if (err < 0)
265         return err;
266 
267     return cbs_read_fragment_content(ctx, frag);
268 }
269 
ff_cbs_read_extradata(CodedBitstreamContext * ctx,CodedBitstreamFragment * frag,const AVCodecParameters * par)270 int ff_cbs_read_extradata(CodedBitstreamContext *ctx,
271                           CodedBitstreamFragment *frag,
272                           const AVCodecParameters *par)
273 {
274     return cbs_read_data(ctx, frag, NULL,
275                          par->extradata,
276                          par->extradata_size, 1);
277 }
278 
ff_cbs_read_extradata_from_codec(CodedBitstreamContext * ctx,CodedBitstreamFragment * frag,const AVCodecContext * avctx)279 int ff_cbs_read_extradata_from_codec(CodedBitstreamContext *ctx,
280                                      CodedBitstreamFragment *frag,
281                                      const AVCodecContext *avctx)
282 {
283     return cbs_read_data(ctx, frag, NULL,
284                          avctx->extradata,
285                          avctx->extradata_size, 1);
286 }
287 
ff_cbs_read_packet(CodedBitstreamContext * ctx,CodedBitstreamFragment * frag,const AVPacket * pkt)288 int ff_cbs_read_packet(CodedBitstreamContext *ctx,
289                        CodedBitstreamFragment *frag,
290                        const AVPacket *pkt)
291 {
292     return cbs_read_data(ctx, frag, pkt->buf,
293                          pkt->data, pkt->size, 0);
294 }
295 
ff_cbs_read(CodedBitstreamContext * ctx,CodedBitstreamFragment * frag,const uint8_t * data,size_t size)296 int ff_cbs_read(CodedBitstreamContext *ctx,
297                 CodedBitstreamFragment *frag,
298                 const uint8_t *data, size_t size)
299 {
300     return cbs_read_data(ctx, frag, NULL,
301                          data, size, 0);
302 }
303 
cbs_write_unit_data(CodedBitstreamContext * ctx,CodedBitstreamUnit * unit)304 static int cbs_write_unit_data(CodedBitstreamContext *ctx,
305                                CodedBitstreamUnit *unit)
306 {
307     PutBitContext pbc;
308     int ret;
309 
310     if (!ctx->write_buffer) {
311         // Initial write buffer size is 1MB.
312         ctx->write_buffer_size = 1024 * 1024;
313 
314     reallocate_and_try_again:
315         ret = av_reallocp(&ctx->write_buffer, ctx->write_buffer_size);
316         if (ret < 0) {
317             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
318                    "sufficiently large write buffer (last attempt "
319                    "%"SIZE_SPECIFIER" bytes).\n", ctx->write_buffer_size);
320             return ret;
321         }
322     }
323 
324     init_put_bits(&pbc, ctx->write_buffer, ctx->write_buffer_size);
325 
326     ret = ctx->codec->write_unit(ctx, unit, &pbc);
327     if (ret < 0) {
328         if (ret == AVERROR(ENOSPC)) {
329             // Overflow.
330             if (ctx->write_buffer_size == INT_MAX / 8)
331                 return AVERROR(ENOMEM);
332             ctx->write_buffer_size = FFMIN(2 * ctx->write_buffer_size, INT_MAX / 8);
333             goto reallocate_and_try_again;
334         }
335         // Write failed for some other reason.
336         return ret;
337     }
338 
339     // Overflow but we didn't notice.
340     av_assert0(put_bits_count(&pbc) <= 8 * ctx->write_buffer_size);
341 
342     if (put_bits_count(&pbc) % 8)
343         unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
344     else
345         unit->data_bit_padding = 0;
346 
347     flush_put_bits(&pbc);
348 
349     ret = ff_cbs_alloc_unit_data(unit, put_bits_count(&pbc) / 8);
350     if (ret < 0)
351         return ret;
352 
353     memcpy(unit->data, ctx->write_buffer, unit->data_size);
354 
355     return 0;
356 }
357 
ff_cbs_write_fragment_data(CodedBitstreamContext * ctx,CodedBitstreamFragment * frag)358 int ff_cbs_write_fragment_data(CodedBitstreamContext *ctx,
359                                CodedBitstreamFragment *frag)
360 {
361     int err, i;
362 
363     for (i = 0; i < frag->nb_units; i++) {
364         CodedBitstreamUnit *unit = &frag->units[i];
365 
366         if (!unit->content)
367             continue;
368 
369         av_buffer_unref(&unit->data_ref);
370         unit->data = NULL;
371 
372         err = cbs_write_unit_data(ctx, unit);
373         if (err < 0) {
374             av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to write unit %d "
375                    "(type %"PRIu32").\n", i, unit->type);
376             return err;
377         }
378         av_assert0(unit->data && unit->data_ref);
379     }
380 
381     av_buffer_unref(&frag->data_ref);
382     frag->data = NULL;
383 
384     err = ctx->codec->assemble_fragment(ctx, frag);
385     if (err < 0) {
386         av_log(ctx->log_ctx, AV_LOG_ERROR, "Failed to assemble fragment.\n");
387         return err;
388     }
389     av_assert0(frag->data && frag->data_ref);
390 
391     return 0;
392 }
393 
ff_cbs_write_extradata(CodedBitstreamContext * ctx,AVCodecParameters * par,CodedBitstreamFragment * frag)394 int ff_cbs_write_extradata(CodedBitstreamContext *ctx,
395                            AVCodecParameters *par,
396                            CodedBitstreamFragment *frag)
397 {
398     int err;
399 
400     err = ff_cbs_write_fragment_data(ctx, frag);
401     if (err < 0)
402         return err;
403 
404     av_freep(&par->extradata);
405 
406     par->extradata = av_malloc(frag->data_size +
407                                AV_INPUT_BUFFER_PADDING_SIZE);
408     if (!par->extradata)
409         return AVERROR(ENOMEM);
410 
411     memcpy(par->extradata, frag->data, frag->data_size);
412     memset(par->extradata + frag->data_size, 0,
413            AV_INPUT_BUFFER_PADDING_SIZE);
414     par->extradata_size = frag->data_size;
415 
416     return 0;
417 }
418 
ff_cbs_write_packet(CodedBitstreamContext * ctx,AVPacket * pkt,CodedBitstreamFragment * frag)419 int ff_cbs_write_packet(CodedBitstreamContext *ctx,
420                         AVPacket *pkt,
421                         CodedBitstreamFragment *frag)
422 {
423     AVBufferRef *buf;
424     int err;
425 
426     err = ff_cbs_write_fragment_data(ctx, frag);
427     if (err < 0)
428         return err;
429 
430     buf = av_buffer_ref(frag->data_ref);
431     if (!buf)
432         return AVERROR(ENOMEM);
433 
434     av_buffer_unref(&pkt->buf);
435 
436     pkt->buf  = buf;
437     pkt->data = frag->data;
438     pkt->size = frag->data_size;
439 
440     return 0;
441 }
442 
443 
ff_cbs_trace_header(CodedBitstreamContext * ctx,const char * name)444 void ff_cbs_trace_header(CodedBitstreamContext *ctx,
445                          const char *name)
446 {
447     if (!ctx->trace_enable)
448         return;
449 
450     av_log(ctx->log_ctx, ctx->trace_level, "%s\n", name);
451 }
452 
ff_cbs_trace_syntax_element(CodedBitstreamContext * ctx,int position,const char * str,const int * subscripts,const char * bits,int64_t value)453 void ff_cbs_trace_syntax_element(CodedBitstreamContext *ctx, int position,
454                                  const char *str, const int *subscripts,
455                                  const char *bits, int64_t value)
456 {
457     char name[256];
458     size_t name_len, bits_len;
459     int pad, subs, i, j, k, n;
460 
461     if (!ctx->trace_enable)
462         return;
463 
464     av_assert0(value >= INT_MIN && value <= UINT32_MAX);
465 
466     subs = subscripts ? subscripts[0] : 0;
467     n = 0;
468     for (i = j = 0; str[i];) {
469         if (str[i] == '[') {
470             if (n < subs) {
471                 ++n;
472                 k = snprintf(name + j, sizeof(name) - j, "[%d", subscripts[n]);
473                 av_assert0(k > 0 && j + k < sizeof(name));
474                 j += k;
475                 for (++i; str[i] && str[i] != ']'; i++);
476                 av_assert0(str[i] == ']');
477             } else {
478                 while (str[i] && str[i] != ']')
479                     name[j++] = str[i++];
480                 av_assert0(str[i] == ']');
481             }
482         } else {
483             av_assert0(j + 1 < sizeof(name));
484             name[j++] = str[i++];
485         }
486     }
487     av_assert0(j + 1 < sizeof(name));
488     name[j] = 0;
489     av_assert0(n == subs);
490 
491     name_len = strlen(name);
492     bits_len = strlen(bits);
493 
494     if (name_len + bits_len > 60)
495         pad = bits_len + 2;
496     else
497         pad = 61 - name_len;
498 
499     av_log(ctx->log_ctx, ctx->trace_level, "%-10d  %s%*s = %"PRId64"\n",
500            position, name, pad, bits, value);
501 }
502 
ff_cbs_read_unsigned(CodedBitstreamContext * ctx,GetBitContext * gbc,int width,const char * name,const int * subscripts,uint32_t * write_to,uint32_t range_min,uint32_t range_max)503 int ff_cbs_read_unsigned(CodedBitstreamContext *ctx, GetBitContext *gbc,
504                          int width, const char *name,
505                          const int *subscripts, uint32_t *write_to,
506                          uint32_t range_min, uint32_t range_max)
507 {
508     uint32_t value;
509     int position;
510 
511     av_assert0(width > 0 && width <= 32);
512 
513     if (get_bits_left(gbc) < width) {
514         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid value at "
515                "%s: bitstream ended.\n", name);
516         return AVERROR_INVALIDDATA;
517     }
518 
519     if (ctx->trace_enable)
520         position = get_bits_count(gbc);
521 
522     value = get_bits_long(gbc, width);
523 
524     if (ctx->trace_enable) {
525         char bits[33];
526         int i;
527         for (i = 0; i < width; i++)
528             bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
529         bits[i] = 0;
530 
531         ff_cbs_trace_syntax_element(ctx, position, name, subscripts,
532                                     bits, value);
533     }
534 
535     if (value < range_min || value > range_max) {
536         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
537                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
538                name, value, range_min, range_max);
539         return AVERROR_INVALIDDATA;
540     }
541 
542     *write_to = value;
543     return 0;
544 }
545 
ff_cbs_write_unsigned(CodedBitstreamContext * ctx,PutBitContext * pbc,int width,const char * name,const int * subscripts,uint32_t value,uint32_t range_min,uint32_t range_max)546 int ff_cbs_write_unsigned(CodedBitstreamContext *ctx, PutBitContext *pbc,
547                           int width, const char *name,
548                           const int *subscripts, uint32_t value,
549                           uint32_t range_min, uint32_t range_max)
550 {
551     av_assert0(width > 0 && width <= 32);
552 
553     if (value < range_min || value > range_max) {
554         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
555                "%"PRIu32", but must be in [%"PRIu32",%"PRIu32"].\n",
556                name, value, range_min, range_max);
557         return AVERROR_INVALIDDATA;
558     }
559 
560     if (put_bits_left(pbc) < width)
561         return AVERROR(ENOSPC);
562 
563     if (ctx->trace_enable) {
564         char bits[33];
565         int i;
566         for (i = 0; i < width; i++)
567             bits[i] = value >> (width - i - 1) & 1 ? '1' : '0';
568         bits[i] = 0;
569 
570         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc),
571                                     name, subscripts, bits, value);
572     }
573 
574     if (width < 32)
575         put_bits(pbc, width, value);
576     else
577         put_bits32(pbc, value);
578 
579     return 0;
580 }
581 
ff_cbs_read_signed(CodedBitstreamContext * ctx,GetBitContext * gbc,int width,const char * name,const int * subscripts,int32_t * write_to,int32_t range_min,int32_t range_max)582 int ff_cbs_read_signed(CodedBitstreamContext *ctx, GetBitContext *gbc,
583                        int width, const char *name,
584                        const int *subscripts, int32_t *write_to,
585                        int32_t range_min, int32_t range_max)
586 {
587     int32_t value;
588     int position;
589 
590     av_assert0(width > 0 && width <= 32);
591 
592     if (get_bits_left(gbc) < width) {
593         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid value at "
594                "%s: bitstream ended.\n", name);
595         return AVERROR_INVALIDDATA;
596     }
597 
598     if (ctx->trace_enable)
599         position = get_bits_count(gbc);
600 
601     value = get_sbits_long(gbc, width);
602 
603     if (ctx->trace_enable) {
604         char bits[33];
605         int i;
606         for (i = 0; i < width; i++)
607             bits[i] = value & (1U << (width - i - 1)) ? '1' : '0';
608         bits[i] = 0;
609 
610         ff_cbs_trace_syntax_element(ctx, position, name, subscripts,
611                                     bits, value);
612     }
613 
614     if (value < range_min || value > range_max) {
615         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
616                "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
617                name, value, range_min, range_max);
618         return AVERROR_INVALIDDATA;
619     }
620 
621     *write_to = value;
622     return 0;
623 }
624 
ff_cbs_write_signed(CodedBitstreamContext * ctx,PutBitContext * pbc,int width,const char * name,const int * subscripts,int32_t value,int32_t range_min,int32_t range_max)625 int ff_cbs_write_signed(CodedBitstreamContext *ctx, PutBitContext *pbc,
626                         int width, const char *name,
627                         const int *subscripts, int32_t value,
628                         int32_t range_min, int32_t range_max)
629 {
630     av_assert0(width > 0 && width <= 32);
631 
632     if (value < range_min || value > range_max) {
633         av_log(ctx->log_ctx, AV_LOG_ERROR, "%s out of range: "
634                "%"PRId32", but must be in [%"PRId32",%"PRId32"].\n",
635                name, value, range_min, range_max);
636         return AVERROR_INVALIDDATA;
637     }
638 
639     if (put_bits_left(pbc) < width)
640         return AVERROR(ENOSPC);
641 
642     if (ctx->trace_enable) {
643         char bits[33];
644         int i;
645         for (i = 0; i < width; i++)
646             bits[i] = value & (1U << (width - i - 1)) ? '1' : '0';
647         bits[i] = 0;
648 
649         ff_cbs_trace_syntax_element(ctx, put_bits_count(pbc),
650                                     name, subscripts, bits, value);
651     }
652 
653     if (width < 32)
654         put_sbits(pbc, width, value);
655     else
656         put_bits32(pbc, value);
657 
658     return 0;
659 }
660 
661 
ff_cbs_alloc_unit_content(CodedBitstreamUnit * unit,size_t size,void (* free)(void * opaque,uint8_t * data))662 int ff_cbs_alloc_unit_content(CodedBitstreamUnit *unit,
663                               size_t size,
664                               void (*free)(void *opaque, uint8_t *data))
665 {
666     av_assert0(!unit->content && !unit->content_ref);
667 
668     unit->content = av_mallocz(size);
669     if (!unit->content)
670         return AVERROR(ENOMEM);
671 
672     unit->content_ref = av_buffer_create(unit->content, size,
673                                          free, NULL, 0);
674     if (!unit->content_ref) {
675         av_freep(&unit->content);
676         return AVERROR(ENOMEM);
677     }
678 
679     return 0;
680 }
681 
ff_cbs_alloc_unit_data(CodedBitstreamUnit * unit,size_t size)682 int ff_cbs_alloc_unit_data(CodedBitstreamUnit *unit,
683                            size_t size)
684 {
685     av_assert0(!unit->data && !unit->data_ref);
686 
687     unit->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
688     if (!unit->data_ref)
689         return AVERROR(ENOMEM);
690 
691     unit->data      = unit->data_ref->data;
692     unit->data_size = size;
693 
694     memset(unit->data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
695 
696     return 0;
697 }
698 
cbs_insert_unit(CodedBitstreamFragment * frag,int position)699 static int cbs_insert_unit(CodedBitstreamFragment *frag,
700                            int position)
701 {
702     CodedBitstreamUnit *units;
703 
704     if (frag->nb_units < frag->nb_units_allocated) {
705         units = frag->units;
706 
707         if (position < frag->nb_units)
708             memmove(units + position + 1, units + position,
709                     (frag->nb_units - position) * sizeof(*units));
710     } else {
711         units = av_malloc_array(frag->nb_units*2 + 1, sizeof(*units));
712         if (!units)
713             return AVERROR(ENOMEM);
714 
715         frag->nb_units_allocated = 2*frag->nb_units_allocated + 1;
716 
717         if (position > 0)
718             memcpy(units, frag->units, position * sizeof(*units));
719 
720         if (position < frag->nb_units)
721             memcpy(units + position + 1, frag->units + position,
722                    (frag->nb_units - position) * sizeof(*units));
723     }
724 
725     memset(units + position, 0, sizeof(*units));
726 
727     if (units != frag->units) {
728         av_free(frag->units);
729         frag->units = units;
730     }
731 
732     ++frag->nb_units;
733 
734     return 0;
735 }
736 
ff_cbs_insert_unit_content(CodedBitstreamFragment * frag,int position,CodedBitstreamUnitType type,void * content,AVBufferRef * content_buf)737 int ff_cbs_insert_unit_content(CodedBitstreamFragment *frag,
738                                int position,
739                                CodedBitstreamUnitType type,
740                                void *content,
741                                AVBufferRef *content_buf)
742 {
743     CodedBitstreamUnit *unit;
744     AVBufferRef *content_ref;
745     int err;
746 
747     if (position == -1)
748         position = frag->nb_units;
749     av_assert0(position >= 0 && position <= frag->nb_units);
750 
751     if (content_buf) {
752         content_ref = av_buffer_ref(content_buf);
753         if (!content_ref)
754             return AVERROR(ENOMEM);
755     } else {
756         content_ref = NULL;
757     }
758 
759     err = cbs_insert_unit(frag, position);
760     if (err < 0) {
761         av_buffer_unref(&content_ref);
762         return err;
763     }
764 
765     unit = &frag->units[position];
766     unit->type        = type;
767     unit->content     = content;
768     unit->content_ref = content_ref;
769 
770     return 0;
771 }
772 
ff_cbs_insert_unit_data(CodedBitstreamFragment * frag,int position,CodedBitstreamUnitType type,uint8_t * data,size_t data_size,AVBufferRef * data_buf)773 int ff_cbs_insert_unit_data(CodedBitstreamFragment *frag,
774                             int position,
775                             CodedBitstreamUnitType type,
776                             uint8_t *data, size_t data_size,
777                             AVBufferRef *data_buf)
778 {
779     CodedBitstreamUnit *unit;
780     AVBufferRef *data_ref;
781     int err;
782 
783     if (position == -1)
784         position = frag->nb_units;
785     av_assert0(position >= 0 && position <= frag->nb_units);
786 
787     if (data_buf)
788         data_ref = av_buffer_ref(data_buf);
789     else
790         data_ref = av_buffer_create(data, data_size, NULL, NULL, 0);
791     if (!data_ref) {
792         if (!data_buf)
793             av_free(data);
794         return AVERROR(ENOMEM);
795     }
796 
797     err = cbs_insert_unit(frag, position);
798     if (err < 0) {
799         av_buffer_unref(&data_ref);
800         return err;
801     }
802 
803     unit = &frag->units[position];
804     unit->type      = type;
805     unit->data      = data;
806     unit->data_size = data_size;
807     unit->data_ref  = data_ref;
808 
809     return 0;
810 }
811 
ff_cbs_delete_unit(CodedBitstreamFragment * frag,int position)812 void ff_cbs_delete_unit(CodedBitstreamFragment *frag,
813                         int position)
814 {
815     av_assert0(0 <= position && position < frag->nb_units
816                              && "Unit to be deleted not in fragment.");
817 
818     cbs_unit_uninit(&frag->units[position]);
819 
820     --frag->nb_units;
821 
822     if (frag->nb_units > 0)
823         memmove(frag->units + position,
824                 frag->units + position + 1,
825                 (frag->nb_units - position) * sizeof(*frag->units));
826 }
827 
cbs_default_free_unit_content(void * opaque,uint8_t * data)828 static void cbs_default_free_unit_content(void *opaque, uint8_t *data)
829 {
830     const CodedBitstreamUnitTypeDescriptor *desc = opaque;
831     if (desc->content_type == CBS_CONTENT_TYPE_INTERNAL_REFS) {
832         int i;
833         for (i = 0; i < desc->nb_ref_offsets; i++) {
834             void **ptr = (void**)(data + desc->ref_offsets[i]);
835             av_buffer_unref((AVBufferRef**)(ptr + 1));
836         }
837     }
838     av_free(data);
839 }
840 
841 static const CodedBitstreamUnitTypeDescriptor
cbs_find_unit_type_desc(CodedBitstreamContext * ctx,CodedBitstreamUnit * unit)842     *cbs_find_unit_type_desc(CodedBitstreamContext *ctx,
843                              CodedBitstreamUnit *unit)
844 {
845     const CodedBitstreamUnitTypeDescriptor *desc;
846     int i, j;
847 
848     if (!ctx->codec->unit_types)
849         return NULL;
850 
851     for (i = 0;; i++) {
852         desc = &ctx->codec->unit_types[i];
853         if (desc->nb_unit_types == 0)
854             break;
855         if (desc->nb_unit_types == CBS_UNIT_TYPE_RANGE) {
856             if (unit->type >= desc->unit_type_range_start &&
857                 unit->type <= desc->unit_type_range_end)
858                 return desc;
859         } else {
860             for (j = 0; j < desc->nb_unit_types; j++) {
861                 if (desc->unit_types[j] == unit->type)
862                     return desc;
863             }
864         }
865     }
866     return NULL;
867 }
868 
ff_cbs_alloc_unit_content2(CodedBitstreamContext * ctx,CodedBitstreamUnit * unit)869 int ff_cbs_alloc_unit_content2(CodedBitstreamContext *ctx,
870                                CodedBitstreamUnit *unit)
871 {
872     const CodedBitstreamUnitTypeDescriptor *desc;
873 
874     av_assert0(!unit->content && !unit->content_ref);
875 
876     desc = cbs_find_unit_type_desc(ctx, unit);
877     if (!desc)
878         return AVERROR(ENOSYS);
879 
880     unit->content = av_mallocz(desc->content_size);
881     if (!unit->content)
882         return AVERROR(ENOMEM);
883 
884     unit->content_ref =
885         av_buffer_create(unit->content, desc->content_size,
886                          desc->content_free ? desc->content_free
887                                             : cbs_default_free_unit_content,
888                          (void*)desc, 0);
889     if (!unit->content_ref) {
890         av_freep(&unit->content);
891         return AVERROR(ENOMEM);
892     }
893 
894     return 0;
895 }
896 
cbs_clone_unit_content(AVBufferRef ** clone_ref,CodedBitstreamUnit * unit,const CodedBitstreamUnitTypeDescriptor * desc)897 static int cbs_clone_unit_content(AVBufferRef **clone_ref,
898                                   CodedBitstreamUnit *unit,
899                                   const CodedBitstreamUnitTypeDescriptor *desc)
900 {
901     uint8_t *src, *copy;
902     uint8_t **src_ptr, **copy_ptr;
903     AVBufferRef **src_buf, **copy_buf;
904     int err, i;
905 
906     av_assert0(unit->content);
907     src = unit->content;
908 
909     copy = av_memdup(src, desc->content_size);
910     if (!copy)
911         return AVERROR(ENOMEM);
912 
913     for (i = 0; i < desc->nb_ref_offsets; i++) {
914         src_ptr  = (uint8_t**)(src + desc->ref_offsets[i]);
915         src_buf  = (AVBufferRef**)(src_ptr + 1);
916         copy_ptr = (uint8_t**)(copy + desc->ref_offsets[i]);
917         copy_buf = (AVBufferRef**)(copy_ptr + 1);
918 
919         if (!*src_ptr) {
920             av_assert0(!*src_buf);
921             continue;
922         }
923         if (!*src_buf) {
924             // We can't handle a non-refcounted pointer here - we don't
925             // have enough information to handle whatever structure lies
926             // at the other end of it.
927             err = AVERROR(EINVAL);
928             goto fail;
929         }
930 
931         // src_ptr is required to point somewhere inside src_buf.  If it
932         // doesn't, there is a bug somewhere.
933         av_assert0(*src_ptr >= (*src_buf)->data &&
934                    *src_ptr <  (*src_buf)->data + (*src_buf)->size);
935 
936         *copy_buf = av_buffer_ref(*src_buf);
937         if (!*copy_buf) {
938             err = AVERROR(ENOMEM);
939             goto fail;
940         }
941         *copy_ptr = (*copy_buf)->data + (*src_ptr - (*src_buf)->data);
942     }
943 
944     *clone_ref = av_buffer_create(copy, desc->content_size,
945                                   desc->content_free ? desc->content_free :
946                                   cbs_default_free_unit_content,
947                                   (void*)desc, 0);
948     if (!*clone_ref) {
949         err = AVERROR(ENOMEM);
950         goto fail;
951     }
952 
953     return 0;
954 
955 fail:
956     for (--i; i >= 0; i--)
957         av_buffer_unref((AVBufferRef**)(copy + desc->ref_offsets[i]));
958     av_freep(&copy);
959     *clone_ref = NULL;
960     return err;
961 }
962 
ff_cbs_make_unit_refcounted(CodedBitstreamContext * ctx,CodedBitstreamUnit * unit)963 int ff_cbs_make_unit_refcounted(CodedBitstreamContext *ctx,
964                                 CodedBitstreamUnit *unit)
965 {
966     const CodedBitstreamUnitTypeDescriptor *desc;
967     AVBufferRef *ref;
968     int err;
969 
970     av_assert0(unit->content);
971     if (unit->content_ref) {
972         // Already refcounted, nothing to do.
973         return 0;
974     }
975 
976     desc = cbs_find_unit_type_desc(ctx, unit);
977     if (!desc)
978         return AVERROR(ENOSYS);
979 
980     switch (desc->content_type) {
981     case CBS_CONTENT_TYPE_POD:
982         ref = av_buffer_alloc(desc->content_size);
983         if (!ref)
984             return AVERROR(ENOMEM);
985         memcpy(ref->data, unit->content, desc->content_size);
986         err = 0;
987         break;
988 
989     case CBS_CONTENT_TYPE_INTERNAL_REFS:
990         err = cbs_clone_unit_content(&ref, unit, desc);
991         break;
992 
993     case CBS_CONTENT_TYPE_COMPLEX:
994         if (!desc->content_clone)
995             return AVERROR_PATCHWELCOME;
996         err = desc->content_clone(&ref, unit);
997         break;
998 
999     default:
1000         av_assert0(0 && "Invalid content type.");
1001     }
1002 
1003     if (err < 0)
1004         return err;
1005 
1006     unit->content_ref = ref;
1007     unit->content     = ref->data;
1008     return 0;
1009 }
1010 
ff_cbs_make_unit_writable(CodedBitstreamContext * ctx,CodedBitstreamUnit * unit)1011 int ff_cbs_make_unit_writable(CodedBitstreamContext *ctx,
1012                               CodedBitstreamUnit *unit)
1013 {
1014     const CodedBitstreamUnitTypeDescriptor *desc;
1015     AVBufferRef *ref;
1016     int err;
1017 
1018     // This can only be applied to refcounted units.
1019     err = ff_cbs_make_unit_refcounted(ctx, unit);
1020     if (err < 0)
1021         return err;
1022     av_assert0(unit->content && unit->content_ref);
1023 
1024     if (av_buffer_is_writable(unit->content_ref))
1025         return 0;
1026 
1027     desc = cbs_find_unit_type_desc(ctx, unit);
1028     if (!desc)
1029         return AVERROR(ENOSYS);
1030 
1031     switch (desc->content_type) {
1032     case CBS_CONTENT_TYPE_POD:
1033         err = av_buffer_make_writable(&unit->content_ref);
1034         break;
1035 
1036     case CBS_CONTENT_TYPE_INTERNAL_REFS:
1037         err = cbs_clone_unit_content(&ref, unit, desc);
1038         break;
1039 
1040     case CBS_CONTENT_TYPE_COMPLEX:
1041         if (!desc->content_clone)
1042             return AVERROR_PATCHWELCOME;
1043         err = desc->content_clone(&ref, unit);
1044         break;
1045 
1046     default:
1047         av_assert0(0 && "Invalid content type.");
1048     }
1049     if (err < 0)
1050         return err;
1051 
1052     if (desc->content_type != CBS_CONTENT_TYPE_POD) {
1053         av_buffer_unref(&unit->content_ref);
1054         unit->content_ref = ref;
1055     }
1056     unit->content = unit->content_ref->data;
1057     return 0;
1058 }
1059