hlsenc: Add encryption support

Partially based on Christian Suloway <csuloway@globaleagleent.com> work.

1 files changed, 198 insertions, 18 deletions

diff --git a/libavformat/hlsenc.c b/libavformat/hlsenc.c
index 96b780fe37..5eebec8511 100644
--- a/libavformat/hlsenc.c
+++ b/libavformat/hlsenc.c
@@ -22,10 +22,20 @@
 #include <float.h>
 #include <stdint.h>
 
+#include <config.h>
+
+#if CONFIG_GCRYPT
+#include <gcrypt.h>
+#elif CONFIG_OPENSSL
+#include <openssl/rand.h>
+#endif
+
 #include "libavutil/mathematics.h"
 #include "libavutil/parseutils.h"
 #include "libavutil/avstring.h"
+#include "libavutil/intreadwrite.h"
 #include "libavutil/opt.h"
+#include "libavutil/random_seed.h"
 #include "libavutil/log.h"
 
 #include "avformat.h"
@@ -60,8 +70,112 @@ typedef struct HLSContext {
     ListEntry *end_list;
     char *basename;
     char *baseurl;
+
+    int encrypt;           // Set by a private option.
+    char *key;             // Set by a private option.
+    int key_len;
+    char *key_url;         // Set by a private option.
+    char *iv;              // Set by a private option.
+    int iv_len;
+
+    char *key_basename;
+
+    AVDictionary *enc_opts;
 } HLSContext;
 
+
+static int randomize(uint8_t *buf, int len)
+{
+#if CONFIG_GCRYPT
+    gcry_randomize(buf, len, GCRY_VERY_STRONG_RANDOM);
+    return 0;
+#elif CONFIG_OPENSSL
+    if (RAND_bytes(buf, len))
+        return 0;
+#else
+    return AVERROR(ENOSYS);
+#endif
+}
+
+static void free_encryption(AVFormatContext *s)
+{
+    HLSContext *hls = s->priv_data;
+
+    av_dict_free(&hls->enc_opts);
+
+    av_freep(&hls->key_basename);
+}
+
+static int dict_set_bin(AVDictionary **dict, const char *key, uint8_t *buf)
+{
+    char hex[33];
+
+    ff_data_to_hex(hex, buf, sizeof(buf), 0);
+    hex[32] = '\0';
+
+    return av_dict_set(dict, key, hex, 0);
+}
+
+static int setup_encryption(AVFormatContext *s)
+{
+    HLSContext *hls = s->priv_data;
+    AVIOContext *out = NULL;
+    int len, ret;
+    uint8_t buf[16];
+    uint8_t *k;
+
+    len = strlen(hls->basename) + 4 + 1;
+    hls->key_basename = av_mallocz(len);
+    if (!hls->key_basename)
+        return AVERROR(ENOMEM);
+
+    av_strlcpy(hls->key_basename, hls->basename + 7, len);
+    av_strlcat(hls->key_basename, ".key", len);
+
+    if (hls->key) {
+        if (hls->key_len != 16) {
+            av_log(s, AV_LOG_ERROR,
+                   "Invalid key size %d, expected 16-bytes hex-coded key\n",
+                   hls->key_len);
+            return AVERROR(EINVAL);
+        }
+
+        if ((ret = dict_set_bin(&hls->enc_opts, "key", hls->key)) < 0)
+            return ret;
+        k = hls->key;
+    } else {
+        if ((ret = randomize(buf, sizeof(buf))) < 0) {
+            av_log(s, AV_LOG_ERROR, "Cannot generate a strong random key\n");
+            return ret;
+        }
+
+        if ((ret = dict_set_bin(&hls->enc_opts, "key", buf)) < 0)
+            return ret;
+        k = buf;
+    }
+
+    if (hls->iv) {
+        if (hls->iv_len != 16) {
+            av_log(s, AV_LOG_ERROR,
+                   "Invalid key size %d, expected 16-bytes hex-coded initialization vector\n",
+                   hls->iv_len);
+            return AVERROR(EINVAL);
+        }
+
+        if ((ret = dict_set_bin(&hls->enc_opts, "iv", hls->iv)) < 0)
+            return ret;
+    }
+
+    if ((ret = s->io_open(s, &out, hls->key_basename, AVIO_FLAG_WRITE, NULL)) < 0)
+        return ret;
+
+    avio_write(out, k, 16);
+
+    avio_close(out);
+
+    return 0;
+}
+
 static int hls_mux_init(AVFormatContext *s)
 {
     HLSContext *hls = s->priv_data;
@@ -165,6 +279,24 @@ static int hls_window(AVFormatContext *s, int last)
            sequence);
 
     for (en = hls->list; en; en = en->next) {
+        if (hls->encrypt) {
+            char *key_url;
+
+            if (hls->key_url)
+                key_url = hls->key_url;
+            else
+                key_url = hls->baseurl;
+
+            avio_printf(out, "#EXT-X-KEY:METHOD=AES-128");
+            avio_printf(out, ",URI=\"");
+            if (key_url)
+                avio_printf(out, "%s", key_url);
+            avio_printf(out, "%s\"", av_basename(hls->key_basename));
+            if (hls->iv)
+                avio_printf(out, ",IV=\"0x%s\"", hls->iv);
+            avio_printf(out, "\n");
+        }
+
         if (hls->version > 2)
             avio_printf(out, "#EXTINF:%f\n",
                         (double)en->duration / AV_TIME_BASE);
@@ -191,18 +323,75 @@ static int hls_start(AVFormatContext *s)
     HLSContext *c = s->priv_data;
     AVFormatContext *oc = c->avf;
     int err = 0;
+    AVDictionary *opts = NULL;
+
 
     if (av_get_frame_filename(oc->filename, sizeof(oc->filename),
                               c->basename, c->wrap ? c->sequence % c->wrap : c->sequence) < 0)
         return AVERROR(EINVAL);
     c->number++;
 
-    if ((err = s->io_open(s, &oc->pb, oc->filename, AVIO_FLAG_WRITE, NULL)) < 0)
+    if (c->encrypt) {
+        if ((err = av_dict_copy(&opts, c->enc_opts, 0)) < 0)
+            return err;
+        if (!c->iv) {
+            uint8_t iv[16] = { 0 };
+            char buf[33];
+
+            AV_WB64(iv + 8, c->sequence);
+            ff_data_to_hex(buf, iv, sizeof(iv), 0);
+            buf[32] = '\0';
+
+            if ((err = av_dict_set(&opts, "iv", buf, 0)) < 0)
+                goto fail;
+        }
+    }
+
+    if ((err = s->io_open(s, &oc->pb, oc->filename, AVIO_FLAG_WRITE, &opts)) < 0)
         return err;
 
     if (oc->oformat->priv_class && oc->priv_data)
         av_opt_set(oc->priv_data, "mpegts_flags", "resend_headers", 0);
 
+fail:
+    av_dict_free(&opts);
+
+    return err;
+}
+
+static int hls_setup(AVFormatContext *s)
+{
+    HLSContext *hls = s->priv_data;
+    const char *pattern = "%d.ts";
+    int basename_size = strlen(s->filename) + strlen(pattern) + 1;
+    char *p;
+
+    if (hls->encrypt)
+        basename_size += 7;
+
+    hls->basename = av_mallocz(basename_size);
+    if (!hls->basename)
+        return AVERROR(ENOMEM);
+
+    // TODO: support protocol nesting?
+    if (hls->encrypt)
+        strcpy(hls->basename, "crypto:");
+
+    av_strlcat(hls->basename, s->filename, basename_size);
+
+    p = strrchr(hls->basename, '.');
+
+    if (p)
+        *p = '\0';
+
+    if (hls->encrypt) {
+        int ret = setup_encryption(s);
+        if (ret < 0)
+            return ret;
+    }
+
+    av_strlcat(hls->basename, pattern, basename_size);
+
     return 0;
 }
 
