avformat/iamfenc: ensure updated extradata is written

Signed-off-by: James Almer <jamrial@gmail.com>

5 files changed, 51 insertions, 23 deletions

diff --git a/libavformat/iamfenc.c b/libavformat/iamfenc.c
index 72a7555c9a..ebf0c2fd67 100644
--- a/libavformat/iamfenc.c
+++ b/libavformat/iamfenc.c
@@ -35,6 +35,9 @@
 typedef struct IAMFMuxContext {
     IAMFContext iamf;
 
+    int64_t descriptors_offset;
+    int update_extradata;
+
     int first_stream_id;
 } IAMFMuxContext;
 
@@ -124,6 +127,7 @@ static int iamf_write_header(AVFormatContext *s)
     IAMFContext *const iamf = &c->iamf;
     int ret;
 
+    c->descriptors_offset = avio_tell(s->pb);
     ret = ff_iamf_write_descriptors(iamf, s->pb, s);
     if (ret < 0)
         return ret;
@@ -135,7 +139,7 @@ static int iamf_write_header(AVFormatContext *s)
 
 static int iamf_write_packet(AVFormatContext *s, AVPacket *pkt)
 {
-    const IAMFMuxContext *const c = s->priv_data;
+    IAMFMuxContext *const c = s->priv_data;
     AVStream *st = s->streams[pkt->stream_index];
     int ret = 0;
 
@@ -143,10 +147,33 @@ static int iamf_write_packet(AVFormatContext *s, AVPacket *pkt)
         ret = ff_iamf_write_parameter_blocks(&c->iamf, s->pb, pkt, s);
     if (!ret)
         ret = ff_iamf_write_audio_frame(&c->iamf, s->pb, st->id, pkt);
+    if (!ret && !pkt->size)
+        c->update_extradata = 1;
 
     return ret;
 }
 
