]> git.sesse.net Git - ffmpeg/blob - libavformat/movenccenc.c
doc/examples/transcode_aac: use av_packet_alloc() to allocate packets
[ffmpeg] / libavformat / movenccenc.c
1 /*
2  * MOV CENC (Common Encryption) writer
3  * Copyright (c) 2015 Eran Kornblau <erankor at gmail dot com>
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 "movenccenc.h"
22 #include "libavutil/intreadwrite.h"
23 #include "avio_internal.h"
24 #include "movenc.h"
25 #include "avc.h"
26
27 static int auxiliary_info_alloc_size(MOVMuxCencContext* ctx, int size)
28 {
29     size_t new_alloc_size;
30
31     if (ctx->auxiliary_info_size + size > ctx->auxiliary_info_alloc_size) {
32         new_alloc_size = FFMAX(ctx->auxiliary_info_size + size, ctx->auxiliary_info_alloc_size * 2);
33         if (av_reallocp(&ctx->auxiliary_info, new_alloc_size)) {
34             return AVERROR(ENOMEM);
35         }
36
37         ctx->auxiliary_info_alloc_size = new_alloc_size;
38     }
39
40     return 0;
41 }
42
43 static int auxiliary_info_write(MOVMuxCencContext* ctx,
44                                          const uint8_t *buf_in, int size)
45 {
46     int ret;
47
48     ret = auxiliary_info_alloc_size(ctx, size);
49     if (ret) {
50         return ret;
51     }
52     memcpy(ctx->auxiliary_info + ctx->auxiliary_info_size, buf_in, size);
53     ctx->auxiliary_info_size += size;
54
55     return 0;
56 }
57
58 static int auxiliary_info_add_subsample(MOVMuxCencContext* ctx,
59     uint16_t clear_bytes, uint32_t encrypted_bytes)
60 {
61     uint8_t* p;
62     int ret;
63
64     if (!ctx->use_subsamples) {
65         return 0;
66     }
67
68     ret = auxiliary_info_alloc_size(ctx, 6);
69     if (ret) {
70         return ret;
71     }
72
73     p = ctx->auxiliary_info + ctx->auxiliary_info_size;
74
75     AV_WB16(p, clear_bytes);
76     p += sizeof(uint16_t);
77
78     AV_WB32(p, encrypted_bytes);
79
80     ctx->auxiliary_info_size += 6;
81     ctx->subsample_count++;
82
83     return 0;
84 }
85
86 /**
87  * Encrypt the input buffer and write using avio_write
88  */
89 static void mov_cenc_write_encrypted(MOVMuxCencContext* ctx, AVIOContext *pb,
90                                      const uint8_t *buf_in, int size)
91 {
92     uint8_t chunk[4096];
93     const uint8_t* cur_pos = buf_in;
94     int size_left = size;
95     int cur_size;
96
97     while (size_left > 0) {
98         cur_size = FFMIN(size_left, sizeof(chunk));
99         av_aes_ctr_crypt(ctx->aes_ctr, chunk, cur_pos, cur_size);
100         avio_write(pb, chunk, cur_size);
101         cur_pos += cur_size;
102         size_left -= cur_size;
103     }
104 }
105
106 /**
107  * Start writing a packet
108  */
109 static int mov_cenc_start_packet(MOVMuxCencContext* ctx)
110 {
111     int ret;
112
113     /* write the iv */
114     ret = auxiliary_info_write(ctx, av_aes_ctr_get_iv(ctx->aes_ctr), AES_CTR_IV_SIZE);
115     if (ret) {
116         return ret;
117     }
118
119     if (!ctx->use_subsamples) {
120         return 0;
121     }
122
123     /* write a zero subsample count */
124     ctx->auxiliary_info_subsample_start = ctx->auxiliary_info_size;
125     ctx->subsample_count = 0;
126     ret = auxiliary_info_write(ctx, (uint8_t*)&ctx->subsample_count, sizeof(ctx->subsample_count));
127     if (ret) {
128         return ret;
129     }
130
131     return 0;
132 }
133
134 /**
135  * Finalize a packet
136  */
137 static int mov_cenc_end_packet(MOVMuxCencContext* ctx)
138 {
139     size_t new_alloc_size;
140
141     av_aes_ctr_increment_iv(ctx->aes_ctr);
142
143     if (!ctx->use_subsamples) {
144         ctx->auxiliary_info_entries++;
145         return 0;
146     }
147
148     /* add the auxiliary info entry size*/
149     if (ctx->auxiliary_info_entries >= ctx->auxiliary_info_sizes_alloc_size) {
150         new_alloc_size = ctx->auxiliary_info_entries * 2 + 1;
151         if (av_reallocp(&ctx->auxiliary_info_sizes, new_alloc_size)) {
152             return AVERROR(ENOMEM);
153         }
154
155         ctx->auxiliary_info_sizes_alloc_size = new_alloc_size;
156     }
157     ctx->auxiliary_info_sizes[ctx->auxiliary_info_entries] =
158         AES_CTR_IV_SIZE + ctx->auxiliary_info_size - ctx->auxiliary_info_subsample_start;
159     ctx->auxiliary_info_entries++;
160
161     /* update the subsample count*/
162     AV_WB16(ctx->auxiliary_info + ctx->auxiliary_info_subsample_start, ctx->subsample_count);
163
164     return 0;
165 }
166
167 int ff_mov_cenc_write_packet(MOVMuxCencContext* ctx, AVIOContext *pb,
168                           const uint8_t *buf_in, int size)
169 {
170     int ret;
171
172     ret = mov_cenc_start_packet(ctx);
173     if (ret) {
174         return ret;
175     }
176
177     ret = auxiliary_info_add_subsample(ctx, 0, size);
178     if (ret) {
179         return ret;
180     }
181
182     mov_cenc_write_encrypted(ctx, pb, buf_in, size);
183
184     ret = mov_cenc_end_packet(ctx);
185     if (ret) {
186         return ret;
187     }
188
189     return 0;
190 }
191
192 int ff_mov_cenc_avc_parse_nal_units(MOVMuxCencContext* ctx, AVIOContext *pb,
193                                  const uint8_t *buf_in, int size)
194 {
195     const uint8_t *p = buf_in;
196     const uint8_t *end = p + size;
197     const uint8_t *nal_start, *nal_end;
198     int ret;
199
200     ret = mov_cenc_start_packet(ctx);
201     if (ret) {
202         return ret;
203     }
204
205     size = 0;
206     nal_start = ff_avc_find_startcode(p, end);
207     for (;;) {
208         while (nal_start < end && !*(nal_start++));
209         if (nal_start == end)
210             break;
211
212         nal_end = ff_avc_find_startcode(nal_start, end);
213
214         avio_wb32(pb, nal_end - nal_start);
215         avio_w8(pb, *nal_start);
216         mov_cenc_write_encrypted(ctx, pb, nal_start + 1, nal_end - nal_start - 1);
217
218         auxiliary_info_add_subsample(ctx, 5, nal_end - nal_start - 1);
219
220         size += 4 + nal_end - nal_start;
221         nal_start = nal_end;
222     }
223
224     ret = mov_cenc_end_packet(ctx);
225     if (ret) {
226         return ret;
227     }
228
229     return size;
230 }
231
232 int ff_mov_cenc_avc_write_nal_units(AVFormatContext *s, MOVMuxCencContext* ctx,
233     int nal_length_size, AVIOContext *pb, const uint8_t *buf_in, int size)
234 {
235     int nalsize;
236     int ret;
237     int j;
238
239     ret = mov_cenc_start_packet(ctx);
240     if (ret) {
241         return ret;
242     }
243
244     while (size > 0) {
245         /* parse the nal size */
246         if (size < nal_length_size + 1) {
247             av_log(s, AV_LOG_ERROR, "CENC-AVC: remaining size %d smaller than nal length+type %d\n",
248                 size, nal_length_size + 1);
249             return -1;
250         }
251
252         avio_write(pb, buf_in, nal_length_size + 1);
253
254         nalsize = 0;
255         for (j = 0; j < nal_length_size; j++) {
256             nalsize = (nalsize << 8) | *buf_in++;
257         }
258         size -= nal_length_size;
259
260         /* encrypt the nal body */
261         if (nalsize <= 0 || nalsize > size) {
262             av_log(s, AV_LOG_ERROR, "CENC-AVC: nal size %d remaining %d\n", nalsize, size);
263             return -1;
264         }
265
266         mov_cenc_write_encrypted(ctx, pb, buf_in + 1, nalsize - 1);
267         buf_in += nalsize;
268         size -= nalsize;
269
270         auxiliary_info_add_subsample(ctx, nal_length_size + 1, nalsize - 1);
271     }
272
273     ret = mov_cenc_end_packet(ctx);
274     if (ret) {
275         return ret;
276     }
277
278     return 0;
279 }
280
281 /* TODO: reuse this function from movenc.c */
282 static int64_t update_size(AVIOContext *pb, int64_t pos)
283 {
284     int64_t curpos = avio_tell(pb);
285     avio_seek(pb, pos, SEEK_SET);
286     avio_wb32(pb, curpos - pos); /* rewrite size */
287     avio_seek(pb, curpos, SEEK_SET);
288
289     return curpos - pos;
290 }
291
292 static int mov_cenc_write_senc_tag(MOVMuxCencContext* ctx, AVIOContext *pb,
293                                    int64_t* auxiliary_info_offset)
294 {
295     int64_t pos = avio_tell(pb);
296
297     avio_wb32(pb, 0); /* size */
298     ffio_wfourcc(pb, "senc");
299     avio_wb32(pb, ctx->use_subsamples ? 0x02 : 0); /* version & flags */
300     avio_wb32(pb, ctx->auxiliary_info_entries); /* entry count */
301     *auxiliary_info_offset = avio_tell(pb);
302     avio_write(pb, ctx->auxiliary_info, ctx->auxiliary_info_size);
303     return update_size(pb, pos);
304 }
305
306 static int mov_cenc_write_saio_tag(AVIOContext *pb, int64_t auxiliary_info_offset)
307 {
308     int64_t pos = avio_tell(pb);
309     uint8_t version;
310
311     avio_wb32(pb, 0); /* size */
312     ffio_wfourcc(pb, "saio");
313     version = auxiliary_info_offset > 0xffffffff ? 1 : 0;
314     avio_w8(pb, version);
315     avio_wb24(pb, 0); /* flags */
316     avio_wb32(pb, 1); /* entry count */
317     if (version) {
318         avio_wb64(pb, auxiliary_info_offset);
319     } else {
320         avio_wb32(pb, auxiliary_info_offset);
321     }
322     return update_size(pb, pos);
323 }
324
325 static int mov_cenc_write_saiz_tag(MOVMuxCencContext* ctx, AVIOContext *pb)
326 {
327     int64_t pos = avio_tell(pb);
328     avio_wb32(pb, 0); /* size */
329     ffio_wfourcc(pb, "saiz");
330     avio_wb32(pb, 0); /* version & flags */
331     avio_w8(pb, ctx->use_subsamples ? 0 : AES_CTR_IV_SIZE);    /* default size*/
332     avio_wb32(pb, ctx->auxiliary_info_entries); /* entry count */
333     if (ctx->use_subsamples) {
334         avio_write(pb, ctx->auxiliary_info_sizes, ctx->auxiliary_info_entries);
335     }
336     return update_size(pb, pos);
337 }
338
339 void ff_mov_cenc_write_stbl_atoms(MOVMuxCencContext* ctx, AVIOContext *pb)
340 {
341     int64_t auxiliary_info_offset;
342
343     mov_cenc_write_senc_tag(ctx, pb, &auxiliary_info_offset);
344     mov_cenc_write_saio_tag(pb, auxiliary_info_offset);
345     mov_cenc_write_saiz_tag(ctx, pb);
346 }
347
348 static int mov_cenc_write_schi_tag(AVIOContext *pb, uint8_t* kid)
349 {
350     int64_t pos = avio_tell(pb);
351     avio_wb32(pb, 0);     /* size */
352     ffio_wfourcc(pb, "schi");
353
354     avio_wb32(pb, 32);    /* size */
355     ffio_wfourcc(pb, "tenc");
356     avio_wb32(pb, 0);     /* version & flags */
357     avio_wb24(pb, 1);     /* is encrypted */
358     avio_w8(pb, AES_CTR_IV_SIZE); /* iv size */
359     avio_write(pb, kid, CENC_KID_SIZE);
360
361     return update_size(pb, pos);
362 }
363
364 int ff_mov_cenc_write_sinf_tag(MOVTrack* track, AVIOContext *pb, uint8_t* kid)
365 {
366     int64_t pos = avio_tell(pb);
367     avio_wb32(pb, 0); /* size */
368     ffio_wfourcc(pb, "sinf");
369
370     /* frma */
371     avio_wb32(pb, 12);    /* size */
372     ffio_wfourcc(pb, "frma");
373     avio_wl32(pb, track->tag);
374
375     /* schm */
376     avio_wb32(pb, 20);    /* size */
377     ffio_wfourcc(pb, "schm");
378     avio_wb32(pb, 0); /* version & flags */
379     ffio_wfourcc(pb, "cenc");    /* scheme type*/
380     avio_wb32(pb, 0x10000); /* scheme version */
381
382     /* schi */
383     mov_cenc_write_schi_tag(pb, kid);
384
385     return update_size(pb, pos);
386 }
387
388 int ff_mov_cenc_init(MOVMuxCencContext* ctx, uint8_t* encryption_key,
389                      int use_subsamples, int bitexact)
390 {
391     int ret;
392
393     ctx->aes_ctr = av_aes_ctr_alloc();
394     if (!ctx->aes_ctr) {
395         return AVERROR(ENOMEM);
396     }
397
398     ret = av_aes_ctr_init(ctx->aes_ctr, encryption_key);
399     if (ret != 0) {
400         return ret;
401     }
402
403     if (!bitexact) {
404         av_aes_ctr_set_random_iv(ctx->aes_ctr);
405     }
406
407     ctx->use_subsamples = use_subsamples;
408
409     return 0;
410 }
411
412 void ff_mov_cenc_free(MOVMuxCencContext* ctx)
413 {
414     av_aes_ctr_free(ctx->aes_ctr);
415     av_freep(&ctx->auxiliary_info);
416     av_freep(&ctx->auxiliary_info_sizes);
417 }