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;