[FFmpeg-devel] [PATCH] avformat/mov: add support for pssh box

[James Almer]

From: Malek Assaad <maleka at axis.com>

* Enable PSSH box writing to mov & moof.
* Enable senc, saio & saiz writing to moof when writing fragments.

Signed-off-by: James Almer <jamrial at gmail.com>
---
 libavformat/movenc.c     | 48 ++++++++++++++++++++++++++++++++++++++--
 libavformat/movenccenc.c |  5 +++--
 libavformat/movenccenc.h |  2 +-
 3 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/libavformat/movenc.c b/libavformat/movenc.c
index 010d2afa13..9cfbff002a 100644
--- a/libavformat/movenc.c
+++ b/libavformat/movenc.c
@@ -3221,8 +3221,8 @@ static int mov_write_stbl_tag(AVFormatContext *s, AVIOContext *pb, MOVMuxContext
     mov_write_stsc_tag(pb, track);
     mov_write_stsz_tag(pb, track);
     mov_write_stco_tag(pb, track);
-    if (track->cenc.aes_ctr) {
-        ff_mov_cenc_write_stbl_atoms(&track->cenc, pb);
+    if (track->cenc.aes_ctr && !(mov->flags & FF_MOV_FLAG_FRAGMENT)) {
+        ff_mov_cenc_write_stbl_atoms(&track->cenc, pb, 0);
     }
     if (track->par->codec_id == AV_CODEC_ID_OPUS || track->par->codec_id == AV_CODEC_ID_AAC) {
         mov_preroll_write_stbl_atoms(pb, track);
@@ -4899,6 +4899,43 @@ static int mov_write_uuidusmt_tag(AVIOContext *pb, AVFormatContext *s)
     return 0;
 }
 
+static int mov_write_pssh_tag(AVIOContext *pb, AVStream *st)
+{
+    AVEncryptionInitInfo *info;
+    const AVPacketSideData *sd = av_packet_side_data_get(st->codecpar->coded_side_data,
+                                                         st->codecpar->nb_coded_side_data,
+                                                         AV_PKT_DATA_ENCRYPTION_INIT_INFO);
+    if (!sd)
+        return 0;
+
+    info = av_encryption_init_info_get_side_data(sd->data, sd->size);
+    for (AVEncryptionInitInfo *copy = info; copy; copy = copy->next) {
+        int64_t pos;
+
+        if (!copy->data_size && !copy->num_key_ids)
+            continue;
+
+        pos = avio_tell(pb);
+        avio_wb32(pb, 0); /* size placeholder */
+        ffio_wfourcc(pb, "pssh");
+        avio_w8(pb, 1); /* version */
+        avio_wb24(pb, 0);
+        for (int i = 0; i < copy->system_id_size; i++)
+            avio_w8(pb, copy->system_id[i]);
+        avio_wb32(pb, copy->num_key_ids);
+        for (int i = 0; i < copy->num_key_ids; i++)
+            for (int j = 0; j < copy->key_id_size; j++)
+                avio_w8(pb, copy->key_ids[i][j]);
+        avio_wb32(pb, copy->data_size);
+        avio_write(pb, copy->data, copy->data_size);
+        update_size(pb, pos);
+    }
+
+    av_encryption_init_info_free(info);
+
+    return 0;
+}
+
 static void build_chunks(MOVTrack *trk)
 {
     int i;
@@ -5046,6 +5083,8 @@ static int mov_write_moov_tag(AVIOContext *pb, MOVMuxContext *mov,
         mov_write_uuidusmt_tag(pb, s);
     else if (mov->mode != MODE_AVIF)
         mov_write_udta_tag(pb, mov, s);
+    for (i = 0; i < mov->nb_streams; i++)
+        mov_write_pssh_tag(pb, mov->tracks[i].st);
 
     return update_size(pb, pos);
 }
@@ -5499,6 +5538,9 @@ static int mov_write_traf_tag(AVIOContext *pb, MOVMuxContext *mov,
         }
     }
 
+    if (track->cenc.aes_ctr && (mov->flags & FF_MOV_FLAG_FRAGMENT))
+        ff_mov_cenc_write_stbl_atoms(&track->cenc, pb, moof_offset);
+
     return update_size(pb, pos);
 }
 
@@ -5519,6 +5561,8 @@ static int mov_write_moof_tag_internal(AVIOContext *pb, MOVMuxContext *mov,
             continue;
         if (!track->entry)
             continue;
+        if (track->cenc.aes_ctr && (mov->flags & FF_MOV_FLAG_FRAGMENT))
+            mov_write_pssh_tag(pb, track->st);
         mov_write_traf_tag(pb, mov, track, pos, moof_size);
     }
 
diff --git a/libavformat/movenccenc.c b/libavformat/movenccenc.c
index 1b7a773808..f54d3bcbca 100644
--- a/libavformat/movenccenc.c
+++ b/libavformat/movenccenc.c
@@ -338,12 +338,13 @@ static int mov_cenc_write_saiz_tag(MOVMuxCencContext* ctx, AVIOContext *pb)
     return update_size(pb, pos);
 }
 
-void ff_mov_cenc_write_stbl_atoms(MOVMuxCencContext* ctx, AVIOContext *pb)
+void ff_mov_cenc_write_stbl_atoms(MOVMuxCencContext* ctx, AVIOContext *pb,
+                                  int64_t moof_offset)
 {
     int64_t auxiliary_info_offset;
 
     mov_cenc_write_senc_tag(ctx, pb, &auxiliary_info_offset);
-    mov_cenc_write_saio_tag(pb, auxiliary_info_offset);
+    mov_cenc_write_saio_tag(pb, auxiliary_info_offset - moof_offset);
     mov_cenc_write_saiz_tag(ctx, pb);
 }
 
diff --git a/libavformat/movenccenc.h b/libavformat/movenccenc.h
index 6f9e70e905..7da5268090 100644
--- a/libavformat/movenccenc.h
+++ b/libavformat/movenccenc.h
@@ -76,7 +76,7 @@ int ff_mov_cenc_avc_write_nal_units(AVFormatContext *s, MOVMuxCencContext* ctx,
 /**
  * Write the cenc atoms that should reside inside stbl
  */
-void ff_mov_cenc_write_stbl_atoms(MOVMuxCencContext* ctx, AVIOContext *pb);
+void ff_mov_cenc_write_stbl_atoms(MOVMuxCencContext* ctx, AVIOContext *pb, int64_t moof_offset);
 
 /**
  * Write the sinf atom, contained inside stsd
-- 
2.47.1