Changes based on internal review

Author: rohitjoins

libraries/extractor/src/main/java/androidx/media3/extractor/ts/AdtsReader.java

@@ -101,7 +101,6 @@ public final class AdtsReader implements ElementaryStreamReader {
 
   // Used when reading the AAC PCE.
   @Nullable private Format pendingOutputFormat;
-  @Nullable private ParsableBitArray pceBuffer;
 
   // Used when reading the samples.
   private long timeUs;
@@ -192,10 +191,11 @@ public void consume(ParsableByteArray data) throws ParserException {
           }
           break;
         case STATE_READING_AAC_PCE:
-          checkNotNull(pceBuffer);
-          if (continueRead(data, pceBuffer.data, pceBuffer.data.length)) {
-            readAacProgramConfigElement();
-          }
+          int offset = data.getPosition();
+          int limit = offset + min(sampleSize, AAC_PCE_MAX_SIZE);
+          ParsableBitArray pceBuffer =
+              new ParsableBitArray(Arrays.copyOfRange(data.getData(), offset, limit));
+          readAacProgramConfigElement(pceBuffer);
           break;
         case STATE_READING_SAMPLE:
           readSample(data);
@@ -301,7 +301,6 @@ private void setReadingAacPceState(
     this.currentOutput = outputToUse;
     this.currentSampleDuration = currentSampleDuration;
     this.sampleSize = sampleSize;
-    pceBuffer = new ParsableBitArray(new byte[min(sampleSize, AAC_PCE_MAX_SIZE)]);
   }
 
   /**
@@ -573,87 +572,82 @@ private void parseAdtsHeader() throws ParserException {
     }
   }
 
-  @RequiresNonNull("currentOutput")
-  void readAacProgramConfigElement() {
-    Format pendingOutputFormat = checkNotNull(this.pendingOutputFormat);
-    ParsableBitArray pceBuffer = checkNotNull(this.pceBuffer);
-
+  private void readAacProgramConfigElement(ParsableBitArray pceBuffer) throws ParserException {
     // See ISO 13818-7 Advanced Audio Coding (2006) Table 36 for PCE tag encoding.
-    if (pceBuffer.readBits(3) == 5 /* PCE tag */) {
-      // See ISO 13818-7 Advanced Audio Coding (2006) Table 25 for syntax of a PCE.
-      pceBuffer.skipBits(10); // element_instance_tag(4), profile(2), element_instance_tag(4)
-
-      int channelBits = 0;
-      // num_front_channel_elements, front_element_is_cpe(1), front_element_tag_select(4)
-      channelBits += pceBuffer.readBits(4) * 5;
-      // num_side_channel_elements, side_element_is_cpe(1), side_element_tag_select(4)
-      channelBits += pceBuffer.readBits(4) * 5;
-      // num_back_channel_elements, back_element_is_cpe(1), back_element_tag_select(4)
-      channelBits += pceBuffer.readBits(4) * 5;
-      // num_lfe_channel_elements, lfe_element_tag_select(4)
-      channelBits += pceBuffer.readBits(2) * 4;
-      // num_assoc_data_elements, assoc_data_element_tag_select(4)
-      channelBits += pceBuffer.readBits(3) * 4;
-      // num_valid_cc_elements, cc_element_is_ind_sw(1), valid_cc_element_tag_select(4)
-      channelBits += pceBuffer.readBits(4) * 5;
-
-      if (pceBuffer.readBit()) { // mono_mixdown_present
-        pceBuffer.skipBits(4); // mono_mixdown_element_number
-      }
+    if (pceBuffer.readBits(3) != 5 /* PCE tag */) {
+      throw ParserException.createForMalformedContainer(/* message= */ null, /* cause= */ null);
+    }
 
-      if (pceBuffer.readBit()) { // stereo_mixdown_present
-        pceBuffer.skipBits(4); // stereo_mixdown_element_number
-      }
+    // See ISO 13818-7 Advanced Audio Coding (2006) Table 25 for syntax of a PCE.
+    pceBuffer.skipBits(10); // element_instance_tag(4), profile(2), element_instance_tag(4)
+
+    int channelBits = 0;
+    // num_front_channel_elements, front_element_is_cpe(1), front_element_tag_select(4)
+    channelBits += pceBuffer.readBits(4) * 5;
+    // num_side_channel_elements, side_element_is_cpe(1), side_element_tag_select(4)
+    channelBits += pceBuffer.readBits(4) * 5;
+    // num_back_channel_elements, back_element_is_cpe(1), back_element_tag_select(4)
+    channelBits += pceBuffer.readBits(4) * 5;
+    // num_lfe_channel_elements, lfe_element_tag_select(4)
+    channelBits += pceBuffer.readBits(2) * 4;
+    // num_assoc_data_elements, assoc_data_element_tag_select(4)
+    channelBits += pceBuffer.readBits(3) * 4;
+    // num_valid_cc_elements, cc_element_is_ind_sw(1), valid_cc_element_tag_select(4)
+    channelBits += pceBuffer.readBits(4) * 5;
+
+    if (pceBuffer.readBit()) { // mono_mixdown_present
+      pceBuffer.skipBits(4); // mono_mixdown_element_number
+    }
 
-      if (pceBuffer.readBit()) { // matrix_mixdown_idx_present
-        pceBuffer.skipBits(3); // matrix_mixdown_idx(2), matrix_mixdown_idx(1)
-      }
+    if (pceBuffer.readBit()) { // stereo_mixdown_present
+      pceBuffer.skipBits(4); // stereo_mixdown_element_number
+    }
 
-      int numAlignmentBits =
-          8 - (pceBuffer.getPosition() + channelBits + 7) % 8 - 1; // byte_alignment
-      int commentSizeBits = 8; // comment_field_bytes
+    if (pceBuffer.readBit()) { // matrix_mixdown_idx_present
+      pceBuffer.skipBits(3); // matrix_mixdown_idx(2), matrix_mixdown_idx(1)
+    }
 
-      // Beyond this point, pceBuffer may be empty, so check before consuming.
-      if (pceBuffer.bitsLeft() >= channelBits + numAlignmentBits + commentSizeBits) {
-        pceBuffer.skipBits(channelBits);
+    int numAlignmentBits =
+        8 - (pceBuffer.getPosition() + channelBits + 7) % 8 - 1; // byte_alignment
+    int commentSizeBits = 8; // comment_field_bytes
 
-        // Store PCE size excluding initial PCE tag, alignment bits and comment for later.
-        int numPceBits = pceBuffer.getPosition() - 3 /* PCE tag */;
+    // Beyond this point, pceBuffer may be empty, so check before consuming.
+    if (pceBuffer.bitsLeft() < channelBits + numAlignmentBits + commentSizeBits) {
+      throw ParserException.createForMalformedContainer(/* message= */ null, /* cause= */ null);
+    }
 
-        pceBuffer.skipBits(numAlignmentBits);
-        int commentSize = pceBuffer.readBits(commentSizeBits);
+    pceBuffer.skipBits(channelBits);
 
-        if (sampleSize >= pceBuffer.getBytePosition() + commentSize) {
-          // Append PCE to format's audio specific config.
-          byte[] oldConfig = pendingOutputFormat.initializationData.get(0);
+    // Store PCE size excluding initial PCE tag, alignment bits and comment for later.
+    int numPceBits = pceBuffer.getPosition() - 3 /* PCE tag */;
 
-          int configSize = oldConfig.length;
-          configSize += (numPceBits + 7) / 8 + 1; // Byte align and add a zero length comment.
-          byte[] newConfig = Arrays.copyOf(oldConfig, configSize);
+    pceBuffer.skipBits(numAlignmentBits);
+    int commentSize = pceBuffer.readBits(commentSizeBits);
 
-          pceBuffer.setPosition(3 /* PCE tag */);
-          pceBuffer.readBits(newConfig, oldConfig.length, numPceBits);
+    if (sampleSize < pceBuffer.getBytePosition() + commentSize) {
+      throw ParserException.createForMalformedContainer(/* message= */ null, /* cause= */ null);
+    }
 
-          pendingOutputFormat =
-              pendingOutputFormat
-                  .buildUpon()
-                  .setInitializationData(ImmutableList.of(newConfig))
-                  .build();
+    Format pendingOutputFormat = checkNotNull(this.pendingOutputFormat);
+    // Append PCE to format's audio specific config.
+    byte[] oldConfig = pendingOutputFormat.initializationData.get(0);
 
-          // Submit PCE-appended output format.
-          currentOutput.format(pendingOutputFormat);
-          hasOutputFormat = true;
-        }
-      }
-    }
+    int configSize = oldConfig.length;
+    configSize += (numPceBits + 7) / 8 + 1; // Byte align and add a zero length comment.
+    byte[] newConfig = Arrays.copyOf(oldConfig, configSize);
 
