Fixes parsing of PSI sections in MPEG2 transport streams

They can span multiple transport stream packets, something the parser
previously did not support.

Change-Id: I78da6ffd2d422cceedb87aad61cba2062419e635
related-to-bug: 6166596

2 files changed, 160 insertions, 13 deletions

diff --git a/media/libstagefright/mpeg2ts/ATSParser.cpp b/media/libstagefright/mpeg2ts/ATSParser.cpp
index 3f4de1f..5f3e300 100644
--- a/media/libstagefright/mpeg2ts/ATSParser.cpp
+++ b/media/libstagefright/mpeg2ts/ATSParser.cpp
@@ -32,6 +32,7 @@
 #include <media/stagefright/MediaDefs.h>
 #include <media/stagefright/MediaErrors.h>
 #include <media/stagefright/MetaData.h>
+#include <media/stagefright/Utils.h>
 #include <media/IStreamSource.h>
 #include <utils/KeyedVector.h>
 
@@ -46,6 +47,9 @@ static const size_t kTSPacketSize = 188;
 struct ATSParser::Program : public RefBase {
     Program(ATSParser *parser, unsigned programNumber, unsigned programMapPID);
 
+    bool parsePSISection(
+            unsigned pid, ABitReader *br, status_t *err);
+
     bool parsePID(
             unsigned pid, unsigned payload_unit_start_indicator,
             ABitReader *br, status_t *err);
@@ -69,6 +73,10 @@ struct ATSParser::Program : public RefBase {
         mProgramMapPID = programMapPID;
     }
 
+    unsigned programMapPID() const {
+        return mProgramMapPID;
+    }
+
 private:
     ATSParser *mParser;
     unsigned mProgramNumber;
@@ -129,6 +137,27 @@ private:
     DISALLOW_EVIL_CONSTRUCTORS(Stream);
 };
 
