Revert "Revert "Fix decoder instantiation during playback"" -- DO NOT MERGE

This reverts commit de7268d8e20b883ec88a7ff19ad560a665373484.

Bug: 36479980
Bug: 36209723
Bug: 36660223

Change-Id: I8d83305a28dc35cee16f8f1068c90fdd0b9effff
(cherry picked from commit 6aa5c0662c86c7f7b2890577a207086eeecbd177)
diff --git a/drm/libmediadrm/CryptoHal.cpp b/drm/libmediadrm/CryptoHal.cpp
index 1466222..9f41403 100644
--- a/drm/libmediadrm/CryptoHal.cpp
+++ b/drm/libmediadrm/CryptoHal.cpp
@@ -228,11 +228,11 @@
 void CryptoHal::setHeapBase(const sp<IMemoryHeap>& heap) {
     native_handle_t* nativeHandle = native_handle_create(1, 0);
     if (!nativeHandle) {
-        ALOGE("setSharedBufferBase(), failed to create native handle");
+        ALOGE("setHeapBase(), failed to create native handle");
         return;
     }
     if (heap == NULL) {
-        ALOGE("setSharedBufferBase(): heap is NULL");
+        ALOGE("setHeapBase(): heap is NULL");
         return;
     }
     int fd = heap->getHeapID();
@@ -244,6 +244,10 @@
     ALOGE_IF(!hResult.isOk(), "setSharedBufferBase(): remote call failed");
 }
 
+void CryptoHal::clearHeapBase(const sp<IMemoryHeap>& heap) {
+    mHeapBases.removeItem(heap->getBase());
+}
+
 status_t CryptoHal::toSharedBuffer(const sp<IMemory>& memory, ::SharedBuffer* buffer) {
     ssize_t offset;
     size_t size;
@@ -257,9 +261,8 @@
         return UNEXPECTED_NULL;
     }
 
-    if (mHeapBases.indexOfKey(heap->getBase()) < 0) {
-        setHeapBase(heap);
-    }
+    // memory must be in the declared heap
+    CHECK(mHeapBases.indexOfKey(heap->getBase()) >= 0);
 
     buffer->bufferId = mHeapBases.valueFor(heap->getBase());
     buffer->offset = offset >= 0 ? offset : 0;
diff --git a/drm/libmediadrm/ICrypto.cpp b/drm/libmediadrm/ICrypto.cpp
index 92c9548..7b70205 100644
--- a/drm/libmediadrm/ICrypto.cpp
+++ b/drm/libmediadrm/ICrypto.cpp
@@ -36,6 +36,8 @@
     DECRYPT,
     NOTIFY_RESOLUTION,
     SET_MEDIADRM_SESSION,
+    SET_HEAP,
+    UNSET_HEAP,
 };
 
 struct BpCrypto : public BpInterface<ICrypto> {
@@ -177,6 +179,23 @@
         return reply.readInt32();
     }
 
+    virtual void setHeap(const sp<IMemoryHeap> &heap) {
+        Parcel data, reply;
+        data.writeInterfaceToken(ICrypto::getInterfaceDescriptor());
+        data.writeStrongBinder(IInterface::asBinder(heap));
+        remote()->transact(SET_HEAP, data, &reply);
+        return;
+    }
+
+    virtual void unsetHeap(const sp<IMemoryHeap>& heap) {
+        Parcel data, reply;
+        data.writeInterfaceToken(ICrypto::getInterfaceDescriptor());
+        data.writeStrongBinder(IInterface::asBinder(heap));
+        remote()->transact(UNSET_HEAP, data, &reply);
+        return;
+    }
+
+
 private:
     void readVector(Parcel &reply, Vector<uint8_t> &vector) const {
         uint32_t size = reply.readInt32();
@@ -404,6 +423,24 @@
             return OK;
         }
 
+        case SET_HEAP:
+        {
+            CHECK_INTERFACE(ICrypto, data, reply);
+            sp<IMemoryHeap> heap =
+                interface_cast<IMemoryHeap>(data.readStrongBinder());
+            setHeap(heap);
+            return OK;
+        }
+
+        case UNSET_HEAP:
+        {
+            CHECK_INTERFACE(ICrypto, data, reply);
+            sp<IMemoryHeap> heap =
+                interface_cast<IMemoryHeap>(data.readStrongBinder());
+            unsetHeap(heap);
+            return OK;
+        }
+
         default:
             return BBinder::onTransact(code, data, reply, flags);
     }
diff --git a/media/libmedia/include/Crypto.h b/media/libmedia/include/Crypto.h
index ce08f98..b68413d 100644
--- a/media/libmedia/include/Crypto.h
+++ b/media/libmedia/include/Crypto.h
@@ -55,6 +55,9 @@
             const CryptoPlugin::SubSample *subSamples, size_t numSubSamples,
             const DestinationBuffer &destination, AString *errorDetailMsg);
 
+    virtual void setHeap(const sp<IMemoryHeap>&) {}
+    virtual void unsetHeap(const sp<IMemoryHeap>&) {}
+
 private:
     mutable Mutex mLock;
 
diff --git a/media/libmedia/include/CryptoHal.h b/media/libmedia/include/CryptoHal.h
index 28ade20..d6214c2 100644
--- a/media/libmedia/include/CryptoHal.h
+++ b/media/libmedia/include/CryptoHal.h
@@ -60,6 +60,9 @@
             const ICrypto::DestinationBuffer &destination,
             AString *errorDetailMsg);
 
