MediaCas: re-send the first scrambled TS packet to decoder

Currently, for scrambled streams that require secure decoder,
we descramble the first TS packet (which includes the PES header)
into the clear, and this TS packet is not descrambled again
at decoding time.

However, descramblers may not be able to descramble the bitstream
into clear buffers past the PES header portion. So we need to
assume that in the first pass, we can only get up to the PES
header portion, and at decoding time, this TS packet should be
descrambled again fully.

bug: 111271614
Test: CTS MediaCasTest; MediaDrmClearKeyTest;
tested with clearkey cas plugin modified to return only the PES
header portion of the first TS packet, instead of the full TS.

Change-Id: I8e77df3c5fa1268282442b7556851d89918c244d
diff --git a/media/libstagefright/ACodecBufferChannel.cpp b/media/libstagefright/ACodecBufferChannel.cpp
index 710ae68..266a240 100644
--- a/media/libstagefright/ACodecBufferChannel.cpp
+++ b/media/libstagefright/ACodecBufferChannel.cpp
@@ -129,6 +129,7 @@
         secureHandle = static_cast<native_handle_t *>(secureData->getDestinationPointer());
     }
     ssize_t result = -1;
+    ssize_t codecDataOffset = 0;
     if (mCrypto != NULL) {
         ICrypto::DestinationBuffer destination;
         if (secure) {
@@ -180,9 +181,16 @@
 
         Status status = Status::OK;
         hidl_string detailedError;
+        ScramblingControl sctrl = ScramblingControl::UNSCRAMBLED;
+
+        if (key != NULL) {
+            sctrl = (ScramblingControl)key[0];
+            // Adjust for the PES offset
+            codecDataOffset = key[2] | (key[3] << 8);
+        }
 
         auto returnVoid = mDescrambler->descramble(
-                key != NULL ? (ScramblingControl)key[0] : ScramblingControl::UNSCRAMBLED,
+                sctrl,
                 hidlSubSamples,
                 srcBuffer,
                 0,
@@ -202,6 +210,11 @@
             return UNKNOWN_ERROR;
         }
 
+        if (result < codecDataOffset) {
+            ALOGD("invalid codec data offset: %zd, result %zd", codecDataOffset, result);
+            return BAD_VALUE;
+        }
+
         ALOGV("descramble succeeded, %zd bytes", result);
 
         if (dstBuffer.type == BufferType::SHARED_MEMORY) {
@@ -210,7 +223,7 @@
         }
     }
 
-    it->mCodecBuffer->setRange(0, result);
+    it->mCodecBuffer->setRange(codecDataOffset, result - codecDataOffset);
 
     // Copy metadata from client to codec buffer.
     it->mCodecBuffer->meta()->clear();
diff --git a/media/libstagefright/mpeg2ts/ATSParser.cpp b/media/libstagefright/mpeg2ts/ATSParser.cpp
index 271d601..cc31815 100644
--- a/media/libstagefright/mpeg2ts/ATSParser.cpp
+++ b/media/libstagefright/mpeg2ts/ATSParser.cpp
@@ -1437,7 +1437,7 @@
     // Perform the 1st pass descrambling if needed
     if (descrambleBytes > 0) {
         memcpy(mDescrambledBuffer->data(), mBuffer->data(), descrambleBytes);
-        mDescrambledBuffer->setRange(0, descrambleBytes);
+        mDescrambledBuffer->setRange(0, mBuffer->size());
 
         hidl_vec<SubSample> subSamples;
         subSamples.resize(descrambleSubSamples);
@@ -1454,10 +1454,9 @@
             }
         }
 
-        uint64_t srcOffset = 0, dstOffset = 0;
-        // If scrambled at PES-level, PES header should be skipped
+        // If scrambled at PES-level, PES header is in the clear
         if (pesScramblingControl != 0) {
-            srcOffset = dstOffset = pesOffset;
+            subSamples[0].numBytesOfClearData = pesOffset;
             subSamples[0].numBytesOfEncryptedData -= pesOffset;
         }
 
@@ -1473,9 +1472,9 @@
                 (ScramblingControl) sctrl,
                 subSamples,
                 mDescramblerSrcBuffer,
-                srcOffset,
+                0 /*srcOffset*/,
                 dstBuffer,
-                dstOffset,
+                0 /*dstOffset*/,
                 [&status, &bytesWritten, &detailedError] (
                         Status _status, uint32_t _bytesWritten,
                         const hidl_string& _detailedError) {
@@ -1492,9 +1491,15 @@
 
         ALOGV("[stream %d] descramble succeeded, %d bytes",
                 mElementaryPID, bytesWritten);
-        memcpy(mBuffer->data(), mDescrambledBuffer->data(), descrambleBytes);
+
+        // Set descrambleBytes to the returned result.
+        // Note that this might be smaller than the total length of input data.
+        // (eg. when we're descrambling the PES header portion of a secure stream,
+        // the plugin might cut it off right after the PES header.)
+        descrambleBytes = bytesWritten;
     }
 
