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; \
  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 WRITE
  79 #undef READWRITE
  80 #undef RWContext
  81 #undef FUNC
  82 #undef xu
  83
  84
  85 static void cbs_jpeg_free_application_data(void *opaque, 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 *opaque, 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 *opaque, 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             next_marker = -1;
 152             end = start;
 153             for (i = start; i + 1 < frag->data_size; i++) {
 154                 if (frag->data[i] != 0xff)
 155                     continue;
 156                 end = i;
 157                 for (++i; i + 1 < frag->data_size &&
 158                           frag->data[i] == 0xff; i++);
 159                 if (i + 1 < frag->data_size) {
 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             if (length > end - start)
 201                 return AVERROR_INVALIDDATA;
 202
 203             data_ref = NULL;
 204             data     = av_malloc(end - start +
 205                                  AV_INPUT_BUFFER_PADDING_SIZE);
 206             if (!data)
 207                 return AVERROR(ENOMEM);
 208
 209             memcpy(data, frag->data + start, length);
 210             for (i = start + length, j = length; i < end; i++, j++) {
 211                 if (frag->data[i] == 0xff) {
 212                     while (frag->data[i] == 0xff)
 213                         ++i;
 214                     data[j] = 0xff;
 215                 } else {
 216                     data[j] = frag->data[i];
 217                 }
 218             }
 219             data_size = j;
 220
 221             memset(data + data_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 222
 223         } else {
 224             data      = frag->data + start;
 225             data_size = end - start;
 226             data_ref  = frag->data_ref;
 227         }
 228
 229         err = ff_cbs_insert_unit_data(frag, unit, marker,
 230                                       data, data_size, data_ref);
 231         if (err < 0)
 232             return err;
 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(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(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(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(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 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         av_assert0(put_bits_count(pbc) % 8 == 0);
 344
 345         flush_put_bits(pbc);
 346
 347         memcpy(put_bits_ptr(pbc), scan->data, scan->data_size);
 348         skip_put_bytes(pbc, scan->data_size);
 349     }
 350
 351     return 0;
 352 }
 353
 354 static int cbs_jpeg_write_segment(CodedBitstreamContext *ctx,
 355                                   CodedBitstreamUnit *unit,
 356                                   PutBitContext *pbc)
 357 {
 358     int err;
 359
 360     if (unit->type >= JPEG_MARKER_SOF0 &&
 361         unit->type <= JPEG_MARKER_SOF3) {
 362         err = cbs_jpeg_write_frame_header(ctx, pbc, unit->content);
 363     } else if (unit->type >= JPEG_MARKER_APPN &&
 364                unit->type <= JPEG_MARKER_APPN + 15) {
 365         err = cbs_jpeg_write_application_data(ctx, pbc, unit->content);
 366     } else {
 367         switch (unit->type) {
 368 #define SEGMENT(marker, func) \
 369             case JPEG_MARKER_ ## marker: \
 370                 err = cbs_jpeg_write_ ## func(ctx, pbc, unit->content); \
 371                 break;
 372             SEGMENT(DQT, dqt);
 373             SEGMENT(DHT, dht);
 374             SEGMENT(COM, comment);
 375         default:
 376             return AVERROR_PATCHWELCOME;
 377         }
 378     }
 379
 380     return err;
 381 }
 382
 383 static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
 384                                CodedBitstreamUnit *unit,
 385                                PutBitContext *pbc)
 386 {
 387     if (unit->type == JPEG_MARKER_SOS)
 388         return cbs_jpeg_write_scan   (ctx, unit, pbc);
 389     else
 390         return cbs_jpeg_write_segment(ctx, unit, pbc);
 391 }
 392
 393 static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
 394                                        CodedBitstreamFragment *frag)
 395 {
 396     const CodedBitstreamUnit *unit;
 397     uint8_t *data;
 398     size_t size, dp, sp;
 399     int i;
 400
 401     size = 4; // SOI + EOI.
 402     for (i = 0; i < frag->nb_units; i++) {
 403         unit = &frag->units[i];
 404         size += 2 + unit->data_size;
 405         if (unit->type == JPEG_MARKER_SOS) {
 406             for (sp = 0; sp < unit->data_size; sp++) {
 407                 if (unit->data[sp] == 0xff)
 408                     ++size;
 409             }
 410         }
 411     }
 412
 413     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
 414     if (!frag->data_ref)
 415         return AVERROR(ENOMEM);
 416     data = frag->data_ref->data;
 417
 418     dp = 0;
 419
 420     data[dp++] = 0xff;
 421     data[dp++] = JPEG_MARKER_SOI;
 422
 423     for (i = 0; i < frag->nb_units; i++) {
 424         unit = &frag->units[i];
 425
 426         data[dp++] = 0xff;
 427         data[dp++] = unit->type;
 428
 429         if (unit->type != JPEG_MARKER_SOS) {
 430             memcpy(data + dp, unit->data, unit->data_size);
 431             dp += unit->data_size;
 432         } else {
 433             sp = AV_RB16(unit->data);
 434             av_assert0(sp <= unit->data_size);
 435             memcpy(data + dp, unit->data, sp);
 436             dp += sp;
 437
 438             for (; sp < unit->data_size; sp++) {
 439                 if (unit->data[sp] == 0xff) {
 440                     data[dp++] = 0xff;
 441                     data[dp++] = 0x00;
 442                 } else {
 443                     data[dp++] = unit->data[sp];
 444                 }
 445             }
 446         }
 447     }
 448
 449     data[dp++] = 0xff;
 450     data[dp++] = JPEG_MARKER_EOI;
 451
 452     av_assert0(dp == size);
 453
 454     memset(data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 455     frag->data      = data;
 456     frag->data_size = size;
 457
 458     return 0;
 459 }
 460
 461 const CodedBitstreamType ff_cbs_type_jpeg = {
 462     .codec_id          = AV_CODEC_ID_MJPEG,
 463
 464     .split_fragment    = &cbs_jpeg_split_fragment,
 465     .read_unit         = &cbs_jpeg_read_unit,
 466     .write_unit        = &cbs_jpeg_write_unit,
 467     .assemble_fragment = &cbs_jpeg_assemble_fragment,
 468 };