+    virtual void setHeap(const sp<IMemoryHeap>& heap) { setHeapBase(heap); }
+    virtual void unsetHeap(const sp<IMemoryHeap>& heap) { clearHeapBase(heap); }
+
 private:
     mutable Mutex mLock;
 
@@ -82,6 +85,7 @@
             const uint8_t uuid[16], const void *initData, size_t size);
 
     void setHeapBase(const sp<IMemoryHeap>& heap);
+    void clearHeapBase(const sp<IMemoryHeap>& heap);
 
     status_t toSharedBuffer(const sp<IMemory>& memory, ::SharedBuffer* buffer);
 
diff --git a/media/libmedia/include/ICrypto.h b/media/libmedia/include/ICrypto.h
index 8990f4b..f83c846 100644
--- a/media/libmedia/include/ICrypto.h
+++ b/media/libmedia/include/ICrypto.h
@@ -27,6 +27,7 @@
 
 struct AString;
 class IMemory;
+class IMemoryHeap;
 
 struct ICrypto : public IInterface {
     DECLARE_META_INTERFACE(Crypto);
@@ -64,6 +65,13 @@
             const CryptoPlugin::SubSample *subSamples, size_t numSubSamples,
             const DestinationBuffer &destination, AString *errorDetailMsg) = 0;
 
+    /**
+     * Declare the heap that the shared memory source buffers passed
+     * to decrypt will be allocated from.
+     */
+    virtual void setHeap(const sp<IMemoryHeap>& heap) = 0;
+    virtual void unsetHeap(const sp<IMemoryHeap>& heap) = 0;
+
 private:
     DISALLOW_EVIL_CONSTRUCTORS(ICrypto);
 };
diff --git a/media/libstagefright/ACodecBufferChannel.cpp b/media/libstagefright/ACodecBufferChannel.cpp
index ecbbd03..7743079 100644
--- a/media/libstagefright/ACodecBufferChannel.cpp
+++ b/media/libstagefright/ACodecBufferChannel.cpp
@@ -39,6 +39,12 @@
 using BufferInfo = ACodecBufferChannel::BufferInfo;
 using BufferInfoIterator = std::vector<const BufferInfo>::const_iterator;
 
+ACodecBufferChannel::~ACodecBufferChannel() {
+    if (mCrypto != nullptr && mDealer != nullptr) {
+        mCrypto->unsetHeap(mDealer->getMemoryHeap());
+    }
+}
+
 static BufferInfoIterator findClientBuffer(
         const std::shared_ptr<const std::vector<const BufferInfo>> &array,
         const sp<MediaCodecBuffer> &buffer) {
@@ -254,6 +260,18 @@
     }
 }
 
+sp<MemoryDealer> ACodecBufferChannel::makeMemoryDealer(size_t heapSize) {
+    sp<MemoryDealer> dealer;
+    if (mDealer != nullptr && mCrypto != nullptr) {
+        mCrypto->unsetHeap(mDealer->getMemoryHeap());
+    }
+    dealer = new MemoryDealer(heapSize, "ACodecBufferChannel");
+    if (mCrypto != nullptr) {
+        mCrypto->setHeap(dealer->getMemoryHeap());
+    }
+    return dealer;
+}
+
 void ACodecBufferChannel::setInputBufferArray(const std::vector<BufferAndId> &array) {
     if (hasCryptoOrDescrambler()) {
         size_t totalSize = std::accumulate(
@@ -268,8 +286,10 @@
                 (size_t max, const BufferAndId& elem) {
                     return std::max(max, align(elem.mBuffer->capacity(), alignment));
                 });
-        mDealer = new MemoryDealer(totalSize + maxSize, "ACodecBufferChannel");
-        mDecryptDestination = mDealer->allocate(maxSize);
+        size_t destinationBufferSize = maxSize;
+        size_t heapSize = totalSize + destinationBufferSize;
+        mDealer = makeMemoryDealer(heapSize);
+        mDecryptDestination = mDealer->allocate(destinationBufferSize);
     }
     std::vector<const BufferInfo> inputBuffers;
     for (const BufferAndId &elem : array) {
diff --git a/media/libstagefright/include/ACodecBufferChannel.h b/media/libstagefright/include/ACodecBufferChannel.h
index 02468c1..b731666 100644
--- a/media/libstagefright/include/ACodecBufferChannel.h
+++ b/media/libstagefright/include/ACodecBufferChannel.h
@@ -60,7 +60,7 @@
 
     ACodecBufferChannel(
             const sp<AMessage> &inputBufferFilled, const sp<AMessage> &outputBufferDrained);
-    virtual ~ACodecBufferChannel() = default;
+    virtual ~ACodecBufferChannel();
 
     // BufferChannelBase interface
     virtual status_t queueInputBuffer(const sp<MediaCodecBuffer> &buffer) override;
@@ -127,6 +127,8 @@
     std::shared_ptr<const std::vector<const BufferInfo>> mInputBuffers;
     std::shared_ptr<const std::vector<const BufferInfo>> mOutputBuffers;
 
+    sp<MemoryDealer> makeMemoryDealer(size_t heapSize);
+
     bool hasCryptoOrDescrambler() {
         return mCrypto != NULL || mDescrambler != NULL;
     }