diff --git a/drm/libmediadrm/CryptoHal.cpp b/drm/libmediadrm/CryptoHal.cpp
index b53c741..5d0f68e 100644
--- a/drm/libmediadrm/CryptoHal.cpp
+++ b/drm/libmediadrm/CryptoHal.cpp
@@ -253,7 +253,7 @@
 
     int32_t seqNum = mHeapSeqNum++;
     sp<HidlMemory> hidlMemory = fromHeap(heap);
-    mHeapBases.add(seqNum, mNextBufferId);
+    mHeapBases.add(seqNum, HeapBase(mNextBufferId, heap->getSize()));
     Return<void> hResult = mPlugin->setSharedBufferBase(*hidlMemory, mNextBufferId++);
     ALOGE_IF(!hResult.isOk(), "setSharedBufferBase(): remote call failed");
     return seqNum;
@@ -278,10 +278,26 @@
         return UNEXPECTED_NULL;
     }
 
-    // memory must be in the declared heap
-    CHECK(mHeapBases.indexOfKey(seqNum) >= 0);
+    // memory must be in one of the heaps that have been set
+    if (mHeapBases.indexOfKey(seqNum) < 0) {
+        return UNKNOWN_ERROR;
+    }
 
-    buffer->bufferId = mHeapBases.valueFor(seqNum);
+    // heap must be the same size as the one that was set in setHeapBase
+    if (mHeapBases.valueFor(seqNum).getSize() != heap->getSize()) {
+        android_errorWriteLog(0x534e4554, "76221123");
+        return UNKNOWN_ERROR;
+     }
+
+    // memory must be within the address space of the heap
+    if (memory->pointer() != static_cast<uint8_t *>(heap->getBase()) + memory->offset()  ||
+            heap->getSize() < memory->offset() + memory->size() ||
+            SIZE_MAX - memory->offset() < memory->size()) {
+        android_errorWriteLog(0x534e4554, "76221123");
+        return UNKNOWN_ERROR;
+    }
+
+    buffer->bufferId = mHeapBases.valueFor(seqNum).getBufferId();
     buffer->offset = offset >= 0 ? offset : 0;
     buffer->size = size;
     return OK;