avformat/iamf_writer: update extradata from packet side data

Some encoders, like flac, propagate updated extradata at the end of encoding
as packet side data. Use it to update the relevant codec_config.

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

7 files changed, 68 insertions, 36 deletions

diff --git a/libavformat/iamf.h b/libavformat/iamf.h
index d8715b7333..d662f6e76b 100644
--- a/libavformat/iamf.h
+++ b/libavformat/iamf.h
@@ -165,6 +165,19 @@ struct IAMFSoundSystemMap {
 extern const AVChannelLayout ff_iamf_scalable_ch_layouts[10];
 extern const struct IAMFSoundSystemMap ff_iamf_sound_system_map[13];
 
+static inline IAMFCodecConfig *ff_iamf_get_codec_config(const IAMFContext *c,
+                                                        unsigned int codec_config_id)
+{
+    IAMFCodecConfig *codec_config = NULL;
+
+    for (int i = 0; i < c->nb_codec_configs; i++) {
+        if (c->codec_configs[i]->codec_config_id == codec_config_id)
+            codec_config = c->codec_configs[i];
+    }
+
+    return codec_config;
+}
+
 static inline IAMFParamDefinition *ff_iamf_get_param_definition(const IAMFContext *iamf,
                                                                 unsigned int parameter_id)
 {
diff --git a/libavformat/iamf_parse.c b/libavformat/iamf_parse.c
index 50b53aa5e9..cb49cf0a57 100644
--- a/libavformat/iamf_parse.c
+++ b/libavformat/iamf_parse.c
@@ -580,16 +580,6 @@ static int param_parse(void *s, IAMFContext *c, AVIOContext *pb,
     return 0;
 }
 
-static IAMFCodecConfig *get_codec_config(IAMFContext *c, unsigned int codec_config_id)
-{
-    for (int i = 0; i < c->nb_codec_configs; i++) {
-        if (c->codec_configs[i]->codec_config_id == codec_config_id)
-            return c->codec_configs[i];
-    }
-
-    return NULL;
-}
-
 static int audio_element_obu(void *s, IAMFContext *c, AVIOContext *pb, int len)
 {
     const IAMFCodecConfig *codec_config;
@@ -627,7 +617,7 @@ static int audio_element_obu(void *s, IAMFContext *c, AVIOContext *pb, int len)
     audio_element_type = avio_r8(pbc) >> 5;
     codec_config_id = ffio_read_leb(pbc);
 
-    codec_config = get_codec_config(c, codec_config_id);
+    codec_config = ff_iamf_get_codec_config(c, codec_config_id);
     if (!codec_config) {
         av_log(s, AV_LOG_ERROR, "Non existant codec config id %d referenced in an audio element\n", codec_config_id);
         ret = AVERROR_INVALIDDATA;
diff --git a/libavformat/iamf_writer.c b/libavformat/iamf_writer.c
index 18dc0e284b..402cb7f230 100644
--- a/libavformat/iamf_writer.c
+++ b/libavformat/iamf_writer.c
@@ -1013,6 +1013,21 @@ int ff_iamf_write_parameter_blocks(const IAMFContext *iamf, AVIOContext *pb,
     return 0;
 }
 
+static IAMFAudioElement *get_audio_element(const IAMFContext *c,
+                                           unsigned int audio_substream_id)
+{
+    for (int i = 0; i < c->nb_audio_elements; i++) {
+        IAMFAudioElement *audio_element = c->audio_elements[i];
+        for (int j = 0; j < audio_element->nb_substreams; j++) {
+            IAMFSubStream *substream = &audio_element->substreams[j];
+            if (substream->audio_substream_id == audio_substream_id)
+                return audio_element;
+        }
+    }
+
+    return NULL;
+}
+
 int ff_iamf_write_audio_frame(const IAMFContext *iamf, AVIOContext *pb,
                               unsigned audio_substream_id, const AVPacket *pkt)
 {
@@ -1027,15 +1042,29 @@ int ff_iamf_write_audio_frame(const IAMFContext *iamf, AVIOContext *pb,
     int ret;
 
     if (!pkt->size) {
+        IAMFAudioElement *audio_element;
+        IAMFCodecConfig *codec_config;
+        size_t new_extradata_size;
         uint8_t *new_extradata = av_packet_get_side_data(pkt,
                                                          AV_PKT_DATA_NEW_EXTRADATA,
-                                                         NULL);
+                                                         &new_extradata_size);
 
         if (!new_extradata)
             return AVERROR_INVALIDDATA;
+        audio_element = get_audio_element(iamf, audio_substream_id);
+        if (!audio_element)
+            return AVERROR(EINVAL);
+        codec_config = ff_iamf_get_codec_config(iamf, audio_element->codec_config_id);
+        if (!codec_config)
+            return AVERROR(EINVAL);
 
-        // TODO: update FLAC Streaminfo on seekable output
-        return 0;
+        av_free(codec_config->extradata);
+        codec_config->extradata = av_memdup(new_extradata, new_extradata_size);
+        if (!codec_config->extradata)
+            return AVERROR(ENOMEM);
+        codec_config->extradata_size = new_extradata_size;
+
+        return update_extradata(codec_config);
     }
 
     side_data = av_packet_get_side_data(pkt, AV_PKT_DATA_SKIP_SAMPLES,
diff --git a/tests/ref/fate/mov-mp4-iamf-5_1_4 b/tests/ref/fate/mov-mp4-iamf-5_1_4
index 2f29a83cf4..1b6db58ee2 100644
--- a/tests/ref/fate/mov-mp4-iamf-5_1_4
+++ b/tests/ref/fate/mov-mp4-iamf-5_1_4
@@ -1,11 +1,11 @@
-65421714a3fd372d2babc9c6400ff00b *tests/data/fate/mov-mp4-iamf-5_1_4.mp4
+5585ed23481b6f28437b3707a1ed632d *tests/data/fate/mov-mp4-iamf-5_1_4.mp4
 86340 tests/data/fate/mov-mp4-iamf-5_1_4.mp4
-#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/mov-mp4-iamf-7_1_4 b/tests/ref/fate/mov-mp4-iamf-7_1_4
index 891803470a..1b88d362bd 100644
--- a/tests/ref/fate/mov-mp4-iamf-7_1_4
+++ b/tests/ref/fate/mov-mp4-iamf-7_1_4
@@ -1,12 +1,12 @@
-9a0cd5cc7ca09da36f08c6ae2aa44a73 *tests/data/fate/mov-mp4-iamf-7_1_4.mp4
+690d2b7a15b5489c59a9148fcd7975be *tests/data/fate/mov-mp4-iamf-7_1_4.mp4
 100588 tests/data/fate/mov-mp4-iamf-7_1_4.mp4
-#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/mov-mp4-iamf-ambisonic_1 b/tests/ref/fate/mov-mp4-iamf-ambisonic_1
index c4af88ce10..30c8edb14c 100644
--- a/tests/ref/fate/mov-mp4-iamf-ambisonic_1
+++ b/tests/ref/fate/mov-mp4-iamf-ambisonic_1
@@ -1,9 +1,9 @@
-fa740a4e2b84453c4e84908190094e28 *tests/data/fate/mov-mp4-iamf-ambisonic_1.mp4
+2b3517591f7bf20e0f74f3ec1381af1e *tests/data/fate/mov-mp4-iamf-ambisonic_1.mp4
 57743 tests/data/fate/mov-mp4-iamf-ambisonic_1.mp4
-#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/mov-mp4-iamf-stereo b/tests/ref/fate/mov-mp4-iamf-stereo
index 5c66c3e188..2fdc31776b 100644
--- a/tests/ref/fate/mov-mp4-iamf-stereo
+++ b/tests/ref/fate/mov-mp4-iamf-stereo
@@ -1,6 +1,6 @@
-7df5cc61ee480e558389051a83040dd2 *tests/data/fate/mov-mp4-iamf-stereo.mp4
+88c2b547f069f2d4a11d24f7f922251a *tests/data/fate/mov-mp4-iamf-stereo.mp4
 15163 tests/data/fate/mov-mp4-iamf-stereo.mp4
-#extradata 0:       34, 0x40a802c6
+#extradata 0:       34, 0xafa70d5e
 #tb 0: 1/44100
 #media_type 0: audio
 #codec_id 0: flac