Fix a ParsableByteArray limit bug in FlacFrameReader

Also clean-up the implementation a bit to avoid allocating two arrays
and copying between them. Instead we just allocate a single (large
enough) array, which also makes the code more concise.

Writing a test for this required `ParsableByteArray` to enforce its
limit, so I added that logic too behind a test-only (for now) gate. We
can use this gate in other tests as part of flushing out more limit
bugs as we work towards b/147657250.

PiperOrigin-RevId: 785501853

Author: icbaker

libraries/common/src/main/java/androidx/media3/common/util/ParsableByteArray.java

@@ -19,6 +19,7 @@
 import static java.nio.ByteOrder.LITTLE_ENDIAN;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.primitives.Chars;
 import com.google.common.primitives.Ints;
@@ -30,6 +31,7 @@
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Wraps a byte array, providing a set of methods for parsing data from it. Numerical values are
@@ -52,6 +54,9 @@ public final class ParsableByteArray {
           StandardCharsets.UTF_16BE,
           StandardCharsets.UTF_16LE);
 
+  // TODO: b/147657250 - Flip this to true
+  private static final AtomicBoolean shouldEnforceLimitOnLegacyMethods = new AtomicBoolean();
+
   private byte[] data;
   private int position;
   // TODO(internal b/147657250): Enforce this limit on all read methods.
@@ -227,6 +232,7 @@ public void readBytes(ParsableBitArray bitArray, int length) {
    * @param length The number of bytes to read.
    */
   public void readBytes(byte[] buffer, int offset, int length) {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(length);
     System.arraycopy(data, position, buffer, offset, length);
     position += length;
   }
@@ -239,12 +245,14 @@ public void readBytes(byte[] buffer, int offset, int length) {
    * @param length The number of bytes to read.
    */
   public void readBytes(ByteBuffer buffer, int length) {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(length);
     buffer.put(data, position, length);
     position += length;
   }
 
   /** Peeks at the next byte as an unsigned value. */
   public int peekUnsignedByte() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(1);
     return (data[position] & 0xFF);
   }
 
@@ -280,6 +288,7 @@ public char peekChar(Charset charset) {
 
   /** Peek the UTF-16 char at {@link #position}{@code + offset}. */
   private char peekChar(ByteOrder byteOrder, int offset) {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(2);
     return byteOrder == BIG_ENDIAN
         ? Chars.fromBytes(data[position + offset], data[position + offset + 1])
         : Chars.fromBytes(data[position + offset + 1], data[position + offset]);
@@ -320,59 +329,69 @@ public int peekCodePoint(Charset charset) {
 
   /** Reads the next byte as an unsigned value. */
   public int readUnsignedByte() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(1);
     return (data[position++] & 0xFF);
   }
 
   /** Reads the next two bytes as an unsigned value. */
   public int readUnsignedShort() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(2);
     return (data[position++] & 0xFF) << 8 | (data[position++] & 0xFF);
   }
 
   /** Reads the next two bytes as an unsigned value. */
   public int readLittleEndianUnsignedShort() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(2);
     return (data[position++] & 0xFF) | (data[position++] & 0xFF) << 8;
   }
 
   /** Reads the next two bytes as a signed value. */
   public short readShort() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(2);
     return (short) ((data[position++] & 0xFF) << 8 | (data[position++] & 0xFF));
   }
 
   /** Reads the next two bytes as a signed value. */
   public short readLittleEndianShort() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(2);
     return (short) ((data[position++] & 0xFF) | (data[position++] & 0xFF) << 8);
   }
 
   /** Reads the next three bytes as an unsigned value. */
   public int readUnsignedInt24() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(3);
     return (data[position++] & 0xFF) << 16
         | (data[position++] & 0xFF) << 8
         | (data[position++] & 0xFF);
   }
 
   /** Reads the next three bytes as a signed value. */
   public int readInt24() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(3);
     return ((data[position++] & 0xFF) << 24) >> 8
         | (data[position++] & 0xFF) << 8
         | (data[position++] & 0xFF);
   }
 
   /** Reads the next three bytes as a signed value in little endian order. */
   public int readLittleEndianInt24() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(3);
     return (data[position++] & 0xFF)
         | (data[position++] & 0xFF) << 8
         | (data[position++] & 0xFF) << 16;
   }
 
   /** Reads the next three bytes as an unsigned value in little endian order. */
   public int readLittleEndianUnsignedInt24() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(3);
     return (data[position++] & 0xFF)
         | (data[position++] & 0xFF) << 8
         | (data[position++] & 0xFF) << 16;
   }
 
   /** Reads the next four bytes as an unsigned value. */
   public long readUnsignedInt() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(4);
     return (data[position++] & 0xFFL) << 24
         | (data[position++] & 0xFFL) << 16
         | (data[position++] & 0xFFL) << 8
