diff --git a/Keymaster.h b/Keymaster.h
index cf692f3..412110c 100644
--- a/Keymaster.h
+++ b/Keymaster.h
@@ -14,13 +14,12 @@
  * limitations under the License.
  */
 
-#ifndef ANDROID_VOLD_KEYMASTER1_H
-#define ANDROID_VOLD_KEYMASTER1_H
+#ifndef ANDROID_VOLD_KEYMASTER_H
+#define ANDROID_VOLD_KEYMASTER_H
 
+#include <memory>
 #include <string>
-
-#include <hardware/hardware.h>
-#include <hardware/keymaster1.h>
+#include <utility>
 
 #include <keymaster/authorization_set.h>
 
@@ -29,19 +28,21 @@
 
 using namespace keymaster;
 
-// C++ wrappers to the keymaster1 C interface.
+// C++ wrappers to the Keymaster C interface.
 // This is tailored to the needs of KeyStorage, but could be extended to be
 // a more general interface.
 
+// Class that wraps a keymaster1_device_t or keymaster2_device_t and provides methods
+// they have in common. Also closes the device on destruction.
+class IKeymasterDevice;
+
 // Wrapper for a keymaster_operation_handle_t representing an
 // ongoing Keymaster operation.  Aborts the operation
 // in the destructor if it is unfinished. Methods log failures
 // to LOG(ERROR).
 class KeymasterOperation {
   public:
-    ~KeymasterOperation() {
-        if (mDevice) mDevice->abort(mDevice, mOpHandle);
-    }
+    ~KeymasterOperation();
     // Is this instance valid? This is false if creation fails, and becomes
     // false on finish or if an update fails.
     explicit operator bool() { return mDevice != nullptr; }
@@ -54,28 +55,24 @@
     bool finishWithOutput(std::string* output);
     // Move constructor
     KeymasterOperation(KeymasterOperation&& rhs) {
-        mOpHandle = rhs.mOpHandle;
-        mDevice = rhs.mDevice;
-        rhs.mDevice = nullptr;
+        mOpHandle = std::move(rhs.mOpHandle);
+        mDevice = std::move(rhs.mDevice);
     }
 
   private:
-    KeymasterOperation(keymaster1_device_t* d, keymaster_operation_handle_t h)
+    KeymasterOperation(std::shared_ptr<IKeymasterDevice> d, keymaster_operation_handle_t h)
         : mDevice{d}, mOpHandle{h} {}
-    keymaster1_device_t* mDevice;
+    std::shared_ptr<IKeymasterDevice> mDevice;
     keymaster_operation_handle_t mOpHandle;
     DISALLOW_COPY_AND_ASSIGN(KeymasterOperation);
     friend class Keymaster;
 };
 
-// Wrapper for a keymaster1_device_t representing an open connection
-// to the keymaster, which is closed in the destructor.
+// Wrapper for a Keymaster device for methods that start a KeymasterOperation or are not
+// part of one.
 class Keymaster {
   public:
     Keymaster();
-    ~Keymaster() {
-        if (mDevice) keymaster1_close(mDevice);
-    }
     // false if we failed to open the keymaster device.
     explicit operator bool() { return mDevice != nullptr; }
     // Generate a key in the keymaster from the given params.
@@ -90,7 +87,7 @@
                              const AuthorizationSet& inParams);
 
   private:
-    keymaster1_device_t* mDevice;
+    std::shared_ptr<IKeymasterDevice> mDevice;
     DISALLOW_COPY_AND_ASSIGN(Keymaster);
 };