+    sp<ABuffer> buffer;
     if (mQueue->isScrambled()) {
         // Queue subSample info for scrambled queue
         sp<ABuffer> clearSizesBuffer = new ABuffer(mSubSamples.size() * 4);
@@ -1506,8 +1511,7 @@
         for (auto it = mSubSamples.begin();
                 it != mSubSamples.end(); it++, i++) {
             if ((it->transport_scrambling_mode == 0
-                    && pesScramblingControl == 0)
-                    || i < descrambleSubSamples) {
+                    && pesScramblingControl == 0)) {
                 clearSizePtr[i] = it->subSampleSize;
                 encSizePtr[i] = 0;
             } else {
@@ -1516,14 +1520,26 @@
             }
             isSync |= it->random_access_indicator;
         }
+
+        // If scrambled at PES-level, PES header is in the clear
+        if (pesScramblingControl != 0) {
+            clearSizePtr[0] = pesOffset;
+            encSizePtr[0] -= pesOffset;
+        }
         // Pass the original TS subsample size now. The PES header adjust
         // will be applied when the scrambled AU is dequeued.
         mQueue->appendScrambledData(
                 mBuffer->data(), mBuffer->size(), sctrl,
                 isSync, clearSizesBuffer, encSizesBuffer);
+
+        buffer = mDescrambledBuffer;
+    } else {
+        memcpy(mBuffer->data(), mDescrambledBuffer->data(), descrambleBytes);
+
+        buffer = mBuffer;
     }
 
-    ABitReader br(mBuffer->data(), mBuffer->size());
+    ABitReader br(buffer->data(), buffer->size());
     status_t err = parsePES(&br, event);
 
     if (err != OK) {
diff --git a/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp b/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp
index ece0692..9e154a3 100644
--- a/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp
+++ b/media/libstagefright/mpeg2ts/AnotherPacketSource.cpp
@@ -226,6 +226,7 @@
         int32_t cryptoMode;
         if (buffer->meta()->findInt32("cryptoMode", &cryptoMode)) {
             int32_t cryptoKey;
+            int32_t pesOffset;
             sp<ABuffer> clearBytesBuffer, encBytesBuffer;
 
             CHECK(buffer->meta()->findInt32("cryptoKey", &cryptoKey));
@@ -233,6 +234,8 @@
                     && clearBytesBuffer != NULL);
             CHECK(buffer->meta()->findBuffer("encBytes", &encBytesBuffer)
                     && encBytesBuffer != NULL);
+            CHECK(buffer->meta()->findInt32("pesOffset", &pesOffset)
+                    && (pesOffset >= 0) && (pesOffset < 65536));
 
             bufmeta.setInt32(kKeyCryptoMode, cryptoMode);
 
@@ -240,6 +243,11 @@
             bufmeta.setData(kKeyCryptoIV, 0, array, 16);
 
             array[0] = (uint8_t) (cryptoKey & 0xff);
+            // array[1] contains PES header flag, which we don't use.
+            // array[2~3] contain the PES offset.
+            array[2] = (uint8_t) (pesOffset & 0xff);
+            array[3] = (uint8_t) ((pesOffset >> 8) & 0xff);
+
             bufmeta.setData(kKeyCryptoKey, 0, array, 16);
 
             bufmeta.setData(kKeyPlainSizes, 0,
diff --git a/media/libstagefright/mpeg2ts/ESQueue.cpp b/media/libstagefright/mpeg2ts/ESQueue.cpp
index 3deee7e..34d0bcc 100644
--- a/media/libstagefright/mpeg2ts/ESQueue.cpp
+++ b/media/libstagefright/mpeg2ts/ESQueue.cpp
@@ -691,25 +691,9 @@
         return NULL;
     }
 
-    // skip the PES header, and copy the rest into scrambled access unit
+    // copy into scrambled access unit
     sp<ABuffer> scrambledAccessUnit = ABuffer::CreateAsCopy(
-            mScrambledBuffer->data() + pesOffset,
-            scrambledLength - pesOffset);
-
-    // fix up first sample size after skipping the PES header
-    if (pesOffset > 0) {
-        int32_t &firstClearSize = *(int32_t*)clearSizes->data();
-        int32_t &firstEncSize = *(int32_t*)encSizes->data();
-        // Cut away the PES header
-        if (firstClearSize >= pesOffset) {
-            // This is for TS-level scrambling, we descrambled the first
-            // (or it was clear to begin with)
-            firstClearSize -= pesOffset;
-        } else if (firstEncSize >= pesOffset) {
-            // This can only be PES-level scrambling
-            firstEncSize -= pesOffset;
-        }
-    }
+            mScrambledBuffer->data(), scrambledLength);
 
     scrambledAccessUnit->meta()->setInt64("timeUs", timeUs);
     if (isSync) {
@@ -723,6 +707,7 @@
     scrambledAccessUnit->meta()->setInt32("cryptoKey", keyId);
     scrambledAccessUnit->meta()->setBuffer("clearBytes", clearSizes);
     scrambledAccessUnit->meta()->setBuffer("encBytes", encSizes);
+    scrambledAccessUnit->meta()->setInt32("pesOffset", pesOffset);
 
     memmove(mScrambledBuffer->data(),
             mScrambledBuffer->data() + scrambledLength,