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 "encryption_info.h"
20#include "mem.h"
21#include "intreadwrite.h"
22
23#define FF_ENCRYPTION_INFO_EXTRA 24
24
25// The format of the AVEncryptionInfo side data:
26// u32be scheme
27// u32be crypt_byte_block
28// u32be skip_byte_block
29// u32be key_id_size
30// u32be iv_size
31// u32be subsample_count
32// u8[key_id_size] key_id
33// u8[iv_size] iv
34// {
35//   u32be bytes_of_clear_data
36//   u32be bytes_of_protected_data
37// }[subsample_count]
38
39AVEncryptionInfo *av_encryption_info_alloc(uint32_t subsample_count, uint32_t key_id_size, uint32_t iv_size)
40{
41    AVEncryptionInfo *info;
42
43    info = av_mallocz(sizeof(*info));
44    if (!info)
45        return NULL;
46
47    info->key_id = av_mallocz(key_id_size);
48    info->key_id_size = key_id_size;
49    info->iv = av_mallocz(iv_size);
50    info->iv_size = iv_size;
51    info->subsamples = av_calloc(subsample_count, sizeof(*info->subsamples));
52    info->subsample_count = subsample_count;
53
54    // Allow info->subsamples to be NULL if there are no subsamples.
55    if (!info->key_id || !info->iv || (!info->subsamples && subsample_count)) {
56        av_encryption_info_free(info);
57        return NULL;
58    }
59
60    return info;
61}
62
63AVEncryptionInfo *av_encryption_info_clone(const AVEncryptionInfo *info)
64{
65    AVEncryptionInfo *ret;
66
67    ret = av_encryption_info_alloc(info->subsample_count, info->key_id_size, info->iv_size);
68    if (!ret)
69        return NULL;
70
71    ret->scheme = info->scheme;
72    ret->crypt_byte_block = info->crypt_byte_block;
73    ret->skip_byte_block = info->skip_byte_block;
74    memcpy(ret->iv, info->iv, info->iv_size);
75    memcpy(ret->key_id, info->key_id, info->key_id_size);
76    memcpy(ret->subsamples, info->subsamples, sizeof(*info->subsamples) * info->subsample_count);
77    return ret;
78}
79
80void av_encryption_info_free(AVEncryptionInfo *info)
81{
82    if (info) {
83        av_free(info->key_id);
84        av_free(info->iv);
85        av_free(info->subsamples);
86        av_free(info);
87    }
88}
89
90AVEncryptionInfo *av_encryption_info_get_side_data(const uint8_t* buffer, size_t size)
91{
92    AVEncryptionInfo *info;
93    uint64_t key_id_size, iv_size, subsample_count, i;
94
95    if (!buffer || size < FF_ENCRYPTION_INFO_EXTRA)
96        return NULL;
97
98    key_id_size = AV_RB32(buffer + 12);
99    iv_size = AV_RB32(buffer + 16);
100    subsample_count = AV_RB32(buffer + 20);
101
102    if (size < FF_ENCRYPTION_INFO_EXTRA + key_id_size + iv_size + subsample_count * 8)
103        return NULL;
104
105    info = av_encryption_info_alloc(subsample_count, key_id_size, iv_size);
106    if (!info)
107        return NULL;
108
109    info->scheme = AV_RB32(buffer);
110    info->crypt_byte_block = AV_RB32(buffer + 4);
111    info->skip_byte_block = AV_RB32(buffer + 8);
112    memcpy(info->key_id, buffer + 24, key_id_size);
113    memcpy(info->iv, buffer + key_id_size + 24, iv_size);
114
115    buffer += key_id_size + iv_size + 24;
116    for (i = 0; i < subsample_count; i++) {
117        info->subsamples[i].bytes_of_clear_data = AV_RB32(buffer);
118        info->subsamples[i].bytes_of_protected_data = AV_RB32(buffer + 4);
119        buffer += 8;
120    }
121
122    return info;
123}
124
125uint8_t *av_encryption_info_add_side_data(const AVEncryptionInfo *info, size_t *size)
126{
127    uint8_t *buffer, *cur_buffer;
128    uint32_t i;
129
130    if (UINT32_MAX - FF_ENCRYPTION_INFO_EXTRA < info->key_id_size ||
131        UINT32_MAX - FF_ENCRYPTION_INFO_EXTRA - info->key_id_size < info->iv_size ||
132        (UINT32_MAX - FF_ENCRYPTION_INFO_EXTRA - info->key_id_size - info->iv_size) / 8 < info->subsample_count) {
133        return NULL;
134    }
135
136    *size = FF_ENCRYPTION_INFO_EXTRA + info->key_id_size + info->iv_size +
137            (info->subsample_count * 8);
138    cur_buffer = buffer = av_malloc(*size);
139    if (!buffer)
140        return NULL;
141
142    AV_WB32(cur_buffer,      info->scheme);
143    AV_WB32(cur_buffer +  4, info->crypt_byte_block);
144    AV_WB32(cur_buffer +  8, info->skip_byte_block);
145    AV_WB32(cur_buffer + 12, info->key_id_size);
146    AV_WB32(cur_buffer + 16, info->iv_size);
147    AV_WB32(cur_buffer + 20, info->subsample_count);
148    cur_buffer += 24;
149    memcpy(cur_buffer, info->key_id, info->key_id_size);
150    cur_buffer += info->key_id_size;
151    memcpy(cur_buffer, info->iv, info->iv_size);
152    cur_buffer += info->iv_size;
153    for (i = 0; i < info->subsample_count; i++) {
154        AV_WB32(cur_buffer, info->subsamples[i].bytes_of_clear_data);
155        AV_WB32(cur_buffer + 4, info->subsamples[i].bytes_of_protected_data);
156        cur_buffer += 8;
157    }
158
159    return buffer;
160}
161
162#ifdef OHOS_DRM
163static void av_encryption_info_set_drm_algo(uint32_t algo, AV_DrmCencInfo *cenc_info)
164{
165    switch (algo) {
166        case MKBETAG('c','e','n','c'):
167        case MKBETAG('c','e','n','s'):
168            cenc_info->algo = AV_DRM_ALG_CENC_AES_CTR;
169            break;
170        case MKBETAG('c','b','c','1'):
171        case MKBETAG('c','b','c','s'):
172            cenc_info->algo = AV_DRM_ALG_CENC_AES_CBC;
173            break;
174        case MKBETAG('s','m','4','c'):
175        case MKBETAG('s','m','4','s'):
176            cenc_info->algo = AV_DRM_ALG_CENC_SM4_CBC;
177            break;
178        case MKBETAG('s','m','4','t'):
179        case MKBETAG('s','m','4','r'):
180            cenc_info->algo = AV_DRM_ALG_CENC_SM4_CTR;
181            break;
182        default:
183            cenc_info->algo = AV_DRM_ALG_CENC_UNENCRYPTED;
184            break;
185    }
186    return;
187}
188
189AV_DrmCencInfo *av_encryption_info_add_side_data_ex(const AVEncryptionInfo *info, size_t *side_data_size,
190    AV_DrmCencInfo *cenc_info, int pkt_data_size)
191{
192    uint32_t i;
193    if ((info == NULL) || (info->key_id_size != AV_DRM_KEY_ID_SIZE) || (info->iv_size > AV_DRM_IV_SIZE) ||
194        (info->iv_size == 0) || (info->subsample_count > AV_DRM_MAX_SUB_SAMPLE_NUM) || (info->key_id == NULL) ||
195        (info->iv == NULL) || (info->subsamples == NULL) || (side_data_size == NULL)) {
196        return NULL;
197    }
198
199    *side_data_size = sizeof(AV_DrmCencInfo);
200    cenc_info = av_mallocz(*side_data_size);
201    if (!cenc_info)
202        return NULL;
203
204    av_encryption_info_set_drm_algo(info->scheme, cenc_info);
205    cenc_info->key_id_len = info->key_id_size;
206    memcpy(cenc_info->key_id, info->key_id, info->key_id_size);
207    cenc_info->iv_len = info->iv_size;
208    memcpy(cenc_info->iv, info->iv, info->iv_size);
209    cenc_info->mode = AV_DRM_CENC_INFO_KEY_IV_SUBSAMPLES_SET;
210    cenc_info->encrypt_blocks = info->crypt_byte_block;
211    cenc_info->skip_blocks = info->skip_byte_block;
212    cenc_info->first_encrypt_offset = 0;
213    cenc_info->sub_sample_num = info->subsample_count;
214
215    for (i = 0; i < cenc_info->sub_sample_num; i++) {
216        cenc_info->sub_samples[i].clear_header_len = info->subsamples[i].bytes_of_clear_data;
217        cenc_info->sub_samples[i].pay_load_len = info->subsamples[i].bytes_of_protected_data;
218    }
219    if ((info->subsample_count == 0) && (info->crypt_byte_block == 0) && (info->skip_byte_block == 0) &&
220        ((info->scheme == MKBETAG('c','e','n','s')) || (info->scheme == MKBETAG('s','m','4','r')))) {
221        cenc_info->sub_sample_num = 1; // 1: sub_sample num
222        cenc_info->sub_samples[0].clear_header_len = 0;
223        cenc_info->sub_samples[0].pay_load_len = (pkt_data_size / 16) * 16; // 16: block size
224        if ((pkt_data_size % 16) != 0) { // 16: block size
225            cenc_info->sub_sample_num = 2; // 2: sub_sample num
226            cenc_info->sub_samples[1].clear_header_len = pkt_data_size % 16; // 16: block size
227            cenc_info->sub_samples[1].pay_load_len = 0;
228        }
229    }
230    return cenc_info;
231}
232#endif
233
234// The format of the AVEncryptionInitInfo side data:
235// u32be init_info_count
236// {
237//   u32be system_id_size
238//   u32be num_key_ids
239//   u32be key_id_size
240//   u32be data_size
241//   u8[system_id_size] system_id
242//   u8[key_id_size][num_key_id] key_ids
243//   u8[data_size] data
244// }[init_info_count]
245
246#define FF_ENCRYPTION_INIT_INFO_EXTRA 16
247
248AVEncryptionInitInfo *av_encryption_init_info_alloc(
249    uint32_t system_id_size, uint32_t num_key_ids, uint32_t key_id_size, uint32_t data_size)
250{
251    AVEncryptionInitInfo *info;
252    uint32_t i;
253
254    info = av_mallocz(sizeof(*info));
255    if (!info)
256        return NULL;
257
258    info->system_id = av_mallocz(system_id_size);
259    info->system_id_size = system_id_size;
260    info->key_ids = key_id_size ? av_calloc(num_key_ids, sizeof(*info->key_ids)) : NULL;
261    info->num_key_ids = num_key_ids;
262    info->key_id_size = key_id_size;
263    info->data = av_mallocz(data_size);
264    info->data_size = data_size;
265
266    // Allow pointers to be NULL if the size is 0.
267    if ((!info->system_id && system_id_size) || (!info->data && data_size) ||
268        (!info->key_ids && num_key_ids && key_id_size)) {
269        av_encryption_init_info_free(info);
270        return NULL;
271    }
272
273    if (key_id_size) {
274        for (i = 0; i < num_key_ids; i++) {
275            info->key_ids[i] = av_mallocz(key_id_size);
276            if (!info->key_ids[i]) {
277                av_encryption_init_info_free(info);
278                return NULL;
279            }
280        }
281    }
282
283    return info;
284}
285
286void av_encryption_init_info_free(AVEncryptionInitInfo *info)
287{
288    uint32_t i;
289    if (info) {
290        for (i = 0; i < info->num_key_ids; i++) {
291            av_free(info->key_ids[i]);
292        }
293        av_encryption_init_info_free(info->next);
294        av_free(info->system_id);
295        av_free(info->key_ids);
296        av_free(info->data);
297        av_free(info);
298    }
299}
300
301AVEncryptionInitInfo *av_encryption_init_info_get_side_data(
302    const uint8_t *side_data, size_t side_data_size)
303{
304    // |ret| tracks the front of the list, |info| tracks the back.
305    AVEncryptionInitInfo *ret = NULL, *info, *temp_info;
306    uint64_t system_id_size, num_key_ids, key_id_size, data_size, i, j;
307    uint64_t init_info_count;
308
309    if (!side_data || side_data_size < 4)
310        return NULL;
311
312    init_info_count = AV_RB32(side_data);
313    side_data += 4;
314    side_data_size -= 4;
315    for (i = 0; i < init_info_count; i++) {
316        if (side_data_size < FF_ENCRYPTION_INIT_INFO_EXTRA) {
317            av_encryption_init_info_free(ret);
318            return NULL;
319        }
320
321        system_id_size = AV_RB32(side_data);
322        num_key_ids = AV_RB32(side_data + 4);
323        key_id_size = AV_RB32(side_data + 8);
324        data_size = AV_RB32(side_data + 12);
325
326        // UINT32_MAX + UINT32_MAX + UINT32_MAX * UINT32_MAX == UINT64_MAX
327        if (side_data_size - FF_ENCRYPTION_INIT_INFO_EXTRA < system_id_size + data_size + num_key_ids * key_id_size) {
328            av_encryption_init_info_free(ret);
329            return NULL;
330        }
331        side_data += FF_ENCRYPTION_INIT_INFO_EXTRA;
332        side_data_size -= FF_ENCRYPTION_INIT_INFO_EXTRA;
333
334        temp_info = av_encryption_init_info_alloc(system_id_size, num_key_ids, key_id_size, data_size);
335        if (!temp_info) {
336            av_encryption_init_info_free(ret);
337            return NULL;
338        }
339        if (i == 0) {
340            info = ret = temp_info;
341        } else {
342            info->next = temp_info;
343            info = temp_info;
344        }
345
346        memcpy(info->system_id, side_data, system_id_size);
347        side_data += system_id_size;
348        side_data_size -= system_id_size;
349        for (j = 0; j < num_key_ids; j++) {
350            memcpy(info->key_ids[j], side_data, key_id_size);
351            side_data += key_id_size;
352            side_data_size -= key_id_size;
353        }
354        memcpy(info->data, side_data, data_size);
355        side_data += data_size;
356        side_data_size -= data_size;
357    }
358
359    return ret;
360}
361
362uint8_t *av_encryption_init_info_add_side_data(const AVEncryptionInitInfo *info, size_t *side_data_size)
363{
364    const AVEncryptionInitInfo *cur_info;
365    uint8_t *buffer, *cur_buffer;
366    uint32_t i, init_info_count;
367    uint64_t temp_side_data_size;
368
369    temp_side_data_size = 4;
370    init_info_count = 0;
371    for (cur_info = info; cur_info; cur_info = cur_info->next) {
372        temp_side_data_size += (uint64_t)FF_ENCRYPTION_INIT_INFO_EXTRA + cur_info->system_id_size + cur_info->data_size;
373        if (init_info_count == UINT32_MAX || temp_side_data_size > UINT32_MAX) {
374            return NULL;
375        }
376        init_info_count++;
377
378        if (cur_info->num_key_ids) {
379            temp_side_data_size += (uint64_t)cur_info->num_key_ids * cur_info->key_id_size;
380            if (temp_side_data_size > UINT32_MAX) {
381                return NULL;
382            }
383        }
384    }
385    *side_data_size = temp_side_data_size;
386
387    cur_buffer = buffer = av_malloc(*side_data_size);
388    if (!buffer)
389        return NULL;
390
391    AV_WB32(cur_buffer, init_info_count);
392    cur_buffer += 4;
393    for (cur_info = info; cur_info; cur_info = cur_info->next) {
394        AV_WB32(cur_buffer,      cur_info->system_id_size);
395        AV_WB32(cur_buffer +  4, cur_info->num_key_ids);
396        AV_WB32(cur_buffer +  8, cur_info->key_id_size);
397        AV_WB32(cur_buffer + 12, cur_info->data_size);
398        cur_buffer += 16;
399
400        memcpy(cur_buffer, cur_info->system_id, cur_info->system_id_size);
401        cur_buffer += cur_info->system_id_size;
402        for (i = 0; i < cur_info->num_key_ids; i++) {
403            memcpy(cur_buffer, cur_info->key_ids[i], cur_info->key_id_size);
404            cur_buffer += cur_info->key_id_size;
405        }
406        if (cur_info->data_size > 0) {
407            memcpy(cur_buffer, cur_info->data, cur_info->data_size);
408            cur_buffer += cur_info->data_size;
409        }
410    }
411
412    return buffer;
413}
414