git.sesse.net Git - ffmpeg/blob - 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 }