dvdsub_parse_extradata: fix memleak
[ffmpeg.git] / libavcodec / dvdsubdec.c
1 /*
2  * DVD subtitle decoding
3  * Copyright (c) 2005 Fabrice Bellard
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * FFmpeg is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with FFmpeg; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21 #include "avcodec.h"
22 #include "get_bits.h"
23 #include "dsputil.h"
24 #include "internal.h"
25
26 #include "libavutil/attributes.h"
27 #include "libavutil/colorspace.h"
28 #include "libavutil/opt.h"
29 #include "libavutil/imgutils.h"
30 #include "libavutil/avstring.h"
31
32 typedef struct DVDSubContext
33 {
34   AVClass *class;
35   uint32_t palette[16];
36   char    *palette_str;
37   int      has_palette;
38   uint8_t  colormap[4];
39   uint8_t  alpha[256];
40   uint8_t *buf;
41   int      buf_size;
42 #ifdef DEBUG
43   int sub_id;
44 #endif
45 } DVDSubContext;
46
47 static void yuv_a_to_rgba(const uint8_t *ycbcr, const uint8_t *alpha, uint32_t *rgba, int num_values)
48 {
49     const uint8_t *cm = ff_cropTbl + MAX_NEG_CROP;
50     uint8_t r, g, b;
51     int i, y, cb, cr;
52     int r_add, g_add, b_add;
53
54     for (i = num_values; i > 0; i--) {
55         y = *ycbcr++;
56         cr = *ycbcr++;
57         cb = *ycbcr++;
58         YUV_TO_RGB1_CCIR(cb, cr);
59         YUV_TO_RGB2_CCIR(r, g, b, y);
60         *rgba++ = (*alpha++ << 24) | (r << 16) | (g << 8) | b;
61     }
62 }
63
64 static int decode_run_2bit(GetBitContext *gb, int *color)
65 {
66     unsigned int v, t;
67
68     v = 0;
69     for (t = 1; v < t && t <= 0x40; t <<= 2)
70         v = (v << 4) | get_bits(gb, 4);
71     *color = v & 3;
72     if (v < 4) { /* Code for fill rest of line */
73         return INT_MAX;
74     }
75     return v >> 2;
76 }
77
78 static int decode_run_8bit(GetBitContext *gb, int *color)
79 {
80     int len;
81     int has_run = get_bits1(gb);
82     if (get_bits1(gb))
83         *color = get_bits(gb, 8);
84     else
85         *color = get_bits(gb, 2);
86     if (has_run) {
87         if (get_bits1(gb)) {
88             len = get_bits(gb, 7);
89             if (len == 0)
90                 len = INT_MAX;
91             else
92                 len += 9;
93         } else
94             len = get_bits(gb, 3) + 2;
95     } else
96         len = 1;
97     return len;
98 }
99
100 static int decode_rle(uint8_t *bitmap, int linesize, int w, int h,
101                       const uint8_t *buf, int start, int buf_size, int is_8bit)
102 {
103     GetBitContext gb;
104     int bit_len;
105     int x, y, len, color;
106     uint8_t *d;
107
108     bit_len = (buf_size - start) * 8;
109     init_get_bits(&gb, buf + start, bit_len);
110
111     x = 0;
112     y = 0;
113     d = bitmap;
114     for(;;) {
115         if (get_bits_count(&gb) > bit_len)
116             return -1;
117         if (is_8bit)
118             len = decode_run_8bit(&gb, &color);
119         else
120             len = decode_run_2bit(&gb, &color);
121         len = FFMIN(len, w - x);
122         memset(d + x, color, len);
123         x += len;
124         if (x >= w) {
125             y++;
126             if (y >= h)
127                 break;
128             d += linesize;
129             x = 0;
130             /* byte align */
131             align_get_bits(&gb);
132         }
133     }
134     return 0;
135 }
136
137 static void guess_palette(DVDSubContext* ctx,
138                           uint32_t *rgba_palette,
139                           uint32_t subtitle_color)
140 {
141     static const uint8_t level_map[4][4] = {
142         // this configuration (full range, lowest to highest) in tests
143         // seemed most common, so assume this
144         {0xff},
145         {0x00, 0xff},
146         {0x00, 0x80, 0xff},
147         {0x00, 0x55, 0xaa, 0xff},
148     };
149     uint8_t color_used[16] = { 0 };
150     int nb_opaque_colors, i, level, j, r, g, b;
151     uint8_t *colormap = ctx->colormap, *alpha = ctx->alpha;
152
153     if(ctx->has_palette) {
154         for(i = 0; i < 4; i++)
155             rgba_palette[i] = (ctx->palette[colormap[i]] & 0x00ffffff)
156                               | ((alpha[i] * 17U) << 24);
157         return;
158     }
159
160     for(i = 0; i < 4; i++)
161         rgba_palette[i] = 0;
162
163     nb_opaque_colors = 0;
164     for(i = 0; i < 4; i++) {
165         if (alpha[i] != 0 && !color_used[colormap[i]]) {
166             color_used[colormap[i]] = 1;
167             nb_opaque_colors++;
168         }
169     }
170
171     if (nb_opaque_colors == 0)
172         return;
173
174     j = 0;
175     memset(color_used, 0, 16);
176     for(i = 0; i < 4; i++) {
177         if (alpha[i] != 0) {
178             if (!color_used[colormap[i]])  {
179                 level = level_map[nb_opaque_colors][j];
180                 r = (((subtitle_color >> 16) & 0xff) * level) >> 8;
181                 g = (((subtitle_color >> 8) & 0xff) * level) >> 8;
182                 b = (((subtitle_color >> 0) & 0xff) * level) >> 8;
183                 rgba_palette[i] = b | (g << 8) | (r << 16) | ((alpha[i] * 17) << 24);
184                 color_used[colormap[i]] = (i + 1);
185                 j++;
186             } else {
187                 rgba_palette[i] = (rgba_palette[color_used[colormap[i]] - 1] & 0x00ffffff) |
188                                     ((alpha[i] * 17) << 24);
189             }
190         }
191     }
192 }
193
194 static void reset_rects(AVSubtitle *sub_header)
195 {
196     int i;
197
198     if (sub_header->rects != NULL) {
199         for (i = 0; i < sub_header->num_rects; i++) {
200             av_freep(&sub_header->rects[i]->pict.data[0]);
201             av_freep(&sub_header->rects[i]->pict.data[1]);
202             av_freep(&sub_header->rects[i]);
203         }
204         av_freep(&sub_header->rects);
205         sub_header->num_rects = 0;
206     }
207 }
208
209 #define READ_OFFSET(a) (big_offsets ? AV_RB32(a) : AV_RB16(a))
210
211 static int decode_dvd_subtitles(DVDSubContext *ctx, AVSubtitle *sub_header,
212                                 const uint8_t *buf, int buf_size)
213 {
214     int cmd_pos, pos, cmd, x1, y1, x2, y2, offset1, offset2, next_cmd_pos;
215     int big_offsets, offset_size, is_8bit = 0;
216     const uint8_t *yuv_palette = 0;
217     uint8_t *colormap = ctx->colormap, *alpha = ctx->alpha;
218     int date;
219     int i;
220     int is_menu = 0;
221
222     if (buf_size < 10)
223         return -1;
224
225     if (AV_RB16(buf) == 0) {   /* HD subpicture with 4-byte offsets */
226         big_offsets = 1;
227         offset_size = 4;
228         cmd_pos = 6;
229     } else {
230         big_offsets = 0;
231         offset_size = 2;
232         cmd_pos = 2;
233     }
234
235     cmd_pos = READ_OFFSET(buf + cmd_pos);
236
237     if (cmd_pos < 0 || cmd_pos > buf_size - 2 - offset_size)
238         return AVERROR(EAGAIN);
239
240     while (cmd_pos > 0 && cmd_pos < buf_size - 2 - offset_size) {
241         date = AV_RB16(buf + cmd_pos);
242         next_cmd_pos = READ_OFFSET(buf + cmd_pos + 2);
243         av_dlog(NULL, "cmd_pos=0x%04x next=0x%04x date=%d\n",
244                 cmd_pos, next_cmd_pos, date);
245         pos = cmd_pos + 2 + offset_size;
246         offset1 = -1;
247         offset2 = -1;
248         x1 = y1 = x2 = y2 = 0;
249         while (pos < buf_size) {
250             cmd = buf[pos++];
251             av_dlog(NULL, "cmd=%02x\n", cmd);
252             switch(cmd) {
253             case 0x00:
254                 /* menu subpicture */
255                 is_menu = 1;
256                 break;
257             case 0x01:
258                 /* set start date */
259                 sub_header->start_display_time = (date << 10) / 90;
260                 break;
261             case 0x02:
262                 /* set end date */
263                 sub_header->end_display_time = (date << 10) / 90;
264                 break;
265             case 0x03:
266                 /* set colormap */
267                 if ((buf_size - pos) < 2)
268                     goto fail;
269                 colormap[3] = buf[pos] >> 4;
270                 colormap[2] = buf[pos] & 0x0f;
271                 colormap[1] = buf[pos + 1] >> 4;
272                 colormap[0] = buf[pos + 1] & 0x0f;
273                 pos += 2;
274                 break;
275             case 0x04:
276                 /* set alpha */
277                 if ((buf_size - pos) < 2)
278                     goto fail;
279                 alpha[3] = buf[pos] >> 4;
280                 alpha[2] = buf[pos] & 0x0f;
281                 alpha[1] = buf[pos + 1] >> 4;
282                 alpha[0] = buf[pos + 1] & 0x0f;
283                 pos += 2;
284             av_dlog(NULL, "alpha=%x%x%x%x\n", alpha[0],alpha[1],alpha[2],alpha[3]);
285                 break;
286             case 0x05:
287             case 0x85:
288                 if ((buf_size - pos) < 6)
289                     goto fail;
290                 x1 = (buf[pos] << 4) | (buf[pos + 1] >> 4);
291                 x2 = ((buf[pos + 1] & 0x0f) << 8) | buf[pos + 2];
292                 y1 = (buf[pos + 3] << 4) | (buf[pos + 4] >> 4);
293                 y2 = ((buf[pos + 4] & 0x0f) << 8) | buf[pos + 5];
294                 if (cmd & 0x80)
295                     is_8bit = 1;
296                 av_dlog(NULL, "x1=%d x2=%d y1=%d y2=%d\n", x1, x2, y1, y2);
297                 pos += 6;
298                 break;
299             case 0x06:
300                 if ((buf_size - pos) < 4)
301                     goto fail;
302                 offset1 = AV_RB16(buf + pos);
303                 offset2 = AV_RB16(buf + pos + 2);
304                 av_dlog(NULL, "offset1=0x%04x offset2=0x%04x\n", offset1, offset2);
305                 pos += 4;
306                 break;
307             case 0x86:
308                 if ((buf_size - pos) < 8)
309                     goto fail;
310                 offset1 = AV_RB32(buf + pos);
311                 offset2 = AV_RB32(buf + pos + 4);
312                 av_dlog(NULL, "offset1=0x%04x offset2=0x%04x\n", offset1, offset2);
313                 pos += 8;
314                 break;
315
316             case 0x83:
317                 /* HD set palette */
318                 if ((buf_size - pos) < 768)
319                     goto fail;
320                 yuv_palette = buf + pos;
321                 pos += 768;
322                 break;
323             case 0x84:
324                 /* HD set contrast (alpha) */
325                 if ((buf_size - pos) < 256)
326                     goto fail;
327                 for (i = 0; i < 256; i++)
328                     alpha[i] = 0xFF - buf[pos+i];
329                 pos += 256;
330                 break;
331
332             case 0xff:
333                 goto the_end;
334             default:
335                 av_dlog(NULL, "unrecognised subpicture command 0x%x\n", cmd);
336                 goto the_end;
337             }
338         }
339     the_end:
340         if (offset1 >= 0) {
341             int w, h;
342             uint8_t *bitmap;
343
344             /* decode the bitmap */
345             w = x2 - x1 + 1;
346             if (w < 0)
347                 w = 0;
348             h = y2 - y1;
349             if (h < 0)
350                 h = 0;
351             if (w > 0 && h > 0) {
352                 reset_rects(sub_header);
353
354                 bitmap = av_malloc(w * h);
355                 sub_header->rects = av_mallocz(sizeof(*sub_header->rects));
356                 sub_header->rects[0] = av_mallocz(sizeof(AVSubtitleRect));
357                 sub_header->num_rects = 1;
358                 sub_header->rects[0]->pict.data[0] = bitmap;
359                 decode_rle(bitmap, w * 2, w, (h + 1) / 2,
360                            buf, offset1, buf_size, is_8bit);
361                 decode_rle(bitmap + w, w * 2, w, h / 2,
362                            buf, offset2, buf_size, is_8bit);
363                 sub_header->rects[0]->pict.data[1] = av_mallocz(AVPALETTE_SIZE);
364                 if (is_8bit) {
365                     if (yuv_palette == 0)
366                         goto fail;
367                     sub_header->rects[0]->nb_colors = 256;
368                     yuv_a_to_rgba(yuv_palette, alpha, (uint32_t*)sub_header->rects[0]->pict.data[1], 256);
369                 } else {
370                     sub_header->rects[0]->nb_colors = 4;
371                     guess_palette(ctx, (uint32_t*)sub_header->rects[0]->pict.data[1],
372                                   0xffff00);
373                 }
374                 sub_header->rects[0]->x = x1;
375                 sub_header->rects[0]->y = y1;
376                 sub_header->rects[0]->w = w;
377                 sub_header->rects[0]->h = h;
378                 sub_header->rects[0]->type = SUBTITLE_BITMAP;
379                 sub_header->rects[0]->pict.linesize[0] = w;
380                 sub_header->rects[0]->flags = is_menu ? AV_SUBTITLE_FLAG_FORCED : 0;
381             }
382         }
383         if (next_cmd_pos < cmd_pos) {
384             av_log(NULL, AV_LOG_ERROR, "Invalid command offset\n");
385             break;
386         }
387         if (next_cmd_pos == cmd_pos)
388             break;
389         cmd_pos = next_cmd_pos;
390     }
391     if (sub_header->num_rects > 0)
392         return is_menu;
393  fail:
394     reset_rects(sub_header);
395     return -1;
396 }
397
398 static int is_transp(const uint8_t *buf, int pitch, int n,
399                      const uint8_t *transp_color)
400 {
401     int i;
402     for(i = 0; i < n; i++) {
403         if (!transp_color[*buf])
404             return 0;
405         buf += pitch;
406     }
407     return 1;
408 }
409
410 /* return 0 if empty rectangle, 1 if non empty */
411 static int find_smallest_bounding_rectangle(AVSubtitle *s)
412 {
413     uint8_t transp_color[256] = { 0 };
414     int y1, y2, x1, x2, y, w, h, i;
415     uint8_t *bitmap;
416
417     if (s->num_rects == 0 || s->rects == NULL || s->rects[0]->w <= 0 || s->rects[0]->h <= 0)
418         return 0;
419
420     for(i = 0; i < s->rects[0]->nb_colors; i++) {
421         if ((((uint32_t*)s->rects[0]->pict.data[1])[i] >> 24) == 0)
422             transp_color[i] = 1;
423     }
424     y1 = 0;
425     while (y1 < s->rects[0]->h && is_transp(s->rects[0]->pict.data[0] + y1 * s->rects[0]->pict.linesize[0],
426                                   1, s->rects[0]->w, transp_color))
427         y1++;
428     if (y1 == s->rects[0]->h) {
429         av_freep(&s->rects[0]->pict.data[0]);
430         s->rects[0]->w = s->rects[0]->h = 0;
431         return 0;
432     }
433
434     y2 = s->rects[0]->h - 1;
435     while (y2 > 0 && is_transp(s->rects[0]->pict.data[0] + y2 * s->rects[0]->pict.linesize[0], 1,
436                                s->rects[0]->w, transp_color))
437         y2--;
438     x1 = 0;
439     while (x1 < (s->rects[0]->w - 1) && is_transp(s->rects[0]->pict.data[0] + x1, s->rects[0]->pict.linesize[0],
440                                         s->rects[0]->h, transp_color))
441         x1++;
442     x2 = s->rects[0]->w - 1;
443     while (x2 > 0 && is_transp(s->rects[0]->pict.data[0] + x2, s->rects[0]->pict.linesize[0], s->rects[0]->h,
444                                   transp_color))
445         x2--;
446     w = x2 - x1 + 1;
447     h = y2 - y1 + 1;
448     bitmap = av_malloc(w * h);
449     if (!bitmap)
450         return 1;
451     for(y = 0; y < h; y++) {
452         memcpy(bitmap + w * y, s->rects[0]->pict.data[0] + x1 + (y1 + y) * s->rects[0]->pict.linesize[0], w);
453     }
454     av_freep(&s->rects[0]->pict.data[0]);
455     s->rects[0]->pict.data[0] = bitmap;
456     s->rects[0]->pict.linesize[0] = w;
457     s->rects[0]->w = w;
458     s->rects[0]->h = h;
459     s->rects[0]->x += x1;
460     s->rects[0]->y += y1;
461     return 1;
462 }
463
464 #ifdef DEBUG
465 static void ppm_save(const char *filename, uint8_t *bitmap, int w, int h,
466                      uint32_t *rgba_palette)
467 {
468     int x, y, v;
469     FILE *f;
470
471     f = fopen(filename, "w");
472     if (!f) {
473         perror(filename);
474         return;
475     }
476     fprintf(f, "P6\n"
477             "%d %d\n"
478             "%d\n",
479             w, h, 255);
480     for(y = 0; y < h; y++) {
481         for(x = 0; x < w; x++) {
482             v = rgba_palette[bitmap[y * w + x]];
483             putc((v >> 16) & 0xff, f);
484             putc((v >> 8) & 0xff, f);
485             putc((v >> 0) & 0xff, f);
486         }
487     }
488     fclose(f);
489 }
490 #endif
491
492 static int append_to_cached_buf(AVCodecContext *avctx,
493                                 const uint8_t *buf, int buf_size)
494 {
495     DVDSubContext *ctx = avctx->priv_data;
496
497     if (ctx->buf_size > 0xffff - buf_size) {
498         av_log(avctx, AV_LOG_WARNING, "Attempt to reconstruct "
499                "too large SPU packets aborted.\n");
500         av_freep(&ctx->buf);
501         return AVERROR_INVALIDDATA;
502     }
503     ctx->buf = av_realloc(ctx->buf, ctx->buf_size + buf_size);
504     if (!ctx->buf)
505         return AVERROR(ENOMEM);
506     memcpy(ctx->buf + ctx->buf_size, buf, buf_size);
507     ctx->buf_size += buf_size;
508     return 0;
509 }
510
511 static int dvdsub_decode(AVCodecContext *avctx,
512                          void *data, int *data_size,
513                          AVPacket *avpkt)
514 {
515     DVDSubContext *ctx = avctx->priv_data;
516     const uint8_t *buf = avpkt->data;
517     int buf_size = avpkt->size;
518     AVSubtitle *sub = data;
519     int is_menu;
520
521     if (ctx->buf) {
522         int ret = append_to_cached_buf(avctx, buf, buf_size);
523         if (ret < 0) {
524             *data_size = 0;
525             return ret;
526         }
527         buf = ctx->buf;
528         buf_size = ctx->buf_size;
529     }
530
531     is_menu = decode_dvd_subtitles(ctx, sub, buf, buf_size);
532     if (is_menu == AVERROR(EAGAIN)) {
533         *data_size = 0;
534         return append_to_cached_buf(avctx, buf, buf_size);
535     }
536
537     if (is_menu < 0) {
538     no_subtitle:
539         *data_size = 0;
540
541         return buf_size;
542     }
543     if (!is_menu && find_smallest_bounding_rectangle(sub) == 0)
544         goto no_subtitle;
545
546 #if defined(DEBUG)
547     {
548     char ppm_name[32];
549
550     snprintf(ppm_name, sizeof(ppm_name), "/tmp/%05d.ppm", ctx->sub_id++);
551     av_dlog(NULL, "start=%d ms end =%d ms\n",
552             sub->start_display_time,
553             sub->end_display_time);
554     ppm_save(ppm_name, sub->rects[0]->pict.data[0],
555              sub->rects[0]->w, sub->rects[0]->h, sub->rects[0]->pict.data[1]);
556     }
557 #endif
558
559     av_freep(&ctx->buf);
560     ctx->buf_size = 0;
561     *data_size = 1;
562     return buf_size;
563 }
564
565 static void parse_palette(DVDSubContext *ctx, char *p)
566 {
567     int i;
568
569     ctx->has_palette = 1;
570     for(i=0;i<16;i++) {
571         ctx->palette[i] = strtoul(p, &p, 16);
572         while(*p == ',' || av_isspace(*p))
573             p++;
574     }
575 }
576
577 static int dvdsub_parse_extradata(AVCodecContext *avctx)
578 {
579     DVDSubContext *ctx = (DVDSubContext*) avctx->priv_data;
580     char *dataorig, *data;
581
582     if (!avctx->extradata || !avctx->extradata_size)
583         return 1;
584
585     dataorig = data = av_malloc(avctx->extradata_size+1);
586     if (!data)
587         return AVERROR(ENOMEM);
588     memcpy(data, avctx->extradata, avctx->extradata_size);
589     data[avctx->extradata_size] = '\0';
590
591     for(;;) {
592         int pos = strcspn(data, "\n\r");
593         if (pos==0 && *data==0)
594             break;
595
596         if (strncmp("palette:", data, 8) == 0) {
597             parse_palette(ctx, data + 8);
598         } else if (strncmp("size:", data, 5) == 0) {
599             int w, h;
600             if (sscanf(data + 5, "%dx%d", &w, &h) == 2) {
601                int ret = ff_set_dimensions(avctx, w, h);
602                if (ret < 0) {
603                    av_free(dataorig);
604                    return ret;
605                }
606             }
607         }
608
609         data += pos;
610         data += strspn(data, "\n\r");
611     }
612
613     av_free(dataorig);
614     return 1;
615 }
616
617 static av_cold int dvdsub_init(AVCodecContext *avctx)
618 {
619     DVDSubContext *ctx = avctx->priv_data;
620     int ret;
621
622     if ((ret = dvdsub_parse_extradata(avctx)) < 0)
623         return ret;
624
625     if (ctx->palette_str)
626         parse_palette(ctx, ctx->palette_str);
627     if (ctx->has_palette) {
628         int i;
629         av_log(avctx, AV_LOG_DEBUG, "palette:");
630         for(i=0;i<16;i++)
631             av_log(avctx, AV_LOG_DEBUG, " 0x%06x", ctx->palette[i]);
632         av_log(avctx, AV_LOG_DEBUG, "\n");
633     }
634
635     return 1;
636 }
637
638 static av_cold int dvdsub_close(AVCodecContext *avctx)
639 {
640     DVDSubContext *ctx = avctx->priv_data;
641     av_freep(&ctx->buf);
642     ctx->buf_size = 0;
643     return 0;
644 }
645
646 #define OFFSET(field) offsetof(DVDSubContext, field)
647 #define VD AV_OPT_FLAG_SUBTITLE_PARAM | AV_OPT_FLAG_DECODING_PARAM
648 static const AVOption options[] = {
649     { "palette", "set the global palette", OFFSET(palette_str), AV_OPT_TYPE_STRING, { .str = NULL }, 0, 0, VD },
650     { NULL }
651 };
652 static const AVClass dvdsub_class = {
653     .class_name = "dvdsubdec",
654     .item_name  = av_default_item_name,
655     .option     = options,
656     .version    = LIBAVUTIL_VERSION_INT,
657 };
658
659 AVCodec ff_dvdsub_decoder = {
660     .name           = "dvdsub",
661     .long_name      = NULL_IF_CONFIG_SMALL("DVD subtitles"),
662     .type           = AVMEDIA_TYPE_SUBTITLE,
663     .id             = AV_CODEC_ID_DVD_SUBTITLE,
664     .priv_data_size = sizeof(DVDSubContext),
665     .init           = dvdsub_init,
666     .decode         = dvdsub_decode,
667     .close          = dvdsub_close,
668     .priv_class     = &dvdsub_class,
669 };