git.sesse.net Git - ffmpeg/blob - libavcodec/cbs_jpeg.c

   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 "cbs.h"
  20 #include "cbs_internal.h"
  21 #include "cbs_jpeg.h"
  22
  23
  24 #define HEADER(name) do { \
  25         ff_cbs_trace_header(ctx, name); \
  26     } while (0)
  27
  28 #define CHECK(call) do { \
  29         err = (call); \
  30         if (err < 0) \
  31             return err; \
  32     } while (0)
  33
  34 #define SUBSCRIPTS(subs, ...) (subs > 0 ? ((int[subs + 1]){ subs, __VA_ARGS__ }) : NULL)
  35
  36 #define u(width, name, range_min, range_max) \
  37     xu(width, name, range_min, range_max, 0)
  38 #define us(width, name, sub, range_min, range_max) \
  39     xu(width, name, range_min, range_max, 1, sub)
  40
  41
  42 #define READ
  43 #define READWRITE read
  44 #define RWContext GetBitContext
  45 #define FUNC(name) cbs_jpeg_read_ ## name
  46
  47 #define xu(width, name, range_min, range_max, subs, ...) do { \
  48         uint32_t value = range_min; \
  49         CHECK(ff_cbs_read_unsigned(ctx, rw, width, #name, \
  50                                    SUBSCRIPTS(subs, __VA_ARGS__), \
  51                                    &value, range_min, range_max)); \
  52         current->name = value; \
  53     } while (0)
  54
  55 #include "cbs_jpeg_syntax_template.c"
  56
  57 #undef READ
  58 #undef READWRITE
  59 #undef RWContext
  60 #undef FUNC
  61 #undef xu
  62
  63 #define WRITE
  64 #define READWRITE write
  65 #define RWContext PutBitContext
  66 #define FUNC(name) cbs_jpeg_write_ ## name
  67
  68 #define xu(width, name, range_min, range_max, subs, ...) do { \
  69         uint32_t value = current->name; \
  70         CHECK(ff_cbs_write_unsigned(ctx, rw, width, #name, \
  71                                     SUBSCRIPTS(subs, __VA_ARGS__), \
  72                                     value, range_min, range_max)); \
  73     } while (0)
  74
  75
  76 #include "cbs_jpeg_syntax_template.c"
  77
  78 #undef READ
  79 #undef READWRITE
  80 #undef RWContext
  81 #undef FUNC
  82 #undef xu
  83
  84
  85 static void cbs_jpeg_free_application_data(void *unit, uint8_t *content)
  86 {
  87     JPEGRawApplicationData *ad = (JPEGRawApplicationData*)content;
  88     av_buffer_unref(&ad->Ap_ref);
  89     av_freep(&content);
  90 }
  91
  92 static void cbs_jpeg_free_comment(void *unit, uint8_t *content)
  93 {
  94     JPEGRawComment *comment = (JPEGRawComment*)content;
  95     av_buffer_unref(&comment->Cm_ref);
  96     av_freep(&content);
  97 }
  98
  99 static void cbs_jpeg_free_scan(void *unit, uint8_t *content)
 100 {
 101     JPEGRawScan *scan = (JPEGRawScan*)content;
 102     av_buffer_unref(&scan->data_ref);
 103     av_freep(&content);
 104 }
 105
 106 static int cbs_jpeg_split_fragment(CodedBitstreamContext *ctx,
 107                                    CodedBitstreamFragment *frag,
 108                                    int header)
 109 {
 110     AVBufferRef *data_ref;
 111     uint8_t *data;
 112     size_t data_size;
 113     int unit, start, end, marker, next_start, next_marker;
 114     int err, i, j, length;
 115
 116     if (frag->data_size < 4) {
 117         // Definitely too short to be meaningful.
 118         return AVERROR_INVALIDDATA;
 119     }
 120
 121     for (i = 0; i + 1 < frag->data_size && frag->data[i] != 0xff; i++);
 122     if (i > 0) {
 123         av_log(ctx->log_ctx, AV_LOG_WARNING, "Discarding %d bytes at "
 124                "beginning of image.\n", i);
 125     }
 126     for (++i; i + 1 < frag->data_size && frag->data[i] == 0xff; i++);
 127     if (i + 1 >= frag->data_size && frag->data[i]) {
 128         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
 129                "no SOI marker found.\n");
 130         return AVERROR_INVALIDDATA;
 131     }
 132     marker = frag->data[i];
 133     if (marker != JPEG_MARKER_SOI) {
 134         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: first "
 135                "marker is %02x, should be SOI.\n", marker);
 136         return AVERROR_INVALIDDATA;
 137     }
 138     for (++i; i + 1 < frag->data_size && frag->data[i] == 0xff; i++);
 139     if (i + 1 >= frag->data_size) {
 140         av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
 141                "no image content found.\n");
 142         return AVERROR_INVALIDDATA;
 143     }
 144     marker = frag->data[i];
 145     start  = i + 1;
 146
 147     for (unit = 0;; unit++) {
 148         if (marker == JPEG_MARKER_EOI) {
 149             break;
 150         } else if (marker == JPEG_MARKER_SOS) {
 151             for (i = start; i + 1 < frag->data_size; i++) {
 152                 if (frag->data[i] != 0xff)
 153                     continue;
 154                 end = i;
 155                 for (++i; i + 1 < frag->data_size &&
 156                           frag->data[i] == 0xff; i++);
 157                 if (i + 1 >= frag->data_size) {
 158                     next_marker = -1;
 159                 } else {
 160                     if (frag->data[i] == 0x00)
 161                         continue;
 162                     next_marker = frag->data[i];
 163                     next_start  = i + 1;
 164                 }
 165                 break;
 166             }
 167         } else {
 168             i = start;
 169             if (i + 2 > frag->data_size) {
 170                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
 171                        "truncated at %02x marker.\n", marker);
 172                 return AVERROR_INVALIDDATA;
 173             }
 174             length = AV_RB16(frag->data + i);
 175             if (i + length > frag->data_size) {
 176                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
 177                        "truncated at %02x marker segment.\n", marker);
 178                 return AVERROR_INVALIDDATA;
 179             }
 180             end = start + length;
 181
 182             i = end;
 183             if (frag->data[i] != 0xff) {
 184                 next_marker = -1;
 185             } else {
 186                 for (++i; i + 1 < frag->data_size &&
 187                           frag->data[i] == 0xff; i++);
 188                 if (i + 1 >= frag->data_size) {
 189                     next_marker = -1;
 190                 } else {
 191                     next_marker = frag->data[i];
 192                     next_start  = i + 1;
 193                 }
 194             }
 195         }
 196
 197         if (marker == JPEG_MARKER_SOS) {
 198             length = AV_RB16(frag->data + start);
 199
 200             data_ref = NULL;
 201             data     = av_malloc(end - start +
 202                                  AV_INPUT_BUFFER_PADDING_SIZE);
 203             if (!data)
 204                 return AVERROR(ENOMEM);
 205
 206             memcpy(data, frag->data + start, length);
 207             for (i = start + length, j = length; i < end; i++, j++) {
 208                 if (frag->data[i] == 0xff) {
 209                     while (frag->data[i] == 0xff)
 210                         ++i;
 211                     data[j] = 0xff;
 212                 } else {
 213                     data[j] = frag->data[i];
 214                 }
 215             }
 216             data_size = j;
 217
 218             memset(data + data_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 219
 220         } else {
 221             data      = frag->data + start;
 222             data_size = end - start;
 223             data_ref  = frag->data_ref;
 224         }
 225
 226         err = ff_cbs_insert_unit_data(ctx, frag, unit, marker,
 227                                       data, data_size, data_ref);
 228         if (err < 0) {
 229             if (!data_ref)
 230                 av_freep(&data);
 231             return err;
 232         }
 233
 234         if (next_marker == -1)
 235             break;
 236         marker = next_marker;
 237         start  = next_start;
 238     }
 239
 240     return 0;
 241 }
 242
 243 static int cbs_jpeg_read_unit(CodedBitstreamContext *ctx,
 244                               CodedBitstreamUnit *unit)
 245 {
 246     GetBitContext gbc;
 247     int err;
 248
 249     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 250     if (err < 0)
 251         return err;
 252
 253     if (unit->type >= JPEG_MARKER_SOF0 &&
 254         unit->type <= JPEG_MARKER_SOF3) {
 255         err = ff_cbs_alloc_unit_content(ctx, unit,
 256                                         sizeof(JPEGRawFrameHeader),
 257                                         NULL);
 258         if (err < 0)
 259             return err;
 260
 261         err = cbs_jpeg_read_frame_header(ctx, &gbc, unit->content);
 262         if (err < 0)
 263             return err;
 264
 265     } else if (unit->type >= JPEG_MARKER_APPN &&
 266                unit->type <= JPEG_MARKER_APPN + 15) {
 267         err = ff_cbs_alloc_unit_content(ctx, unit,
 268                                         sizeof(JPEGRawApplicationData),
 269                                         &cbs_jpeg_free_application_data);
 270         if (err < 0)
 271             return err;
 272
 273         err = cbs_jpeg_read_application_data(ctx, &gbc, unit->content);
 274         if (err < 0)
 275             return err;
 276
 277     } else if (unit->type == JPEG_MARKER_SOS) {
 278         JPEGRawScan *scan;
 279         int pos;
 280
 281         err = ff_cbs_alloc_unit_content(ctx, unit,
 282                                         sizeof(JPEGRawScan),
 283                                         &cbs_jpeg_free_scan);
 284         if (err < 0)
 285             return err;
 286         scan = unit->content;
 287
 288         err = cbs_jpeg_read_scan_header(ctx, &gbc, &scan->header);
 289         if (err < 0)
 290             return err;
 291
 292         pos = get_bits_count(&gbc);
 293         av_assert0(pos % 8 == 0);
 294         if (pos > 0) {
 295             scan->data_size = unit->data_size - pos / 8;
 296             scan->data_ref  = av_buffer_ref(unit->data_ref);
 297             if (!scan->data_ref)
 298                 return AVERROR(ENOMEM);
 299             scan->data = unit->data + pos / 8;
 300         }
 301
 302     } else {
 303         switch (unit->type) {
 304 #define SEGMENT(marker, type, func, free) \
 305         case JPEG_MARKER_ ## marker: \
 306             { \
 307                 err = ff_cbs_alloc_unit_content(ctx, unit, \
 308                                                 sizeof(type), free); \
 309                 if (err < 0) \
 310                     return err; \
 311                 err = cbs_jpeg_read_ ## func(ctx, &gbc, unit->content); \
 312                 if (err < 0) \
 313                     return err; \
 314             } \
 315             break
 316             SEGMENT(DQT, JPEGRawQuantisationTableSpecification, dqt, NULL);
 317             SEGMENT(DHT, JPEGRawHuffmanTableSpecification,      dht, NULL);
 318             SEGMENT(COM, JPEGRawComment,  comment, &cbs_jpeg_free_comment);
 319 #undef SEGMENT
 320         default:
 321             return AVERROR(ENOSYS);
 322         }
 323     }
 324
 325     return 0;
 326 }
 327
 328 static int cbs_jpeg_write_scan(CodedBitstreamContext *ctx,
 329                                CodedBitstreamUnit *unit,
 330                                PutBitContext *pbc)
 331 {
 332     JPEGRawScan *scan = unit->content;
 333     int i, err;
 334
 335     err = cbs_jpeg_write_scan_header(ctx, pbc, &scan->header);
 336     if (err < 0)
 337         return err;
 338
 339     if (scan->data) {
 340         if (scan->data_size * 8 > put_bits_left(pbc))
 341             return AVERROR(ENOSPC);
 342
 343         for (i = 0; i < scan->data_size; i++)
 344             put_bits(pbc, 8, scan->data[i]);
 345     }
 346
 347     return 0;
 348 }
 349
 350 static int cbs_jpeg_write_segment(CodedBitstreamContext *ctx,
 351                                   CodedBitstreamUnit *unit,
 352                                   PutBitContext *pbc)
 353 {
 354     int err;
 355
 356     if (unit->type >= JPEG_MARKER_SOF0 &&
 357         unit->type <= JPEG_MARKER_SOF3) {
 358         err = cbs_jpeg_write_frame_header(ctx, pbc, unit->content);
 359     } else if (unit->type >= JPEG_MARKER_APPN &&
 360                unit->type <= JPEG_MARKER_APPN + 15) {
 361         err = cbs_jpeg_write_application_data(ctx, pbc, unit->content);
 362     } else {
 363         switch (unit->type) {
 364 #define SEGMENT(marker, func) \
 365             case JPEG_MARKER_ ## marker: \
 366                 err = cbs_jpeg_write_ ## func(ctx, pbc, unit->content); \
 367                 break;
 368             SEGMENT(DQT, dqt);
 369             SEGMENT(DHT, dht);
 370             SEGMENT(COM, comment);
 371         default:
 372             return AVERROR_PATCHWELCOME;
 373         }
 374     }
 375
 376     return err;
 377 }
 378
 379 static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
 380                                 CodedBitstreamUnit *unit)
 381 {
 382     CodedBitstreamJPEGContext *priv = ctx->priv_data;
 383     PutBitContext pbc;
 384     int err;
 385
 386     if (!priv->write_buffer) {
 387         // Initial write buffer size is 1MB.
 388         priv->write_buffer_size = 1024 * 1024;
 389
 390     reallocate_and_try_again:
 391         err = av_reallocp(&priv->write_buffer, priv->write_buffer_size);
 392         if (err < 0) {
 393             av_log(ctx->log_ctx, AV_LOG_ERROR, "Unable to allocate a "
 394                    "sufficiently large write buffer (last attempt "
 395                    "%"SIZE_SPECIFIER" bytes).\n", priv->write_buffer_size);
 396             return err;
 397         }
 398     }
 399
 400     init_put_bits(&pbc, priv->write_buffer, priv->write_buffer_size);
 401
 402     if (unit->type == JPEG_MARKER_SOS)
 403         err = cbs_jpeg_write_scan(ctx, unit, &pbc);
 404     else
 405         err = cbs_jpeg_write_segment(ctx, unit, &pbc);
 406
 407     if (err == AVERROR(ENOSPC)) {
 408         // Overflow.
 409         priv->write_buffer_size *= 2;
 410         goto reallocate_and_try_again;
 411     }
 412     if (err < 0) {
 413         // Write failed for some other reason.
 414         return err;
 415     }
 416
 417     if (put_bits_count(&pbc) % 8)
 418         unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
 419     else
 420         unit->data_bit_padding = 0;
 421
 422     unit->data_size = (put_bits_count(&pbc) + 7) / 8;
 423     flush_put_bits(&pbc);
 424
 425     err = ff_cbs_alloc_unit_data(ctx, unit, unit->data_size);
 426     if (err < 0)
 427         return err;
 428
 429     memcpy(unit->data, priv->write_buffer, unit->data_size);
 430
 431     return 0;
 432 }
 433
 434 static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
 435                                        CodedBitstreamFragment *frag)
 436 {
 437     const CodedBitstreamUnit *unit;
 438     uint8_t *data;
 439     size_t size, dp, sp;
 440     int i;
 441
 442     size = 4; // SOI + EOI.
 443     for (i = 0; i < frag->nb_units; i++) {
 444         unit = &frag->units[i];
 445         size += 2 + unit->data_size;
 446         if (unit->type == JPEG_MARKER_SOS) {
 447             for (sp = 0; sp < unit->data_size; sp++) {
 448                 if (unit->data[sp] == 0xff)
 449                     ++size;
 450             }
 451         }
 452     }
 453
 454     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
 455     if (!frag->data_ref)
 456         return AVERROR(ENOMEM);
 457     data = frag->data_ref->data;
 458
 459     dp = 0;
 460
 461     data[dp++] = 0xff;
 462     data[dp++] = JPEG_MARKER_SOI;
 463
 464     for (i = 0; i < frag->nb_units; i++) {
 465         unit = &frag->units[i];
 466
 467         data[dp++] = 0xff;
 468         data[dp++] = unit->type;
 469
 470         if (unit->type != JPEG_MARKER_SOS) {
 471             memcpy(data + dp, unit->data, unit->data_size);
 472             dp += unit->data_size;
 473         } else {
 474             sp = AV_RB16(unit->data);
 475             av_assert0(sp <= unit->data_size);
 476             memcpy(data + dp, unit->data, sp);
 477             dp += sp;
 478
 479             for (; sp < unit->data_size; sp++) {
 480                 if (unit->data[sp] == 0xff) {
 481                     data[dp++] = 0xff;
 482                     data[dp++] = 0x00;
 483                 } else {
 484                     data[dp++] = unit->data[sp];
 485                 }
 486             }
 487         }
 488     }
 489
 490     data[dp++] = 0xff;
 491     data[dp++] = JPEG_MARKER_EOI;
 492
 493     av_assert0(dp == size);
 494
 495     memset(data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 496     frag->data      = data;
 497     frag->data_size = size;
 498
 499     return 0;
 500 }
 501
 502 static void cbs_jpeg_close(CodedBitstreamContext *ctx)
 503 {
 504     CodedBitstreamJPEGContext *priv = ctx->priv_data;
 505
 506     av_freep(&priv->write_buffer);
 507 }
 508
 509 const CodedBitstreamType ff_cbs_type_jpeg = {
 510     .codec_id          = AV_CODEC_ID_MJPEG,
 511
 512     .priv_data_size    = sizeof(CodedBitstreamJPEGContext),
 513
 514     .split_fragment    = &cbs_jpeg_split_fragment,
 515     .read_unit         = &cbs_jpeg_read_unit,
 516     .write_unit        = &cbs_jpeg_write_unit,
 517     .assemble_fragment = &cbs_jpeg_assemble_fragment,
 518
 519     .close             = &cbs_jpeg_close,
 520 };