From 10e29be8b65f6c5fb8e2677d89d817ab1e3dba63 Mon Sep 17 00:00:00 2001
From: sheenachhabra <sheenachhabra@google.com>
Date: Fri, 3 May 2024 02:44:21 -0700
Subject: [PATCH] Allocate correct size for keys and ilst boxes in Mp4Muxer

The keys and ilst boxes in the MP4 muxer were allocated
with a fixed size of 200 bytes. This was not enough to
store the keys and values of large metadata entries, which
could cause the muxer to throw an exception. This CL allocates
the correct size for the keys and ilst boxes based on the size of the metadata entries.

PiperOrigin-RevId: 630331680
---
 .../java/androidx/media3/muxer/Boxes.java     | 19 +++++++++++++---
 .../media3/muxer/Mp4MuxerMetadataTest.java    | 22 +++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java b/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
index 572781309b..01fc9a0beb 100644
--- a/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
+++ b/libraries/muxer/src/main/java/androidx/media3/muxer/Boxes.java
@@ -415,7 +415,12 @@ import java.util.Locale;
    * <p>This box contains a list of metadata keys.
    */
   public static ByteBuffer keys(List<MdtaMetadataEntry> mdtaMetadataEntries) {
-    ByteBuffer contents = ByteBuffer.allocate(MAX_FIXED_LEAF_BOX_SIZE);
+    int totalSizeToStoreKeys = 0;
+    for (int i = 0; i < mdtaMetadataEntries.size(); i++) {
+      // Add header size to wrap each key into a "mdta" box.
+      totalSizeToStoreKeys += mdtaMetadataEntries.get(i).key.length() + BOX_HEADER_SIZE;
+    }
+    ByteBuffer contents = ByteBuffer.allocate(2 * BYTES_PER_INTEGER + totalSizeToStoreKeys);
     contents.putInt(0x0); // version and flags
     contents.putInt(mdtaMetadataEntries.size()); // Entry count
 
@@ -434,7 +439,15 @@ import java.util.Locale;
    * <p>This box contains a list of metadata values.
    */
   public static ByteBuffer ilst(List<MdtaMetadataEntry> mdtaMetadataEntries) {
-    ByteBuffer contents = ByteBuffer.allocate(MAX_FIXED_LEAF_BOX_SIZE);
+    int totalSizeToStoreValues = 0;
+    for (int i = 0; i < mdtaMetadataEntries.size(); i++) {
+      // Add additional 16 bytes for writing metadata associated to each value.
+      // Add header size to wrap each value into a "data" box.
+      totalSizeToStoreValues +=
+          mdtaMetadataEntries.get(i).value.length + 4 * BYTES_PER_INTEGER + BOX_HEADER_SIZE;
+    }
+
+    ByteBuffer contents = ByteBuffer.allocate(totalSizeToStoreValues);
 
     for (int i = 0; i < mdtaMetadataEntries.size(); i++) {
       int keyId = i + 1;
@@ -448,7 +461,7 @@ import java.util.Locale;
 
       valueContents.flip();
       ByteBuffer valueBox = BoxUtils.wrapIntoBox("data", valueContents);
-      contents.putInt(valueBox.remaining() + 8);
+      contents.putInt(valueBox.remaining() + BOX_HEADER_SIZE);
       contents.putInt(keyId);
       contents.put(valueBox);
     }
diff --git a/libraries/muxer/src/test/java/androidx/media3/muxer/Mp4MuxerMetadataTest.java b/libraries/muxer/src/test/java/androidx/media3/muxer/Mp4MuxerMetadataTest.java
index e71cf8bfac..1419c892f3 100644
--- a/libraries/muxer/src/test/java/androidx/media3/muxer/Mp4MuxerMetadataTest.java
+++ b/libraries/muxer/src/test/java/androidx/media3/muxer/Mp4MuxerMetadataTest.java
@@ -20,6 +20,7 @@ import static androidx.media3.container.MdtaMetadataEntry.TYPE_INDICATOR_FLOAT32
 import static androidx.media3.container.MdtaMetadataEntry.TYPE_INDICATOR_STRING;
 import static androidx.media3.muxer.MuxerTestUtil.FAKE_VIDEO_FORMAT;
 import static androidx.media3.muxer.MuxerTestUtil.XMP_SAMPLE_DATA;
+import static com.google.common.truth.Truth.assertThat;
 
 import android.content.Context;
 import android.media.MediaCodec.BufferInfo;
@@ -277,6 +278,27 @@ public class Mp4MuxerMetadataTest {
         MuxerTestUtil.getExpectedDumpFilePath("mp4_with_string_metadata.mp4"));
   }
 
+  @Test
+  public void writeMp4File_addManyLargeStringMetadata_doesNotThrow() throws Exception {
+    String outputFilePath = temporaryFolder.newFile().getPath();
+    Mp4Muxer muxer = new Mp4Muxer.Builder(new FileOutputStream(outputFilePath)).build();
+
+    String metadataKey = "SomeStringKey";
+    byte[] metadataValue = Util.getUtf8Bytes("AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA");
+    for (int i = 0; i < 100; i++) {
+      muxer.addMetadataEntry(
+          new MdtaMetadataEntry(metadataKey, metadataValue, TYPE_INDICATOR_STRING));
+    }
+    TrackToken token = muxer.addTrack(FAKE_VIDEO_FORMAT);
+
+    try {
+      muxer.writeSampleData(token, sampleAndSampleInfo.first, sampleAndSampleInfo.second);
+      assertThat(sampleAndSampleInfo.first.remaining()).isEqualTo(0);
+    } finally {
+      muxer.close();
+    }
+  }
+
   @Test
   public void writeMp4File_addFloatMetadata_matchesExpected() throws Exception {
     String outputFilePath = temporaryFolder.newFile().getPath();