+struct ATSParser::PSISection : public RefBase {
+    PSISection();
+
+    status_t append(const void *data, size_t size);
+    void clear();
+
+    bool isComplete() const;
+    bool isEmpty() const;
+
+    const uint8_t *data() const;
+    size_t size() const;
+
+protected:
+    virtual ~PSISection();
+
+private:
+    sp<ABuffer> mBuffer;
+
+    DISALLOW_EVIL_CONSTRUCTORS(PSISection);
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 
 ATSParser::Program::Program(
@@ -141,21 +170,23 @@ ATSParser::Program::Program(
     ALOGV("new program number %u", programNumber);
 }
 
-bool ATSParser::Program::parsePID(
-        unsigned pid, unsigned payload_unit_start_indicator,
-        ABitReader *br, status_t *err) {
+bool ATSParser::Program::parsePSISection(
+        unsigned pid, ABitReader *br, status_t *err) {
     *err = OK;
 
-    if (pid == mProgramMapPID) {
-        if (payload_unit_start_indicator) {
-            unsigned skip = br->getBits(8);
-            br->skipBits(skip * 8);
-        }
+    if (pid != mProgramMapPID) {
+        return false;
+    }
 
-        *err = parseProgramMap(br);
+    *err = parseProgramMap(br);
 
-        return true;
-    }
+    return true;
+}
+
+bool ATSParser::Program::parsePID(
+        unsigned pid, unsigned payload_unit_start_indicator,
+        ABitReader *br, status_t *err) {
+    *err = OK;
 
     ssize_t index = mStreams.indexOfKey(pid);
     if (index < 0) {
@@ -817,6 +848,7 @@ sp<MediaSource> ATSParser::Stream::getSource(SourceType type) {
 
 ATSParser::ATSParser(uint32_t flags)
     : mFlags(flags) {
+    mPSISections.add(0 /* PID */, new PSISection);
 }
 
 ATSParser::~ATSParser() {
@@ -898,6 +930,10 @@ void ATSParser::parseProgramAssociationTable(ABitReader *br) {
                 mPrograms.push(
                         new Program(this, program_number, programMapPID));
             }
+
+            if (mPSISections.indexOfKey(programMapPID) < 0) {
+                mPSISections.add(programMapPID, new PSISection);
+            }
         }
     }
 
@@ -907,12 +943,58 @@ void ATSParser::parseProgramAssociationTable(ABitReader *br) {
 status_t ATSParser::parsePID(
         ABitReader *br, unsigned PID,
         unsigned payload_unit_start_indicator) {
-    if (PID == 0) {
+    ssize_t sectionIndex = mPSISections.indexOfKey(PID);
+
+    if (sectionIndex >= 0) {
+        const sp<PSISection> &section = mPSISections.valueAt(sectionIndex);
+
         if (payload_unit_start_indicator) {
+            CHECK(section->isEmpty());
+
             unsigned skip = br->getBits(8);
             br->skipBits(skip * 8);
         }
-        parseProgramAssociationTable(br);
+
+
+        CHECK((br->numBitsLeft() % 8) == 0);
+        status_t err = section->append(br->data(), br->numBitsLeft() / 8);
+
+        if (err != OK) {
+            return err;
+        }
+
+        if (!section->isComplete()) {
+            return OK;
+        }
+
+        ABitReader sectionBits(section->data(), section->size());
+
+        if (PID == 0) {
+            parseProgramAssociationTable(&sectionBits);
+        } else {
+            bool handled = false;
+            for (size_t i = 0; i < mPrograms.size(); ++i) {
+                status_t err;
+                if (!mPrograms.editItemAt(i)->parsePSISection(
+                            PID, &sectionBits, &err)) {
+                    continue;
+                }
+
+                if (err != OK) {
+                    return err;
+                }
+
+                handled = true;
+                break;
+            }
+
+            if (!handled) {
+                mPSISections.removeItem(PID);
+            }
+        }
+
+        section->clear();
+
         return OK;
     }
 
@@ -1009,4 +1091,64 @@ bool ATSParser::PTSTimeDeltaEstablished() {
     return mPrograms.editItemAt(0)->PTSTimeDeltaEstablished();
 }
 
+////////////////////////////////////////////////////////////////////////////////
+
+ATSParser::PSISection::PSISection() {
+}
+
+ATSParser::PSISection::~PSISection() {
+}
+
+status_t ATSParser::PSISection::append(const void *data, size_t size) {
+    if (mBuffer == NULL || mBuffer->size() + size > mBuffer->capacity()) {
+        size_t newCapacity =
+            (mBuffer == NULL) ? size : mBuffer->capacity() + size;
+
+        newCapacity = (newCapacity + 1023) & ~1023;
+
+        sp<ABuffer> newBuffer = new ABuffer(newCapacity);
+
+        if (mBuffer != NULL) {
+            memcpy(newBuffer->data(), mBuffer->data(), mBuffer->size());
+            newBuffer->setRange(0, mBuffer->size());
+        } else {
+            newBuffer->setRange(0, 0);
+        }
+
+        mBuffer = newBuffer;
+    }
+
+    memcpy(mBuffer->data() + mBuffer->size(), data, size);
+    mBuffer->setRange(0, mBuffer->size() + size);
+
+    return OK;
+}
+
+void ATSParser::PSISection::clear() {
+    if (mBuffer != NULL) {
+        mBuffer->setRange(0, 0);
+    }
+}
+
+bool ATSParser::PSISection::isComplete() const {
+    if (mBuffer == NULL || mBuffer->size() < 3) {
+        return false;
+    }
+
+    unsigned sectionLength = U16_AT(mBuffer->data() + 1) & 0xfff;
+    return mBuffer->size() >= sectionLength + 3;
+}
+
+bool ATSParser::PSISection::isEmpty() const {
+    return mBuffer == NULL || mBuffer->size() == 0;
+}
+
+const uint8_t *ATSParser::PSISection::data() const {
+    return mBuffer == NULL ? NULL : mBuffer->data();
+}
+
+size_t ATSParser::PSISection::size() const {
+    return mBuffer == NULL ? 0 : mBuffer->size();
+}
+
 }  // namespace android
diff --git a/media/libstagefright/mpeg2ts/ATSParser.h b/media/libstagefright/mpeg2ts/ATSParser.h
index c8038d1..9ef2939 100644
--- a/media/libstagefright/mpeg2ts/ATSParser.h
+++ b/media/libstagefright/mpeg2ts/ATSParser.h
@@ -22,6 +22,7 @@
 
 #include <media/stagefright/foundation/ABase.h>
 #include <media/stagefright/foundation/AMessage.h>
+#include <utils/KeyedVector.h>
 #include <utils/Vector.h>
 #include <utils/RefBase.h>
 
@@ -91,10 +92,14 @@ protected:
 private:
     struct Program;
     struct Stream;
+    struct PSISection;
 
     uint32_t mFlags;
     Vector<sp<Program> > mPrograms;
 
+    // Keyed by PID
+    KeyedVector<unsigned, sp<PSISection> > mPSISections;
+
     void parseProgramAssociationTable(ABitReader *br);
     void parseProgramMap(ABitReader *br);
     void parsePES(ABitReader *br);