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             for (i = start; i + 1 < frag->data_size; i++) {
 153                 if (frag->data[i] != 0xff)
 154                     continue;
 155                 end = i;
 156                 for (++i; i + 1 < frag->data_size &&
 157                           frag->data[i] == 0xff; i++);
 158                 if (i + 1 < frag->data_size) {
 159                     if (frag->data[i] == 0x00)
 160                         continue;
 161                     next_marker = frag->data[i];
 162                     next_start  = i + 1;
 163                 }
 164                 break;
 165             }
 166         } else {
 167             i = start;
 168             if (i + 2 > frag->data_size) {
 169                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
 170                        "truncated at %02x marker.\n", marker);
 171                 return AVERROR_INVALIDDATA;
 172             }
 173             length = AV_RB16(frag->data + i);
 174             if (i + length > frag->data_size) {
 175                 av_log(ctx->log_ctx, AV_LOG_ERROR, "Invalid JPEG image: "
 176                        "truncated at %02x marker segment.\n", marker);
 177                 return AVERROR_INVALIDDATA;
 178             }
 179             end = start + length;
 180
 181             i = end;
 182             if (frag->data[i] != 0xff) {
 183                 next_marker = -1;
 184             } else {
 185                 for (++i; i + 1 < frag->data_size &&
 186                           frag->data[i] == 0xff; i++);
 187                 if (i + 1 >= frag->data_size) {
 188                     next_marker = -1;
 189                 } else {
 190                     next_marker = frag->data[i];
 191                     next_start  = i + 1;
 192                 }
 193             }
 194         }
 195
 196         if (marker == JPEG_MARKER_SOS) {
 197             length = AV_RB16(frag->data + start);
 198
 199             if (length > end - start)
 200                 return AVERROR_INVALIDDATA;
 201
 202             data_ref = NULL;
 203             data     = av_malloc(end - start +
 204                                  AV_INPUT_BUFFER_PADDING_SIZE);
 205             if (!data)
 206                 return AVERROR(ENOMEM);
 207
 208             memcpy(data, frag->data + start, length);
 209             for (i = start + length, j = length; i < end; i++, j++) {
 210                 if (frag->data[i] == 0xff) {
 211                     while (frag->data[i] == 0xff)
 212                         ++i;
 213                     data[j] = 0xff;
 214                 } else {
 215                     data[j] = frag->data[i];
 216                 }
 217             }
 218             data_size = j;
 219
 220             memset(data + data_size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 221
 222         } else {
 223             data      = frag->data + start;
 224             data_size = end - start;
 225             data_ref  = frag->data_ref;
 226         }
 227
 228         err = ff_cbs_insert_unit_data(frag, unit, marker,
 229                                       data, data_size, data_ref);
 230         if (err < 0)
 231             return err;
 232
 233         if (next_marker == -1)
 234             break;
 235         marker = next_marker;
 236         start  = next_start;
 237     }
 238
 239     return 0;
 240 }
 241
 242 static int cbs_jpeg_read_unit(CodedBitstreamContext *ctx,
 243                               CodedBitstreamUnit *unit)
 244 {
 245     GetBitContext gbc;
 246     int err;
 247
 248     err = init_get_bits(&gbc, unit->data, 8 * unit->data_size);
 249     if (err < 0)
 250         return err;
 251
 252     if (unit->type >= JPEG_MARKER_SOF0 &&
 253         unit->type <= JPEG_MARKER_SOF3) {
 254         err = ff_cbs_alloc_unit_content(unit,
 255                                         sizeof(JPEGRawFrameHeader),
 256                                         NULL);
 257         if (err < 0)
 258             return err;
 259
 260         err = cbs_jpeg_read_frame_header(ctx, &gbc, unit->content);
 261         if (err < 0)
 262             return err;
 263
 264     } else if (unit->type >= JPEG_MARKER_APPN &&
 265                unit->type <= JPEG_MARKER_APPN + 15) {
 266         err = ff_cbs_alloc_unit_content(unit,
 267                                         sizeof(JPEGRawApplicationData),
 268                                         &cbs_jpeg_free_application_data);
 269         if (err < 0)
 270             return err;
 271
 272         err = cbs_jpeg_read_application_data(ctx, &gbc, unit->content);
 273         if (err < 0)
 274             return err;
 275
 276     } else if (unit->type == JPEG_MARKER_SOS) {
 277         JPEGRawScan *scan;
 278         int pos;
 279
 280         err = ff_cbs_alloc_unit_content(unit,
 281                                         sizeof(JPEGRawScan),
 282                                         &cbs_jpeg_free_scan);
 283         if (err < 0)
 284             return err;
 285         scan = unit->content;
 286
 287         err = cbs_jpeg_read_scan_header(ctx, &gbc, &scan->header);
 288         if (err < 0)
 289             return err;
 290
 291         pos = get_bits_count(&gbc);
 292         av_assert0(pos % 8 == 0);
 293         if (pos > 0) {
 294             scan->data_size = unit->data_size - pos / 8;
 295             scan->data_ref  = av_buffer_ref(unit->data_ref);
 296             if (!scan->data_ref)
 297                 return AVERROR(ENOMEM);
 298             scan->data = unit->data + pos / 8;
 299         }
 300
 301     } else {
 302         switch (unit->type) {
 303 #define SEGMENT(marker, type, func, free) \
 304         case JPEG_MARKER_ ## marker: \
 305             { \
 306                 err = ff_cbs_alloc_unit_content(unit, \
 307                                                 sizeof(type), free); \
 308                 if (err < 0) \
 309                     return err; \
 310                 err = cbs_jpeg_read_ ## func(ctx, &gbc, unit->content); \
 311                 if (err < 0) \
 312                     return err; \
 313             } \
 314             break
 315             SEGMENT(DQT, JPEGRawQuantisationTableSpecification, dqt, NULL);
 316             SEGMENT(DHT, JPEGRawHuffmanTableSpecification,      dht, NULL);
 317             SEGMENT(COM, JPEGRawComment,  comment, &cbs_jpeg_free_comment);
 318 #undef SEGMENT
 319         default:
 320             return AVERROR(ENOSYS);
 321         }
 322     }
 323
 324     return 0;
 325 }
 326
 327 static int cbs_jpeg_write_scan(CodedBitstreamContext *ctx,
 328                                CodedBitstreamUnit *unit,
 329                                PutBitContext *pbc)
 330 {
 331     JPEGRawScan *scan = unit->content;
 332     int err;
 333
 334     err = cbs_jpeg_write_scan_header(ctx, pbc, &scan->header);
 335     if (err < 0)
 336         return err;
 337
 338     if (scan->data) {
 339         if (scan->data_size * 8 > put_bits_left(pbc))
 340             return AVERROR(ENOSPC);
 341
 342         av_assert0(put_bits_count(pbc) % 8 == 0);
 343
 344         flush_put_bits(pbc);
 345
 346         memcpy(put_bits_ptr(pbc), scan->data, scan->data_size);
 347         skip_put_bytes(pbc, scan->data_size);
 348     }
 349
 350     return 0;
 351 }
 352
 353 static int cbs_jpeg_write_segment(CodedBitstreamContext *ctx,
 354                                   CodedBitstreamUnit *unit,
 355                                   PutBitContext *pbc)
 356 {
 357     int err;
 358
 359     if (unit->type >= JPEG_MARKER_SOF0 &&
 360         unit->type <= JPEG_MARKER_SOF3) {
 361         err = cbs_jpeg_write_frame_header(ctx, pbc, unit->content);
 362     } else if (unit->type >= JPEG_MARKER_APPN &&
 363                unit->type <= JPEG_MARKER_APPN + 15) {
 364         err = cbs_jpeg_write_application_data(ctx, pbc, unit->content);
 365     } else {
 366         switch (unit->type) {
 367 #define SEGMENT(marker, func) \
 368             case JPEG_MARKER_ ## marker: \
 369                 err = cbs_jpeg_write_ ## func(ctx, pbc, unit->content); \
 370                 break;
 371             SEGMENT(DQT, dqt);
 372             SEGMENT(DHT, dht);
 373             SEGMENT(COM, comment);
 374         default:
 375             return AVERROR_PATCHWELCOME;
 376         }
 377     }
 378
 379     return err;
 380 }
 381
 382 static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
 383                                CodedBitstreamUnit *unit,
 384                                PutBitContext *pbc)
 385 {
 386     if (unit->type == JPEG_MARKER_SOS)
 387         return cbs_jpeg_write_scan   (ctx, unit, pbc);
 388     else
 389         return cbs_jpeg_write_segment(ctx, unit, pbc);
 390 }
 391
 392 static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
 393                                        CodedBitstreamFragment *frag)
 394 {
 395     const CodedBitstreamUnit *unit;
 396     uint8_t *data;
 397     size_t size, dp, sp;
 398     int i;
 399
 400     size = 4; // SOI + EOI.
 401     for (i = 0; i < frag->nb_units; i++) {
 402         unit = &frag->units[i];
 403         size += 2 + unit->data_size;
 404         if (unit->type == JPEG_MARKER_SOS) {
 405             for (sp = 0; sp < unit->data_size; sp++) {
 406                 if (unit->data[sp] == 0xff)
 407                     ++size;
 408             }
 409         }
 410     }
 411
 412     frag->data_ref = av_buffer_alloc(size + AV_INPUT_BUFFER_PADDING_SIZE);
 413     if (!frag->data_ref)
 414         return AVERROR(ENOMEM);
 415     data = frag->data_ref->data;
 416
 417     dp = 0;
 418
 419     data[dp++] = 0xff;
 420     data[dp++] = JPEG_MARKER_SOI;
 421
 422     for (i = 0; i < frag->nb_units; i++) {
 423         unit = &frag->units[i];
 424
 425         data[dp++] = 0xff;
 426         data[dp++] = unit->type;
 427
 428         if (unit->type != JPEG_MARKER_SOS) {
 429             memcpy(data + dp, unit->data, unit->data_size);
 430             dp += unit->data_size;
 431         } else {
 432             sp = AV_RB16(unit->data);
 433             av_assert0(sp <= unit->data_size);
 434             memcpy(data + dp, unit->data, sp);
 435             dp += sp;
 436
 437             for (; sp < unit->data_size; sp++) {
 438                 if (unit->data[sp] == 0xff) {
 439                     data[dp++] = 0xff;
 440                     data[dp++] = 0x00;
 441                 } else {
 442                     data[dp++] = unit->data[sp];
 443                 }
 444             }
 445         }
 446     }
 447
 448     data[dp++] = 0xff;
 449     data[dp++] = JPEG_MARKER_EOI;
 450
 451     av_assert0(dp == size);
 452
 453     memset(data + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 454     frag->data      = data;
 455     frag->data_size = size;
 456
 457     return 0;
 458 }
 459
 460 const CodedBitstreamType ff_cbs_type_jpeg = {
 461     .codec_id          = AV_CODEC_ID_MJPEG,
 462
 463     .split_fragment    = &cbs_jpeg_split_fragment,
 464     .read_unit         = &cbs_jpeg_read_unit,
 465     .write_unit        = &cbs_jpeg_write_unit,
 466     .assemble_fragment = &cbs_jpeg_assemble_fragment,
 467 };