@@ -210,9 +399,6 @@ static int hls_write_header(AVFormatContext *s)
 {
     HLSContext *hls = s->priv_data;
     int ret, i;
-    char *p;
-    const char *pattern = "%d.ts";
-    int basename_size = strlen(s->filename) + strlen(pattern) + 1;
 
     hls->sequence       = hls->start_sequence;
     hls->recording_time = hls->time * AV_TIME_BASE;
@@ -234,21 +420,8 @@ static int hls_write_header(AVFormatContext *s)
         goto fail;
     }
 
-    hls->basename = av_malloc(basename_size);
-
-    if (!hls->basename) {
-        ret = AVERROR(ENOMEM);
+    if ((ret = hls_setup(s)) < 0)
         goto fail;
-    }
-
-    strcpy(hls->basename, s->filename);
-
-    p = strrchr(hls->basename, '.');
-
-    if (p)
-        *p = '\0';
-
-    av_strlcat(hls->basename, pattern, basename_size);
 
     if ((ret = hls_mux_init(s)) < 0)
         goto fail;
@@ -265,6 +438,8 @@ fail:
         av_free(hls->basename);
         if (hls->avf)
             avformat_free_context(hls->avf);
+
+        free_encryption(s);
     }
     return ret;
 }
@@ -332,6 +507,7 @@ static int hls_write_trailer(struct AVFormatContext *s)
     hls_window(s, 1);
 
     free_entries(hls);
+    free_encryption(s);
     return 0;
 }
 
@@ -345,6 +521,10 @@ static const AVOption options[] = {
     {"hls_allow_cache", "explicitly set whether the client MAY (1) or MUST NOT (0) cache media segments", OFFSET(allowcache), AV_OPT_TYPE_INT, {.i64 = -1}, INT_MIN, INT_MAX, E},
     {"hls_base_url",  "url to prepend to each playlist entry",   OFFSET(baseurl), AV_OPT_TYPE_STRING, {.str = NULL},  0, 0,       E},
     {"hls_version",   "protocol version",                        OFFSET(version), AV_OPT_TYPE_INT,    {.i64 = 3},     2, 3, E},
+    {"hls_enc",       "AES128 encryption support",               OFFSET(encrypt), AV_OPT_TYPE_INT,    {.i64 = 0},     0, 1, E},
+    {"hls_enc_key",   "use the specified hex-coded 16byte key to encrypt the segments",  OFFSET(key), AV_OPT_TYPE_BINARY, .flags = E},
+    {"hls_enc_key_url", "url to access the key to decrypt the segments",    OFFSET(key_url), AV_OPT_TYPE_STRING, {.str = NULL},  0, 0, E},
+    {"hls_enc_iv",     "use the specified hex-coded 16byte initialization vector",  OFFSET(iv), AV_OPT_TYPE_BINARY, .flags = E},
     { NULL },
 };