avcodec/wmaenc: Check input for finiteness
[ffmpeg.git] / libavcodec / wmaenc.c
1 /*
2  * WMA compatible encoder
3  * Copyright (c) 2007 Michael Niedermayer
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
22 #include "libavutil/attributes.h"
23 #include "libavutil/internal.h"
24
25 #include "avcodec.h"
26 #include "internal.h"
27 #include "wma.h"
28 #include "libavutil/avassert.h"
29
30
31 static av_cold int encode_init(AVCodecContext *avctx)
32 {
33     WMACodecContext *s = avctx->priv_data;
34     int i, flags1, flags2, block_align;
35     uint8_t *extradata;
36
37     s->avctx = avctx;
38
39     if (avctx->channels > MAX_CHANNELS) {
40         av_log(avctx, AV_LOG_ERROR,
41                "too many channels: got %i, need %i or fewer\n",
42                avctx->channels, MAX_CHANNELS);
43         return AVERROR(EINVAL);
44     }
45
46     if (avctx->sample_rate > 48000) {
47         av_log(avctx, AV_LOG_ERROR, "sample rate is too high: %d > 48kHz\n",
48                avctx->sample_rate);
49         return AVERROR(EINVAL);
50     }
51
52     if (avctx->bit_rate < 24 * 1000) {
53         av_log(avctx, AV_LOG_ERROR,
54                "bitrate too low: got %"PRId64", need 24000 or higher\n",
55                (int64_t)avctx->bit_rate);
56         return AVERROR(EINVAL);
57     }
58
59     /* extract flag infos */
60     flags1 = 0;
61     flags2 = 1;
62     if (avctx->codec->id == AV_CODEC_ID_WMAV1) {
63         extradata             = av_malloc(4);
64         if (!extradata)
65             return AVERROR(ENOMEM);
66         avctx->extradata_size = 4;
67         AV_WL16(extradata, flags1);
68         AV_WL16(extradata + 2, flags2);
69     } else if (avctx->codec->id == AV_CODEC_ID_WMAV2) {
70         extradata             = av_mallocz(10);
71         if (!extradata)
72             return AVERROR(ENOMEM);
73         avctx->extradata_size = 10;
74         AV_WL32(extradata, flags1);
75         AV_WL16(extradata + 4, flags2);
76     } else {
77         av_assert0(0);
78     }
79     avctx->extradata          = extradata;
80     s->use_exp_vlc            = flags2 & 0x0001;
81     s->use_bit_reservoir      = flags2 & 0x0002;
82     s->use_variable_block_len = flags2 & 0x0004;
83     if (avctx->channels == 2)
84         s->ms_stereo = 1;
85
86     ff_wma_init(avctx, flags2);
87
88     /* init MDCT */
89     for (i = 0; i < s->nb_block_sizes; i++)
90         ff_mdct_init(&s->mdct_ctx[i], s->frame_len_bits - i + 1, 0, 1.0);
91
92     block_align        = avctx->bit_rate * (int64_t) s->frame_len /
93                          (avctx->sample_rate * 8);
94     block_align        = FFMIN(block_align, MAX_CODED_SUPERFRAME_SIZE);
95     avctx->block_align = block_align;
96     avctx->frame_size = avctx->initial_padding = s->frame_len;
97
98     return 0;
99 }
100
101 static int apply_window_and_mdct(AVCodecContext *avctx, const AVFrame *frame)
102 {
103     WMACodecContext *s = avctx->priv_data;
104     float **audio      = (float **) frame->extended_data;
105     int len            = frame->nb_samples;
106     int window_index   = s->frame_len_bits - s->block_len_bits;
107     FFTContext *mdct   = &s->mdct_ctx[window_index];
108     int ch;
109     const float *win   = s->windows[window_index];
110     int window_len     = 1 << s->block_len_bits;
111     float n            = 2.0 * 32768.0 / window_len;
112
113     for (ch = 0; ch < avctx->channels; ch++) {
114         memcpy(s->output, s->frame_out[ch], window_len * sizeof(*s->output));
115         s->fdsp->vector_fmul_scalar(s->frame_out[ch], audio[ch], n, len);
116         s->fdsp->vector_fmul_reverse(&s->output[window_len], s->frame_out[ch],
117                                     win, len);
118         s->fdsp->vector_fmul(s->frame_out[ch], s->frame_out[ch], win, len);
119         mdct->mdct_calc(mdct, s->coefs[ch], s->output);
120         if (!isfinite(s->coefs[ch][0])) {
121             av_log(avctx, AV_LOG_ERROR, "Input contains NaN/+-Inf\n");
122             return AVERROR(EINVAL);
123         }
124     }
125
126     return 0;
127 }
128
129 // FIXME use for decoding too
130 static void init_exp(WMACodecContext *s, int ch, const int *exp_param)
131 {
132     int n;
133     const uint16_t *ptr;
134     float v, *q, max_scale, *q_end;
135
136     ptr       = s->exponent_bands[s->frame_len_bits - s->block_len_bits];
137     q         = s->exponents[ch];
138     q_end     = q + s->block_len;
139     max_scale = 0;
140     while (q < q_end) {
141         /* XXX: use a table */
142         v         = ff_exp10(*exp_param++ *(1.0 / 16.0));
143         max_scale = FFMAX(max_scale, v);
144         n         = *ptr++;
145         do {
146             *q++ = v;
147         } while (--n);
148     }
149     s->max_exponent[ch] = max_scale;
150 }
151
152 static void encode_exp_vlc(WMACodecContext *s, int ch, const int *exp_param)
153 {
154     int last_exp;
155     const uint16_t *ptr;
156     float *q, *q_end;
157
158     ptr   = s->exponent_bands[s->frame_len_bits - s->block_len_bits];
159     q     = s->exponents[ch];
160     q_end = q + s->block_len;
161     if (s->version == 1) {
162         last_exp = *exp_param++;
163         av_assert0(last_exp - 10 >= 0 && last_exp - 10 < 32);
164         put_bits(&s->pb, 5, last_exp - 10);
165         q += *ptr++;
166     } else
167         last_exp = 36;
168     while (q < q_end) {
169         int exp  = *exp_param++;
170         int code = exp - last_exp + 60;
171         av_assert1(code >= 0 && code < 120);
172         put_bits(&s->pb, ff_aac_scalefactor_bits[code],
173                  ff_aac_scalefactor_code[code]);
174         /* XXX: use a table */
175         q       += *ptr++;
176         last_exp = exp;
177     }
178 }
179
180 static int encode_block(WMACodecContext *s, float (*src_coefs)[BLOCK_MAX_SIZE],
181                         int total_gain)
182 {
183     int v, bsize, ch, coef_nb_bits, parse_exponents;
184     float mdct_norm;
185     int nb_coefs[MAX_CHANNELS];
186     static const int fixed_exp[25] = {
187         20, 20, 20, 20, 20,
188         20, 20, 20, 20, 20,
189         20, 20, 20, 20, 20,
190         20, 20, 20, 20, 20,
191         20, 20, 20, 20, 20
192     };
193
194     // FIXME remove duplication relative to decoder
195     if (s->use_variable_block_len) {
196         av_assert0(0); // FIXME not implemented
197     } else {
198         /* fixed block len */
199         s->next_block_len_bits = s->frame_len_bits;
200         s->prev_block_len_bits = s->frame_len_bits;
201         s->block_len_bits      = s->frame_len_bits;
202     }
203
204     s->block_len = 1 << s->block_len_bits;
205 //     av_assert0((s->block_pos + s->block_len) <= s->frame_len);
206     bsize = s->frame_len_bits - s->block_len_bits;
207
208     // FIXME factor
209     v = s->coefs_end[bsize] - s->coefs_start;
210     for (ch = 0; ch < s->avctx->channels; ch++)
211         nb_coefs[ch] = v;
212     {
213         int n4 = s->block_len / 2;
214         mdct_norm = 1.0 / (float) n4;
215         if (s->version == 1)
216             mdct_norm *= sqrt(n4);
217     }
218
219     if (s->avctx->channels == 2)
220         put_bits(&s->pb, 1, !!s->ms_stereo);
221
222     for (ch = 0; ch < s->avctx->channels; ch++) {
223         // FIXME only set channel_coded when needed, instead of always
224         s->channel_coded[ch] = 1;
225         if (s->channel_coded[ch])
226             init_exp(s, ch, fixed_exp);
227     }
228
229     for (ch = 0; ch < s->avctx->channels; ch++) {
230         if (s->channel_coded[ch]) {
231             WMACoef *coefs1;
232             float *coefs, *exponents, mult;
233             int i, n;
234
235             coefs1    = s->coefs1[ch];
236             exponents = s->exponents[ch];
237             mult      = ff_exp10(total_gain * 0.05) / s->max_exponent[ch];
238             mult     *= mdct_norm;
239             coefs     = src_coefs[ch];
240             if (s->use_noise_coding && 0) {
241                 av_assert0(0); // FIXME not implemented
242             } else {
243                 coefs += s->coefs_start;
244                 n      = nb_coefs[ch];
245                 for (i = 0; i < n; i++) {
246                     double t = *coefs++ / (exponents[i] * mult);
247                     if (t < -32768 || t > 32767)
248                         return -1;
249
250                     coefs1[i] = lrint(t);
251                 }
252             }
253         }
254     }
255
256     v = 0;
257     for (ch = 0; ch < s->avctx->channels; ch++) {
258         int a = s->channel_coded[ch];
259         put_bits(&s->pb, 1, a);
260         v |= a;
261     }
262
263     if (!v)
264         return 1;
265
266     for (v = total_gain - 1; v >= 127; v -= 127)
267         put_bits(&s->pb, 7, 127);
268     put_bits(&s->pb, 7, v);
269
270     coef_nb_bits = ff_wma_total_gain_to_bits(total_gain);
271
272     if (s->use_noise_coding) {
273         for (ch = 0; ch < s->avctx->channels; ch++) {
274             if (s->channel_coded[ch]) {
275                 int i, n;
276                 n = s->exponent_high_sizes[bsize];
277                 for (i = 0; i < n; i++) {
278                     put_bits(&s->pb, 1, s->high_band_coded[ch][i] = 0);
279                     if (0)
280                         nb_coefs[ch] -= s->exponent_high_bands[bsize][i];
281                 }
282             }
283         }
284     }
285
286     parse_exponents = 1;
287     if (s->block_len_bits != s->frame_len_bits)
288         put_bits(&s->pb, 1, parse_exponents);
289
290     if (parse_exponents) {
291         for (ch = 0; ch < s->avctx->channels; ch++) {
292             if (s->channel_coded[ch]) {
293                 if (s->use_exp_vlc) {
294                     encode_exp_vlc(s, ch, fixed_exp);
295                 } else {
296                     av_assert0(0); // FIXME not implemented
297 //                    encode_exp_lsp(s, ch);
298                 }
299             }
300         }
301     } else
302         av_assert0(0); // FIXME not implemented
303
304     for (ch = 0; ch < s->avctx->channels; ch++) {
305         if (s->channel_coded[ch]) {
306             int run, tindex;
307             WMACoef *ptr, *eptr;
308             tindex = (ch == 1 && s->ms_stereo);
309             ptr    = &s->coefs1[ch][0];
310             eptr   = ptr + nb_coefs[ch];
311
312             run = 0;
313             for (; ptr < eptr; ptr++) {
314                 if (*ptr) {
315                     int level     = *ptr;
316                     int abs_level = FFABS(level);
317                     int code      = 0;
318                     if (abs_level <= s->coef_vlcs[tindex]->max_level)
319                         if (run < s->coef_vlcs[tindex]->levels[abs_level - 1])
320                             code = run + s->int_table[tindex][abs_level - 1];
321
322                     av_assert2(code < s->coef_vlcs[tindex]->n);
323                     put_bits(&s->pb, s->coef_vlcs[tindex]->huffbits[code],
324                              s->coef_vlcs[tindex]->huffcodes[code]);
325
326                     if (code == 0) {
327                         if (1 << coef_nb_bits <= abs_level)
328                             return -1;
329
330                         put_bits(&s->pb, coef_nb_bits, abs_level);
331                         put_bits(&s->pb, s->frame_len_bits, run);
332                     }
333                     // FIXME the sign is flipped somewhere
334                     put_bits(&s->pb, 1, level < 0);
335                     run = 0;
336                 } else
337                     run++;
338             }
339             if (run)
340                 put_bits(&s->pb, s->coef_vlcs[tindex]->huffbits[1],
341                          s->coef_vlcs[tindex]->huffcodes[1]);
342         }
343         if (s->version == 1 && s->avctx->channels >= 2)
344             avpriv_align_put_bits(&s->pb);
345     }
346     return 0;
347 }
348
349 static int encode_frame(WMACodecContext *s, float (*src_coefs)[BLOCK_MAX_SIZE],
350                         uint8_t *buf, int buf_size, int total_gain)
351 {
352     init_put_bits(&s->pb, buf, buf_size);
353
354     if (s->use_bit_reservoir)
355         av_assert0(0); // FIXME not implemented
356     else if (encode_block(s, src_coefs, total_gain) < 0)
357         return INT_MAX;
358
359     avpriv_align_put_bits(&s->pb);
360
361     return put_bits_count(&s->pb) / 8 - s->avctx->block_align;
362 }
363
364 static int encode_superframe(AVCodecContext *avctx, AVPacket *avpkt,
365                              const AVFrame *frame, int *got_packet_ptr)
366 {
367     WMACodecContext *s = avctx->priv_data;
368     int i, total_gain, ret, error;
369
370     s->block_len_bits = s->frame_len_bits; // required by non variable block len
371     s->block_len      = 1 << s->block_len_bits;
372
373     ret = apply_window_and_mdct(avctx, frame);
374
375     if (ret < 0)
376         return ret;
377
378     if (s->ms_stereo) {
379         float a, b;
380         int i;
381
382         for (i = 0; i < s->block_len; i++) {
383             a              = s->coefs[0][i] * 0.5;
384             b              = s->coefs[1][i] * 0.5;
385             s->coefs[0][i] = a + b;
386             s->coefs[1][i] = a - b;
387         }
388     }
389
390     if ((ret = ff_alloc_packet2(avctx, avpkt, 2 * MAX_CODED_SUPERFRAME_SIZE, 0)) < 0)
391         return ret;
392
393     total_gain = 128;
394     for (i = 64; i; i >>= 1) {
395         error = encode_frame(s, s->coefs, avpkt->data, avpkt->size,
396                                  total_gain - i);
397         if (error <= 0)
398             total_gain -= i;
399     }
400
401     while(total_gain <= 128 && error > 0)
402         error = encode_frame(s, s->coefs, avpkt->data, avpkt->size, total_gain++);
403     if (error > 0) {
404         av_log(avctx, AV_LOG_ERROR, "Invalid input data or requested bitrate too low, cannot encode\n");
405         avpkt->size = 0;
406         return AVERROR(EINVAL);
407     }
408     av_assert0((put_bits_count(&s->pb) & 7) == 0);
409     i= avctx->block_align - (put_bits_count(&s->pb)+7)/8;
410     av_assert0(i>=0);
411     while(i--)
412         put_bits(&s->pb, 8, 'N');
413
414     flush_put_bits(&s->pb);
415     av_assert0(put_bits_ptr(&s->pb) - s->pb.buf == avctx->block_align);
416
417     if (frame->pts != AV_NOPTS_VALUE)
418         avpkt->pts = frame->pts - ff_samples_to_time_base(avctx, avctx->initial_padding);
419
420     avpkt->size     = avctx->block_align;
421     *got_packet_ptr = 1;
422     return 0;
423 }
424
425 #if CONFIG_WMAV1_ENCODER
426 AVCodec ff_wmav1_encoder = {
427     .name           = "wmav1",
428     .long_name      = NULL_IF_CONFIG_SMALL("Windows Media Audio 1"),
429     .type           = AVMEDIA_TYPE_AUDIO,
430     .id             = AV_CODEC_ID_WMAV1,
431     .priv_data_size = sizeof(WMACodecContext),
432     .init           = encode_init,
433     .encode2        = encode_superframe,
434     .close          = ff_wma_end,
435     .sample_fmts    = (const enum AVSampleFormat[]) { AV_SAMPLE_FMT_FLTP,
436                                                       AV_SAMPLE_FMT_NONE },
437 };
438 #endif
439 #if CONFIG_WMAV2_ENCODER
440 AVCodec ff_wmav2_encoder = {
441     .name           = "wmav2",
442     .long_name      = NULL_IF_CONFIG_SMALL("Windows Media Audio 2"),
443     .type           = AVMEDIA_TYPE_AUDIO,
444     .id             = AV_CODEC_ID_WMAV2,
445     .priv_data_size = sizeof(WMACodecContext),
446     .init           = encode_init,
447     .encode2        = encode_superframe,
448     .close          = ff_wma_end,
449     .sample_fmts    = (const enum AVSampleFormat[]) { AV_SAMPLE_FMT_FLTP,
450                                                       AV_SAMPLE_FMT_NONE },
451 };
452 #endif