-    // Pass through all accumulated data as sample data.
-    ParsableByteArray data = new ParsableByteArray(pceBuffer.data);
-    setReadingSampleState(currentOutput, currentSampleDuration, 0, sampleSize);
-    readSample(data);
+    pceBuffer.setPosition(3 /* PCE tag */);
+    pceBuffer.readBits(newConfig, oldConfig.length, numPceBits);
+
+    pendingOutputFormat =
+        pendingOutputFormat.buildUpon().setInitializationData(ImmutableList.of(newConfig)).build();
 
+    // Submit PCE-appended output format.
+    checkNotNull(currentOutput).format(pendingOutputFormat);
+    this.hasOutputFormat = true;
     this.pendingOutputFormat = null;
-    this.pceBuffer = null;
+
+    setReadingSampleState(currentOutput, currentSampleDuration, 0, sampleSize);
   }
 
   /** Reads the rest of the sample */

libraries/extractor/src/test/java/androidx/media3/extractor/ts/AdtsReaderTest.java

@@ -219,7 +219,7 @@ public void aacPceDataFail() throws ParserException {
     bytes[AAC_PCE_ADTS_HEADER.length] |= 0x20;
 
     // Should throw as FakeTrackOutput expects a format before sampleMetadata.
-    assertThrows(IllegalStateException.class, this::feed);
+    assertThrows(ParserException.class, this::feed);
   }
 
   private void feedLimited(int limit) throws ParserException {