PesReader.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 androidx.media3.extractor.ts;
import static java.lang.Math.min;
import androidx.annotation.Nullable;
import androidx.media3.common.C;
import androidx.media3.common.ParserException;
import androidx.media3.common.util.Assertions;
import androidx.media3.common.util.Log;
import androidx.media3.common.util.ParsableBitArray;
import androidx.media3.common.util.ParsableByteArray;
import androidx.media3.common.util.TimestampAdjuster;
import androidx.media3.common.util.UnstableApi;
import androidx.media3.extractor.ExtractorOutput;
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
import org.checkerframework.checker.nullness.qual.RequiresNonNull;
/** Parses PES packet data and extracts samples. */
@UnstableApi
public final class PesReader implements TsPayloadReader {
private static final String TAG = "PesReader";
private static final int STATE_FINDING_HEADER = 0;
private static final int STATE_READING_HEADER = 1;
private static final int STATE_READING_HEADER_EXTENSION = 2;
private static final int STATE_READING_BODY = 3;
private static final int HEADER_SIZE = 9;
private static final int MAX_HEADER_EXTENSION_SIZE = 10;
private static final int PES_SCRATCH_SIZE = 10; // max(HEADER_SIZE, MAX_HEADER_EXTENSION_SIZE)
private final ElementaryStreamReader reader;
private final ParsableBitArray pesScratch;
private int state;
private int bytesRead;
private @MonotonicNonNull TimestampAdjuster timestampAdjuster;
private boolean ptsFlag;
private boolean dtsFlag;
private boolean seenFirstDts;
private int extendedHeaderLength;
private int payloadSize;
private boolean dataAlignmentIndicator;
private long timeUs;
public PesReader(ElementaryStreamReader reader) {
this.reader = reader;
pesScratch = new ParsableBitArray(new byte[PES_SCRATCH_SIZE]);
state = STATE_FINDING_HEADER;
}
@Override
public void init(
TimestampAdjuster timestampAdjuster,
ExtractorOutput extractorOutput,
TrackIdGenerator idGenerator) {
this.timestampAdjuster = timestampAdjuster;
reader.createTracks(extractorOutput, idGenerator);
}
// TsPayloadReader implementation.
@Override
public final void seek() {
state = STATE_FINDING_HEADER;
bytesRead = 0;
seenFirstDts = false;
reader.seek();
}
@Override
public final void consume(ParsableByteArray data, @Flags int flags) throws ParserException {
Assertions.checkStateNotNull(timestampAdjuster); // Asserts init has been called.
if ((flags & FLAG_PAYLOAD_UNIT_START_INDICATOR) != 0) {
switch (state) {
case STATE_FINDING_HEADER:
case STATE_READING_HEADER:
// Expected.
break;
case STATE_READING_HEADER_EXTENSION:
Log.w(TAG, "Unexpected start indicator reading extended header");
break;
case STATE_READING_BODY:
// If payloadSize is unset then the length of the previous packet was unspecified, and so
// we only know that it's finished now that we've seen the start of the next one. This is
// expected. If payloadSize is set, then the length of the previous packet was known, but
// we didn't receive that amount of data. This is not expected.
if (payloadSize != C.LENGTH_UNSET) {
Log.w(TAG, "Unexpected start indicator: expected " + payloadSize + " more bytes");
}
// Either way, notify the reader that it has now finished.
reader.packetFinished();
break;
default:
throw new IllegalStateException();
}
setState(STATE_READING_HEADER);
}
while (data.bytesLeft() > 0) {
switch (state) {
case STATE_FINDING_HEADER:
data.skipBytes(data.bytesLeft());
break;
case STATE_READING_HEADER:
if (continueRead(data, pesScratch.data, HEADER_SIZE)) {
setState(parseHeader() ? STATE_READING_HEADER_EXTENSION : STATE_FINDING_HEADER);
}
break;
case STATE_READING_HEADER_EXTENSION:
int readLength = min(MAX_HEADER_EXTENSION_SIZE, extendedHeaderLength);
// Read as much of the extended header as we're interested in, and skip the rest.
if (continueRead(data, pesScratch.data, readLength)
&& continueRead(data, /* target= */ null, extendedHeaderLength)) {
parseHeaderExtension();
flags |= dataAlignmentIndicator ? FLAG_DATA_ALIGNMENT_INDICATOR : 0;
reader.packetStarted(timeUs, flags);
setState(STATE_READING_BODY);
}
break;
case STATE_READING_BODY:
readLength = data.bytesLeft();
int padding = payloadSize == C.LENGTH_UNSET ? 0 : readLength - payloadSize;
if (padding > 0) {
readLength -= padding;
data.setLimit(data.getPosition() + readLength);
}
reader.consume(data);
if (payloadSize != C.LENGTH_UNSET) {
payloadSize -= readLength;
if (payloadSize == 0) {
reader.packetFinished();
setState(STATE_READING_HEADER);
}
}
break;
default:
throw new IllegalStateException();
}
}
}
private void setState(int state) {
this.state = state;
bytesRead = 0;
}
/**
* Continues a read from the provided {@code source} into a given {@code target}. It's assumed
* that the data should be written into {@code target} starting from an offset of zero.
*
* @param source The source from which to read.
* @param target The target into which data is to be read, or {@code null} to skip.
* @param targetLength The target length of the read.
* @return Whether the target length has been reached.
*/
private boolean continueRead(
ParsableByteArray source, @Nullable byte[] target, int targetLength) {
int bytesToRead = min(source.bytesLeft(), targetLength - bytesRead);
if (bytesToRead <= 0) {
return true;
} else if (target == null) {
source.skipBytes(bytesToRead);
} else {
source.readBytes(target, bytesRead, bytesToRead);
}
bytesRead += bytesToRead;
return bytesRead == targetLength;
}
private boolean parseHeader() {
// Note: see ISO/IEC 13818-1, section 2.4.3.6 for detailed information on the format of
// the header.
pesScratch.setPosition(0);
int startCodePrefix = pesScratch.readBits(24);
if (startCodePrefix != 0x000001) {
Log.w(TAG, "Unexpected start code prefix: " + startCodePrefix);
payloadSize = C.LENGTH_UNSET;
return false;
}
pesScratch.skipBits(8); // stream_id.
int packetLength = pesScratch.readBits(16);
pesScratch.skipBits(5); // '10' (2), PES_scrambling_control (2), PES_priority (1)
dataAlignmentIndicator = pesScratch.readBit();
pesScratch.skipBits(2); // copyright (1), original_or_copy (1)
ptsFlag = pesScratch.readBit();
dtsFlag = pesScratch.readBit();
// ESCR_flag (1), ES_rate_flag (1), DSM_trick_mode_flag (1),
// additional_copy_info_flag (1), PES_CRC_flag (1), PES_extension_flag (1)
pesScratch.skipBits(6);
extendedHeaderLength = pesScratch.readBits(8);
if (packetLength == 0) {
payloadSize = C.LENGTH_UNSET;
} else {
payloadSize =
packetLength
+ 6 /* packetLength does not include the first 6 bytes */
- HEADER_SIZE
- extendedHeaderLength;
if (payloadSize < 0) {
Log.w(TAG, "Found negative packet payload size: " + payloadSize);
payloadSize = C.LENGTH_UNSET;
}
}
return true;
}
@RequiresNonNull("timestampAdjuster")
private void parseHeaderExtension() {
pesScratch.setPosition(0);
timeUs = C.TIME_UNSET;
if (ptsFlag) {
pesScratch.skipBits(4); // '0010' or '0011'
long pts = (long) pesScratch.readBits(3) << 30;
pesScratch.skipBits(1); // marker_bit
pts |= pesScratch.readBits(15) << 15;
pesScratch.skipBits(1); // marker_bit
pts |= pesScratch.readBits(15);
pesScratch.skipBits(1); // marker_bit
if (!seenFirstDts && dtsFlag) {
pesScratch.skipBits(4); // '0011'
long dts = (long) pesScratch.readBits(3) << 30;
pesScratch.skipBits(1); // marker_bit
dts |= pesScratch.readBits(15) << 15;
pesScratch.skipBits(1); // marker_bit
dts |= pesScratch.readBits(15);
pesScratch.skipBits(1); // marker_bit
// Subsequent PES packets may have earlier presentation timestamps than this one, but they
// should all be greater than or equal to this packet's decode timestamp. We feed the
// decode timestamp to the adjuster here so that in the case that this is the first to be
// fed, the adjuster will be able to compute an offset to apply such that the adjusted
// presentation timestamps of all future packets are non-negative.
timestampAdjuster.adjustTsTimestamp(dts);
seenFirstDts = true;
}
timeUs = timestampAdjuster.adjustTsTimestamp(pts);
}
}
}