Add ElementaryStreamReader's factory to inject custom readers in TSExtractor

Issue:#726

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=134433012

library/src/androidTest/java/com/google/android/exoplayer2/extractor/ts/TsExtractorTest.java

@@ -16,8 +16,16 @@
 package com.google.android.exoplayer2.extractor.ts;
 
 import android.test.InstrumentationTestCase;
+import com.google.android.exoplayer2.Format;
 import com.google.android.exoplayer2.extractor.Extractor;
+import com.google.android.exoplayer2.extractor.ExtractorOutput;
+import com.google.android.exoplayer2.extractor.PositionHolder;
+import com.google.android.exoplayer2.extractor.TrackOutput;
+import com.google.android.exoplayer2.testutil.FakeExtractorInput;
+import com.google.android.exoplayer2.testutil.FakeExtractorOutput;
+import com.google.android.exoplayer2.testutil.FakeTrackOutput;
 import com.google.android.exoplayer2.testutil.TestUtil;
+import com.google.android.exoplayer2.util.ParsableByteArray;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.Random;
@@ -61,6 +69,31 @@ public final class TsExtractorTest extends InstrumentationTestCase {
     }, "ts/sample.ts", fileData, getInstrumentation());
   }
 
+  public void testCustomPesReader() throws Exception {
+    CustomEsReaderFactory factory = new CustomEsReaderFactory();
+    TsExtractor tsExtractor = new TsExtractor(new TimestampAdjuster(0), factory);
+    FakeExtractorInput input = new FakeExtractorInput.Builder()
+        .setData(TestUtil.getByteArray(getInstrumentation(), "ts/sample.ts"))
+        .setSimulateIOErrors(false)
+        .setSimulateUnknownLength(false)
+        .setSimulatePartialReads(false).build();
+    FakeExtractorOutput output = new FakeExtractorOutput();
+    tsExtractor.init(output);
+    tsExtractor.seek(input.getPosition());
+    PositionHolder seekPositionHolder = new PositionHolder();
+    int readResult = Extractor.RESULT_CONTINUE;
+    while (readResult != Extractor.RESULT_END_OF_INPUT) {
+      readResult = tsExtractor.read(input, seekPositionHolder);
+    }
+    CustomEsReader reader = factory.reader;
+    assertEquals(2, reader.packetsRead);
+    TrackOutput trackOutput = reader.getTrackOutput();
+    assertTrue(trackOutput == output.trackOutputs.get(257 /* PID of audio track. */));
+    assertEquals(
+        Format.createTextSampleFormat("Overriding format", "mime", null, 0, 0, "und", null, 0),
+        ((FakeTrackOutput) trackOutput).format);
+  }
+
   private static void writeJunkData(ByteArrayOutputStream out, int length) throws IOException {
     for (int i = 0; i < length; i++) {
       if (((byte) i) == TS_SYNC_BYTE) {
@@ -71,4 +104,62 @@ public final class TsExtractorTest extends InstrumentationTestCase {
     }
   }
 
+  private static final class CustomEsReader extends ElementaryStreamReader {
+
+    public int packetsRead = 0;
+
+    public CustomEsReader(TrackOutput output, String language) {
+      super(output);
+      output.format(Format.createTextSampleFormat("Overriding format", "mime", null, 0, 0,
+          language, null, 0));
+    }
+
+    @Override
+    public void seek() {
+    }
+
+    @Override
+    public void packetStarted(long pesTimeUs, boolean dataAlignmentIndicator) {
+    }
+
+    @Override
+    public void consume(ParsableByteArray data) {
+    }
+
+    @Override
+    public void packetFinished() {
+      packetsRead++;
+    }
+
+    public TrackOutput getTrackOutput() {
+      return output;
+    }
+
+  }
+
+  private static final class CustomEsReaderFactory implements ElementaryStreamReader.Factory {
+
+    private final ElementaryStreamReader.Factory defaultFactory;
+    private CustomEsReader reader;
+
+    public CustomEsReaderFactory() {
+      defaultFactory = new DefaultStreamReaderFactory();
+    }
+
+    @Override
+    public ElementaryStreamReader onPmtEntry(int pid, int streamType,
+        ElementaryStreamReader.EsInfo esInfo, ExtractorOutput output) {
+      if (streamType == 3) {
+        // We need to manually avoid a duplicate custom reader creation.
+        if (reader == null) {
+          reader = new CustomEsReader(output.track(pid), esInfo.language);
+        }
+        return reader;
+      } else {
+        return defaultFactory.onPmtEntry(pid, streamType, esInfo, output);
+      }
+    }
+
+  }
+
 }

library/src/main/java/com/google/android/exoplayer2/extractor/ts/DefaultStreamReaderFactory.java

+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.android.exoplayer2.extractor.ts;
+
+import android.support.annotation.IntDef;
+import android.util.SparseBooleanArray;
+import com.google.android.exoplayer2.extractor.DummyTrackOutput;
+import com.google.android.exoplayer2.extractor.ExtractorOutput;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/**
+ * Default implementation for {@link ElementaryStreamReader.Factory}.
+ */
+public final class DefaultStreamReaderFactory implements ElementaryStreamReader.Factory {
+
+  /**
+   * Flags controlling what workarounds are enabled for elementary stream readers.
+   */
+  @Retention(RetentionPolicy.SOURCE)
+  @IntDef(flag = true, value = {WORKAROUND_ALLOW_NON_IDR_KEYFRAMES, WORKAROUND_IGNORE_AAC_STREAM,
+      WORKAROUND_IGNORE_H264_STREAM, WORKAROUND_DETECT_ACCESS_UNITS, WORKAROUND_MAP_BY_TYPE})
+  public @interface WorkaroundFlags {
+  }
+  public static final int WORKAROUND_ALLOW_NON_IDR_KEYFRAMES = 1;
+  public static final int WORKAROUND_IGNORE_AAC_STREAM = 2;
+  public static final int WORKAROUND_IGNORE_H264_STREAM = 4;
+  public static final int WORKAROUND_DETECT_ACCESS_UNITS = 8;
+  public static final int WORKAROUND_MAP_BY_TYPE = 16;
+
+  private static final int BASE_EMBEDDED_TRACK_ID = 0x2000; // 0xFF + 1.
+
+  private final SparseBooleanArray trackIds;
+  @WorkaroundFlags
+  private final int workaroundFlags;
+  private Id3Reader id3Reader;
+  private int nextEmbeddedTrackId = BASE_EMBEDDED_TRACK_ID;
+
+  public DefaultStreamReaderFactory() {
+    this(0);
+  }
+
+  public DefaultStreamReaderFactory(int workaroundFlags) {
+    trackIds = new SparseBooleanArray();
+    this.workaroundFlags = workaroundFlags;
+  }
+
+  @Override
+  public ElementaryStreamReader onPmtEntry(int pid, int streamType,
+      ElementaryStreamReader.EsInfo esInfo, ExtractorOutput output) {
+
+    if ((workaroundFlags & WORKAROUND_MAP_BY_TYPE) != 0 && id3Reader == null) {
+      // Setup an ID3 track regardless of whether there's a corresponding entry, in case one
+      // appears intermittently during playback. See b/20261500.
+      id3Reader = new Id3Reader(output.track(TsExtractor.TS_STREAM_TYPE_ID3));
+    }
+    int trackId = (workaroundFlags & WORKAROUND_MAP_BY_TYPE) != 0 ? streamType : pid;
+    if (trackIds.get(trackId)) {
+      return null;
+    }
+    trackIds.put(trackId, true);
+    switch (streamType) {
+      case TsExtractor.TS_STREAM_TYPE_MPA:
+      case TsExtractor.TS_STREAM_TYPE_MPA_LSF:
+        return new MpegAudioReader(output.track(trackId), esInfo.language);
+      case TsExtractor.TS_STREAM_TYPE_AAC:
+        return (workaroundFlags & WORKAROUND_IGNORE_AAC_STREAM) != 0 ? null
+            : new AdtsReader(output.track(trackId), new DummyTrackOutput(), esInfo.language);
+      case TsExtractor.TS_STREAM_TYPE_AC3:
+      case TsExtractor.TS_STREAM_TYPE_E_AC3:
+        return new Ac3Reader(output.track(trackId), esInfo.language);
+      case TsExtractor.TS_STREAM_TYPE_DTS:
+      case TsExtractor.TS_STREAM_TYPE_HDMV_DTS:
+        return new DtsReader(output.track(trackId), esInfo.language);
+      case TsExtractor.TS_STREAM_TYPE_H262:
+        return new H262Reader(output.track(trackId));
+      case TsExtractor.TS_STREAM_TYPE_H264:
+        return (workaroundFlags & WORKAROUND_IGNORE_H264_STREAM) != 0
+            ? null : new H264Reader(output.track(trackId),
+                new SeiReader(output.track(nextEmbeddedTrackId++)),
+                (workaroundFlags & WORKAROUND_ALLOW_NON_IDR_KEYFRAMES) != 0,
+                (workaroundFlags & WORKAROUND_DETECT_ACCESS_UNITS) != 0);
+      case TsExtractor.TS_STREAM_TYPE_H265:
+        return new H265Reader(output.track(trackId),
+            new SeiReader(output.track(nextEmbeddedTrackId++)));
+      case TsExtractor.TS_STREAM_TYPE_ID3:
+        if ((workaroundFlags & WORKAROUND_MAP_BY_TYPE) != 0) {
+          return id3Reader;
+        } else {
+          return new Id3Reader(output.track(nextEmbeddedTrackId++));
+        }
+      default:
+        return null;
+    }
+  }
+
+}

library/src/main/java/com/google/android/exoplayer2/extractor/ts/ElementaryStreamReader.java

@@ -15,13 +15,60 @@
  */
 package com.google.android.exoplayer2.extractor.ts;
 
+import com.google.android.exoplayer2.extractor.ExtractorOutput;
 import com.google.android.exoplayer2.extractor.TrackOutput;
 import com.google.android.exoplayer2.util.ParsableByteArray;
 
 /**
  * Extracts individual samples from an elementary media stream, preserving original order.
  */
-/* package */ abstract class ElementaryStreamReader {
+public abstract class ElementaryStreamReader {
+
+  /**
+   * Factory of {@link ElementaryStreamReader} instances.
+   */
+  public interface Factory {
+
+    /**
+     * Returns an {@link ElementaryStreamReader} for a given PMT entry. May return null if the
+     * stream type is not supported or if the stream already has a reader assigned to it.
+     *
+     * @param pid The pid for the PMT entry.
+     * @param streamType One of the {@link TsExtractor}{@code .TS_STREAM_TYPE_*} constants defining
+     *     the type of the stream.
+     * @param esInfo The descriptor information linked to the elementary stream.
+     * @param output The {@link ExtractorOutput} that provides the {@link TrackOutput}s for the
+     *     created readers.
+     * @return An {@link ElementaryStreamReader} for the elementary streams carried by the provided
+     *     pid. {@code null} if the stream is not supported or if it should be ignored.
+     */
+    ElementaryStreamReader onPmtEntry(int pid, int streamType, EsInfo esInfo,
+        ExtractorOutput output);
+
+  }
+
+  /**
+   * Holds descriptor information associated with an elementary stream.
+   */
+  public static final class EsInfo {
+
+    public final int streamType;
+    public String language;
+    public byte[] descriptorBytes;
+
+    /**
+     * @param streamType The type of the stream as defined by the
+     *     {@link TsExtractor}{@code .TS_STREAM_TYPE_*}.
+     * @param language The language of the stream, as defined by ISO/IEC 13818-1, section 2.6.18.
+     * @param descriptorBytes The descriptor bytes associated to the stream.
+     */
+    public EsInfo(int streamType, String language, byte[] descriptorBytes) {
+      this.streamType = streamType;
+      this.language = language;
+      this.descriptorBytes = descriptorBytes;
+    }
+
+  }
 
   protected final TrackOutput output;
 

library/src/main/java/com/google/android/exoplayer2/source/hls/HlsChunkSource.java

@@ -24,6 +24,7 @@ import com.google.android.exoplayer2.extractor.Extractor;
 import com.google.android.exoplayer2.extractor.mp3.Mp3Extractor;
 import com.google.android.exoplayer2.extractor.ts.Ac3Extractor;
 import com.google.android.exoplayer2.extractor.ts.AdtsExtractor;
+import com.google.android.exoplayer2.extractor.ts.DefaultStreamReaderFactory;
 import com.google.android.exoplayer2.extractor.ts.TimestampAdjuster;
 import com.google.android.exoplayer2.extractor.ts.TsExtractor;
 import com.google.android.exoplayer2.source.BehindLiveWindowException;
@@ -355,20 +356,22 @@ import java.util.Locale;
       timestampAdjuster = timestampAdjusterProvider.getAdjuster(segment.discontinuitySequenceNumber,
           startTimeUs);
       // This flag ensures the change of pid between streams does not affect the sample queues.
-      @TsExtractor.WorkaroundFlags int workaroundFlags = TsExtractor.WORKAROUND_MAP_BY_TYPE;
+      @DefaultStreamReaderFactory.WorkaroundFlags
+      int workaroundFlags = DefaultStreamReaderFactory.WORKAROUND_MAP_BY_TYPE;
       String codecs = variants[newVariantIndex].format.codecs;
       if (!TextUtils.isEmpty(codecs)) {
         // Sometimes AAC and H264 streams are declared in TS chunks even though they don't really
         // exist. If we know from the codec attribute that they don't exist, then we can explicitly
         // ignore them even if they're declared.
         if (!MimeTypes.AUDIO_AAC.equals(MimeTypes.getAudioMediaMimeType(codecs))) {
-          workaroundFlags |= TsExtractor.WORKAROUND_IGNORE_AAC_STREAM;
+          workaroundFlags |= DefaultStreamReaderFactory.WORKAROUND_IGNORE_AAC_STREAM;
         }
         if (!MimeTypes.VIDEO_H264.equals(MimeTypes.getVideoMediaMimeType(codecs))) {
-          workaroundFlags |= TsExtractor.WORKAROUND_IGNORE_H264_STREAM;
+          workaroundFlags |= DefaultStreamReaderFactory.WORKAROUND_IGNORE_H264_STREAM;
         }
       }
-      extractor = new TsExtractor(timestampAdjuster, workaroundFlags);
+      extractor = new TsExtractor(timestampAdjuster,
+          new DefaultStreamReaderFactory(workaroundFlags));
     } else {
       // MPEG-2 TS segments, and we need to continue using the same extractor.
       extractor = previous.extractor;