@@ -381,6 +400,7 @@ public long readUnsignedInt() {
 
   /** Reads the next four bytes as an unsigned value in little endian order. */
   public long readLittleEndianUnsignedInt() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(4);
     return (data[position++] & 0xFFL)
         | (data[position++] & 0xFFL) << 8
         | (data[position++] & 0xFFL) << 16
@@ -389,6 +409,7 @@ public long readLittleEndianUnsignedInt() {
 
   /** Reads the next four bytes as a signed value */
   public int readInt() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(4);
     return (data[position++] & 0xFF) << 24
         | (data[position++] & 0xFF) << 16
         | (data[position++] & 0xFF) << 8
@@ -397,6 +418,7 @@ public int readInt() {
 
   /** Reads the next four bytes as a signed value in little endian order. */
   public int readLittleEndianInt() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(4);
     return (data[position++] & 0xFF)
         | (data[position++] & 0xFF) << 8
         | (data[position++] & 0xFF) << 16
@@ -405,6 +427,7 @@ public int readLittleEndianInt() {
 
   /** Reads the next eight bytes as a signed value. */
   public long readLong() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(8);
     return (data[position++] & 0xFFL) << 56
         | (data[position++] & 0xFFL) << 48
         | (data[position++] & 0xFFL) << 40
@@ -417,6 +440,7 @@ public long readLong() {
 
   /** Reads the next eight bytes as a signed value in little endian order. */
   public long readLittleEndianLong() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(8);
     return (data[position++] & 0xFFL)
         | (data[position++] & 0xFFL) << 8
         | (data[position++] & 0xFFL) << 16
@@ -429,6 +453,7 @@ public long readLittleEndianLong() {
 
   /** Reads the next four bytes, returning the integer portion of the fixed point 16.16 integer. */
   public int readUnsignedFixedPoint1616() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(4);
     int result = (data[position++] & 0xFF) << 8 | (data[position++] & 0xFF);
     position += 2; // Skip the non-integer portion.
     return result;
@@ -518,6 +543,7 @@ public String readString(int length) {
    * @return The string encoded by the bytes in the specified character set.
    */
   public String readString(int length, Charset charset) {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(length);
     String result = new String(data, position, length, charset);
     position += length;
     return result;
@@ -531,6 +557,7 @@ public String readString(int length, Charset charset) {
    * @return The string, not including any terminating NUL byte.
    */
   public String readNullTerminatedString(int length) {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(length);
     if (length == 0) {
       return "";
     }
@@ -630,6 +657,7 @@ public String readLine(Charset charset) {
    * @return Decoded long value
    */
   public long readUtf8EncodedLong() {
+    maybeAssertAtLeastBytesLeftForLegacyMethod(1);
     int length = 0;
     long value = data[position];
     // find the high most 0 bit
@@ -647,6 +675,7 @@ public long readUtf8EncodedLong() {
     if (length == 0) {
       throw new NumberFormatException("Invalid UTF-8 sequence first byte: " + value);
     }
+    maybeAssertAtLeastBytesLeftForLegacyMethod(length);
     for (int i = 1; i < length; i++) {
       int x = data[position + i];
       if ((x & 0xC0) != 0x80) { // if the high most 0 bit not 7th
@@ -724,6 +753,23 @@ public Charset readUtfCharsetFromBom() {
     return null;
   }
 
+  /**
+   * Sets whether all read/peek methods should enforce that {@link #getPosition()} never exceeds
+   * {@link #limit()}.
+   *
+   * <p>Setting this to {@code true} in tests can help catch cases of accidentally reading beyond
+   * {@link #limit()} but still within the bounds of the underlying {@link #getData()}.
+   *
+   * <p>Some (newer) methods will always enforce the invariant, even when this is set to {@code
+   * false}.
+   *
+   * <p>Defaults to false (this may change in a later release).
+   */
+  @VisibleForTesting
+  public static void setShouldEnforceLimitOnLegacyMethods(boolean enforceLimit) {
+    ParsableByteArray.shouldEnforceLimitOnLegacyMethods.set(enforceLimit);
+  }
+
   /**
    * Returns the index of the next occurrence of '\n' or '\r', or {@link #limit} if none is found.
    */
@@ -898,6 +944,22 @@ && isUtf8ContinuationByte(data[position + 3])) {
     }
   }
 
+  /**
+   * Enforces that {@link #bytesLeft()} is at least {@code bytesNeeded} if {@link
+   * #shouldEnforceLimitOnLegacyMethods} is set to {@code true}.
+   *
+   * <p>This should only be called from methods that previously didn't enforce the limit. All new
+   * methods added to this class should unconditionally enforce the limit.
+   */
+  private void maybeAssertAtLeastBytesLeftForLegacyMethod(int bytesNeeded) {
+    if (shouldEnforceLimitOnLegacyMethods.get()) {
+      if (bytesLeft() < bytesNeeded) {
+        throw new IndexOutOfBoundsException(
+            "bytesNeeded= " + bytesNeeded + ", bytesLeft=" + bytesLeft());
+      }
+    }
+  }
+
   private static boolean isUtf8ContinuationByte(byte b) {
     return (b & 0xC0) == 0x80;
   }

libraries/extractor/src/main/java/androidx/media3/extractor/FlacFrameReader.java

@@ -106,22 +106,21 @@ public static boolean checkFrameHeaderFromPeek(
       throws IOException {
     long originalPeekPosition = input.getPeekPosition();
 
-    byte[] frameStartBytes = new byte[2];
-    input.peekFully(frameStartBytes, 0, 2);
-    int frameStart = (frameStartBytes[0] & 0xFF) << 8 | (frameStartBytes[1] & 0xFF);
+    // We will try and read the first subframe header following the frame header as well.
+    int dataToCheck = FlacConstants.MAX_FRAME_HEADER_SIZE + 1;
+    ParsableByteArray scratch = new ParsableByteArray(dataToCheck);
+
+    // Check the start marker first, before peeking the rest of the frame header.
+    input.peekFully(scratch.getData(), 0, 2);
+    int frameStart = scratch.peekChar();
     if (frameStart != frameStartMarker) {
       input.resetPeekPosition();
       input.advancePeekPosition((int) (originalPeekPosition - input.getPosition()));
       return false;
     }
 
-    // Try and read the first subframe header following the frame header as well.
-    int dataToCheck = FlacConstants.MAX_FRAME_HEADER_SIZE + 1;
-    ParsableByteArray scratch = new ParsableByteArray(dataToCheck);
-    System.arraycopy(
-        frameStartBytes, /* srcPos= */ 0, scratch.getData(), /* destPos= */ 0, /* length= */ 2);
-
-    int totalBytesPeeked = ExtractorUtil.peekToLength(input, scratch.getData(), 2, dataToCheck - 2);
+    int totalBytesPeeked = 2;
+    totalBytesPeeked += ExtractorUtil.peekToLength(input, scratch.getData(), 2, dataToCheck - 2);
     scratch.setLimit(totalBytesPeeked);
 
     input.resetPeekPosition();

libraries/extractor/src/main/java/androidx/media3/extractor/FlacStreamMetadata.java

@@ -18,6 +18,7 @@
 import static androidx.media3.extractor.VorbisUtil.parseVorbisComments;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import androidx.media3.common.C;
 import androidx.media3.common.Format;
 import androidx.media3.common.Metadata;
@@ -165,7 +166,8 @@ public FlacStreamMetadata(
         concatenateVorbisMetadata(vorbisComments, pictureFrames));
   }
 
-  private FlacStreamMetadata(
+  @VisibleForTesting
+  /* package */ FlacStreamMetadata(
       int minBlockSizeSamples,
       int maxBlockSizeSamples,
       int minFrameSize,

libraries/extractor/src/test/java/androidx/media3/extractor/FlacFrameReaderTest.java

@@ -15,19 +15,25 @@
  */
 package androidx.media3.extractor;
 
+import static androidx.media3.test.utils.TestUtil.createByteArray;
 import static com.google.common.truth.Truth.assertThat;
 
+import androidx.media3.common.Metadata;
 import androidx.media3.common.util.ParsableByteArray;
 import androidx.media3.common.util.Util;
 import androidx.media3.extractor.FlacFrameReader.SampleNumberHolder;
 import androidx.media3.extractor.FlacMetadataReader.FlacStreamMetadataHolder;
+import androidx.media3.extractor.FlacStreamMetadata.SeekTable;
 import androidx.media3.extractor.flac.FlacConstants;
 import androidx.media3.test.utils.FakeExtractorInput;
 import androidx.media3.test.utils.TestUtil;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.common.primitives.Bytes;
 import com.google.common.primitives.UnsignedBytes;
 import java.io.IOException;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -40,6 +46,16 @@
 @RunWith(AndroidJUnit4.class)
 public class FlacFrameReaderTest {
 
+  @BeforeClass
+  public static void enableParsableByteArrayLimitEnforcement() {
+    ParsableByteArray.setShouldEnforceLimitOnLegacyMethods(true);
+  }
+
+  @AfterClass
+  public static void disableParsableByteArrayLimitEnforcement() {
+    ParsableByteArray.setShouldEnforceLimitOnLegacyMethods(false);
+  }
+
   @Test
   public void checkAndReadFrameHeader_validData_updatesPosition() throws Exception {
     FlacStreamMetadataHolder streamMetadataHolder =
@@ -294,6 +310,47 @@ public void checkFrameHeaderFromPeek_validData_isTrue() throws Exception {
     assertThat(result).isTrue();
   }
 
+  @Test
+  public void checkFrameHeaderFromPeek_validData_maxHeaderSize() throws Exception {
+    byte[] frameHeader =
+        Bytes.concat(
+            // The sync code and blocking strategy bit (variable block size)
+            createByteArray(0xFF, 0xF9),
+            // Indicates 16-bit uncommon block size and sample rate, and mono channel count.
+            createByteArray(0b0111_1101, 0x00),
+            // The smallest 7 byte coded number: 0x8000_0000
+            createByteArray(0xFE, 0x82, 0x80, 0x80, 0x80, 0x80, 0x80),
+            // The 16-bit uncommon block size (257 - 1) & sample rate (256kHz), as indicated above.
+            createByteArray(0x01, 0x00, 0x01, 0x00),
+            // The CRC placeholder (updated below).
+            createByteArray(0x00));
+    frameHeader[15] =
+        UnsignedBytes.checkedCast(
+            Util.crc8(frameHeader, /* start= */ 0, /* end= */ 15, /* initialValue= */ 0));
+
+    ExtractorInput input = new FakeExtractorInput.Builder().setData(frameHeader).build();
+    FlacStreamMetadata flacStreamMetadata =
+        new FlacStreamMetadata(
+            /* minBlockSizeSamples= */ 16,
+            /* maxBlockSizeSamples= */ 257,
+            /* minFrameSize= */ 0, /* unknown */
+            /* maxFrameSize= */ 0 /* unknown */,
+            /* sampleRate= */ 256,
+            /* channels= */ 1,
+            /* bitsPerSample= */ 16,
+            /* totalSamples= */ 0 /* unknown */,
+            (SeekTable) null,
+            (Metadata) null);
+
+    int frameStartMarker = FlacMetadataReader.getFrameStartMarker(input);
+
+    boolean result =
+        FlacFrameReader.checkFrameHeaderFromPeek(
+            input, flacStreamMetadata, frameStartMarker, new SampleNumberHolder());
+
+    assertThat(result).isTrue();
+  }
+
   @Test
   public void checkFrameHeaderFromPeek_validData_writesSampleNumber() throws Exception {
     FlacStreamMetadataHolder streamMetadataHolder =