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);
     }