+static int iamf_write_trailer(AVFormatContext *s)
+{
+    const IAMFMuxContext *const c = s->priv_data;
+    const IAMFContext *const iamf = &c->iamf;
+    int64_t pos;
+    int ret;
+
+    if (!c->update_extradata || !(s->pb->seekable & AVIO_SEEKABLE_NORMAL))
+        return 0;
+
+    pos = avio_tell(s->pb);
+    avio_seek(s->pb, c->descriptors_offset, SEEK_SET);
+    ret = ff_iamf_write_descriptors(iamf, s->pb, s);
+    if (ret < 0)
+        return ret;
+
+    avio_seek(s->pb, pos, SEEK_SET);
+
+    return 0;
+}
+
 static void iamf_deinit(AVFormatContext *s)
 {
     IAMFMuxContext *const c = s->priv_data;
@@ -178,6 +205,7 @@ const FFOutputFormat ff_iamf_muxer = {
     .deinit            = iamf_deinit,
     .write_header      = iamf_write_header,
     .write_packet      = iamf_write_packet,
+    .write_trailer     = iamf_write_trailer,
     .p.codec_tag       = (const AVCodecTag* const []){ iamf_codec_tags, NULL },
     .p.flags           = AVFMT_GLOBALHEADER | AVFMT_NOTIMESTAMPS,
 };
diff --git a/tests/ref/fate/iamf-5_1_4 b/tests/ref/fate/iamf-5_1_4
index e6eb356ff0..030765ecee 100644
--- a/tests/ref/fate/iamf-5_1_4
+++ b/tests/ref/fate/iamf-5_1_4
@@ -1,11 +1,11 @@
-c447cbbc8943cfb751fdf1145a094250 *tests/data/fate/iamf-5_1_4.iamf
+7749af0b8f37c91e4cc0dbc6f31d7f67 *tests/data/fate/iamf-5_1_4.iamf
 85603 tests/data/fate/iamf-5_1_4.iamf
-#extradata 0:       34, 0x40a802c6
-#extradata 1:       34, 0x40a802c6
-#extradata 2:       34, 0x407c02c4
-#extradata 3:       34, 0x407c02c4
-#extradata 4:       34, 0x40a802c6
-#extradata 5:       34, 0x40a802c6
+#extradata 0:       34, 0xafa70d5e
+#extradata 1:       34, 0xafa70d5e
+#extradata 2:       34, 0xaf7b0d5c
+#extradata 3:       34, 0xaf7b0d5c
+#extradata 4:       34, 0xafa70d5e
+#extradata 5:       34, 0xafa70d5e
 #tb 0: 1/44100
 #media_type 0: audio
 #codec_id 0: flac
diff --git a/tests/ref/fate/iamf-7_1_4 b/tests/ref/fate/iamf-7_1_4
index c176f4a1cd..39cbd6d840 100644
--- a/tests/ref/fate/iamf-7_1_4
+++ b/tests/ref/fate/iamf-7_1_4
@@ -1,12 +1,12 @@
-157c3185684e12cc8385ee7c3ef2fb4c *tests/data/fate/iamf-7_1_4.iamf
+5d9fcee2b9b2ad3c802c40bb1147016e *tests/data/fate/iamf-7_1_4.iamf
 99851 tests/data/fate/iamf-7_1_4.iamf
-#extradata 0:       34, 0x40a802c6
-#extradata 1:       34, 0x40a802c6
-#extradata 2:       34, 0x407c02c4
-#extradata 3:       34, 0x407c02c4
-#extradata 4:       34, 0x40a802c6
-#extradata 5:       34, 0x40a802c6
-#extradata 6:       34, 0x40a802c6
+#extradata 0:       34, 0xafa70d5e
+#extradata 1:       34, 0xafa70d5e
+#extradata 2:       34, 0xaf7b0d5c
+#extradata 3:       34, 0xaf7b0d5c
+#extradata 4:       34, 0xafa70d5e
+#extradata 5:       34, 0xafa70d5e
+#extradata 6:       34, 0xafa70d5e
 #tb 0: 1/44100
 #media_type 0: audio
 #codec_id 0: flac
diff --git a/tests/ref/fate/iamf-ambisonic_1 b/tests/ref/fate/iamf-ambisonic_1
index 928a9c6cf2..ed73dd4d09 100644
--- a/tests/ref/fate/iamf-ambisonic_1
+++ b/tests/ref/fate/iamf-ambisonic_1
@@ -1,9 +1,9 @@
-0b66877e65e3e5bae46887aced977593 *tests/data/fate/iamf-ambisonic_1.iamf
+c3751e6a2d85bd6a8318fe93b5754309 *tests/data/fate/iamf-ambisonic_1.iamf
 57006 tests/data/fate/iamf-ambisonic_1.iamf
-#extradata 0:       34, 0x3615025b
-#extradata 1:       34, 0x3615025b
-#extradata 2:       34, 0x3615025b
-#extradata 3:       34, 0x3615025b
+#extradata 0:       34, 0xad120cfe
+#extradata 1:       34, 0xad120cfe
+#extradata 2:       34, 0xad120cfe
+#extradata 3:       34, 0xad120cfe
 #tb 0: 1/44100
 #media_type 0: audio
 #codec_id 0: flac
diff --git a/tests/ref/fate/iamf-stereo b/tests/ref/fate/iamf-stereo
index 65d6b506d4..82206f84bf 100644
--- a/tests/ref/fate/iamf-stereo
+++ b/tests/ref/fate/iamf-stereo
@@ -1,6 +1,6 @@
-ace731a4fbc302e24498d6b64daa16e7 *tests/data/fate/iamf-stereo.iamf
+4000d96f7363212c5a75c6f7fa5756f5 *tests/data/fate/iamf-stereo.iamf
 14426 tests/data/fate/iamf-stereo.iamf
-#extradata 0:       34, 0x40a802c6
+#extradata 0:       34, 0xafa70d5e
 #tb 0: 1/44100
 #media_type 0: audio
 #codec_id 0: flac