Merge "Use accessors for certificates and RSA keys."
diff --git a/keystore/keystore_cli_v2.cpp b/keystore/keystore_cli_v2.cpp
index 6e45ee2..43f72a9 100644
--- a/keystore/keystore_cli_v2.cpp
+++ b/keystore/keystore_cli_v2.cpp
@@ -56,7 +56,7 @@
     keymint::AuthorizationSet parameters;
 };
 
-constexpr const char keystore2_service_name[] = "android.system.keystore2";
+constexpr const char keystore2_service_name[] = "android.system.keystore2.IKeystoreService/default";
 
 int unwrapError(const ndk::ScopedAStatus& status) {
     if (status.isOk()) return 0;
@@ -769,7 +769,7 @@
         sec_level->generateKey(keyDescriptor(name), {} /* attestationKey */, params.vector_data(),
                                0 /* flags */, {} /* entropy */, &keyMetadata);
 
-    if (rc.isOk()) {
+    if (!rc.isOk()) {
         std::cerr << "GenerateKey failed: " << rc.getDescription() << std::endl;
         return unwrapError(rc);
     }
diff --git a/keystore/tests/Android.bp b/keystore/tests/Android.bp
index 249cb77..39601eb 100644
--- a/keystore/tests/Android.bp
+++ b/keystore/tests/Android.bp
@@ -62,9 +62,9 @@
         "libgtest_main",
         "libutils",
         "liblog",
+        "android.security.apc-ndk_platform",
     ],
     shared_libs: [
-        "android.security.apc-ndk_platform",
         "libbinder_ndk",
     ],
    sanitize: {
diff --git a/keystore2/Android.bp b/keystore2/Android.bp
index 8d0f4e7..32493c0 100644
--- a/keystore2/Android.bp
+++ b/keystore2/Android.bp
@@ -62,6 +62,9 @@
     shared_libs: [
         "libcutils",
     ],
+    features: [
+        "watchdog",
+    ],
 }
 
 rust_library {
@@ -76,7 +79,7 @@
     rustlibs: [
         "liblog_rust",
         "librand",
-    ]
+    ],
 }
 
 rust_test {
@@ -91,6 +94,10 @@
         "libkeystore2_test_utils",
         "libnix",
     ],
+    // The test should always include watchdog.
+    features: [
+        "watchdog",
+    ],
 }
 
 rust_binary {
@@ -105,5 +112,39 @@
     ],
     init_rc: ["keystore2.rc"],
 
+    // In S, keystore2 is the only process using dynamically linked Rust from
+    // /system. As a result, the usual savings from sharing libraries don't
+    // apply.
+    // Remove `prefer_rlib: true` once we have several processes, once a space
+    // calculation shows net RAM savings, or once we have automatic variant
+    // selection available in the build system.
+    prefer_rlib: true,
+
+    // TODO(b/187412695)
+    // This is a hack to work around the build system not installing
+    // dynamic dependencies of rlibs to the device. This section should
+    // be removed once that works correctly.
+    shared_libs: [
+        "android.hardware.confirmationui@1.0",
+        "android.hardware.security.sharedsecret-V1-ndk_platform",
+        "android.security.compat-ndk_platform",
+        "libc",
+        "libdl_android",
+        "libdl",
+        "libandroidicu",
+        "libkeymint",
+        "libkeystore2_aaid",
+        "libkeystore2_apc_compat",
+        "libkeystore2_crypto",
+        "libkeystore2_vintf_cpp",
+        "libkm_compat_service",
+        "libkm_compat",
+        "libm",
+        "libstatspull",
+        "libstatssocket",
+    ],
+
     vintf_fragments: ["android.system.keystore2-service.xml"],
+
+    required: ["keystore_cli_v2"],
 }
diff --git a/keystore2/selinux/Android.bp b/keystore2/selinux/Android.bp
index 18063d3..254f95e 100644
--- a/keystore2/selinux/Android.bp
+++ b/keystore2/selinux/Android.bp
@@ -34,6 +34,7 @@
 
     rustlibs: [
         "libanyhow",
+        "liblazy_static",
         "liblog_rust",
         "libselinux_bindgen",
         "libthiserror",
@@ -56,8 +57,30 @@
     rustlibs: [
         "libandroid_logger",
         "libanyhow",
+        "liblazy_static",
         "liblog_rust",
         "libselinux_bindgen",
         "libthiserror",
     ],
 }
+
+rust_test {
+    name: "keystore2_selinux_concurrency_test",
+    srcs: [
+        "src/concurrency_test.rs",
+    ],
+    crate_name: "keystore2_selinux_concurrency_test",
+    test_suites: ["general-tests"],
+    auto_gen_config: true,
+
+    rustlibs: [
+        "libandroid_logger",
+        "libanyhow",
+        "libkeystore2_selinux",
+        "liblazy_static",
+        "liblog_rust",
+        "libnix",
+        "libnum_cpus",
+        "libthiserror",
+    ],
+}
diff --git a/keystore2/selinux/src/concurrency_test.rs b/keystore2/selinux/src/concurrency_test.rs
new file mode 100644
index 0000000..a5d2df2
--- /dev/null
+++ b/keystore2/selinux/src/concurrency_test.rs
@@ -0,0 +1,190 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use keystore2_selinux::{check_access, Context};
+use nix::sched::sched_setaffinity;
+use nix::sched::CpuSet;
+use nix::unistd::getpid;
+use std::thread;
+use std::{
+    sync::{atomic::AtomicU8, atomic::Ordering, Arc},
+    time::{Duration, Instant},
+};
+
+#[derive(Clone, Copy)]
+struct CatCount(u8, u8, u8, u8);
+
+impl CatCount {
+    fn next(&mut self) -> CatCount {
+        let result = *self;
+        if self.3 == 255 {
+            if self.2 == 254 {
+                if self.1 == 253 {
+                    if self.0 == 252 {
+                        self.0 = 255;
+                    }
+                    self.0 += 1;
+                    self.1 = self.0;
+                }
+                self.1 += 1;
+                self.2 = self.1;
+            }
+            self.2 += 1;
+            self.3 = self.2;
+        }
+        self.3 += 1;
+        result
+    }
+
+    fn make_string(&self) -> String {
+        format!("c{},c{},c{},c{}", self.0, self.1, self.2, self.3)
+    }
+}
+
+impl Default for CatCount {
+    fn default() -> Self {
+        Self(0, 1, 2, 3)
+    }
+}
+
+/// This test calls selinux_check_access concurrently causing access vector cache misses
+/// in libselinux avc. The test then checks if any of the threads fails to report back
+/// after a burst of access checks. The purpose of the test is to draw out a specific
+/// access vector cache corruption that sends a calling thread into an infinite loop.
+/// This was observed when keystore2 used libselinux concurrently in a non thread safe
+/// way. See b/184006658.
+#[test]
+fn test_concurrent_check_access() {
+    android_logger::init_once(
+        android_logger::Config::default()
+            .with_tag("keystore2_selinux_concurrency_test")
+            .with_min_level(log::Level::Debug),
+    );
+
+    let cpus = num_cpus::get();
+    let turnpike = Arc::new(AtomicU8::new(0));
+    let complete_count = Arc::new(AtomicU8::new(0));
+    let mut threads: Vec<thread::JoinHandle<()>> = Vec::new();
+
+    for i in 0..cpus {
+        log::info!("Spawning thread {}", i);
+        let turnpike_clone = turnpike.clone();
+        let complete_count_clone = complete_count.clone();
+        threads.push(thread::spawn(move || {
+            let mut cpu_set = CpuSet::new();
+            cpu_set.set(i).unwrap();
+            sched_setaffinity(getpid(), &cpu_set).unwrap();
+            let mut cat_count: CatCount = Default::default();
+
+            log::info!("Thread 0 reached turnpike");
+            loop {
+                turnpike_clone.fetch_add(1, Ordering::Relaxed);
+                loop {
+                    match turnpike_clone.load(Ordering::Relaxed) {
+                        0 => break,
+                        255 => return,
+                        _ => {}
+                    }
+                }
+
+                for _ in 0..250 {
+                    let (tctx, sctx, perm, class) = (
+                        Context::new("u:object_r:keystore:s0").unwrap(),
+                        Context::new(&format!(
+                            "u:r:untrusted_app:s0:{}",
+                            cat_count.next().make_string()
+                        ))
+                        .unwrap(),
+                        "use",
+                        "keystore2_key",
+                    );
+
+                    check_access(&sctx, &tctx, class, perm).unwrap();
+                }
+
+                complete_count_clone.fetch_add(1, Ordering::Relaxed);
+                while complete_count_clone.load(Ordering::Relaxed) as usize != cpus {
+                    thread::sleep(Duration::from_millis(5));
+                }
+            }
+        }));
+    }
+
+    let mut i = 0;
+    let run_time = Instant::now();
+
+    loop {
+        const TEST_ITERATIONS: u32 = 500;
+        const MAX_SLEEPS: u64 = 500;
+        const SLEEP_MILLISECONDS: u64 = 5;
+        let mut sleep_count: u64 = 0;
+        while turnpike.load(Ordering::Relaxed) as usize != cpus {
+            thread::sleep(Duration::from_millis(SLEEP_MILLISECONDS));
+            sleep_count += 1;
+            assert!(
+                sleep_count < MAX_SLEEPS,
+                "Waited too long to go ready on iteration {}, only {} are ready",
+                i,
+                turnpike.load(Ordering::Relaxed)
+            );
+        }
+
+        if i % 100 == 0 {
+            let elapsed = run_time.elapsed().as_secs();
+            println!("{:02}:{:02}: Iteration {}", elapsed / 60, elapsed % 60, i);
+        }
+
+        // Give the threads some time to reach and spin on the turn pike.
+        assert_eq!(turnpike.load(Ordering::Relaxed) as usize, cpus, "i = {}", i);
+        if i >= TEST_ITERATIONS {
+            turnpike.store(255, Ordering::Relaxed);
+            break;
+        }
+
+        // Now go.
+        complete_count.store(0, Ordering::Relaxed);
+        turnpike.store(0, Ordering::Relaxed);
+        i += 1;
+
+        // Wait for them to all complete.
+        sleep_count = 0;
+        while complete_count.load(Ordering::Relaxed) as usize != cpus {
+            thread::sleep(Duration::from_millis(SLEEP_MILLISECONDS));
+            sleep_count += 1;
+            if sleep_count >= MAX_SLEEPS {
+                // Enable the following block to park the thread to allow attaching a debugger.
+                if false {
+                    println!(
+                        "Waited {} seconds and we seem stuck. Going to sleep forever.",
+                        (MAX_SLEEPS * SLEEP_MILLISECONDS) as f32 / 1000.0
+                    );
+                    loop {
+                        thread::park();
+                    }
+                } else {
+                    assert!(
+                        sleep_count < MAX_SLEEPS,
+                        "Waited too long to complete on iteration {}, only {} are complete",
+                        i,
+                        complete_count.load(Ordering::Relaxed)
+                    );
+                }
+            }
+        }
+    }
+
+    for t in threads {
+        t.join().unwrap();
+    }
+}
diff --git a/keystore2/selinux/src/lib.rs b/keystore2/selinux/src/lib.rs
index cc707e7..5197cf6 100644
--- a/keystore2/selinux/src/lib.rs
+++ b/keystore2/selinux/src/lib.rs
@@ -20,6 +20,13 @@
 //!  * selabel_lookup for the keystore2_key backend.
 //! And it provides an owning wrapper around context strings `Context`.
 
+use anyhow::Context as AnyhowContext;
+use anyhow::{anyhow, Result};
+use lazy_static::lazy_static;
+pub use selinux::pid_t;
+use selinux::SELABEL_CTX_ANDROID_KEYSTORE2_KEY;
+use selinux::SELINUX_CB_LOG;
+use selinux_bindgen as selinux;
 use std::ffi::{CStr, CString};
 use std::fmt;
 use std::io;
@@ -29,18 +36,18 @@
 use std::ptr;
 use std::sync;
 
-use selinux_bindgen as selinux;
-
-use anyhow::Context as AnyhowContext;
-use anyhow::{anyhow, Result};
-
-use selinux::SELABEL_CTX_ANDROID_KEYSTORE2_KEY;
-use selinux::SELINUX_CB_LOG;
-
-pub use selinux::pid_t;
-
 static SELINUX_LOG_INIT: sync::Once = sync::Once::new();
 
+lazy_static! {
+    /// `selinux_check_access` is only thread safe if avc_init was called with lock callbacks.
+    /// However, avc_init is deprecated and not exported by androids version of libselinux.
+    /// `selinux_set_callbacks` does not allow setting lock callbacks. So the only option
+    /// that remains right now is to put a big lock around calls into libselinux.
+    /// TODO b/188079221 It should suffice to protect `selinux_check_access` but until we are
+    /// certain of that, we leave the extra locks in place
+    static ref LIB_SELINUX_LOCK: sync::Mutex<()> = Default::default();
+}
+
 fn redirect_selinux_logs_to_logcat() {
     // `selinux_set_callback` assigns the static lifetime function pointer
     // `selinux_log_callback` to a static lifetime variable.
@@ -164,6 +171,8 @@
     /// `selinux_android_keystore2_key_context_handle`.
     pub fn new() -> Result<Self> {
         init_logger_once();
+        let _lock = LIB_SELINUX_LOCK.lock().unwrap();
+
         let handle = unsafe { selinux::selinux_android_keystore2_key_context_handle() };
         if handle.is_null() {
             return Err(anyhow!(Error::sys("Failed to open KeystoreKeyBackend")));
@@ -192,6 +201,8 @@
         match unsafe {
             // No need to initialize the logger here because it cannot run unless
             // KeystoreKeyBackend::new has run.
+            let _lock = LIB_SELINUX_LOCK.lock().unwrap();
+
             selinux::selabel_lookup(self.handle, &mut con, c_key.as_ptr(), Self::BACKEND_TYPE)
         } {
             0 => {
@@ -219,6 +230,8 @@
 ///  * Err(io::Error::last_os_error()) if getcon failed.
 pub fn getcon() -> Result<Context> {
     init_logger_once();
+    let _lock = LIB_SELINUX_LOCK.lock().unwrap();
+
     let mut con: *mut c_char = ptr::null_mut();
     match unsafe { selinux::getcon(&mut con) } {
         0 => {
@@ -241,6 +254,8 @@
 ///  * Err(io::Error::last_os_error()) if getpidcon failed.
 pub fn getpidcon(pid: selinux::pid_t) -> Result<Context> {
     init_logger_once();
+    let _lock = LIB_SELINUX_LOCK.lock().unwrap();
+
     let mut con: *mut c_char = ptr::null_mut();
     match unsafe { selinux::getpidcon(pid, &mut con) } {
         0 => {
@@ -267,6 +282,7 @@
 ///            the access check.
 pub fn check_access(source: &CStr, target: &CStr, tclass: &str, perm: &str) -> Result<()> {
     init_logger_once();
+
     let c_tclass = CString::new(tclass).with_context(|| {
         format!("check_access: Failed to convert tclass \"{}\" to CString.", tclass)
     })?;
@@ -275,6 +291,8 @@
     })?;
 
     match unsafe {
+        let _lock = LIB_SELINUX_LOCK.lock().unwrap();
+
         selinux::selinux_check_access(
             source.as_ptr(),
             target.as_ptr(),
diff --git a/keystore2/src/apc.rs b/keystore2/src/apc.rs
index 848b770..0096686 100644
--- a/keystore2/src/apc.rs
+++ b/keystore2/src/apc.rs
@@ -21,7 +21,7 @@
     sync::{mpsc::Sender, Arc, Mutex},
 };
 
-use crate::utils::{compat_2_response_code, ui_opts_2_compat};
+use crate::utils::{compat_2_response_code, ui_opts_2_compat, watchdog as wd};
 use android_security_apc::aidl::android::security::apc::{
     IConfirmationCallback::IConfirmationCallback,
     IProtectedConfirmation::{BnProtectedConfirmation, IProtectedConfirmation},
@@ -363,6 +363,8 @@
         locale: &str,
         ui_option_flags: i32,
     ) -> BinderResult<()> {
+        // presentPrompt can take more time than other operations.
+        let _wp = wd::watch_millis("IProtectedConfirmation::presentPrompt", 3000);
         map_or_log_err(
             self.present_prompt(listener, prompt_text, extra_data, locale, ui_option_flags),
             Ok,
@@ -372,9 +374,11 @@
         &self,
         listener: &binder::Strong<dyn IConfirmationCallback>,
     ) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IProtectedConfirmation::cancelPrompt", 500);
         map_or_log_err(self.cancel_prompt(listener), Ok)
     }
     fn isSupported(&self) -> BinderResult<bool> {
+        let _wp = wd::watch_millis("IProtectedConfirmation::isSupported", 500);
         map_or_log_err(Self::is_supported(), Ok)
     }
 }
diff --git a/keystore2/src/async_task.rs b/keystore2/src/async_task.rs
index 4d0034a..45f0274 100644
--- a/keystore2/src/async_task.rs
+++ b/keystore2/src/async_task.rs
@@ -19,6 +19,7 @@
 //! processed all tasks before it terminates.
 //! Note that low priority tasks are processed only when the high priority queue is empty.
 
+use crate::utils::watchdog as wd;
 use std::{any::Any, any::TypeId, time::Duration};
 use std::{
     collections::{HashMap, VecDeque},
@@ -170,6 +171,7 @@
     {
         let (ref condvar, ref state) = *self.state;
         let mut state = state.lock().unwrap();
+
         if hi_prio {
             state.hi_prio_req.push_back(Box::new(f));
         } else {
@@ -239,11 +241,14 @@
                     // Now that the lock has been dropped, perform the action.
                     match action {
                         Action::QueuedFn(f) => {
+                            let _wd = wd::watch_millis("async_task thread: calling queued fn", 500);
                             f(&mut shelf);
                             done_idle = false;
                         }
                         Action::IdleFns(idle_fns) => {
                             for idle_fn in idle_fns {
+                                let _wd =
+                                    wd::watch_millis("async_task thread: calling idle_fn", 500);
                                 idle_fn(&mut shelf);
                             }
                             done_idle = true;
diff --git a/keystore2/src/audit_log.rs b/keystore2/src/audit_log.rs
index 30fc155..3d7d26e 100644
--- a/keystore2/src/audit_log.rs
+++ b/keystore2/src/audit_log.rs
@@ -25,14 +25,15 @@
 const TAG_KEY_GENERATED: u32 = 210024;
 const TAG_KEY_IMPORTED: u32 = 210025;
 const TAG_KEY_DESTROYED: u32 = 210026;
+const TAG_KEY_INTEGRITY_VIOLATION: u32 = 210032;
 
-const NAMESPACE_MASK: i64 = 0x80000000;
+const FLAG_NAMESPACE: i64 = 0x80000000;
 
-/// For app domain returns calling app uid, for SELinux domain returns masked namespace.
-fn key_owner(key: &KeyDescriptor, calling_app: uid_t) -> i32 {
-    match key.domain {
-        Domain::APP => calling_app as i32,
-        Domain::SELINUX => (key.nspace | NAMESPACE_MASK) as i32,
+/// Encode key owner as either uid or namespace with a flag.
+fn key_owner(domain: Domain, nspace: i64, uid: i32) -> i32 {
+    match domain {
+        Domain::APP => uid,
+        Domain::SELINUX => (nspace | FLAG_NAMESPACE) as i32,
         _ => {
             log::info!("Not logging audit event for key with unexpected domain");
             0
@@ -55,12 +56,29 @@
     log_key_event(TAG_KEY_DESTROYED, key, calling_app, success);
 }
 
+/// Logs key integrity violation to NIAP audit log.
+pub fn log_key_integrity_violation(key: &KeyDescriptor) {
+    with_log_context(TAG_KEY_INTEGRITY_VIOLATION, |ctx| {
+        let owner = key_owner(key.domain, key.nspace, key.nspace as i32);
+        ctx.append_str(key.alias.as_ref().map_or("none", String::as_str)).append_i32(owner)
+    })
+}
+
 fn log_key_event(tag: u32, key: &KeyDescriptor, calling_app: uid_t, success: bool) {
-    if let Some(ctx) = LogContext::new(LogIdSecurity, tag) {
-        let event = ctx
-            .append_i32(if success { 1 } else { 0 })
+    with_log_context(tag, |ctx| {
+        let owner = key_owner(key.domain, key.nspace, calling_app as i32);
+        ctx.append_i32(if success { 1 } else { 0 })
             .append_str(key.alias.as_ref().map_or("none", String::as_str))
-            .append_i32(key_owner(key, calling_app));
+            .append_i32(owner)
+    })
+}
+
+fn with_log_context<F>(tag: u32, f: F)
+where
+    F: Fn(LogContext) -> LogContext,
+{
+    if let Some(ctx) = LogContext::new(LogIdSecurity, tag) {
+        let event = f(ctx);
         LOGS_HANDLER.queue_lo(move |_| {
             event.write();
         });
diff --git a/keystore2/src/authorization.rs b/keystore2/src/authorization.rs
index cac75c0..777089f 100644
--- a/keystore2/src/authorization.rs
+++ b/keystore2/src/authorization.rs
@@ -18,7 +18,7 @@
 use crate::globals::{ENFORCEMENTS, SUPER_KEY, DB, LEGACY_MIGRATOR};
 use crate::permission::KeystorePerm;
 use crate::super_key::UserState;
-use crate::utils::check_keystore_permission;
+use crate::utils::{check_keystore_permission, watchdog as wd};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     HardwareAuthToken::HardwareAuthToken,
 };
@@ -121,7 +121,7 @@
         // Check keystore permission.
         check_keystore_permission(KeystorePerm::add_auth()).context("In add_auth_token.")?;
 
-        ENFORCEMENTS.add_auth_token(auth_token.clone())?;
+        ENFORCEMENTS.add_auth_token(auth_token.clone());
         Ok(())
     }
 
@@ -234,6 +234,7 @@
 
 impl IKeystoreAuthorization for AuthorizationManager {
     fn addAuthToken(&self, auth_token: &HardwareAuthToken) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreAuthorization::addAuthToken", 500);
         map_or_log_err(self.add_auth_token(auth_token), Ok)
     }
 
@@ -244,6 +245,10 @@
         password: Option<&[u8]>,
         unlocking_sids: Option<&[i64]>,
     ) -> BinderResult<()> {
+        let _wp =
+            wd::watch_millis_with("IKeystoreAuthorization::onLockScreenEvent", 500, move || {
+                format!("lock event: {}", lock_screen_event.0)
+            });
         map_or_log_err(
             self.on_lock_screen_event(
                 lock_screen_event,
@@ -261,6 +266,7 @@
         secure_user_id: i64,
         auth_token_max_age_millis: i64,
     ) -> binder::public_api::Result<AuthorizationTokens> {
+        let _wp = wd::watch_millis("IKeystoreAuthorization::getAuthTokensForCredStore", 500);
         map_or_log_err(
             self.get_auth_tokens_for_credstore(
                 challenge,
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 32e2c98..2930162 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -41,12 +41,12 @@
 //! from the database module these functions take permission check
 //! callbacks.
 
-#![allow(clippy::needless_question_mark)]
+mod perboot;
 
 use crate::impl_metadata; // This is in db_utils.rs
 use crate::key_parameter::{KeyParameter, Tag};
 use crate::permission::KeyPermSet;
-use crate::utils::{get_current_time_in_seconds, AID_USER_OFFSET};
+use crate::utils::{get_current_time_in_milliseconds, watchdog as wd, AID_USER_OFFSET};
 use crate::{
     db_utils::{self, SqlField},
     gc::Gc,
@@ -63,9 +63,6 @@
     HardwareAuthToken::HardwareAuthToken,
     HardwareAuthenticatorType::HardwareAuthenticatorType, SecurityLevel::SecurityLevel,
 };
-use android_hardware_security_secureclock::aidl::android::hardware::security::secureclock::{
-    Timestamp::Timestamp,
-};
 use android_system_keystore2::aidl::android::system::keystore2::{
     Domain::Domain, KeyDescriptor::KeyDescriptor,
 };
@@ -93,7 +90,7 @@
 use std::{
     collections::{HashMap, HashSet},
     path::Path,
-    sync::{Condvar, Mutex},
+    sync::{Arc, Condvar, Mutex},
     time::{Duration, SystemTime},
 };
 
@@ -735,33 +732,29 @@
 /// ownership. It also implements all of Keystore 2.0's database functionality.
 pub struct KeystoreDB {
     conn: Connection,
-    gc: Option<Gc>,
+    gc: Option<Arc<Gc>>,
+    perboot: Arc<perboot::PerbootDB>,
 }
 
 /// Database representation of the monotonic time retrieved from the system call clock_gettime with
-/// CLOCK_MONOTONIC_RAW. Stores monotonic time as i64 in seconds.
+/// CLOCK_MONOTONIC_RAW. Stores monotonic time as i64 in milliseconds.
 #[derive(Debug, Copy, Clone, Default, Eq, PartialEq, Ord, PartialOrd)]
 pub struct MonotonicRawTime(i64);
 
 impl MonotonicRawTime {
     /// Constructs a new MonotonicRawTime
     pub fn now() -> Self {
-        Self(get_current_time_in_seconds())
+        Self(get_current_time_in_milliseconds())
     }
 
-    /// Constructs a new MonotonicRawTime from a given number of seconds.
-    pub fn from_secs(val: i64) -> Self {
-        Self(val)
+    /// Returns the value of MonotonicRawTime in milliseconds as i64
+    pub fn milliseconds(&self) -> i64 {
+        self.0
     }
 
     /// Returns the integer value of MonotonicRawTime as i64
     pub fn seconds(&self) -> i64 {
-        self.0
-    }
-
-    /// Returns the value of MonotonicRawTime in milli seconds as i64
-    pub fn milli_seconds(&self) -> i64 {
-        self.0 * 1000
+        self.0 / 1000
     }
 
     /// Like i64::checked_sub.
@@ -784,8 +777,10 @@
 
 /// This struct encapsulates the information to be stored in the database about the auth tokens
 /// received by keystore.
+#[derive(Clone)]
 pub struct AuthTokenEntry {
     auth_token: HardwareAuthToken,
+    // Time received in milliseconds
     time_received: MonotonicRawTime,
 }
 
@@ -830,19 +825,22 @@
 
 impl KeystoreDB {
     const UNASSIGNED_KEY_ID: i64 = -1i64;
-    const PERBOOT_DB_FILE_NAME: &'static str = &"file:perboot.sqlite?mode=memory&cache=shared";
 
     /// Name of the file that holds the cross-boot persistent database.
     pub const PERSISTENT_DB_FILENAME: &'static str = &"persistent.sqlite";
 
-    /// This creates a PerBootDbKeepAlive object to keep the per boot database alive.
-    pub fn keep_perboot_db_alive() -> Result<PerBootDbKeepAlive> {
-        let conn = Connection::open_in_memory()
-            .context("In keep_perboot_db_alive: Failed to initialize SQLite connection.")?;
-
-        conn.execute("ATTACH DATABASE ? as perboot;", params![Self::PERBOOT_DB_FILE_NAME])
-            .context("In keep_perboot_db_alive: Failed to attach database perboot.")?;
-        Ok(PerBootDbKeepAlive(conn))
+    /// Set write-ahead logging mode on the persistent database found in `db_root`.
+    pub fn set_wal_mode(db_root: &Path) -> Result<()> {
+        let path = Self::make_persistent_path(&db_root)?;
+        let conn =
+            Connection::open(path).context("In KeystoreDB::set_wal_mode: Failed to open DB")?;
+        let mode: String = conn
+            .pragma_update_and_check(None, "journal_mode", &"WAL", |row| row.get(0))
+            .context("In KeystoreDB::set_wal_mode: Failed to set journal_mode")?;
+        match mode.as_str() {
+            "wal" => Ok(()),
+            _ => Err(anyhow!("Unable to set WAL mode, db is still in {} mode.", mode)),
+        }
     }
 
     /// This will create a new database connection connecting the two
@@ -850,21 +848,13 @@
     /// It also attempts to initialize all of the tables.
     /// KeystoreDB cannot be used by multiple threads.
     /// Each thread should open their own connection using `thread_local!`.
-    pub fn new(db_root: &Path, gc: Option<Gc>) -> Result<Self> {
-        // Build the path to the sqlite file.
-        let mut persistent_path = db_root.to_path_buf();
-        persistent_path.push(Self::PERSISTENT_DB_FILENAME);
+    pub fn new(db_root: &Path, gc: Option<Arc<Gc>>) -> Result<Self> {
+        let _wp = wd::watch_millis("KeystoreDB::new", 500);
 
-        // Now convert them to strings prefixed with "file:"
-        let mut persistent_path_str = "file:".to_owned();
-        persistent_path_str.push_str(&persistent_path.to_string_lossy());
+        let persistent_path = Self::make_persistent_path(&db_root)?;
+        let conn = Self::make_connection(&persistent_path)?;
 
-        let conn = Self::make_connection(&persistent_path_str, &Self::PERBOOT_DB_FILE_NAME)?;
-
-        // On busy fail Immediately. It is unlikely to succeed given a bug in sqlite.
-        conn.busy_handler(None).context("In KeystoreDB::new: Failed to set busy handler.")?;
-
-        let mut db = Self { conn, gc };
+        let mut db = Self { conn, gc, perboot: perboot::PERBOOT_DB.clone() };
         db.with_transaction(TransactionBehavior::Immediate, |tx| {
             Self::init_tables(tx).context("Trying to initialize tables.").no_gc()
         })?;
@@ -978,41 +968,22 @@
         )
         .context("Failed to initialize \"grant\" table.")?;
 
-        //TODO: only drop the following two perboot tables if this is the first start up
-        //during the boot (b/175716626).
-        // tx.execute("DROP TABLE IF EXISTS perboot.authtoken;", NO_PARAMS)
-        //     .context("Failed to drop perboot.authtoken table")?;
-        tx.execute(
-            "CREATE TABLE IF NOT EXISTS perboot.authtoken (
-                        id INTEGER PRIMARY KEY,
-                        challenge INTEGER,
-                        user_id INTEGER,
-                        auth_id INTEGER,
-                        authenticator_type INTEGER,
-                        timestamp INTEGER,
-                        mac BLOB,
-                        time_received INTEGER,
-                        UNIQUE(user_id, auth_id, authenticator_type));",
-            NO_PARAMS,
-        )
-        .context("Failed to initialize \"authtoken\" table.")?;
-
-        // tx.execute("DROP TABLE IF EXISTS perboot.metadata;", NO_PARAMS)
-        //     .context("Failed to drop perboot.metadata table")?;
-        // metadata table stores certain miscellaneous information required for keystore functioning
-        // during a boot cycle, as key-value pairs.
-        tx.execute(
-            "CREATE TABLE IF NOT EXISTS perboot.metadata (
-                        key TEXT,
-                        value BLOB,
-                        UNIQUE(key));",
-            NO_PARAMS,
-        )
-        .context("Failed to initialize \"metadata\" table.")?;
         Ok(())
     }
 
-    fn make_connection(persistent_file: &str, perboot_file: &str) -> Result<Connection> {
+    fn make_persistent_path(db_root: &Path) -> Result<String> {
+        // Build the path to the sqlite file.
+        let mut persistent_path = db_root.to_path_buf();
+        persistent_path.push(Self::PERSISTENT_DB_FILENAME);
+
+        // Now convert them to strings prefixed with "file:"
+        let mut persistent_path_str = "file:".to_owned();
+        persistent_path_str.push_str(&persistent_path.to_string_lossy());
+
+        Ok(persistent_path_str)
+    }
+
+    fn make_connection(persistent_file: &str) -> Result<Connection> {
         let conn =
             Connection::open_in_memory().context("Failed to initialize SQLite connection.")?;
 
@@ -1030,20 +1001,10 @@
             }
             break;
         }
-        loop {
-            if let Err(e) = conn
-                .execute("ATTACH DATABASE ? as perboot;", params![perboot_file])
-                .context("Failed to attach database perboot.")
-            {
-                if Self::is_locked_error(&e) {
-                    std::thread::sleep(std::time::Duration::from_micros(500));
-                    continue;
-                } else {
-                    return Err(e);
-                }
-            }
-            break;
-        }
+
+        // Drop the cache size from default (2M) to 0.5M
+        conn.execute("PRAGMA persistent.cache_size = -500;", params![])
+            .context("Failed to decrease cache size for persistent db")?;
 
         Ok(conn)
     }
@@ -1096,6 +1057,8 @@
         &mut self,
         storage_type: StatsdStorageType,
     ) -> Result<Keystore2StorageStats> {
+        let _wp = wd::watch_millis("KeystoreDB::get_storage_stat", 500);
+
         match storage_type {
             StatsdStorageType::Database => self.get_total_size(),
             StatsdStorageType::KeyEntry => {
@@ -1127,7 +1090,15 @@
             }
             StatsdStorageType::Grant => self.get_table_size(storage_type, "persistent", "grant"),
             StatsdStorageType::AuthToken => {
-                self.get_table_size(storage_type, "perboot", "authtoken")
+                // Since the table is actually a BTreeMap now, unused_size is not meaningfully
+                // reportable
+                // Size provided is only an approximation
+                Ok(Keystore2StorageStats {
+                    storage_type,
+                    size: (self.perboot.auth_tokens_len() * std::mem::size_of::<AuthTokenEntry>())
+                        as i64,
+                    unused_size: 0,
+                })
             }
             StatsdStorageType::BlobMetadata => {
                 self.get_table_size(storage_type, "persistent", "blobmetadata")
@@ -1143,51 +1114,71 @@
     }
 
     /// This function is intended to be used by the garbage collector.
-    /// It deletes the blob given by `blob_id_to_delete`. It then tries to find a superseded
-    /// key blob that might need special handling by the garbage collector.
+    /// It deletes the blobs given by `blob_ids_to_delete`. It then tries to find up to `max_blobs`
+    /// superseded key blobs that might need special handling by the garbage collector.
     /// If no further superseded blobs can be found it deletes all other superseded blobs that don't
     /// need special handling and returns None.
-    pub fn handle_next_superseded_blob(
+    pub fn handle_next_superseded_blobs(
         &mut self,
-        blob_id_to_delete: Option<i64>,
-    ) -> Result<Option<(i64, Vec<u8>, BlobMetaData)>> {
+        blob_ids_to_delete: &[i64],
+        max_blobs: usize,
+    ) -> Result<Vec<(i64, Vec<u8>, BlobMetaData)>> {
+        let _wp = wd::watch_millis("KeystoreDB::handle_next_superseded_blob", 500);
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            // Delete the given blob if one was given.
-            if let Some(blob_id_to_delete) = blob_id_to_delete {
+            // Delete the given blobs.
+            for blob_id in blob_ids_to_delete {
                 tx.execute(
                     "DELETE FROM persistent.blobmetadata WHERE blobentryid = ?;",
-                    params![blob_id_to_delete],
+                    params![blob_id],
                 )
                 .context("Trying to delete blob metadata.")?;
-                tx.execute(
-                    "DELETE FROM persistent.blobentry WHERE id = ?;",
-                    params![blob_id_to_delete],
-                )
-                .context("Trying to blob.")?;
+                tx.execute("DELETE FROM persistent.blobentry WHERE id = ?;", params![blob_id])
+                    .context("Trying to blob.")?;
             }
 
-            // Find another superseded keyblob load its metadata and return it.
-            if let Some((blob_id, blob)) = tx
-                .query_row(
-                    "SELECT id, blob FROM persistent.blobentry
-                     WHERE subcomponent_type = ?
-                     AND (
-                         id NOT IN (
-                             SELECT MAX(id) FROM persistent.blobentry
-                             WHERE subcomponent_type = ?
-                             GROUP BY keyentryid, subcomponent_type
-                         )
-                     OR keyentryid NOT IN (SELECT id FROM persistent.keyentry)
-                 );",
-                    params![SubComponentType::KEY_BLOB, SubComponentType::KEY_BLOB],
-                    |row| Ok((row.get(0)?, row.get(1)?)),
-                )
-                .optional()
-                .context("Trying to query superseded blob.")?
-            {
-                let blob_metadata = BlobMetaData::load_from_db(blob_id, tx)
-                    .context("Trying to load blob metadata.")?;
-                return Ok(Some((blob_id, blob, blob_metadata))).no_gc();
+            Self::cleanup_unreferenced(tx).context("Trying to cleanup unreferenced.")?;
+
+            // Find up to max_blobx more superseded key blobs, load their metadata and return it.
+            let result: Vec<(i64, Vec<u8>)> = {
+                let mut stmt = tx
+                    .prepare(
+                        "SELECT id, blob FROM persistent.blobentry
+                        WHERE subcomponent_type = ?
+                        AND (
+                            id NOT IN (
+                                SELECT MAX(id) FROM persistent.blobentry
+                                WHERE subcomponent_type = ?
+                                GROUP BY keyentryid, subcomponent_type
+                            )
+                        OR keyentryid NOT IN (SELECT id FROM persistent.keyentry)
+                    ) LIMIT ?;",
+                    )
+                    .context("Trying to prepare query for superseded blobs.")?;
+
+                let rows = stmt
+                    .query_map(
+                        params![
+                            SubComponentType::KEY_BLOB,
+                            SubComponentType::KEY_BLOB,
+                            max_blobs as i64,
+                        ],
+                        |row| Ok((row.get(0)?, row.get(1)?)),
+                    )
+                    .context("Trying to query superseded blob.")?;
+
+                rows.collect::<Result<Vec<(i64, Vec<u8>)>, rusqlite::Error>>()
+                    .context("Trying to extract superseded blobs.")?
+            };
+
+            let result = result
+                .into_iter()
+                .map(|(blob_id, blob)| {
+                    Ok((blob_id, blob, BlobMetaData::load_from_db(blob_id, tx)?))
+                })
+                .collect::<Result<Vec<(i64, Vec<u8>, BlobMetaData)>>>()
+                .context("Trying to load blob metadata.")?;
+            if !result.is_empty() {
+                return Ok(result).no_gc();
             }
 
             // We did not find any superseded key blob, so let's remove other superseded blob in
@@ -1206,9 +1197,9 @@
             )
             .context("Trying to purge superseded blobs.")?;
 
-            Ok(None).no_gc()
+            Ok(vec![]).no_gc()
         })
-        .context("In handle_next_superseded_blob.")
+        .context("In handle_next_superseded_blobs.")
     }
 
     /// This maintenance function should be called only once before the database is used for the
@@ -1220,6 +1211,8 @@
     /// Unlike with `mark_unreferenced`, we don't need to purge grants, because only keys that made
     /// it to `KeyLifeCycle::Live` may have grants.
     pub fn cleanup_leftovers(&mut self) -> Result<usize> {
+        let _wp = wd::watch_millis("KeystoreDB::cleanup_leftovers", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             tx.execute(
                 "UPDATE persistent.keyentry SET state = ? WHERE state = ?;",
@@ -1239,6 +1232,8 @@
         alias: &str,
         key_type: KeyType,
     ) -> Result<bool> {
+        let _wp = wd::watch_millis("KeystoreDB::key_exists", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let key_descriptor =
                 KeyDescriptor { domain, nspace, alias: Some(alias.to_string()), blob: None };
@@ -1264,6 +1259,8 @@
         blob_metadata: &BlobMetaData,
         key_metadata: &KeyMetaData,
     ) -> Result<KeyEntry> {
+        let _wp = wd::watch_millis("KeystoreDB::store_super_key", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let key_id = Self::insert_with_retry(|id| {
                 tx.execute(
@@ -1307,6 +1304,8 @@
         key_type: &SuperKeyType,
         user_id: u32,
     ) -> Result<Option<(KeyIdGuard, KeyEntry)>> {
+        let _wp = wd::watch_millis("KeystoreDB::load_super_key", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let key_descriptor = KeyDescriptor {
                 domain: Domain::APP,
@@ -1346,6 +1345,8 @@
     where
         F: Fn() -> Result<(Vec<u8>, BlobMetaData)>,
     {
+        let _wp = wd::watch_millis("KeystoreDB::get_or_create_key_with", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let id = {
                 let mut stmt = tx
@@ -1424,18 +1425,6 @@
         .context("In get_or_create_key_with.")
     }
 
-    /// SQLite3 seems to hold a shared mutex while running the busy handler when
-    /// waiting for the database file to become available. This makes it
-    /// impossible to successfully recover from a locked database when the
-    /// transaction holding the device busy is in the same process on a
-    /// different connection. As a result the busy handler has to time out and
-    /// fail in order to make progress.
-    ///
-    /// Instead, we set the busy handler to None (return immediately). And catch
-    /// Busy and Locked errors (the latter occur on in memory databases with
-    /// shared cache, e.g., the per-boot database.) and restart the transaction
-    /// after a grace period of half a millisecond.
-    ///
     /// Creates a transaction with the given behavior and executes f with the new transaction.
     /// The transaction is committed only if f returns Ok and retried if DatabaseBusy
     /// or DatabaseLocked is encountered.
@@ -1494,6 +1483,8 @@
         namespace: &i64,
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
+        let _wp = wd::watch_millis("KeystoreDB::create_key_entry", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             Self::create_key_entry_internal(tx, domain, namespace, km_uuid).no_gc()
         })
@@ -1545,6 +1536,8 @@
         private_key: &[u8],
         km_uuid: &Uuid,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::create_attestation_key_entry", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let key_id = KEY_ID_LOCK.get(
                 Self::insert_with_retry(|id| {
@@ -1587,6 +1580,8 @@
         blob: Option<&[u8]>,
         blob_metadata: Option<&BlobMetaData>,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::set_blob", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             Self::set_blob_internal(&tx, key_id.0, sc_type, blob, blob_metadata).need_gc()
         })
@@ -1598,6 +1593,8 @@
     /// We use this to insert key blobs into the database which can then be garbage collected
     /// lazily by the key garbage collector.
     pub fn set_deleted_blob(&mut self, blob: &[u8], blob_metadata: &BlobMetaData) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::set_deleted_blob", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             Self::set_blob_internal(
                 &tx,
@@ -1708,6 +1705,8 @@
         expiration_date: i64,
         km_uuid: &Uuid,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::store_signed_attestation_certificate_chain", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let mut stmt = tx
                 .prepare(
@@ -1777,6 +1776,8 @@
         namespace: i64,
         km_uuid: &Uuid,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::assign_attestation_key", 500);
+
         match domain {
             Domain::APP | Domain::SELINUX => {}
             _ => {
@@ -1839,6 +1840,8 @@
         num_keys: i32,
         km_uuid: &Uuid,
     ) -> Result<Vec<Vec<u8>>> {
+        let _wp = wd::watch_millis("KeystoreDB::fetch_unsigned_attestation_keys", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let mut stmt = tx
                 .prepare(
@@ -1864,7 +1867,7 @@
                         km_uuid,
                         num_keys
                     ],
-                    |row| Ok(row.get(0)?),
+                    |row| row.get(0),
                 )?
                 .collect::<rusqlite::Result<Vec<Vec<u8>>>>()
                 .context("Failed to execute statement")?;
@@ -1876,6 +1879,8 @@
     /// Removes any keys that have expired as of the current time. Returns the number of keys
     /// marked unreferenced that are bound to be garbage collected.
     pub fn delete_expired_attestation_keys(&mut self) -> Result<i32> {
+        let _wp = wd::watch_millis("KeystoreDB::delete_expired_attestation_keys", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let mut stmt = tx
                 .prepare(
@@ -1911,6 +1916,8 @@
     /// Deletes all remotely provisioned attestation keys in the system, regardless of the state
     /// they are in. This is useful primarily as a testing mechanism.
     pub fn delete_all_attestation_keys(&mut self) -> Result<i64> {
+        let _wp = wd::watch_millis("KeystoreDB::delete_all_attestation_keys", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let mut stmt = tx
                 .prepare(
@@ -1919,7 +1926,7 @@
                 )
                 .context("Failed to prepare statement")?;
             let keys_to_delete = stmt
-                .query_map(params![KeyType::Attestation], |row| Ok(row.get(0)?))?
+                .query_map(params![KeyType::Attestation], |row| row.get(0))?
                 .collect::<rusqlite::Result<Vec<i64>>>()
                 .context("Failed to execute statement")?;
             let num_deleted = keys_to_delete
@@ -1942,6 +1949,8 @@
         date: i64,
         km_uuid: &Uuid,
     ) -> Result<AttestationPoolStatus> {
+        let _wp = wd::watch_millis("KeystoreDB::get_attestation_pool_status", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let mut stmt = tx.prepare(
                 "SELECT data
@@ -1962,7 +1971,7 @@
                         km_uuid,
                         KeyLifeCycle::Live
                     ],
-                    |row| Ok(row.get(0)?),
+                    |row| row.get(0),
                 )?
                 .collect::<rusqlite::Result<Vec<DateTime>>>()
                 .context("Failed to execute metadata statement")?;
@@ -2009,6 +2018,8 @@
         namespace: i64,
         km_uuid: &Uuid,
     ) -> Result<Option<CertificateChain>> {
+        let _wp = wd::watch_millis("KeystoreDB::retrieve_attestation_key_and_cert_chain", 500);
+
         match domain {
             Domain::APP | Domain::SELINUX => {}
             _ => {
@@ -2143,6 +2154,8 @@
         caller_uid: u32,
         check_permission: impl Fn(&KeyDescriptor) -> Result<()>,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::migrate_key_namespace", 500);
+
         let destination = match destination.domain {
             Domain::APP => KeyDescriptor { nspace: caller_uid as i64, ..(*destination).clone() },
             Domain::SELINUX => (*destination).clone(),
@@ -2211,6 +2224,8 @@
         metadata: &KeyMetaData,
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
+        let _wp = wd::watch_millis("KeystoreDB::store_new_key", 500);
+
         let (alias, domain, namespace) = match key {
             KeyDescriptor { alias: Some(alias), domain: Domain::APP, nspace, blob: None }
             | KeyDescriptor { alias: Some(alias), domain: Domain::SELINUX, nspace, blob: None } => {
@@ -2266,6 +2281,8 @@
         cert: &[u8],
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
+        let _wp = wd::watch_millis("KeystoreDB::store_new_certificate", 500);
+
         let (alias, domain, namespace) = match key {
             KeyDescriptor { alias: Some(alias), domain: Domain::APP, nspace, blob: None }
             | KeyDescriptor { alias: Some(alias), domain: Domain::SELINUX, nspace, blob: None } => {
@@ -2378,11 +2395,12 @@
                 let mut stmt = tx
                     .prepare(
                         "SELECT keyentryid, access_vector FROM persistent.grant
-                            WHERE grantee = ? AND id = ?;",
+                            WHERE grantee = ? AND id = ? AND
+                            (SELECT state FROM persistent.keyentry WHERE id = keyentryid) = ?;",
                     )
                     .context("Domain::GRANT prepare statement failed")?;
                 let mut rows = stmt
-                    .query(params![caller_uid as i64, key.nspace])
+                    .query(params![caller_uid as i64, key.nspace, KeyLifeCycle::Live])
                     .context("Domain:Grant: query failed.")?;
                 let (key_id, access_vector): (i64, i32) =
                     db_utils::with_rows_extract_one(&mut rows, |row| {
@@ -2545,6 +2563,8 @@
     /// zero, the key also gets marked unreferenced and scheduled for deletion.
     /// Returns Ok(true) if the key was marked unreferenced as a hint to the garbage collector.
     pub fn check_and_update_key_usage_count(&mut self, key_id: i64) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::check_and_update_key_usage_count", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let limit: Option<i32> = tx
                 .query_row(
@@ -2591,6 +2611,8 @@
         caller_uid: u32,
         check_permission: impl Fn(&KeyDescriptor, Option<KeyPermSet>) -> Result<()>,
     ) -> Result<(KeyIdGuard, KeyEntry)> {
+        let _wp = wd::watch_millis("KeystoreDB::load_key_entry", 500);
+
         loop {
             match self.load_key_entry_internal(
                 key,
@@ -2718,6 +2740,8 @@
         caller_uid: u32,
         check_permission: impl Fn(&KeyDescriptor, Option<KeyPermSet>) -> Result<()>,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::unbind_key", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let (key_id, access_key_descriptor, access_vector) =
                 Self::load_access_tuple(tx, key, key_type, caller_uid)
@@ -2747,6 +2771,8 @@
     /// Delete all artifacts belonging to the namespace given by the domain-namespace tuple.
     /// This leaves all of the blob entries orphaned for subsequent garbage collection.
     pub fn unbind_keys_for_namespace(&mut self, domain: Domain, namespace: i64) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::unbind_keys_for_namespace", 500);
+
         if !(domain == Domain::APP || domain == Domain::SELINUX) {
             return Err(KsError::Rc(ResponseCode::INVALID_ARGUMENT))
                 .context("In unbind_keys_for_namespace.");
@@ -2756,32 +2782,33 @@
                 "DELETE FROM persistent.keymetadata
                 WHERE keyentryid IN (
                     SELECT id FROM persistent.keyentry
-                    WHERE domain = ? AND namespace = ?
+                    WHERE domain = ? AND namespace = ? AND key_type = ?
                 );",
-                params![domain.0, namespace],
+                params![domain.0, namespace, KeyType::Client],
             )
             .context("Trying to delete keymetadata.")?;
             tx.execute(
                 "DELETE FROM persistent.keyparameter
                 WHERE keyentryid IN (
                     SELECT id FROM persistent.keyentry
-                    WHERE domain = ? AND namespace = ?
+                    WHERE domain = ? AND namespace = ? AND key_type = ?
                 );",
-                params![domain.0, namespace],
+                params![domain.0, namespace, KeyType::Client],
             )
             .context("Trying to delete keyparameters.")?;
             tx.execute(
                 "DELETE FROM persistent.grant
                 WHERE keyentryid IN (
                     SELECT id FROM persistent.keyentry
-                    WHERE domain = ? AND namespace = ?
+                    WHERE domain = ? AND namespace = ? AND key_type = ?
                 );",
-                params![domain.0, namespace],
+                params![domain.0, namespace, KeyType::Client],
             )
             .context("Trying to delete grants.")?;
             tx.execute(
-                "DELETE FROM persistent.keyentry WHERE domain = ? AND namespace = ?;",
-                params![domain.0, namespace],
+                "DELETE FROM persistent.keyentry
+                 WHERE domain = ? AND namespace = ? AND key_type = ?;",
+                params![domain.0, namespace, KeyType::Client],
             )
             .context("Trying to delete keyentry.")?;
             Ok(()).need_gc()
@@ -2789,6 +2816,47 @@
         .context("In unbind_keys_for_namespace")
     }
 
+    fn cleanup_unreferenced(tx: &Transaction) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::cleanup_unreferenced", 500);
+        {
+            tx.execute(
+                "DELETE FROM persistent.keymetadata
+            WHERE keyentryid IN (
+                SELECT id FROM persistent.keyentry
+                WHERE state = ?
+            );",
+                params![KeyLifeCycle::Unreferenced],
+            )
+            .context("Trying to delete keymetadata.")?;
+            tx.execute(
+                "DELETE FROM persistent.keyparameter
+            WHERE keyentryid IN (
+                SELECT id FROM persistent.keyentry
+                WHERE state = ?
+            );",
+                params![KeyLifeCycle::Unreferenced],
+            )
+            .context("Trying to delete keyparameters.")?;
+            tx.execute(
+                "DELETE FROM persistent.grant
+            WHERE keyentryid IN (
+                SELECT id FROM persistent.keyentry
+                WHERE state = ?
+            );",
+                params![KeyLifeCycle::Unreferenced],
+            )
+            .context("Trying to delete grants.")?;
+            tx.execute(
+                "DELETE FROM persistent.keyentry
+                WHERE state = ?;",
+                params![KeyLifeCycle::Unreferenced],
+            )
+            .context("Trying to delete keyentry.")?;
+            Result::<()>::Ok(())
+        }
+        .context("In cleanup_unreferenced")
+    }
+
     /// Delete the keys created on behalf of the user, denoted by the user id.
     /// Delete all the keys unless 'keep_non_super_encrypted_keys' set to true.
     /// Returned boolean is to hint the garbage collector to delete the unbound keys.
@@ -2798,6 +2866,8 @@
         user_id: u32,
         keep_non_super_encrypted_keys: bool,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::unbind_keys_for_user", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             let mut stmt = tx
                 .prepare(&format!(
@@ -2897,17 +2967,28 @@
     /// Returns a list of KeyDescriptors in the selected domain/namespace.
     /// The key descriptors will have the domain, nspace, and alias field set.
     /// Domain must be APP or SELINUX, the caller must make sure of that.
-    pub fn list(&mut self, domain: Domain, namespace: i64) -> Result<Vec<KeyDescriptor>> {
+    pub fn list(
+        &mut self,
+        domain: Domain,
+        namespace: i64,
+        key_type: KeyType,
+    ) -> Result<Vec<KeyDescriptor>> {
+        let _wp = wd::watch_millis("KeystoreDB::list", 500);
+
         self.with_transaction(TransactionBehavior::Deferred, |tx| {
             let mut stmt = tx
                 .prepare(
                     "SELECT alias FROM persistent.keyentry
-             WHERE domain = ? AND namespace = ? AND alias IS NOT NULL AND state = ?;",
+                     WHERE domain = ?
+                     AND namespace = ?
+                     AND alias IS NOT NULL
+                     AND state = ?
+                     AND key_type = ?;",
                 )
                 .context("In list: Failed to prepare.")?;
 
             let mut rows = stmt
-                .query(params![domain.0 as u32, namespace, KeyLifeCycle::Live])
+                .query(params![domain.0 as u32, namespace, KeyLifeCycle::Live, key_type])
                 .context("In list: Failed to query.")?;
 
             let mut descriptors: Vec<KeyDescriptor> = Vec::new();
@@ -2939,6 +3020,8 @@
         access_vector: KeyPermSet,
         check_permission: impl Fn(&KeyDescriptor, &KeyPermSet) -> Result<()>,
     ) -> Result<KeyDescriptor> {
+        let _wp = wd::watch_millis("KeystoreDB::grant", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             // Load the key_id and complete the access control tuple.
             // We ignore the access vector here because grants cannot be granted.
@@ -3004,6 +3087,8 @@
         grantee_uid: u32,
         check_permission: impl Fn(&KeyDescriptor) -> Result<()>,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("KeystoreDB::ungrant", 500);
+
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
             // Load the key_id and complete the access control tuple.
             // We ignore the access vector here because grants cannot be granted.
@@ -3053,100 +3138,59 @@
         }
     }
 
-    /// Insert or replace the auth token based on the UNIQUE constraint of the auth token table
-    pub fn insert_auth_token(&mut self, auth_token: &HardwareAuthToken) -> Result<()> {
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            tx.execute(
-                "INSERT OR REPLACE INTO perboot.authtoken (challenge, user_id, auth_id,
-            authenticator_type, timestamp, mac, time_received) VALUES(?, ?, ?, ?, ?, ?, ?);",
-                params![
-                    auth_token.challenge,
-                    auth_token.userId,
-                    auth_token.authenticatorId,
-                    auth_token.authenticatorType.0 as i32,
-                    auth_token.timestamp.milliSeconds as i64,
-                    auth_token.mac,
-                    MonotonicRawTime::now(),
-                ],
-            )
-            .context("In insert_auth_token: failed to insert auth token into the database")?;
-            Ok(()).no_gc()
-        })
+    /// Insert or replace the auth token based on (user_id, auth_id, auth_type)
+    pub fn insert_auth_token(&mut self, auth_token: &HardwareAuthToken) {
+        self.perboot.insert_auth_token_entry(AuthTokenEntry::new(
+            auth_token.clone(),
+            MonotonicRawTime::now(),
+        ))
     }
 
     /// Find the newest auth token matching the given predicate.
-    pub fn find_auth_token_entry<F>(
-        &mut self,
-        p: F,
-    ) -> Result<Option<(AuthTokenEntry, MonotonicRawTime)>>
+    pub fn find_auth_token_entry<F>(&self, p: F) -> Option<(AuthTokenEntry, MonotonicRawTime)>
     where
         F: Fn(&AuthTokenEntry) -> bool,
     {
-        self.with_transaction(TransactionBehavior::Deferred, |tx| {
-            let mut stmt = tx
-                .prepare("SELECT * from perboot.authtoken ORDER BY time_received DESC;")
-                .context("Prepare statement failed.")?;
-
-            let mut rows = stmt.query(NO_PARAMS).context("Failed to query.")?;
-
-            while let Some(row) = rows.next().context("Failed to get next row.")? {
-                let entry = AuthTokenEntry::new(
-                    HardwareAuthToken {
-                        challenge: row.get(1)?,
-                        userId: row.get(2)?,
-                        authenticatorId: row.get(3)?,
-                        authenticatorType: HardwareAuthenticatorType(row.get(4)?),
-                        timestamp: Timestamp { milliSeconds: row.get(5)? },
-                        mac: row.get(6)?,
-                    },
-                    row.get(7)?,
-                );
-                if p(&entry) {
-                    return Ok(Some((
-                        entry,
-                        Self::get_last_off_body(tx)
-                            .context("In find_auth_token_entry: Trying to get last off body")?,
-                    )))
-                    .no_gc();
-                }
-            }
-            Ok(None).no_gc()
-        })
-        .context("In find_auth_token_entry.")
+        self.perboot.find_auth_token_entry(p).map(|entry| (entry, self.get_last_off_body()))
     }
 
     /// Insert last_off_body into the metadata table at the initialization of auth token table
-    pub fn insert_last_off_body(&mut self, last_off_body: MonotonicRawTime) -> Result<()> {
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            tx.execute(
-                "INSERT OR REPLACE INTO perboot.metadata (key, value) VALUES (?, ?);",
-                params!["last_off_body", last_off_body],
-            )
-            .context("In insert_last_off_body: failed to insert.")?;
-            Ok(()).no_gc()
-        })
+    pub fn insert_last_off_body(&self, last_off_body: MonotonicRawTime) {
+        self.perboot.set_last_off_body(last_off_body)
     }
 
     /// Update last_off_body when on_device_off_body is called
-    pub fn update_last_off_body(&mut self, last_off_body: MonotonicRawTime) -> Result<()> {
-        self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            tx.execute(
-                "UPDATE perboot.metadata SET value = ? WHERE key = ?;",
-                params![last_off_body, "last_off_body"],
-            )
-            .context("In update_last_off_body: failed to update.")?;
-            Ok(()).no_gc()
-        })
+    pub fn update_last_off_body(&self, last_off_body: MonotonicRawTime) {
+        self.perboot.set_last_off_body(last_off_body)
     }
 
     /// Get last_off_body time when finding auth tokens
-    fn get_last_off_body(tx: &Transaction) -> Result<MonotonicRawTime> {
-        tx.query_row(
-            "SELECT value from perboot.metadata WHERE key = ?;",
-            params!["last_off_body"],
-            |row| Ok(row.get(0)?),
-        )
-        .context("In get_last_off_body: query_row failed.")
+    fn get_last_off_body(&self) -> MonotonicRawTime {
+        self.perboot.get_last_off_body()
+    }
+
+    /// Load descriptor of a key by key id
+    pub fn load_key_descriptor(&mut self, key_id: i64) -> Result<Option<KeyDescriptor>> {
+        let _wp = wd::watch_millis("KeystoreDB::load_key_descriptor", 500);
+
+        self.with_transaction(TransactionBehavior::Deferred, |tx| {
+            tx.query_row(
+                "SELECT domain, namespace, alias FROM persistent.keyentry WHERE id = ?;",
+                params![key_id],
+                |row| {
+                    Ok(KeyDescriptor {
+                        domain: Domain(row.get(0)?),
+                        nspace: row.get(1)?,
+                        alias: row.get(2)?,
+                        blob: None,
+                    })
+                },
+            )
+            .optional()
+            .context("Trying to load key descriptor")
+            .no_gc()
+        })
+        .context("In load_key_descriptor.")
     }
 }
 
@@ -3169,8 +3213,9 @@
     use android_hardware_security_secureclock::aidl::android::hardware::security::secureclock::{
         Timestamp::Timestamp,
     };
+    use rusqlite::DatabaseName::Attached;
     use rusqlite::NO_PARAMS;
-    use rusqlite::{Error, TransactionBehavior};
+    use rusqlite::TransactionBehavior;
     use std::cell::RefCell;
     use std::collections::BTreeMap;
     use std::fmt::Write;
@@ -3182,9 +3227,9 @@
     use std::time::Instant;
 
     fn new_test_db() -> Result<KeystoreDB> {
-        let conn = KeystoreDB::make_connection("file::memory:", "file::memory:")?;
+        let conn = KeystoreDB::make_connection("file::memory:")?;
 
-        let mut db = KeystoreDB { conn, gc: None };
+        let mut db = KeystoreDB { conn, gc: None, perboot: Arc::new(perboot::PerbootDB::new()) };
         db.with_transaction(TransactionBehavior::Immediate, |tx| {
             KeystoreDB::init_tables(tx).context("Failed to initialize tables.").no_gc()
         })?;
@@ -3200,7 +3245,7 @@
         let gc_db = KeystoreDB::new(path, None).expect("Failed to open test gc db_connection.");
         let gc = Gc::new_init_with(Default::default(), move || (Box::new(cb), gc_db, super_key));
 
-        KeystoreDB::new(path, Some(gc))
+        KeystoreDB::new(path, Some(Arc::new(gc)))
     }
 
     fn rebind_alias(
@@ -3270,15 +3315,6 @@
         assert_eq!(tables[3], "keyentry");
         assert_eq!(tables[4], "keymetadata");
         assert_eq!(tables[5], "keyparameter");
-        let tables = db
-            .conn
-            .prepare("SELECT name from perboot.sqlite_master WHERE type='table' ORDER BY name;")?
-            .query_map(params![], |row| row.get(0))?
-            .collect::<rusqlite::Result<Vec<String>>>()?;
-
-        assert_eq!(tables.len(), 2);
-        assert_eq!(tables[0], "authtoken");
-        assert_eq!(tables[1], "metadata");
         Ok(())
     }
 
@@ -3293,8 +3329,8 @@
             timestamp: Timestamp { milliSeconds: 500 },
             mac: String::from("mac").into_bytes(),
         };
-        db.insert_auth_token(&auth_token1)?;
-        let auth_tokens_returned = get_auth_tokens(&mut db)?;
+        db.insert_auth_token(&auth_token1);
+        let auth_tokens_returned = get_auth_tokens(&db);
         assert_eq!(auth_tokens_returned.len(), 1);
 
         // insert another auth token with the same values for the columns in the UNIQUE constraint
@@ -3308,8 +3344,8 @@
             mac: String::from("mac").into_bytes(),
         };
 
-        db.insert_auth_token(&auth_token2)?;
-        let mut auth_tokens_returned = get_auth_tokens(&mut db)?;
+        db.insert_auth_token(&auth_token2);
+        let mut auth_tokens_returned = get_auth_tokens(&db);
         assert_eq!(auth_tokens_returned.len(), 1);
 
         if let Some(auth_token) = auth_tokens_returned.pop() {
@@ -3327,33 +3363,16 @@
             mac: String::from("mac").into_bytes(),
         };
 
-        db.insert_auth_token(&auth_token3)?;
-        let auth_tokens_returned = get_auth_tokens(&mut db)?;
+        db.insert_auth_token(&auth_token3);
+        let auth_tokens_returned = get_auth_tokens(&db);
         assert_eq!(auth_tokens_returned.len(), 2);
 
         Ok(())
     }
 
     // utility function for test_auth_token_table_invariant()
-    fn get_auth_tokens(db: &mut KeystoreDB) -> Result<Vec<AuthTokenEntry>> {
-        let mut stmt = db.conn.prepare("SELECT * from perboot.authtoken;")?;
-
-        let auth_token_entries: Vec<AuthTokenEntry> = stmt
-            .query_map(NO_PARAMS, |row| {
-                Ok(AuthTokenEntry::new(
-                    HardwareAuthToken {
-                        challenge: row.get(1)?,
-                        userId: row.get(2)?,
-                        authenticatorId: row.get(3)?,
-                        authenticatorType: HardwareAuthenticatorType(row.get(4)?),
-                        timestamp: Timestamp { milliSeconds: row.get(5)? },
-                        mac: row.get(6)?,
-                    },
-                    row.get(7)?,
-                ))
-            })?
-            .collect::<Result<Vec<AuthTokenEntry>, Error>>()?;
-        Ok(auth_token_entries)
+    fn get_auth_tokens(db: &KeystoreDB) -> Vec<AuthTokenEntry> {
+        db.perboot.get_all_auth_token_entries()
     }
 
     #[test]
@@ -4737,7 +4756,7 @@
                 })
                 .collect();
             list_o_descriptors.sort();
-            let mut list_result = db.list(*domain, *namespace)?;
+            let mut list_result = db.list(*domain, *namespace, KeyType::Client)?;
             list_result.sort();
             assert_eq!(list_o_descriptors, list_result);
 
@@ -4767,7 +4786,7 @@
             loaded_entries.sort_unstable();
             assert_eq!(list_o_ids, loaded_entries);
         }
-        assert_eq!(Vec::<KeyDescriptor>::new(), db.list(Domain::SELINUX, 101)?);
+        assert_eq!(Vec::<KeyDescriptor>::new(), db.list(Domain::SELINUX, 101, KeyType::Client)?);
 
         Ok(())
     }
@@ -5207,17 +5226,17 @@
     #[test]
     fn test_last_off_body() -> Result<()> {
         let mut db = new_test_db()?;
-        db.insert_last_off_body(MonotonicRawTime::now())?;
+        db.insert_last_off_body(MonotonicRawTime::now());
         let tx = db.conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
-        let last_off_body_1 = KeystoreDB::get_last_off_body(&tx)?;
         tx.commit()?;
+        let last_off_body_1 = db.get_last_off_body();
         let one_second = Duration::from_secs(1);
         thread::sleep(one_second);
-        db.update_last_off_body(MonotonicRawTime::now())?;
+        db.update_last_off_body(MonotonicRawTime::now());
         let tx2 = db.conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
-        let last_off_body_2 = KeystoreDB::get_last_off_body(&tx2)?;
         tx2.commit()?;
-        assert!(last_off_body_1.seconds() < last_off_body_2.seconds());
+        let last_off_body_2 = db.get_last_off_body();
+        assert!(last_off_body_1 < last_off_body_2);
         Ok(())
     }
 
@@ -5230,11 +5249,11 @@
         make_test_key_entry(&mut db, Domain::APP, 110000, TEST_ALIAS, None)?;
         db.unbind_keys_for_user(2, false)?;
 
-        assert_eq!(1, db.list(Domain::APP, 110000)?.len());
-        assert_eq!(0, db.list(Domain::APP, 210000)?.len());
+        assert_eq!(1, db.list(Domain::APP, 110000, KeyType::Client)?.len());
+        assert_eq!(0, db.list(Domain::APP, 210000, KeyType::Client)?.len());
 
         db.unbind_keys_for_user(1, true)?;
-        assert_eq!(0, db.list(Domain::APP, 110000)?.len());
+        assert_eq!(0, db.list(Domain::APP, 110000, KeyType::Client)?.len());
 
         Ok(())
     }
@@ -5303,7 +5322,12 @@
 
         for t in get_valid_statsd_storage_types() {
             let stat = db.get_storage_stat(t)?;
-            assert!(stat.size >= PAGE_SIZE);
+            // AuthToken can be less than a page since it's in a btree, not sqlite
+            // TODO(b/187474736) stop using if-let here
+            if let StatsdStorageType::AuthToken = t {
+            } else {
+                assert!(stat.size >= PAGE_SIZE);
+            }
             assert!(stat.size >= stat.unused_size);
         }
 
@@ -5433,7 +5457,7 @@
             authenticatorType: kmhw_authenticator_type::ANY,
             timestamp: Timestamp { milliSeconds: 10 },
             mac: b"mac".to_vec(),
-        })?;
+        });
         assert_storage_increased(&mut db, vec![StatsdStorageType::AuthToken], &mut working_stats);
         Ok(())
     }
@@ -5462,4 +5486,78 @@
 
         Ok(())
     }
+
+    #[test]
+    fn find_auth_token_entry_returns_latest() -> Result<()> {
+        let mut db = new_test_db()?;
+        db.insert_auth_token(&HardwareAuthToken {
+            challenge: 123,
+            userId: 456,
+            authenticatorId: 789,
+            authenticatorType: kmhw_authenticator_type::ANY,
+            timestamp: Timestamp { milliSeconds: 10 },
+            mac: b"mac0".to_vec(),
+        });
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        db.insert_auth_token(&HardwareAuthToken {
+            challenge: 123,
+            userId: 457,
+            authenticatorId: 789,
+            authenticatorType: kmhw_authenticator_type::ANY,
+            timestamp: Timestamp { milliSeconds: 12 },
+            mac: b"mac1".to_vec(),
+        });
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        db.insert_auth_token(&HardwareAuthToken {
+            challenge: 123,
+            userId: 458,
+            authenticatorId: 789,
+            authenticatorType: kmhw_authenticator_type::ANY,
+            timestamp: Timestamp { milliSeconds: 3 },
+            mac: b"mac2".to_vec(),
+        });
+        // All three entries are in the database
+        assert_eq!(db.perboot.auth_tokens_len(), 3);
+        // It selected the most recent timestamp
+        assert_eq!(db.find_auth_token_entry(|_| true).unwrap().0.auth_token.mac, b"mac2".to_vec());
+        Ok(())
+    }
+
+    #[test]
+    fn test_set_wal_mode() -> Result<()> {
+        let temp_dir = TempDir::new("test_set_wal_mode")?;
+        let mut db = KeystoreDB::new(temp_dir.path(), None)?;
+        let mode: String =
+            db.conn.pragma_query_value(Some(Attached("persistent")), "journal_mode", |row| {
+                row.get(0)
+            })?;
+        assert_eq!(mode, "delete");
+        db.conn.close().expect("Close didn't work");
+
+        KeystoreDB::set_wal_mode(temp_dir.path())?;
+
+        db = KeystoreDB::new(temp_dir.path(), None)?;
+        let mode: String =
+            db.conn.pragma_query_value(Some(Attached("persistent")), "journal_mode", |row| {
+                row.get(0)
+            })?;
+        assert_eq!(mode, "wal");
+        Ok(())
+    }
+
+    #[test]
+    fn test_load_key_descriptor() -> Result<()> {
+        let mut db = new_test_db()?;
+        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS, None)?.0;
+
+        let key = db.load_key_descriptor(key_id)?.unwrap();
+
+        assert_eq!(key.domain, Domain::APP);
+        assert_eq!(key.nspace, 1);
+        assert_eq!(key.alias, Some(TEST_ALIAS.to_string()));
+
+        // No such id
+        assert_eq!(db.load_key_descriptor(key_id + 1)?, None);
+        Ok(())
+    }
 }
diff --git a/keystore2/src/database/perboot.rs b/keystore2/src/database/perboot.rs
new file mode 100644
index 0000000..7ff35fa
--- /dev/null
+++ b/keystore2/src/database/perboot.rs
@@ -0,0 +1,122 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! This module implements a per-boot, shared, in-memory storage of auth tokens
+//! and last-time-on-body for the main Keystore 2.0 database module.
+
+use super::{AuthTokenEntry, MonotonicRawTime};
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    HardwareAuthToken::HardwareAuthToken, HardwareAuthenticatorType::HardwareAuthenticatorType,
+};
+use lazy_static::lazy_static;
+use std::collections::HashSet;
+use std::sync::atomic::{AtomicI64, Ordering};
+use std::sync::Arc;
+use std::sync::RwLock;
+
+#[derive(PartialEq, PartialOrd, Ord, Eq, Hash)]
+struct AuthTokenId {
+    user_id: i64,
+    auth_id: i64,
+    authenticator_type: HardwareAuthenticatorType,
+}
+
+impl AuthTokenId {
+    fn from_auth_token(tok: &HardwareAuthToken) -> Self {
+        AuthTokenId {
+            user_id: tok.userId,
+            auth_id: tok.authenticatorId,
+            authenticator_type: tok.authenticatorType,
+        }
+    }
+}
+
+//Implements Eq/Hash to only operate on the AuthTokenId portion
+//of the AuthTokenEntry. This allows a HashSet to DTRT.
+#[derive(Clone)]
+struct AuthTokenEntryWrap(AuthTokenEntry);
+
+impl std::hash::Hash for AuthTokenEntryWrap {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        AuthTokenId::from_auth_token(&self.0.auth_token).hash(state)
+    }
+}
+
+impl PartialEq<AuthTokenEntryWrap> for AuthTokenEntryWrap {
+    fn eq(&self, other: &AuthTokenEntryWrap) -> bool {
+        AuthTokenId::from_auth_token(&self.0.auth_token)
+            == AuthTokenId::from_auth_token(&other.0.auth_token)
+    }
+}
+
+impl Eq for AuthTokenEntryWrap {}
+
+/// Per-boot state structure. Currently only used to track auth tokens and
+/// last-off-body.
+#[derive(Default)]
+pub struct PerbootDB {
+    // We can use a .unwrap() discipline on this lock, because only panicking
+    // while holding a .write() lock will poison it. The only write usage is
+    // an insert call which inserts a pre-constructed pair.
+    auth_tokens: RwLock<HashSet<AuthTokenEntryWrap>>,
+    // Ordering::Relaxed is appropriate for accessing this atomic, since it
+    // does not currently need to be synchronized with anything else.
+    last_off_body: AtomicI64,
+}
+
+lazy_static! {
+    /// The global instance of the perboot DB. Located here rather than in globals
+    /// in order to restrict access to the database module.
+    pub static ref PERBOOT_DB: Arc<PerbootDB> = Arc::new(PerbootDB::new());
+}
+
+impl PerbootDB {
+    /// Construct a new perboot database. Currently just uses default values.
+    pub fn new() -> Self {
+        Default::default()
+    }
+    /// Add a new auth token + timestamp to the database, replacing any which
+    /// match all of user_id, auth_id, and auth_type.
+    pub fn insert_auth_token_entry(&self, entry: AuthTokenEntry) {
+        self.auth_tokens.write().unwrap().replace(AuthTokenEntryWrap(entry));
+    }
+    /// Locate an auth token entry which matches the predicate with the most
+    /// recent update time.
+    pub fn find_auth_token_entry<P: Fn(&AuthTokenEntry) -> bool>(
+        &self,
+        p: P,
+    ) -> Option<AuthTokenEntry> {
+        let reader = self.auth_tokens.read().unwrap();
+        let mut matches: Vec<_> = reader.iter().filter(|x| p(&x.0)).collect();
+        matches.sort_by_key(|x| x.0.time_received);
+        matches.last().map(|x| x.0.clone())
+    }
+    /// Get the last time the device was off the user's body
+    pub fn get_last_off_body(&self) -> MonotonicRawTime {
+        MonotonicRawTime(self.last_off_body.load(Ordering::Relaxed))
+    }
+    /// Set the last time the device was off the user's body
+    pub fn set_last_off_body(&self, last_off_body: MonotonicRawTime) {
+        self.last_off_body.store(last_off_body.0, Ordering::Relaxed)
+    }
+    /// Return how many auth tokens are currently tracked.
+    pub fn auth_tokens_len(&self) -> usize {
+        self.auth_tokens.read().unwrap().len()
+    }
+    #[cfg(test)]
+    /// For testing, return all auth tokens currently tracked.
+    pub fn get_all_auth_token_entries(&self) -> Vec<AuthTokenEntry> {
+        self.auth_tokens.read().unwrap().iter().cloned().map(|x| x.0).collect()
+    }
+}
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 04d1f77..29a3f0b 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -638,8 +638,7 @@
                 } else {
                     unlocked_device_required
                 }
-            })
-            .context("In authorize_create: Trying to get required auth token.")?;
+            });
             Some(
                 hat_and_last_off_body
                     .ok_or(Error::Km(Ec::KEY_USER_NOT_AUTHENTICATED))
@@ -700,15 +699,11 @@
         })
     }
 
-    fn find_auth_token<F>(p: F) -> Result<Option<(AuthTokenEntry, MonotonicRawTime)>>
+    fn find_auth_token<F>(p: F) -> Option<(AuthTokenEntry, MonotonicRawTime)>
     where
         F: Fn(&AuthTokenEntry) -> bool,
     {
-        DB.with(|db| {
-            let mut db = db.borrow_mut();
-            db.find_auth_token_entry(p).context("Trying to find auth token.")
-        })
-        .context("In find_auth_token.")
+        DB.with(|db| db.borrow().find_auth_token_entry(p))
     }
 
     /// Checks if the time now since epoch is greater than (or equal, if is_given_time_inclusive is
@@ -752,11 +747,9 @@
     /// Add this auth token to the database.
     /// Then present the auth token to the op auth map. If an operation is waiting for this
     /// auth token this fulfills the request and removes the receiver from the map.
-    pub fn add_auth_token(&self, hat: HardwareAuthToken) -> Result<()> {
-        DB.with(|db| db.borrow_mut().insert_auth_token(&hat)).context("In add_auth_token.")?;
-
+    pub fn add_auth_token(&self, hat: HardwareAuthToken) {
+        DB.with(|db| db.borrow_mut().insert_auth_token(&hat));
         self.op_auth_map.add_auth_token(hat);
-        Ok(())
     }
 
     /// This allows adding an entry to the op_auth_map, indexed by the operation challenge.
@@ -824,28 +817,22 @@
         // Filter the matching auth tokens by challenge
         let result = Self::find_auth_token(|hat: &AuthTokenEntry| {
             (challenge == hat.challenge()) && hat.satisfies(&sids, auth_type)
-        })
-        .context(
-            "In get_auth_tokens: Failed to get a matching auth token filtered by challenge.",
-        )?;
+        });
 
         let auth_token = if let Some((auth_token_entry, _)) = result {
             auth_token_entry.take_auth_token()
         } else {
             // Filter the matching auth tokens by age.
             if auth_token_max_age_millis != 0 {
-                let now_in_millis = MonotonicRawTime::now().milli_seconds();
+                let now_in_millis = MonotonicRawTime::now();
                 let result = Self::find_auth_token(|auth_token_entry: &AuthTokenEntry| {
                     let token_valid = now_in_millis
-                        .checked_sub(auth_token_entry.time_received().milli_seconds())
+                        .checked_sub(&auth_token_entry.time_received())
                         .map_or(false, |token_age_in_millis| {
-                            auth_token_max_age_millis > token_age_in_millis
+                            auth_token_max_age_millis > token_age_in_millis.milliseconds()
                         });
                     token_valid && auth_token_entry.satisfies(&sids, auth_type)
-                })
-                .context(
-                    "In get_auth_tokens: Failed to get a matching auth token filtered by age.",
-                )?;
+                });
 
                 if let Some((auth_token_entry, _)) = result {
                     auth_token_entry.take_auth_token()
diff --git a/keystore2/src/error.rs b/keystore2/src/error.rs
index d1b2ffb..f969cb6 100644
--- a/keystore2/src/error.rs
+++ b/keystore2/src/error.rs
@@ -30,16 +30,13 @@
 //! Keystore functions should use `anyhow::Result` to return error conditions, and
 //! context should be added every time an error is forwarded.
 
-use std::cmp::PartialEq;
-
 pub use android_hardware_security_keymint::aidl::android::hardware::security::keymint::ErrorCode::ErrorCode;
 pub use android_system_keystore2::aidl::android::system::keystore2::ResponseCode::ResponseCode;
-
-use keystore2_selinux as selinux;
-
 use android_system_keystore2::binder::{
     ExceptionCode, Result as BinderResult, Status as BinderStatus, StatusCode,
 };
+use keystore2_selinux as selinux;
+use std::cmp::PartialEq;
 
 /// This is the main Keystore error type. It wraps the Keystore `ResponseCode` generated
 /// from AIDL in the `Rc` variant and Keymint `ErrorCode` in the Km variant.
diff --git a/keystore2/src/gc.rs b/keystore2/src/gc.rs
index 6cc0f27..2010c79 100644
--- a/keystore2/src/gc.rs
+++ b/keystore2/src/gc.rs
@@ -20,22 +20,28 @@
 
 use crate::{
     async_task,
-    database::{KeystoreDB, Uuid},
+    database::{BlobMetaData, KeystoreDB, Uuid},
     super_key::SuperKeyManager,
 };
 use anyhow::{Context, Result};
 use async_task::AsyncTask;
-use std::sync::Arc;
+use std::sync::{
+    atomic::{AtomicU8, Ordering},
+    Arc,
+};
 
 pub struct Gc {
     async_task: Arc<AsyncTask>,
+    notified: Arc<AtomicU8>,
 }
 
 impl Gc {
     /// Creates a garbage collector using the given async_task.
-    /// The garbage collector needs a function to invalidate key blobs and a database connection.
-    /// Both are obtained from the init function. The function is only called if this is first
-    /// time a garbage collector was initialized with the given AsyncTask instance.
+    /// The garbage collector needs a function to invalidate key blobs, a database connection,
+    /// and a reference to the `SuperKeyManager`. They are obtained from the init function.
+    /// The function is only called if this is first time a garbage collector was initialized
+    /// with the given AsyncTask instance.
+    /// Note: It is a logical error to initialize different Gc instances with the same `AsyncTask`.
     pub fn new_init_with<F>(async_task: Arc<AsyncTask>, init: F) -> Self
     where
         F: FnOnce() -> (
@@ -46,34 +52,43 @@
             + 'static,
     {
         let weak_at = Arc::downgrade(&async_task);
+        let notified = Arc::new(AtomicU8::new(0));
+        let notified_clone = notified.clone();
         // Initialize the task's shelf.
         async_task.queue_hi(move |shelf| {
             let (invalidate_key, db, super_key) = init();
+            let notified = notified_clone;
             shelf.get_or_put_with(|| GcInternal {
-                blob_id_to_delete: None,
+                deleted_blob_ids: vec![],
+                superseded_blobs: vec![],
                 invalidate_key,
                 db,
                 async_task: weak_at,
                 super_key,
+                notified,
             });
         });
-        Self { async_task }
+        Self { async_task, notified }
     }
 
     /// Notifies the key garbage collector to iterate through orphaned and superseded blobs and
     /// attempts their deletion. We only process one key at a time and then schedule another
     /// attempt by queueing it in the async_task (low priority) queue.
     pub fn notify_gc(&self) {
-        self.async_task.queue_lo(|shelf| shelf.get_downcast_mut::<GcInternal>().unwrap().step())
+        if let Ok(0) = self.notified.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed) {
+            self.async_task.queue_lo(|shelf| shelf.get_downcast_mut::<GcInternal>().unwrap().step())
+        }
     }
 }
 
 struct GcInternal {
-    blob_id_to_delete: Option<i64>,
+    deleted_blob_ids: Vec<i64>,
+    superseded_blobs: Vec<(i64, Vec<u8>, BlobMetaData)>,
     invalidate_key: Box<dyn Fn(&Uuid, &[u8]) -> Result<()> + Send + 'static>,
     db: KeystoreDB,
     async_task: std::sync::Weak<AsyncTask>,
     super_key: Arc<SuperKeyManager>,
+    notified: Arc<AtomicU8>,
 }
 
 impl GcInternal {
@@ -81,16 +96,23 @@
     /// We process one key at a time, because deleting a key is a time consuming process which
     /// may involve calling into the KeyMint backend and we don't want to hog neither the backend
     /// nor the database for extended periods of time.
+    /// To limit the number of database transactions, which are also expensive and competing
+    /// with threads on the critical path, deleted blobs are loaded in batches.
     fn process_one_key(&mut self) -> Result<()> {
-        if let Some((blob_id, blob, blob_metadata)) = self
-            .db
-            .handle_next_superseded_blob(self.blob_id_to_delete.take())
-            .context("In process_one_key: Trying to handle superseded blob.")?
-        {
-            // Set the blob_id as the next to be deleted blob. So it will be
+        if self.superseded_blobs.is_empty() {
+            let blobs = self
+                .db
+                .handle_next_superseded_blobs(&self.deleted_blob_ids, 20)
+                .context("In process_one_key: Trying to handle superseded blob.")?;
+            self.deleted_blob_ids = vec![];
+            self.superseded_blobs = blobs;
+        }
+
+        if let Some((blob_id, blob, blob_metadata)) = self.superseded_blobs.pop() {
+            // Add the next blob_id to the deleted blob ids list. So it will be
             // removed from the database regardless of whether the following
             // succeeds or not.
-            self.blob_id_to_delete = Some(blob_id);
+            self.deleted_blob_ids.push(blob_id);
 
             // If the key has a km_uuid we try to get the corresponding device
             // and delete the key, unwrapping if necessary and possible.
@@ -110,13 +132,20 @@
 
     /// Processes one key and then schedules another attempt until it runs out of blobs to delete.
     fn step(&mut self) {
+        self.notified.store(0, Ordering::Relaxed);
         if let Err(e) = self.process_one_key() {
             log::error!("Error trying to delete blob entry. {:?}", e);
         }
         // Schedule the next step. This gives high priority requests a chance to interleave.
-        if self.blob_id_to_delete.is_some() {
+        if !self.deleted_blob_ids.is_empty() {
             if let Some(at) = self.async_task.upgrade() {
-                at.queue_lo(move |shelf| shelf.get_downcast_mut::<GcInternal>().unwrap().step());
+                if let Ok(0) =
+                    self.notified.compare_exchange(0, 1, Ordering::Relaxed, Ordering::Relaxed)
+                {
+                    at.queue_lo(move |shelf| {
+                        shelf.get_downcast_mut::<GcInternal>().unwrap().step()
+                    });
+                }
             }
         }
     }
diff --git a/keystore2/src/globals.rs b/keystore2/src/globals.rs
index bd28ca6..89114a6 100644
--- a/keystore2/src/globals.rs
+++ b/keystore2/src/globals.rs
@@ -20,6 +20,7 @@
 use crate::legacy_blob::LegacyBlobLoader;
 use crate::legacy_migrator::LegacyMigrator;
 use crate::super_key::SuperKeyManager;
+use crate::utils::watchdog as wd;
 use crate::utils::Asp;
 use crate::{async_task::AsyncTask, database::MonotonicRawTime};
 use crate::{
@@ -38,11 +39,12 @@
 use binder::FromIBinder;
 use keystore2_vintf::get_aidl_instances;
 use lazy_static::lazy_static;
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, RwLock};
 use std::{cell::RefCell, sync::Once};
 use std::{collections::HashMap, path::Path, path::PathBuf};
 
 static DB_INIT: Once = Once::new();
+static DB_SET_WAL_MODE: Once = Once::new();
 
 /// Open a connection to the Keystore 2.0 database. This is called during the initialization of
 /// the thread local DB field. It should never be called directly. The first time this is called
@@ -53,27 +55,19 @@
 /// is run only once, as long as the ASYNC_TASK instance is the same. So only one additional
 /// database connection is created for the garbage collector worker.
 pub fn create_thread_local_db() -> KeystoreDB {
-    let gc = Gc::new_init_with(ASYNC_TASK.clone(), || {
-        (
-            Box::new(|uuid, blob| {
-                let km_dev: Strong<dyn IKeyMintDevice> =
-                    get_keymint_dev_by_uuid(uuid).map(|(dev, _)| dev)?.get_interface()?;
-                map_km_error(km_dev.deleteKey(&*blob))
-                    .context("In invalidate key closure: Trying to invalidate key blob.")
-            }),
-            KeystoreDB::new(&DB_PATH.lock().expect("Could not get the database directory."), None)
-                .expect("Failed to open database."),
-            SUPER_KEY.clone(),
-        )
+    let db_path = DB_PATH.read().expect("Could not get the database directory.");
+
+    DB_SET_WAL_MODE.call_once(|| {
+        log::info!("Setting Keystore 2.0 database to WAL mode first time since boot.");
+        KeystoreDB::set_wal_mode(&db_path)
+            .expect("In create_thread_local_db: Could not set WAL mode.");
     });
 
-    let mut db =
-        KeystoreDB::new(&DB_PATH.lock().expect("Could not get the database directory."), Some(gc))
-            .expect("Failed to open database.");
+    let mut db = KeystoreDB::new(&db_path, Some(GC.clone())).expect("Failed to open database.");
+
     DB_INIT.call_once(|| {
         log::info!("Touching Keystore 2.0 database for this first time since boot.");
-        db.insert_last_off_body(MonotonicRawTime::now())
-            .expect("Could not initialize database with last off body.");
+        db.insert_last_off_body(MonotonicRawTime::now());
         log::info!("Calling cleanup leftovers.");
         let n = db.cleanup_leftovers().expect("Failed to cleanup database on startup.");
         if n != 0 {
@@ -151,7 +145,7 @@
 
 lazy_static! {
     /// The path where keystore stores all its keys.
-    pub static ref DB_PATH: Mutex<PathBuf> = Mutex::new(
+    pub static ref DB_PATH: RwLock<PathBuf> = RwLock::new(
         Path::new("/data/misc/keystore").to_path_buf());
     /// Runtime database of unwrapped super keys.
     pub static ref SUPER_KEY: Arc<SuperKeyManager> = Default::default();
@@ -169,12 +163,27 @@
     /// LegacyBlobLoader is initialized and exists globally.
     /// The same directory used by the database is used by the LegacyBlobLoader as well.
     pub static ref LEGACY_BLOB_LOADER: Arc<LegacyBlobLoader> = Arc::new(LegacyBlobLoader::new(
-        &DB_PATH.lock().expect("Could not get the database path for legacy blob loader.")));
+        &DB_PATH.read().expect("Could not get the database path for legacy blob loader.")));
     /// Legacy migrator. Atomically migrates legacy blobs to the database.
     pub static ref LEGACY_MIGRATOR: Arc<LegacyMigrator> =
         Arc::new(LegacyMigrator::new(Arc::new(Default::default())));
     /// Background thread which handles logging via statsd and logd
     pub static ref LOGS_HANDLER: Arc<AsyncTask> = Default::default();
+
+    static ref GC: Arc<Gc> = Arc::new(Gc::new_init_with(ASYNC_TASK.clone(), || {
+        (
+            Box::new(|uuid, blob| {
+                let km_dev: Strong<dyn IKeyMintDevice> =
+                    get_keymint_dev_by_uuid(uuid).map(|(dev, _)| dev)?.get_interface()?;
+                let _wp = wd::watch_millis("In invalidate key closure: calling deleteKey", 500);
+                map_km_error(km_dev.deleteKey(&*blob))
+                    .context("In invalidate key closure: Trying to invalidate key blob.")
+            }),
+            KeystoreDB::new(&DB_PATH.read().expect("Could not get the database directory."), None)
+                .expect("Failed to open database."),
+            SUPER_KEY.clone(),
+        )
+    }));
 }
 
 static KEYMINT_SERVICE_NAME: &str = "android.hardware.security.keymint.IKeyMintDevice";
@@ -227,8 +236,10 @@
             .context("In connect_keymint: Trying to get Legacy wrapper.")
     }?;
 
+    let wp = wd::watch_millis("In connect_keymint: calling getHardwareInfo()", 500);
     let hw_info = map_km_error(keymint.getHardwareInfo())
         .context("In connect_keymint: Failed to get hardware info.")?;
+    drop(wp);
 
     Ok((Asp::new(keymint.as_binder()), hw_info))
 }
diff --git a/keystore2/src/key_parameter.rs b/keystore2/src/key_parameter.rs
index 74a9b23..549f574 100644
--- a/keystore2/src/key_parameter.rs
+++ b/keystore2/src/key_parameter.rs
@@ -90,8 +90,6 @@
 //!  * The termination condition which has an empty in list.
 //!  * The public interface, which does not have @marker and calls itself with an empty out list.
 
-#![allow(clippy::from_over_into, clippy::needless_question_mark)]
-
 use std::convert::TryInto;
 
 use crate::db_utils::SqlField;
@@ -601,9 +599,9 @@
         ], [$($in)*]
     }};
     (@into $enum_name:ident, [$($out:tt)*], []) => {
-        impl Into<KmKeyParameter> for $enum_name {
-            fn into(self) -> KmKeyParameter {
-                match self {
+        impl From<$enum_name> for KmKeyParameter {
+            fn from(x: $enum_name) -> Self {
+                match x {
                     $($out)*
                 }
             }
@@ -1389,11 +1387,11 @@
             db.prepare("SELECT tag, data, security_level FROM persistent.keyparameter")?;
         let mut rows = stmt.query(NO_PARAMS)?;
         let row = rows.next()?.unwrap();
-        Ok(KeyParameter::new_from_sql(
+        KeyParameter::new_from_sql(
             Tag(row.get(0)?),
             &SqlField::new(1, row),
             SecurityLevel(row.get(2)?),
-        )?)
+        )
     }
 }
 
diff --git a/keystore2/src/keystore2_main.rs b/keystore2/src/keystore2_main.rs
index 4d4a718..53461da 100644
--- a/keystore2/src/keystore2_main.rs
+++ b/keystore2/src/keystore2_main.rs
@@ -47,10 +47,6 @@
     // Saying hi.
     info!("Keystore2 is starting.");
 
-    // Initialize the per boot database.
-    let _keep_me_alive = keystore2::database::KeystoreDB::keep_perboot_db_alive()
-        .expect("Failed to initialize the perboot database.");
-
     let mut args = std::env::args();
     args.next().expect("That's odd. How is there not even a first argument?");
 
@@ -60,7 +56,7 @@
     // For the ground truth check the service startup rule for init (typically in keystore2.rc).
     let id_rotation_state = if let Some(dir) = args.next() {
         let db_path = Path::new(&dir);
-        *keystore2::globals::DB_PATH.lock().expect("Could not lock DB_PATH.") =
+        *keystore2::globals::DB_PATH.write().expect("Could not lock DB_PATH.") =
             db_path.to_path_buf();
         IdRotationState::new(&db_path)
     } else {
@@ -125,7 +121,7 @@
     }
 
     let vpnprofilestore = VpnProfileStore::new_native_binder(
-        &keystore2::globals::DB_PATH.lock().expect("Could not get DB_PATH."),
+        &keystore2::globals::DB_PATH.read().expect("Could not get DB_PATH."),
     );
     binder::add_service(VPNPROFILESTORE_SERVICE_NAME, vpnprofilestore.as_binder()).unwrap_or_else(
         |e| {
diff --git a/keystore2/src/legacy_blob.rs b/keystore2/src/legacy_blob.rs
index 29d46ad..9eebb36 100644
--- a/keystore2/src/legacy_blob.rs
+++ b/keystore2/src/legacy_blob.rs
@@ -14,8 +14,6 @@
 
 //! This module implements methods to load legacy keystore key blob files.
 
-#![allow(clippy::redundant_slicing)]
-
 use crate::{
     error::{Error as KsError, ResponseCode},
     key_parameter::{KeyParameter, KeyParameterValue},
@@ -484,7 +482,7 @@
         let element_size =
             read_ne_u32(stream).context("In read_key_parameters: While reading element size.")?;
 
-        let elements_buffer = stream
+        let mut element_stream = stream
             .get(0..element_size as usize)
             .ok_or(KsError::Rc(ResponseCode::VALUE_CORRUPTED))
             .context("In read_key_parameters: While reading elements buffer.")?;
@@ -492,8 +490,6 @@
         // update the stream position.
         *stream = &stream[element_size as usize..];
 
-        let mut element_stream = &elements_buffer[..];
-
         let mut params: Vec<KeyParameterValue> = Vec::new();
         for _ in 0..element_count {
             let tag = Tag(read_ne_i32(&mut element_stream).context("In read_key_parameters.")?);
diff --git a/keystore2/src/legacy_migrator.rs b/keystore2/src/legacy_migrator.rs
index 5e0d573..d5647cd 100644
--- a/keystore2/src/legacy_migrator.rs
+++ b/keystore2/src/legacy_migrator.rs
@@ -17,7 +17,7 @@
 use crate::error::Error;
 use crate::key_parameter::KeyParameterValue;
 use crate::legacy_blob::BlobValue;
-use crate::utils::uid_to_android_user;
+use crate::utils::{uid_to_android_user, watchdog as wd};
 use crate::{async_task::AsyncTask, legacy_blob::LegacyBlobLoader};
 use crate::{
     database::{
@@ -196,6 +196,8 @@
 
     /// List all aliases for uid in the legacy database.
     pub fn list_uid(&self, domain: Domain, namespace: i64) -> Result<Vec<KeyDescriptor>> {
+        let _wp = wd::watch_millis("LegacyMigrator::list_uid", 500);
+
         let uid = match (domain, namespace) {
             (Domain::APP, namespace) => namespace as u32,
             (Domain::SELINUX, Self::WIFI_NAMESPACE) => Self::AID_WIFI,
@@ -290,6 +292,8 @@
     where
         F: Fn() -> Result<T>,
     {
+        let _wp = wd::watch_millis("LegacyMigrator::with_try_migrate", 500);
+
         // Access the key and return on success.
         match key_accessor() {
             Ok(result) => return Ok(result),
@@ -342,6 +346,8 @@
     where
         F: FnMut() -> Result<Option<T>>,
     {
+        let _wp = wd::watch_millis("LegacyMigrator::with_try_migrate_super_key", 500);
+
         match key_accessor() {
             Ok(Some(result)) => return Ok(Some(result)),
             Ok(None) => {}
@@ -364,6 +370,8 @@
     /// Deletes all keys belonging to the given namespace, migrating them into the database
     /// for subsequent garbage collection if necessary.
     pub fn bulk_delete_uid(&self, domain: Domain, nspace: i64) -> Result<()> {
+        let _wp = wd::watch_millis("LegacyMigrator::bulk_delete_uid", 500);
+
         let uid = match (domain, nspace) {
             (Domain::APP, nspace) => nspace as u32,
             (Domain::SELINUX, Self::WIFI_NAMESPACE) => Self::AID_WIFI,
@@ -385,6 +393,8 @@
         user_id: u32,
         keep_non_super_encrypted_keys: bool,
     ) -> Result<()> {
+        let _wp = wd::watch_millis("LegacyMigrator::bulk_delete_user", 500);
+
         let result = self.do_serialized(move |migrator_state| {
             migrator_state
                 .bulk_delete(BulkDeleteRequest::User(user_id), keep_non_super_encrypted_keys)
diff --git a/keystore2/src/lib.rs b/keystore2/src/lib.rs
index 11af9df..c04c4b0 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -47,4 +47,6 @@
 mod db_utils;
 mod gc;
 mod super_key;
+
+#[cfg(feature = "watchdog")]
 mod watchdog;
diff --git a/keystore2/src/maintenance.rs b/keystore2/src/maintenance.rs
index 9e7576e..0633bc1 100644
--- a/keystore2/src/maintenance.rs
+++ b/keystore2/src/maintenance.rs
@@ -22,7 +22,7 @@
 use crate::globals::{DB, LEGACY_MIGRATOR, SUPER_KEY};
 use crate::permission::{KeyPerm, KeystorePerm};
 use crate::super_key::UserState;
-use crate::utils::{check_key_permission, check_keystore_permission};
+use crate::utils::{check_key_permission, check_keystore_permission, watchdog as wd};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::IKeyMintDevice::IKeyMintDevice;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::SecurityLevel::SecurityLevel;
 use android_security_maintenance::aidl::android::security::maintenance::{
@@ -133,11 +133,17 @@
         }
     }
 
-    fn early_boot_ended_help(sec_level: &SecurityLevel) -> Result<()> {
-        let (dev, _, _) =
-            get_keymint_device(sec_level).context("In early_boot_ended: getting keymint device")?;
+    fn early_boot_ended_help(sec_level: SecurityLevel) -> Result<()> {
+        let (dev, _, _) = get_keymint_device(&sec_level)
+            .context("In early_boot_ended: getting keymint device")?;
         let km_dev: Strong<dyn IKeyMintDevice> =
             dev.get_interface().context("In early_boot_ended: getting keymint device interface")?;
+
+        let _wp = wd::watch_millis_with(
+            "In early_boot_ended_help: calling earlyBootEnded()",
+            500,
+            move || format!("Seclevel: {:?}", sec_level),
+        );
         map_km_error(km_dev.earlyBootEnded())
             .context("In keymint device: calling earlyBootEnded")?;
         Ok(())
@@ -157,7 +163,7 @@
             (SecurityLevel::STRONGBOX, "STRONGBOX"),
         ];
         sec_levels.iter().fold(Ok(()), |result, (sec_level, sec_level_string)| {
-            let curr_result = Maintenance::early_boot_ended_help(sec_level);
+            let curr_result = Maintenance::early_boot_ended_help(*sec_level);
             if curr_result.is_err() {
                 log::error!(
                     "Call to earlyBootEnded failed for security level {}.",
@@ -173,8 +179,8 @@
         check_keystore_permission(KeystorePerm::report_off_body())
             .context("In on_device_off_body.")?;
 
-        DB.with(|db| db.borrow_mut().update_last_off_body(MonotonicRawTime::now()))
-            .context("In on_device_off_body: Trying to update last off body time.")
+        DB.with(|db| db.borrow_mut().update_last_off_body(MonotonicRawTime::now()));
+        Ok(())
     }
 
     fn migrate_key_namespace(source: &KeyDescriptor, destination: &KeyDescriptor) -> Result<()> {
@@ -219,30 +225,37 @@
 
 impl IKeystoreMaintenance for Maintenance {
     fn onUserPasswordChanged(&self, user_id: i32, password: Option<&[u8]>) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::onUserPasswordChanged", 500);
         map_or_log_err(Self::on_user_password_changed(user_id, password.map(|pw| pw.into())), Ok)
     }
 
     fn onUserAdded(&self, user_id: i32) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::onUserAdded", 500);
         map_or_log_err(Self::add_or_remove_user(user_id), Ok)
     }
 
     fn onUserRemoved(&self, user_id: i32) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::onUserRemoved", 500);
         map_or_log_err(Self::add_or_remove_user(user_id), Ok)
     }
 
     fn clearNamespace(&self, domain: Domain, nspace: i64) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::clearNamespace", 500);
         map_or_log_err(Self::clear_namespace(domain, nspace), Ok)
     }
 
     fn getState(&self, user_id: i32) -> BinderResult<AidlUserState> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::getState", 500);
         map_or_log_err(Self::get_state(user_id), Ok)
     }
 
     fn earlyBootEnded(&self) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::earlyBootEnded", 500);
         map_or_log_err(Self::early_boot_ended(), Ok)
     }
 
     fn onDeviceOffBody(&self) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::onDeviceOffBody", 500);
         map_or_log_err(Self::on_device_off_body(), Ok)
     }
 
@@ -251,6 +264,7 @@
         source: &KeyDescriptor,
         destination: &KeyDescriptor,
     ) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IKeystoreMaintenance::migrateKeyNamespace", 500);
         map_or_log_err(Self::migrate_key_namespace(source, destination), Ok)
     }
 }
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 0b5c77a..8d7ad0a 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -128,7 +128,7 @@
 use crate::enforcements::AuthInfo;
 use crate::error::{map_err_with, map_km_error, map_or_log_err, Error, ErrorCode, ResponseCode};
 use crate::metrics::log_key_operation_event_stats;
-use crate::utils::Asp;
+use crate::utils::{watchdog as wd, Asp};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     IKeyMintOperation::IKeyMintOperation, KeyParameter::KeyParameter, KeyPurpose::KeyPurpose,
     SecurityLevel::SecurityLevel,
@@ -291,6 +291,8 @@
                 }
             };
 
+        let _wp = wd::watch_millis("In Operation::prune: calling abort()", 500);
+
         // We abort the operation. If there was an error we log it but ignore it.
         if let Err(e) = map_km_error(km_op.abort()) {
             log::error!("In prune: KeyMint::abort failed with {:?}.", e);
@@ -370,10 +372,10 @@
             .before_update()
             .context("In update_aad: Trying to get auth tokens.")?;
 
-        self.update_outcome(
-            &mut *outcome,
-            map_km_error(km_op.updateAad(aad_input, hat.as_ref(), tst.as_ref())),
-        )
+        self.update_outcome(&mut *outcome, {
+            let _wp = wd::watch_millis("Operation::update_aad: calling updateAad", 500);
+            map_km_error(km_op.updateAad(aad_input, hat.as_ref(), tst.as_ref()))
+        })
         .context("In update_aad: KeyMint::update failed.")?;
 
         Ok(())
@@ -397,10 +399,10 @@
             .context("In update: Trying to get auth tokens.")?;
 
         let output = self
-            .update_outcome(
-                &mut *outcome,
-                map_km_error(km_op.update(input, hat.as_ref(), tst.as_ref())),
-            )
+            .update_outcome(&mut *outcome, {
+                let _wp = wd::watch_millis("Operation::update: calling update", 500);
+                map_km_error(km_op.update(input, hat.as_ref(), tst.as_ref()))
+            })
             .context("In update: KeyMint::update failed.")?;
 
         if output.is_empty() {
@@ -430,16 +432,16 @@
             .context("In finish: Trying to get auth tokens.")?;
 
         let output = self
-            .update_outcome(
-                &mut *outcome,
+            .update_outcome(&mut *outcome, {
+                let _wp = wd::watch_millis("Operation::finish: calling finish", 500);
                 map_km_error(km_op.finish(
                     input,
                     signature,
                     hat.as_ref(),
                     tst.as_ref(),
                     confirmation_token.as_deref(),
-                )),
-            )
+                ))
+            })
             .context("In finish: KeyMint::finish failed.")?;
 
         self.auth_info.lock().unwrap().after_finish().context("In finish.")?;
@@ -463,7 +465,10 @@
         let km_op: binder::public_api::Strong<dyn IKeyMintOperation> =
             self.km_op.get_interface().context("In abort: Failed to get KeyMintOperation.")?;
 
-        map_km_error(km_op.abort()).context("In abort: KeyMint::abort failed.")
+        {
+            let _wp = wd::watch_millis("Operation::abort: calling abort", 500);
+            map_km_error(km_op.abort()).context("In abort: KeyMint::abort failed.")
+        }
     }
 }
 
@@ -837,6 +842,7 @@
 
 impl IKeystoreOperation for KeystoreOperation {
     fn updateAad(&self, aad_input: &[u8]) -> binder::public_api::Result<()> {
+        let _wp = wd::watch_millis("IKeystoreOperation::updateAad", 500);
         map_or_log_err(
             self.with_locked_operation(
                 |op| op.update_aad(aad_input).context("In KeystoreOperation::updateAad"),
@@ -847,6 +853,7 @@
     }
 
     fn update(&self, input: &[u8]) -> binder::public_api::Result<Option<Vec<u8>>> {
+        let _wp = wd::watch_millis("IKeystoreOperation::update", 500);
         map_or_log_err(
             self.with_locked_operation(
                 |op| op.update(input).context("In KeystoreOperation::update"),
@@ -860,6 +867,7 @@
         input: Option<&[u8]>,
         signature: Option<&[u8]>,
     ) -> binder::public_api::Result<Option<Vec<u8>>> {
+        let _wp = wd::watch_millis("IKeystoreOperation::finish", 500);
         map_or_log_err(
             self.with_locked_operation(
                 |op| op.finish(input, signature).context("In KeystoreOperation::finish"),
@@ -870,6 +878,7 @@
     }
 
     fn abort(&self) -> binder::public_api::Result<()> {
+        let _wp = wd::watch_millis("IKeystoreOperation::abort", 500);
         map_err_with(
             self.with_locked_operation(
                 |op| op.abort(Outcome::Abort).context("In KeystoreOperation::abort"),
diff --git a/keystore2/src/permission.rs b/keystore2/src/permission.rs
index 726c2ec..e7999bc 100644
--- a/keystore2/src/permission.rs
+++ b/keystore2/src/permission.rs
@@ -18,8 +18,6 @@
 //! It also provides KeystorePerm and KeyPerm as convenience wrappers for the SELinux permission
 //! defined by keystore2 and keystore2_key respectively.
 
-#![allow(clippy::from_over_into)]
-
 use android_system_keystore2::aidl::android::system::keystore2::{
     Domain::Domain, KeyDescriptor::KeyDescriptor, KeyPermission::KeyPermission,
 };
@@ -151,9 +149,9 @@
             }
         }
 
-        impl Into<$aidl_name> for $name {
-            fn into(self) -> $aidl_name {
-                self.0
+        impl From<$name> for $aidl_name {
+            fn from(p: $name) -> $aidl_name {
+                p.0
             }
         }
 
@@ -259,9 +257,9 @@
             }
         }
 
-        impl Into<i32> for $name {
-            fn into(self) -> i32 {
-                self as i32
+        impl From<$name> for i32 {
+            fn from(p: $name) -> i32 {
+                p as i32
             }
         }
 
diff --git a/keystore2/src/raw_device.rs b/keystore2/src/raw_device.rs
index 06432fe..9e6ef41 100644
--- a/keystore2/src/raw_device.rs
+++ b/keystore2/src/raw_device.rs
@@ -22,7 +22,7 @@
     error::{map_km_error, Error},
     globals::get_keymint_device,
     super_key::KeyBlob,
-    utils::{key_characteristics_to_internal, Asp, AID_KEYSTORE},
+    utils::{key_characteristics_to_internal, watchdog as wd, Asp, AID_KEYSTORE},
 };
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     BeginResult::BeginResult, ErrorCode::ErrorCode, HardwareAuthToken::HardwareAuthToken,
@@ -151,7 +151,7 @@
             self.create_and_store_key(db, &key_desc, |km_dev| km_dev.generateKey(&params, None))
                 .context("In lookup_or_generate_key: generate_and_store_key failed")?;
             Self::lookup_from_desc(db, key_desc)
-                .context("In lookup_or_generate_key: secpnd lookup failed")
+                .context("In lookup_or_generate_key: second lookup failed")
         }
     }
 
@@ -170,8 +170,14 @@
     {
         match f(key_blob) {
             Err(Error::Km(ErrorCode::KEY_REQUIRES_UPGRADE)) => {
-                let upgraded_blob = map_km_error(km_dev.upgradeKey(key_blob, &[]))
-                    .context("In upgrade_keyblob_if_required_with: Upgrade failed")?;
+                let upgraded_blob = map_km_error({
+                    let _wp = wd::watch_millis(
+                        "In KeyMintDevice::upgrade_keyblob_if_required_with: calling upgradeKey.",
+                        500,
+                    );
+                    km_dev.upgradeKey(key_blob, &[])
+                })
+                .context("In upgrade_keyblob_if_required_with: Upgrade failed")?;
 
                 let mut new_blob_metadata = BlobMetaData::new();
                 new_blob_metadata.add(BlobMetaEntry::KmUuid(self.km_uuid));
@@ -223,14 +229,20 @@
 
         let begin_result: BeginResult = self
             .upgrade_keyblob_if_required_with(db, &km_dev, key_id_guard, &key_blob, |blob| {
-                map_km_error(km_dev.begin(purpose, blob, operation_parameters, auth_token))
+                map_km_error({
+                    let _wp = wd::watch_millis("In use_key_in_one_step: calling: begin", 500);
+                    km_dev.begin(purpose, blob, operation_parameters, auth_token)
+                })
             })
             .context("In use_key_in_one_step: Failed to begin operation.")?;
         let operation: Strong<dyn IKeyMintOperation> = begin_result
             .operation
             .ok_or_else(Error::sys)
             .context("In use_key_in_one_step: Operation missing")?;
-        map_km_error(operation.finish(Some(input), None, None, None, None))
-            .context("In use_key_in_one_step: Failed to finish operation.")
+        map_km_error({
+            let _wp = wd::watch_millis("In use_key_in_one_step: calling: finish", 500);
+            operation.finish(Some(input), None, None, None, None)
+        })
+        .context("In use_key_in_one_step: Failed to finish operation.")
     }
 }
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index f8ee369..1f3f8e8 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -19,8 +19,6 @@
 //! certificate chains signed by some root authority and stored in a keystore SQLite
 //! DB.
 
-#![allow(clippy::from_over_into, clippy::needless_question_mark, clippy::vec_init_then_push)]
-
 use std::collections::HashMap;
 
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
@@ -45,7 +43,7 @@
 use crate::database::{CertificateChain, KeystoreDB, Uuid};
 use crate::error::{self, map_or_log_err, map_rem_prov_error, Error};
 use crate::globals::{get_keymint_device, get_remotely_provisioned_component, DB};
-use crate::utils::Asp;
+use crate::utils::{watchdog as wd, Asp};
 
 /// Contains helper functions to check if remote provisioning is enabled on the system and, if so,
 /// to assign and retrieve attestation keys and certificate chains.
@@ -252,7 +250,7 @@
             // attestation keys unless the pool status is checked first, so this call should be
             // enough to routinely clean out expired keys.
             db.delete_expired_attestation_keys()?;
-            Ok(db.get_attestation_pool_status(expired_by, &uuid)?)
+            db.get_attestation_pool_status(expired_by, &uuid)
         })
     }
 
@@ -294,14 +292,15 @@
             protected_data,
         ))
         .context("In generate_csr: Failed to generate csr")?;
-        let mut cose_mac_0 = Vec::<u8>::new();
         // TODO(b/180392379): Replace this manual CBOR generation with the cbor-serde crate as well.
         //                    This generates an array consisting of the mac and the public key Maps.
         //                    Just generate the actual MacedPublicKeys structure when the crate is
         //                    available.
-        cose_mac_0.push((0b100_00000 | (keys_to_sign.len() + 1)) as u8);
-        cose_mac_0.push(0b010_11000); //push mac
-        cose_mac_0.push(mac.len() as u8);
+        let mut cose_mac_0: Vec<u8> = vec![
+            (0b100_00000 | (keys_to_sign.len() + 1)) as u8,
+            0b010_11000, // mac
+            (mac.len() as u8),
+        ];
         cose_mac_0.append(&mut mac);
         for maced_public_key in keys_to_sign {
             if maced_public_key.macedKey.len() > 83 + 8 {
@@ -327,13 +326,13 @@
         DB.with::<_, Result<()>>(|db| {
             let mut db = db.borrow_mut();
             let (_, _, uuid) = get_keymint_device(&sec_level)?;
-            Ok(db.store_signed_attestation_certificate_chain(
+            db.store_signed_attestation_certificate_chain(
                 public_key,
                 batch_cert,
                 certs, /* DER encoded certificate chain */
                 expiration_date,
                 &uuid,
-            )?)
+            )
         })
     }
 
@@ -362,7 +361,7 @@
         raw_key[32..64].clone_from_slice(&data[53..53 + 32]);
         DB.with::<_, Result<()>>(|db| {
             let mut db = db.borrow_mut();
-            Ok(db.create_attestation_key_entry(&maced_key.macedKey, &raw_key, &priv_key, &uuid)?)
+            db.create_attestation_key_entry(&maced_key.macedKey, &raw_key, &priv_key, &uuid)
         })
     }
 
@@ -377,7 +376,7 @@
     pub fn delete_all_keys(&self) -> Result<i64> {
         DB.with::<_, Result<i64>>(|db| {
             let mut db = db.borrow_mut();
-            Ok(db.delete_all_attestation_keys()?)
+            db.delete_all_attestation_keys()
         })
     }
 }
@@ -392,6 +391,7 @@
         expired_by: i64,
         sec_level: SecurityLevel,
     ) -> binder::public_api::Result<AttestationPoolStatus> {
+        let _wp = wd::watch_millis("IRemoteProvisioning::getPoolStatus", 500);
         map_or_log_err(self.get_pool_status(expired_by, sec_level), Ok)
     }
 
@@ -405,6 +405,7 @@
         protected_data: &mut ProtectedData,
         device_info: &mut DeviceInfo,
     ) -> binder::public_api::Result<Vec<u8>> {
+        let _wp = wd::watch_millis("IRemoteProvisioning::generateCsr", 500);
         map_or_log_err(
             self.generate_csr(
                 test_mode,
@@ -427,6 +428,7 @@
         expiration_date: i64,
         sec_level: SecurityLevel,
     ) -> binder::public_api::Result<()> {
+        let _wp = wd::watch_millis("IRemoteProvisioning::provisionCertChain", 500);
         map_or_log_err(
             self.provision_cert_chain(public_key, batch_cert, certs, expiration_date, sec_level),
             Ok,
@@ -438,14 +440,17 @@
         is_test_mode: bool,
         sec_level: SecurityLevel,
     ) -> binder::public_api::Result<()> {
+        let _wp = wd::watch_millis("IRemoteProvisioning::generateKeyPair", 500);
         map_or_log_err(self.generate_key_pair(is_test_mode, sec_level), Ok)
     }
 
     fn getSecurityLevels(&self) -> binder::public_api::Result<Vec<SecurityLevel>> {
+        let _wp = wd::watch_millis("IRemoteProvisioning::getSecurityLevels", 500);
         map_or_log_err(self.get_security_levels(), Ok)
     }
 
     fn deleteAllKeys(&self) -> binder::public_api::Result<i64> {
+        let _wp = wd::watch_millis("IRemoteProvisioning::deleteAllKeys", 500);
         map_or_log_err(self.delete_all_keys(), Ok)
     }
 }
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 53880a1..00b26e4 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -14,6 +14,32 @@
 
 //! This crate implements the IKeystoreSecurityLevel interface.
 
+use crate::attestation_key_utils::{get_attest_key_info, AttestationKeyInfo};
+use crate::audit_log::{
+    log_key_deleted, log_key_generated, log_key_imported, log_key_integrity_violation,
+};
+use crate::database::{CertificateInfo, KeyIdGuard};
+use crate::error::{self, map_km_error, map_or_log_err, Error, ErrorCode};
+use crate::globals::{DB, ENFORCEMENTS, LEGACY_MIGRATOR, SUPER_KEY};
+use crate::key_parameter::KeyParameter as KsKeyParam;
+use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
+use crate::metrics::log_key_creation_event_stats;
+use crate::remote_provisioning::RemProvState;
+use crate::super_key::{KeyBlob, SuperKeyManager};
+use crate::utils::{
+    check_device_attestation_permissions, check_key_permission, is_device_id_attestation_tag,
+    key_characteristics_to_internal, uid_to_android_user, watchdog as wd, Asp,
+};
+use crate::{
+    database::{
+        BlobMetaData, BlobMetaEntry, DateTime, KeyEntry, KeyEntryLoadBits, KeyMetaData,
+        KeyMetaEntry, KeyType, SubComponentType, Uuid,
+    },
+    operation::KeystoreOperation,
+    operation::LoggingInfo,
+    operation::OperationDb,
+    permission::KeyPerm,
+};
 use crate::{globals::get_keymint_device, id_rotation::IdRotationState};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     Algorithm::Algorithm, AttestationKey::AttestationKey,
@@ -30,34 +56,6 @@
     IKeystoreSecurityLevel::IKeystoreSecurityLevel, KeyDescriptor::KeyDescriptor,
     KeyMetadata::KeyMetadata, KeyParameters::KeyParameters,
 };
-
-use crate::attestation_key_utils::{get_attest_key_info, AttestationKeyInfo};
-use crate::audit_log::{log_key_deleted, log_key_generated, log_key_imported};
-use crate::database::{CertificateInfo, KeyIdGuard};
-use crate::globals::{DB, ENFORCEMENTS, LEGACY_MIGRATOR, SUPER_KEY};
-use crate::key_parameter::KeyParameter as KsKeyParam;
-use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
-use crate::metrics::log_key_creation_event_stats;
-use crate::remote_provisioning::RemProvState;
-use crate::super_key::{KeyBlob, SuperKeyManager};
-use crate::utils::{
-    check_device_attestation_permissions, check_key_permission, is_device_id_attestation_tag,
-    uid_to_android_user, Asp,
-};
-use crate::{
-    database::{
-        BlobMetaData, BlobMetaEntry, DateTime, KeyEntry, KeyEntryLoadBits, KeyMetaData,
-        KeyMetaEntry, KeyType, SubComponentType, Uuid,
-    },
-    operation::KeystoreOperation,
-    operation::LoggingInfo,
-    operation::OperationDb,
-    permission::KeyPerm,
-};
-use crate::{
-    error::{self, map_km_error, map_or_log_err, Error, ErrorCode},
-    utils::key_characteristics_to_internal,
-};
 use anyhow::{anyhow, Context, Result};
 
 /// Implementation of the IKeystoreSecurityLevel Interface.
@@ -104,6 +102,11 @@
         Ok((result, km_uuid))
     }
 
+    fn watch_millis(&self, id: &'static str, millis: u64) -> Option<wd::WatchPoint> {
+        let sec_level = self.security_level;
+        wd::watch_millis_with(id, millis, move || format!("SecurityLevel {:?}", sec_level))
+    }
+
     fn store_new_key(
         &self,
         key: KeyDescriptor,
@@ -313,16 +316,29 @@
                 &blob_metadata,
                 &operation_parameters,
                 |blob| loop {
-                    match map_km_error(km_dev.begin(
-                        purpose,
-                        blob,
-                        &operation_parameters,
-                        immediate_hat.as_ref(),
-                    )) {
+                    match map_km_error({
+                        let _wp = self.watch_millis(
+                            "In KeystoreSecurityLevel::create_operation: calling begin",
+                            500,
+                        );
+                        km_dev.begin(purpose, blob, &operation_parameters, immediate_hat.as_ref())
+                    }) {
                         Err(Error::Km(ErrorCode::TOO_MANY_OPERATIONS)) => {
                             self.operation_db.prune(caller_uid, forced)?;
                             continue;
                         }
+                        v @ Err(Error::Km(ErrorCode::INVALID_KEY_BLOB)) => {
+                            if let Some((key_id, _)) = key_properties {
+                                if let Ok(Some(key)) =
+                                    DB.with(|db| db.borrow_mut().load_key_descriptor(key_id))
+                                {
+                                    log_key_integrity_violation(&key);
+                                } else {
+                                    log::error!("Failed to load key descriptor for audit log");
+                                }
+                            }
+                            return v;
+                        }
                         v => return v,
                     }
                 },
@@ -334,12 +350,19 @@
         let op_params: Vec<KeyParameter> = operation_parameters.to_vec();
 
         let operation = match begin_result.operation {
-            Some(km_op) => {
-                self.operation_db.create_operation(km_op, caller_uid, auth_info, forced,
-                    LoggingInfo::new(self.security_level, purpose, op_params,
-                         upgraded_blob.is_some()))
-            },
-            None => return Err(Error::sys()).context("In create_operation: Begin operation returned successfully, but did not return a valid operation."),
+            Some(km_op) => self.operation_db.create_operation(
+                km_op,
+                caller_uid,
+                auth_info,
+                forced,
+                LoggingInfo::new(self.security_level, purpose, op_params, upgraded_blob.is_some()),
+            ),
+            None => {
+                return Err(Error::sys()).context(concat!(
+                    "In create_operation: Begin operation returned successfully, ",
+                    "but did not return a valid operation."
+                ))
+            }
         };
 
         let op_binder: binder::public_api::Strong<dyn IKeystoreOperation> =
@@ -371,9 +394,19 @@
         let mut result = params.to_vec();
         // If there is an attestation challenge we need to get an application id.
         if params.iter().any(|kp| kp.tag == Tag::ATTESTATION_CHALLENGE) {
-            let aaid = keystore2_aaid::get_aaid(uid).map_err(|e| {
-                anyhow!(format!("In add_certificate_parameters: get_aaid returned status {}.", e))
-            })?;
+            let aaid = {
+                let _wp = self.watch_millis(
+                    "In KeystoreSecurityLevel::add_certificate_parameters calling: get_aaid",
+                    500,
+                );
+                keystore2_aaid::get_aaid(uid).map_err(|e| {
+                    anyhow!(format!(
+                        "In add_certificate_parameters: get_aaid returned status {}.",
+                        e
+                    ))
+                })
+            }?;
+
             result.push(KeyParameter {
                 tag: Tag::ATTESTATION_APPLICATION_ID,
                 value: KeyParameterValue::Blob(aaid),
@@ -495,21 +528,48 @@
                             attestKeyParams: vec![],
                             issuerSubjectName: issuer_subject.clone(),
                         });
-                        map_km_error(km_dev.generateKey(&params, attest_key.as_ref()))
+                        map_km_error({
+                            let _wp = self.watch_millis(
+                                concat!(
+                                    "In KeystoreSecurityLevel::generate_key (UserGenerated): ",
+                                    "calling generate_key."
+                                ),
+                                5000, // Generate can take a little longer.
+                            );
+                            km_dev.generateKey(&params, attest_key.as_ref())
+                        })
                     },
                 )
                 .context("In generate_key: Using user generated attestation key.")
                 .map(|(result, _)| result),
             Some(AttestationKeyInfo::RemoteProvisioned { attestation_key, attestation_certs }) => {
-                map_km_error(km_dev.generateKey(&params, Some(&attestation_key)))
-                    .context("While generating Key with remote provisioned attestation key.")
-                    .map(|mut creation_result| {
-                        creation_result.certificateChain.push(attestation_certs);
-                        creation_result
-                    })
+                map_km_error({
+                    let _wp = self.watch_millis(
+                        concat!(
+                            "In KeystoreSecurityLevel::generate_key (RemoteProvisioned): ",
+                            "calling generate_key.",
+                        ),
+                        5000, // Generate can take a little longer.
+                    );
+                    km_dev.generateKey(&params, Some(&attestation_key))
+                })
+                .context("While generating Key with remote provisioned attestation key.")
+                .map(|mut creation_result| {
+                    creation_result.certificateChain.push(attestation_certs);
+                    creation_result
+                })
             }
-            None => map_km_error(km_dev.generateKey(&params, None))
-                .context("While generating Key without explicit attestation key."),
+            None => map_km_error({
+                let _wp = self.watch_millis(
+                    concat!(
+                        "In KeystoreSecurityLevel::generate_key (No attestation): ",
+                        "calling generate_key.",
+                    ),
+                    5000, // Generate can take a little longer.
+                );
+                km_dev.generateKey(&params, None)
+            })
+            .context("While generating Key without explicit attestation key."),
         }
         .context("In generate_key.")?;
 
@@ -566,9 +626,12 @@
 
         let km_dev: Strong<dyn IKeyMintDevice> =
             self.keymint.get_interface().context("In import_key: Trying to get the KM device")?;
-        let creation_result =
-            map_km_error(km_dev.importKey(&params, format, key_data, None /* attestKey */))
-                .context("In import_key: Trying to call importKey")?;
+        let creation_result = map_km_error({
+            let _wp =
+                self.watch_millis("In KeystoreSecurityLevel::import_key: calling importKey.", 500);
+            km_dev.importKey(&params, format, key_data, None /* attestKey */)
+        })
+        .context("In import_key: Trying to call importKey")?;
 
         let user_id = uid_to_android_user(caller_uid);
         self.store_new_key(key, creation_result, user_id, Some(flags)).context("In import_key.")
@@ -681,6 +744,10 @@
                 &wrapping_blob_metadata,
                 &[],
                 |wrapping_blob| {
+                    let _wp = self.watch_millis(
+                        "In KeystoreSecurityLevel::import_wrapped_key: calling importWrappedKey.",
+                        500,
+                    );
                     let creation_result = map_km_error(km_dev.importWrappedKey(
                         wrapped_data,
                         wrapping_blob,
@@ -739,8 +806,17 @@
     {
         match f(key_blob) {
             Err(Error::Km(ErrorCode::KEY_REQUIRES_UPGRADE)) => {
-                let upgraded_blob = map_km_error(km_dev.upgradeKey(key_blob, params))
-                    .context("In upgrade_keyblob_if_required_with: Upgrade failed.")?;
+                let upgraded_blob = {
+                    let _wp = self.watch_millis(
+                        concat!(
+                            "In KeystoreSecurityLevel::upgrade_keyblob_if_required_with: ",
+                            "calling upgradeKey."
+                        ),
+                        500,
+                    );
+                    map_km_error(km_dev.upgradeKey(key_blob, params))
+                }
+                .context("In upgrade_keyblob_if_required_with: Upgrade failed.")?;
 
                 if let Some(kid) = key_id_guard {
                     Self::store_upgraded_keyblob(
@@ -810,14 +886,35 @@
             "In IKeystoreSecurityLevel convert_storage_key_to_ephemeral: ",
             "Getting keymint device interface"
         ))?;
-        match map_km_error(km_dev.convertStorageKeyToEphemeral(key_blob)) {
+        match {
+            let _wp = self.watch_millis(
+                concat!(
+                    "In IKeystoreSecurityLevel::convert_storage_key_to_ephemeral: ",
+                    "calling convertStorageKeyToEphemeral (1)"
+                ),
+                500,
+            );
+            map_km_error(km_dev.convertStorageKeyToEphemeral(key_blob))
+        } {
             Ok(result) => {
                 Ok(EphemeralStorageKeyResponse { ephemeralKey: result, upgradedBlob: None })
             }
             Err(error::Error::Km(ErrorCode::KEY_REQUIRES_UPGRADE)) => {
-                let upgraded_blob = map_km_error(km_dev.upgradeKey(key_blob, &[]))
-                    .context("In convert_storage_key_to_ephemeral: Failed to upgrade key blob.")?;
-                let ephemeral_key = map_km_error(km_dev.convertStorageKeyToEphemeral(key_blob))
+                let upgraded_blob = {
+                    let _wp = self.watch_millis(
+                        "In convert_storage_key_to_ephemeral: calling upgradeKey",
+                        500,
+                    );
+                    map_km_error(km_dev.upgradeKey(key_blob, &[]))
+                }
+                .context("In convert_storage_key_to_ephemeral: Failed to upgrade key blob.")?;
+                let ephemeral_key = {
+                    let _wp = self.watch_millis(
+                        "In convert_storage_key_to_ephemeral: calling convertStorageKeyToEphemeral (2)",
+                        500,
+                    );
+                    map_km_error(km_dev.convertStorageKeyToEphemeral(key_blob))
+                }
                     .context(concat!(
                         "In convert_storage_key_to_ephemeral: ",
                         "Failed to retrieve ephemeral key (after upgrade)."
@@ -851,7 +948,11 @@
             .keymint
             .get_interface()
             .context("In IKeystoreSecurityLevel delete_key: Getting keymint device interface")?;
-        map_km_error(km_dev.deleteKey(&key_blob)).context("In keymint device deleteKey")
+        {
+            let _wp =
+                self.watch_millis("In KeystoreSecuritylevel::delete_key: calling deleteKey", 500);
+            map_km_error(km_dev.deleteKey(&key_blob)).context("In keymint device deleteKey")
+        }
     }
 }
 
@@ -864,6 +965,7 @@
         operation_parameters: &[KeyParameter],
         forced: bool,
     ) -> binder::public_api::Result<CreateOperationResponse> {
+        let _wp = self.watch_millis("IKeystoreSecurityLevel::createOperation", 500);
         map_or_log_err(self.create_operation(key, operation_parameters, forced), Ok)
     }
     fn generateKey(
@@ -874,6 +976,9 @@
         flags: i32,
         entropy: &[u8],
     ) -> binder::public_api::Result<KeyMetadata> {
+        // Duration is set to 5 seconds, because generateKey - especially for RSA keys, takes more
+        // time than other operations
+        let _wp = self.watch_millis("IKeystoreSecurityLevel::generateKey", 5000);
         let result = self.generate_key(key, attestation_key, params, flags, entropy);
         log_key_creation_event_stats(self.security_level, params, &result);
         log_key_generated(key, ThreadState::get_calling_uid(), result.is_ok());
@@ -887,6 +992,7 @@
         flags: i32,
         key_data: &[u8],
     ) -> binder::public_api::Result<KeyMetadata> {
+        let _wp = self.watch_millis("IKeystoreSecurityLevel::importKey", 500);
         let result = self.import_key(key, attestation_key, params, flags, key_data);
         log_key_creation_event_stats(self.security_level, params, &result);
         log_key_imported(key, ThreadState::get_calling_uid(), result.is_ok());
@@ -900,6 +1006,7 @@
         params: &[KeyParameter],
         authenticators: &[AuthenticatorSpec],
     ) -> binder::public_api::Result<KeyMetadata> {
+        let _wp = self.watch_millis("IKeystoreSecurityLevel::importWrappedKey", 500);
         let result =
             self.import_wrapped_key(key, wrapping_key, masking_key, params, authenticators);
         log_key_creation_event_stats(self.security_level, params, &result);
@@ -910,9 +1017,11 @@
         &self,
         storage_key: &KeyDescriptor,
     ) -> binder::public_api::Result<EphemeralStorageKeyResponse> {
+        let _wp = self.watch_millis("IKeystoreSecurityLevel::convertStorageKeyToEphemeral", 500);
         map_or_log_err(self.convert_storage_key_to_ephemeral(storage_key), Ok)
     }
     fn deleteKey(&self, key: &KeyDescriptor) -> binder::public_api::Result<()> {
+        let _wp = self.watch_millis("IKeystoreSecurityLevel::deleteKey", 500);
         let result = self.delete_key(key);
         log_key_deleted(key, ThreadState::get_calling_uid(), result.is_ok());
         map_or_log_err(result, Ok)
diff --git a/keystore2/src/service.rs b/keystore2/src/service.rs
index b8ea244..1f61729 100644
--- a/keystore2/src/service.rs
+++ b/keystore2/src/service.rs
@@ -22,7 +22,7 @@
 use crate::security_level::KeystoreSecurityLevel;
 use crate::utils::{
     check_grant_permission, check_key_permission, check_keystore_permission,
-    key_parameters_to_authorizations, Asp,
+    key_parameters_to_authorizations, watchdog as wd, Asp,
 };
 use crate::{
     database::Uuid,
@@ -291,7 +291,7 @@
             &mut DB
                 .with(|db| {
                     let mut db = db.borrow_mut();
-                    db.list(k.domain, k.nspace)
+                    db.list(k.domain, k.nspace, KeyType::Client)
                 })
                 .context("In list_entries: Trying to list keystore database.")?,
         );
@@ -354,9 +354,13 @@
         &self,
         security_level: SecurityLevel,
     ) -> binder::public_api::Result<Strong<dyn IKeystoreSecurityLevel>> {
+        let _wp = wd::watch_millis_with("IKeystoreService::getSecurityLevel", 500, move || {
+            format!("security_level: {}", security_level.0)
+        });
         map_or_log_err(self.get_security_level(security_level), Ok)
     }
     fn getKeyEntry(&self, key: &KeyDescriptor) -> binder::public_api::Result<KeyEntryResponse> {
+        let _wp = wd::watch_millis("IKeystoreService::get_key_entry", 500);
         map_or_log_err(self.get_key_entry(key), Ok)
     }
     fn updateSubcomponent(
@@ -365,6 +369,7 @@
         public_cert: Option<&[u8]>,
         certificate_chain: Option<&[u8]>,
     ) -> binder::public_api::Result<()> {
+        let _wp = wd::watch_millis("IKeystoreService::updateSubcomponent", 500);
         map_or_log_err(self.update_subcomponent(key, public_cert, certificate_chain), Ok)
     }
     fn listEntries(
@@ -372,9 +377,11 @@
         domain: Domain,
         namespace: i64,
     ) -> binder::public_api::Result<Vec<KeyDescriptor>> {
+        let _wp = wd::watch_millis("IKeystoreService::listEntries", 500);
         map_or_log_err(self.list_entries(domain, namespace), Ok)
     }
     fn deleteKey(&self, key: &KeyDescriptor) -> binder::public_api::Result<()> {
+        let _wp = wd::watch_millis("IKeystoreService::deleteKey", 500);
         let result = self.delete_key(key);
         log_key_deleted(key, ThreadState::get_calling_uid(), result.is_ok());
         map_or_log_err(result, Ok)
@@ -385,9 +392,11 @@
         grantee_uid: i32,
         access_vector: i32,
     ) -> binder::public_api::Result<KeyDescriptor> {
+        let _wp = wd::watch_millis("IKeystoreService::grant", 500);
         map_or_log_err(self.grant(key, grantee_uid, access_vector.into()), Ok)
     }
     fn ungrant(&self, key: &KeyDescriptor, grantee_uid: i32) -> binder::public_api::Result<()> {
+        let _wp = wd::watch_millis("IKeystoreService::ungrant", 500);
         map_or_log_err(self.ungrant(key, grantee_uid), Ok)
     }
 }
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index 50a5f31..7a8b9be 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -29,6 +29,7 @@
     legacy_migrator::LegacyMigrator,
     raw_device::KeyMintDevice,
     try_insert::TryInsert,
+    utils::watchdog as wd,
 };
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     Algorithm::Algorithm, BlockMode::BlockMode, HardwareAuthToken::HardwareAuthToken,
@@ -949,6 +950,10 @@
                     let key_params: Vec<KmKeyParameter> =
                         key_params.into_iter().map(|x| x.into()).collect();
                     km_dev.create_and_store_key(db, &key_desc, |dev| {
+                        let _wp = wd::watch_millis(
+                            "In lock_screen_lock_bound_key: calling importKey.",
+                            500,
+                        );
                         dev.importKey(key_params.as_slice(), KeyFormat::RAW, &encrypting_key, None)
                     })?;
                     entry.biometric_unlock = Some(BiometricUnlock {
@@ -987,7 +992,7 @@
             for sid in &biometric.sids {
                 if let Some((auth_token_entry, _)) = db.find_auth_token_entry(|entry| {
                     entry.auth_token().userId == *sid || entry.auth_token().authenticatorId == *sid
-                })? {
+                }) {
                     let res: Result<(Arc<SuperKey>, Arc<SuperKey>)> = (|| {
                         let slb = biometric.screen_lock_bound.decrypt(
                             db,
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index bca27d1..a110c64 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -36,7 +36,6 @@
     APC_COMPAT_ERROR_IGNORED, APC_COMPAT_ERROR_OK, APC_COMPAT_ERROR_OPERATION_PENDING,
     APC_COMPAT_ERROR_SYSTEM_ERROR,
 };
-use std::convert::TryFrom;
 use std::sync::Mutex;
 
 /// This function uses its namesake in the permission module and in
@@ -107,11 +106,17 @@
     let permission_controller: binder::Strong<dyn IPermissionController::IPermissionController> =
         binder::get_interface("permission")?;
 
-    let binder_result = permission_controller.checkPermission(
-        "android.permission.READ_PRIVILEGED_PHONE_STATE",
-        ThreadState::get_calling_pid(),
-        ThreadState::get_calling_uid() as i32,
-    );
+    let binder_result = {
+        let _wp = watchdog::watch_millis(
+            "In check_device_attestation_permissions: calling checkPermission.",
+            500,
+        );
+        permission_controller.checkPermission(
+            "android.permission.READ_PRIVILEGED_PHONE_STATE",
+            ThreadState::get_calling_pid(),
+            ThreadState::get_calling_uid() as i32,
+        )
+    };
     let has_permissions = map_binder_status(binder_result)
         .context("In check_device_attestation_permissions: checkPermission failed")?;
     match has_permissions {
@@ -180,19 +185,15 @@
     parameters.into_iter().map(|p| p.into_authorization()).collect()
 }
 
-/// This returns the current time (in seconds) as an instance of a monotonic clock, by invoking the
-/// system call since Rust does not support getting monotonic time instance as an integer.
-pub fn get_current_time_in_seconds() -> i64 {
+/// This returns the current time (in milliseconds) as an instance of a monotonic clock,
+/// by invoking the system call since Rust does not support getting monotonic time instance
+/// as an integer.
+pub fn get_current_time_in_milliseconds() -> i64 {
     let mut current_time = libc::timespec { tv_sec: 0, tv_nsec: 0 };
     // Following unsafe block includes one system call to get monotonic time.
     // Therefore, it is not considered harmful.
     unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC_RAW, &mut current_time) };
-    // It is safe to unwrap here because try_from() returns std::convert::Infallible, which is
-    // defined to be an error that can never happen (i.e. the result is always ok).
-    // This suppresses the compiler's complaint about converting tv_sec to i64 in method
-    // get_current_time_in_seconds.
-    #[allow(clippy::useless_conversion)]
-    i64::try_from(current_time.tv_sec).unwrap()
+    current_time.tv_sec as i64 * 1000 + (current_time.tv_nsec as i64 / 1_000_000)
 }
 
 /// Converts a response code as returned by the Android Protected Confirmation HIDL compatibility
@@ -233,6 +234,55 @@
     unsafe { cutils_bindgen::multiuser_get_user_id(uid) }
 }
 
+/// This module provides helpers for simplified use of the watchdog module.
+#[cfg(feature = "watchdog")]
+pub mod watchdog {
+    pub use crate::watchdog::WatchPoint;
+    use crate::watchdog::Watchdog;
+    use lazy_static::lazy_static;
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    lazy_static! {
+        /// A Watchdog thread, that can be used to create watch points.
+        static ref WD: Arc<Watchdog> = Watchdog::new(Duration::from_secs(10));
+    }
+
+    /// Sets a watch point with `id` and a timeout of `millis` milliseconds.
+    pub fn watch_millis(id: &'static str, millis: u64) -> Option<WatchPoint> {
+        Watchdog::watch(&WD, id, Duration::from_millis(millis))
+    }
+
+    /// Like `watch_millis` but with a callback that is called every time a report
+    /// is printed about this watch point.
+    pub fn watch_millis_with(
+        id: &'static str,
+        millis: u64,
+        callback: impl Fn() -> String + Send + 'static,
+    ) -> Option<WatchPoint> {
+        Watchdog::watch_with(&WD, id, Duration::from_millis(millis), callback)
+    }
+}
+
+/// This module provides empty/noop implementations of the watch dog utility functions.
+#[cfg(not(feature = "watchdog"))]
+pub mod watchdog {
+    /// Noop watch point.
+    pub struct WatchPoint();
+    /// Sets a Noop watch point.
+    fn watch_millis(_: &'static str, _: u64) -> Option<WatchPoint> {
+        None
+    }
+
+    pub fn watch_millis_with(
+        _: &'static str,
+        _: u64,
+        _: impl Fn() -> String + Send + 'static,
+    ) -> Option<WatchPoint> {
+        None
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/keystore2/src/watchdog.rs b/keystore2/src/watchdog.rs
index 0171901..9cca171 100644
--- a/keystore2/src/watchdog.rs
+++ b/keystore2/src/watchdog.rs
@@ -66,18 +66,19 @@
     thread: Option<thread::JoinHandle<()>>,
     timeout: Duration,
     records: HashMap<Index, Record>,
+    last_report: Instant,
     has_overdue: bool,
 }
 
 impl WatchdogState {
-    fn update_overdue_and_find_next_timeout(&mut self) -> Option<Duration> {
+    fn update_overdue_and_find_next_timeout(&mut self) -> (bool, Option<Duration>) {
         let now = Instant::now();
         let mut next_timeout: Option<Duration> = None;
-        self.has_overdue = false;
+        let mut has_overdue = false;
         for (_, r) in self.records.iter() {
             let timeout = r.deadline.saturating_duration_since(now);
             if timeout == Duration::new(0, 0) {
-                self.has_overdue = true;
+                has_overdue = true;
                 continue;
             }
             next_timeout = match next_timeout {
@@ -91,13 +92,25 @@
                 None => Some(timeout),
             };
         }
-        next_timeout
+        (has_overdue, next_timeout)
     }
 
-    fn log_report(&self) -> bool {
-        if !self.has_overdue {
-            return false;
+    fn log_report(&mut self, has_overdue: bool) -> bool {
+        match (self.has_overdue, has_overdue) {
+            (true, true) => {
+                if self.last_report.elapsed() < Watchdog::NOISY_REPORT_TIMEOUT {
+                    self.has_overdue = false;
+                    return false;
+                }
+            }
+            (_, false) => {
+                self.has_overdue = false;
+                return false;
+            }
+            (false, true) => {}
         }
+        self.last_report = Instant::now();
+        self.has_overdue = has_overdue;
         log::warn!("Keystore Watchdog report:");
         log::warn!("Overdue records:");
         let now = Instant::now();
@@ -164,6 +177,7 @@
                     thread: None,
                     timeout,
                     records: HashMap::new(),
+                    last_report: Instant::now(),
                     has_overdue: false,
                 }),
             )),
@@ -249,8 +263,8 @@
             let mut state = state.lock().unwrap();
 
             loop {
-                let next_timeout = state.update_overdue_and_find_next_timeout();
-                let has_overdue = state.log_report();
+                let (has_overdue, next_timeout) = state.update_overdue_and_find_next_timeout();
+                state.log_report(has_overdue);
                 let (next_timeout, idle) = match (has_overdue, next_timeout) {
                     (true, Some(next_timeout)) => {
                         (min(next_timeout, Self::NOISY_REPORT_TIMEOUT), false)
@@ -268,6 +282,7 @@
                     break;
                 }
             }
+            log::info!("Watchdog thread idle -> terminating. Have a great day.");
         }));
         state.state = State::Running;
     }
diff --git a/keystore2/vpnprofilestore/lib.rs b/keystore2/vpnprofilestore/lib.rs
index d92e045..df2731a 100644
--- a/keystore2/vpnprofilestore/lib.rs
+++ b/keystore2/vpnprofilestore/lib.rs
@@ -22,33 +22,50 @@
     BinderFeatures, ExceptionCode, Result as BinderResult, Status as BinderStatus, Strong,
     ThreadState,
 };
-use anyhow::{Context, Result};
-use keystore2::{async_task::AsyncTask, legacy_blob::LegacyBlobLoader};
+use anyhow::{anyhow, Context, Result};
+use keystore2::{async_task::AsyncTask, legacy_blob::LegacyBlobLoader, utils::watchdog as wd};
 use rusqlite::{
     params, Connection, OptionalExtension, Transaction, TransactionBehavior, NO_PARAMS,
 };
 use std::{
     collections::HashSet,
     path::{Path, PathBuf},
+    sync::Once,
 };
 
+static DB_SET_WAL_MODE: Once = Once::new();
+
 struct DB {
     conn: Connection,
 }
 
 impl DB {
     fn new(db_file: &Path) -> Result<Self> {
+        DB_SET_WAL_MODE.call_once(|| {
+            log::info!("Setting VpnProfileStore database to WAL mode first time since boot.");
+            Self::set_wal_mode(&db_file).expect("In vpnprofilestore: Could not set WAL mode.");
+        });
+
         let mut db = Self {
             conn: Connection::open(db_file).context("Failed to initialize SQLite connection.")?,
         };
 
-        // On busy fail Immediately. It is unlikely to succeed given a bug in sqlite.
-        db.conn.busy_handler(None).context("Failed to set busy handler.")?;
-
         db.init_tables().context("Trying to initialize vpnstore db.")?;
         Ok(db)
     }
 
+    fn set_wal_mode(db_file: &Path) -> Result<()> {
+        let conn = Connection::open(db_file)
+            .context("In VpnProfileStore set_wal_mode: Failed to open DB.")?;
+        let mode: String = conn
+            .pragma_update_and_check(None, "journal_mode", &"WAL", |row| row.get(0))
+            .context("In VpnProfileStore set_wal_mode: Failed to set journal_mode")?;
+        match mode.as_str() {
+            "wal" => Ok(()),
+            _ => Err(anyhow!("Unable to set WAL mode, db is still in {} mode.", mode)),
+        }
+    }
+
     fn with_transaction<T, F>(&mut self, behavior: TransactionBehavior, f: F) -> Result<T>
     where
         F: Fn(&Transaction) -> Result<T>,
@@ -366,15 +383,19 @@
 
 impl IVpnProfileStore for VpnProfileStore {
     fn get(&self, alias: &str) -> BinderResult<Vec<u8>> {
+        let _wp = wd::watch_millis("IVpnProfileStore::get", 500);
         map_or_log_err(self.get(alias), Ok)
     }
     fn put(&self, alias: &str, profile: &[u8]) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IVpnProfileStore::put", 500);
         map_or_log_err(self.put(alias, profile), Ok)
     }
     fn remove(&self, alias: &str) -> BinderResult<()> {
+        let _wp = wd::watch_millis("IVpnProfileStore::remove", 500);
         map_or_log_err(self.remove(alias), Ok)
     }
     fn list(&self, prefix: &str) -> BinderResult<Vec<String>> {
+        let _wp = wd::watch_millis("IVpnProfileStore::list", 500);
         map_or_log_err(self.list(prefix), Ok)
     }
 }
@@ -466,6 +487,9 @@
         const PROFILE_COUNT: u32 = 5000u32;
         const PROFILE_DB_COUNT: u32 = 5000u32;
 
+        let mode: String = db.conn.pragma_query_value(None, "journal_mode", |row| row.get(0))?;
+        assert_eq!(mode, "wal");
+
         let mut actual_profile_count = PROFILE_COUNT;
         // First insert PROFILE_COUNT profiles.
         for count in 0..PROFILE_COUNT {
diff --git a/ondevice-signing/Android.bp b/ondevice-signing/Android.bp
index 2e5e02e..432e585 100644
--- a/ondevice-signing/Android.bp
+++ b/ondevice-signing/Android.bp
@@ -84,22 +84,19 @@
   srcs: [
     "odsign_main.cpp",
     "CertUtils.cpp",
-    "Keymaster.cpp",
-    "KeymasterSigningKey.cpp",
     "KeystoreKey.cpp",
+    "KeystoreHmacKey.cpp",
     "VerityUtils.cpp",
   ],
 
   header_libs: ["odrefresh_headers"],
 
   static_libs: [
-    "libmini_keyctl_static", // TODO need static?
     "libc++fs",
     "lib_odsign_proto",
   ],
 
   shared_libs: [
-    "android.hardware.keymaster@4.1",
     "android.system.keystore2-V1-cpp",
     "android.hardware.security.keymint-V1-cpp",
     "libbase",
@@ -107,11 +104,7 @@
     "libcrypto",
     "libcrypto_utils",
     "libfsverity",
-    "libhidlbase",
     "liblogwrap",
-    "libkeymaster4support", // For authorization_set
-    "libkeymaster4_1support",
-    "libkeyutils",
     "libprotobuf-cpp-full",
     "libutils",
   ],
diff --git a/ondevice-signing/KeyConstants.h b/ondevice-signing/KeyConstants.h
index 9e1a513..ccc9251 100644
--- a/ondevice-signing/KeyConstants.h
+++ b/ondevice-signing/KeyConstants.h
@@ -16,3 +16,6 @@
 
 static constexpr int kRsaKeySize = 2048;
 static constexpr int kRsaKeyExponent = 65537;
+
+static constexpr int kHmacKeySize = 256;
+static constexpr int kHmacMinMacLength = 256;
diff --git a/ondevice-signing/Keymaster.cpp b/ondevice-signing/Keymaster.cpp
deleted file mode 100644
index f9bf9b2..0000000
--- a/ondevice-signing/Keymaster.cpp
+++ /dev/null
@@ -1,293 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <string>
-
-#include <android-base/logging.h>
-#include <keymasterV4_1/Keymaster.h>
-#include <keymasterV4_1/authorization_set.h>
-
-#include <fcntl.h>
-#include <sys/stat.h>
-#include <sys/types.h>
-
-#include "Keymaster.h"
-
-using AuthorizationSet = ::android::hardware::keymaster::V4_0::AuthorizationSet;
-using AuthorizationSetBuilder = ::android::hardware::keymaster::V4_0::AuthorizationSetBuilder;
-using Digest = ::android::hardware::keymaster::V4_0::Digest;
-using ErrorCode = ::android::hardware::keymaster::V4_0::ErrorCode;
-using HardwareAuthToken = ::android::hardware::keymaster::V4_0::HardwareAuthToken;
-using HidlBuf = ::android::hardware::hidl_vec<uint8_t>;
-using KeyCharacteristics = ::android::hardware::keymaster::V4_0::KeyCharacteristics;
-using KeyFormat = ::android::hardware::keymaster::V4_0::KeyFormat;
-using KeyParameter = ::android::hardware::keymaster::V4_0::KeyParameter;
-using KeyPurpose = ::android::hardware::keymaster::V4_0::KeyPurpose;
-using KmSupport = ::android::hardware::keymaster::V4_1::support::Keymaster;
-using KmDevice = ::android::hardware::keymaster::V4_1::IKeymasterDevice;
-using OperationHandle = ::android::hardware::keymaster::V4_0::OperationHandle;
-using PaddingMode = ::android::hardware::keymaster::V4_0::PaddingMode;
-using VerificationToken = ::android::hardware::keymaster::V4_0::VerificationToken;
-
-using android::sp;
-using android::base::Error;
-using android::base::Result;
-using android::hardware::hidl_vec;
-
-Keymaster::Keymaster() {}
-
-bool Keymaster::initialize() {
-    // TODO(b/165630556): Stop using Keymaster directly and migrate to keystore2
-    // (once available).
-    auto devices = KmSupport::enumerateAvailableDevices();
-    sp<KmDevice> devToUse = nullptr;
-    for (const auto& dev : devices) {
-        auto version = dev->halVersion();
-        if (version.majorVersion > 4 || (version.majorVersion == 4 && version.minorVersion >= 1)) {
-            // TODO we probably have a preference for the SE, hoping Keystore2 will provide this
-            LOG(INFO) << "Using keymaster " << version.keymasterName << " "
-                      << (int)version.majorVersion << "." << (int)version.minorVersion;
-            devToUse = dev;
-            break;
-        }
-    }
-
-    if (devToUse == nullptr) {
-        LOG(WARNING) << "Didn't find a keymaster to use.";
-    }
-    mDevice = devToUse;
-
-    return mDevice != nullptr;
-}
-
-std::optional<Keymaster> Keymaster::getInstance() {
-    static Keymaster keymaster;
-
-    if (!keymaster.initialize()) {
-        return {};
-    } else {
-        return {keymaster};
-    }
-}
-
-Result<std::vector<uint8_t>> Keymaster::createKey() const {
-    ErrorCode error;
-    HidlBuf keyBlob;
-
-    auto params = AuthorizationSetBuilder()
-                      .Authorization(::android::hardware::keymaster::V4_0::TAG_NO_AUTH_REQUIRED)
-                      // TODO MAKE SURE WE ADD THE EARLY_BOOT_ONLY FLAG here
-                      // currently doesn't work on cuttlefish (b/173618442)
-                      //.Authorization(::android::hardware::keymaster::V4_1::TAG_EARLY_BOOT_ONLY)
-                      .RsaSigningKey(2048, 65537)
-                      .Digest(Digest::SHA_2_256)
-                      .Padding(PaddingMode::RSA_PKCS1_1_5_SIGN);
-
-    mDevice->generateKey(params.hidl_data(), [&](ErrorCode hidl_error, const HidlBuf& hidl_key_blob,
-                                                 const KeyCharacteristics&
-                                                 /* hidl_key_characteristics */) {
-        error = hidl_error;
-        keyBlob = hidl_key_blob;
-    });
-
-    if (error != ErrorCode::OK) {
-        return Error() << "Error creating keymaster signing key: "
-                       << static_cast<std::underlying_type<ErrorCode>::type>(error);
-    }
-
-    return keyBlob;
-}
-
-static ErrorCode Begin(const sp<KmDevice>& keymaster_, KeyPurpose purpose, const HidlBuf& key_blob,
-                       const AuthorizationSet& in_params, AuthorizationSet* out_params,
-                       OperationHandle* op_handle) {
-    ErrorCode error;
-    OperationHandle saved_handle = *op_handle;
-    CHECK(keymaster_
-              ->begin(purpose, key_blob, in_params.hidl_data(), HardwareAuthToken(),
-                      [&](ErrorCode hidl_error, const hidl_vec<KeyParameter>& hidl_out_params,
-                          uint64_t hidl_op_handle) {
-                          error = hidl_error;
-                          *out_params = hidl_out_params;
-                          *op_handle = hidl_op_handle;
-                      })
-              .isOk());
-    if (error != ErrorCode::OK) {
-        // Some implementations may modify *op_handle on error.
-        *op_handle = saved_handle;
-    }
-    return error;
-}
-
-static ErrorCode Update(const sp<KmDevice>& keymaster_, OperationHandle op_handle,
-                        const AuthorizationSet& in_params, const std::string& input,
-                        AuthorizationSet* out_params, std::string* output, size_t* input_consumed) {
-    ErrorCode error;
-    HidlBuf inputData(input.size());
-    memcpy(inputData.data(), input.c_str(), input.size());
-    CHECK(keymaster_
-              ->update(op_handle, in_params.hidl_data(), inputData, HardwareAuthToken(),
-                       VerificationToken(),
-                       [&](ErrorCode hidl_error, uint32_t hidl_input_consumed,
-                           const hidl_vec<KeyParameter>& hidl_out_params,
-                           const HidlBuf& hidl_output) {
-                           error = hidl_error;
-                           out_params->push_back(AuthorizationSet(hidl_out_params));
-                           std::string retdata(reinterpret_cast<const char*>(hidl_output.data()),
-                                               hidl_output.size());
-                           output->append(retdata);
-                           *input_consumed = hidl_input_consumed;
-                       })
-              .isOk());
-    return error;
-}
-
-static ErrorCode Finish(const sp<KmDevice>& keymaster_, OperationHandle op_handle,
-                        const AuthorizationSet& in_params, const std::string& input,
-                        const std::string& signature, AuthorizationSet* out_params,
-                        std::string* output) {
-    ErrorCode error;
-    HidlBuf inputData(input.size());
-    memcpy(inputData.data(), input.c_str(), input.size());
-    HidlBuf signatureData(signature.size());
-    memcpy(signatureData.data(), signature.c_str(), signature.size());
-    // TODO still need to handle error -62 - key requires upgrade
-    CHECK(keymaster_
-              ->finish(op_handle, in_params.hidl_data(), inputData, signatureData,
-                       HardwareAuthToken(), VerificationToken(),
-                       [&](ErrorCode hidl_error, const hidl_vec<KeyParameter>& hidl_out_params,
-                           const HidlBuf& hidl_output) {
-                           error = hidl_error;
-                           *out_params = hidl_out_params;
-                           std::string retdata(reinterpret_cast<const char*>(hidl_output.data()),
-                                               hidl_output.size());
-                           output->append(retdata);
-                       })
-              .isOk());
-    return error;
-}
-
-static std::string ProcessMessage(const sp<KmDevice>& keymaster_, const HidlBuf& key_blob,
-                                  KeyPurpose operation, const std::string& message,
-                                  const AuthorizationSet& in_params, AuthorizationSet* out_params) {
-    AuthorizationSet begin_out_params;
-    OperationHandle op_handle_;
-    ErrorCode ec =
-        Begin(keymaster_, operation, key_blob, in_params, &begin_out_params, &op_handle_);
-
-    std::string output;
-    size_t consumed = 0;
-    AuthorizationSet update_params;
-    AuthorizationSet update_out_params;
-    ec = Update(keymaster_, op_handle_, update_params, message, &update_out_params, &output,
-                &consumed);
-
-    std::string unused;
-    AuthorizationSet finish_params;
-    AuthorizationSet finish_out_params;
-    ec = Finish(keymaster_, op_handle_, finish_params, message.substr(consumed), unused,
-                &finish_out_params, &output);
-
-    out_params->push_back(begin_out_params);
-    out_params->push_back(finish_out_params);
-    return output;
-}
-
-Result<std::vector<uint8_t>>
-Keymaster::extractPublicKey(const std::vector<uint8_t>& keyBlob) const {
-    std::vector<uint8_t> publicKey;
-    ErrorCode error;
-
-    mDevice->exportKey(KeyFormat::X509, keyBlob, {} /* clientId */, {} /* appData */,
-                       [&](ErrorCode hidl_error, const HidlBuf& keyData) {
-                           error = hidl_error;
-                           publicKey = keyData;
-                       });
-
-    if (error != ErrorCode::OK) {
-        return Error() << "Error extracting public key: "
-                       << static_cast<std::underlying_type<ErrorCode>::type>(error);
-    }
-
-    return publicKey;
-}
-
-Result<KeymasterVerifyResult> Keymaster::verifyKey(const std::vector<uint8_t>& keyBlob) const {
-    ErrorCode error;
-    KeyCharacteristics characteristics;
-
-    mDevice->getKeyCharacteristics(
-        keyBlob, {} /* clientId */, {} /* appData */,
-        [&](ErrorCode hidl_error, const KeyCharacteristics& hidl_characteristics) {
-            error = hidl_error;
-            characteristics = hidl_characteristics;
-        });
-
-    if (error == ErrorCode::KEY_REQUIRES_UPGRADE) {
-        return KeymasterVerifyResult::UPGRADE;
-    }
-
-    if (error != ErrorCode::OK) {
-        return Error() << "Error getting key characteristics: "
-                       << static_cast<std::underlying_type<ErrorCode>::type>(error);
-    }
-
-    // TODO(b/165630556)
-    // Verify this is an early boot key and the other key parameters
-    return KeymasterVerifyResult::OK;
-}
-
-Result<std::vector<uint8_t>> Keymaster::upgradeKey(const std::vector<uint8_t>& keyBlob) const {
-    ErrorCode error;
-    HidlBuf newKeyBlob;
-
-    // TODO deduplicate
-    auto params = AuthorizationSetBuilder()
-                      .Authorization(::android::hardware::keymaster::V4_0::TAG_NO_AUTH_REQUIRED)
-                      // TODO MAKE SURE WE ADD THE EARLY_BOOT_ONLY FLAG here
-                      // currently doesn't work on cuttlefish (b/173618442)
-                      //.Authorization(::android::hardware::keymaster::V4_1::TAG_EARLY_BOOT_ONLY)
-                      .RsaSigningKey(2048, 65537)
-                      .Digest(Digest::SHA_2_256)
-                      .Padding(PaddingMode::RSA_PKCS1_1_5_SIGN);
-
-    mDevice->upgradeKey(keyBlob, params.hidl_data(),
-                        [&](ErrorCode hidl_error, const HidlBuf& hidl_key_blob) {
-                            error = hidl_error;
-                            newKeyBlob = hidl_key_blob;
-                        });
-
-    if (error != ErrorCode::OK) {
-        return Error() << "Error upgrading keymaster signing key: "
-                       << static_cast<std::underlying_type<ErrorCode>::type>(error);
-    }
-
-    return newKeyBlob;
-}
-
-Result<std::string> Keymaster::sign(const std::vector<uint8_t>& keyBlob,
-                                    const std::string& message) const {
-    AuthorizationSet out_params;
-    auto params = AuthorizationSetBuilder()
-                      .Digest(Digest::SHA_2_256)
-                      .Padding(PaddingMode::RSA_PKCS1_1_5_SIGN);
-    std::string signature =
-        ProcessMessage(mDevice, keyBlob, KeyPurpose::SIGN, message, params, &out_params);
-    if (!out_params.empty()) {
-        return Error() << "Error signing key: expected empty out params.";
-    }
-    return signature;
-}
diff --git a/ondevice-signing/Keymaster.h b/ondevice-signing/Keymaster.h
deleted file mode 100644
index 455289f..0000000
--- a/ondevice-signing/Keymaster.h
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#pragma once
-
-#include <optional>
-
-#include <android-base/macros.h>
-#include <android-base/result.h>
-#include <android-base/unique_fd.h>
-
-#include <keymasterV4_1/Keymaster.h>
-
-#include <utils/StrongPointer.h>
-
-enum class KeymasterVerifyResult {
-    OK = 0,
-    UPGRADE = -1,
-};
-
-class Keymaster {
-    using KmDevice = ::android::hardware::keymaster::V4_1::IKeymasterDevice;
-
-  public:
-    static std::optional<Keymaster> getInstance();
-
-    android::base::Result<std::vector<uint8_t>> createKey() const;
-
-    android::base::Result<std::vector<uint8_t>>
-    extractPublicKey(const std::vector<uint8_t>& keyBlob) const;
-
-    android::base::Result<KeymasterVerifyResult>
-    verifyKey(const std::vector<uint8_t>& keyBlob) const;
-
-    android::base::Result<std::vector<uint8_t>>
-    upgradeKey(const std::vector<uint8_t>& keyBlob) const;
-
-    /* Sign a message with an initialized signing key */
-    android::base::Result<std::string> sign(const std::vector<uint8_t>& keyBlob,
-                                            const std::string& message) const;
-
-  private:
-    Keymaster();
-    bool initialize();
-
-    android::sp<KmDevice> mDevice;
-};
diff --git a/ondevice-signing/KeymasterSigningKey.cpp b/ondevice-signing/KeymasterSigningKey.cpp
deleted file mode 100644
index dc3ef8a..0000000
--- a/ondevice-signing/KeymasterSigningKey.cpp
+++ /dev/null
@@ -1,168 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <string>
-
-#include <android-base/file.h>
-#include <android-base/logging.h>
-
-#include <fcntl.h>
-#include <sys/stat.h>
-#include <sys/types.h>
-
-#include "CertUtils.h"
-#include "Keymaster.h"
-#include "KeymasterSigningKey.h"
-
-using android::base::ErrnoError;
-using android::base::Error;
-using android::base::ReadFileToString;
-using android::base::Result;
-using android::base::unique_fd;
-
-const std::string kSigningKeyBlob = "/data/misc/odsign/key.blob";
-
-KeymasterSigningKey::KeymasterSigningKey() {}
-
-Result<std::unique_ptr<KeymasterSigningKey>>
-KeymasterSigningKey::loadFromBlobAndVerify(const std::string& path) {
-    auto signingKey = std::make_unique<KeymasterSigningKey>();
-
-    auto status = signingKey->initializeFromKeyblob(path);
-
-    if (!status.ok()) {
-        return status.error();
-    }
-
-    return signingKey;
-}
-
-Result<void> KeymasterSigningKey::saveKeyblob(const std::string& path) const {
-    int flags = O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC;
-
-    unique_fd fd(TEMP_FAILURE_RETRY(open(path.c_str(), flags, 0600)));
-    if (fd == -1) {
-        return ErrnoError() << "Error creating key blob file " << path;
-    }
-
-    if (!android::base::WriteFully(fd, mVerifiedKeyBlob.data(), mVerifiedKeyBlob.size())) {
-        return ErrnoError() << "Error writing key blob file " << path;
-    } else {
-        return {};
-    }
-}
-
-Result<void> KeymasterSigningKey::createSigningKey() {
-    KeymasterSigningKey signingKey;
-    auto keymaster = Keymaster::getInstance();
-    if (!keymaster.has_value()) {
-        return Error() << "Failed to initialize keymaster.";
-    }
-    mKeymaster = keymaster;
-
-    auto keyBlob = mKeymaster->createKey();
-
-    if (!keyBlob.ok()) {
-        return keyBlob.error();
-    }
-
-    mVerifiedKeyBlob.assign(keyBlob->begin(), keyBlob->end());
-
-    return {};
-}
-
-Result<std::unique_ptr<KeymasterSigningKey>> KeymasterSigningKey::createAndPersistNewKey() {
-    auto signingKey = std::make_unique<KeymasterSigningKey>();
-
-    auto status = signingKey->createSigningKey();
-
-    if (!status.ok()) {
-        return status.error();
-    }
-
-    status = signingKey->saveKeyblob(kSigningKeyBlob);
-    if (!status.ok()) {
-        return status.error();
-    }
-
-    return signingKey;
-}
-
-Result<SigningKey*> KeymasterSigningKey::getInstance() {
-    auto key = loadFromBlobAndVerify(kSigningKeyBlob);
-
-    if (!key.ok()) {
-        key = createAndPersistNewKey();
-        if (!key.ok()) {
-            return key.error();
-        }
-    }
-
-    return key->release();
-}
-
-Result<std::vector<uint8_t>> KeymasterSigningKey::getPublicKey() const {
-    auto publicKey = mKeymaster->extractPublicKey(mVerifiedKeyBlob);
-    if (!publicKey.ok()) {
-        return publicKey.error();
-    }
-
-    // Keymaster returns the public key not in a full X509 cert, but just the
-    // "SubjectPublicKeyInfo"
-    return extractPublicKeyFromSubjectPublicKeyInfo(publicKey.value());
-}
-
-Result<void> KeymasterSigningKey::initializeFromKeyblob(const std::string& path) {
-    std::string keyBlobData;
-    auto keymaster = Keymaster::getInstance();
-    if (!keymaster.has_value()) {
-        return Error() << "Failed to initialize keymaster.";
-    }
-    mKeymaster = keymaster;
-
-    bool result = ReadFileToString(path, &keyBlobData);
-    if (!result) {
-        return ErrnoError() << "Failed to read " << path;
-    }
-
-    std::vector<uint8_t> keyBlob = {keyBlobData.begin(), keyBlobData.end()};
-
-    auto verifyResult = mKeymaster->verifyKey(keyBlob);
-    if (!verifyResult.ok()) {
-        return Error() << "Failed to verify key: " << verifyResult.error().message();
-    }
-
-    if (*verifyResult == KeymasterVerifyResult::UPGRADE) {
-        auto upgradeResult = mKeymaster->upgradeKey(keyBlob);
-        if (!upgradeResult.ok()) {
-            return Error() << "Failed to upgrade key: " << upgradeResult.error().message();
-        }
-        mVerifiedKeyBlob = *upgradeResult;
-        // Make sure we persist the new blob
-        auto saveResult = saveKeyblob(path);
-        if (!saveResult.ok()) {
-            return Error() << "Failed to store upgraded key";
-        }
-    } else {
-        mVerifiedKeyBlob = keyBlob;
-    }
-
-    return {};
-}
-
-Result<std::string> KeymasterSigningKey::sign(const std::string& message) const {
-    return mKeymaster->sign(mVerifiedKeyBlob, message);
-}
diff --git a/ondevice-signing/KeymasterSigningKey.h b/ondevice-signing/KeymasterSigningKey.h
deleted file mode 100644
index e66781f..0000000
--- a/ondevice-signing/KeymasterSigningKey.h
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#pragma once
-
-#include <android-base/macros.h>
-#include <android-base/result.h>
-#include <android-base/unique_fd.h>
-
-#include <utils/StrongPointer.h>
-
-#include "Keymaster.h"
-#include "SigningKey.h"
-
-class KeymasterSigningKey : public SigningKey {
-    using KmDevice = ::android::hardware::keymaster::V4_1::IKeymasterDevice;
-
-  public:
-    friend std::unique_ptr<KeymasterSigningKey> std::make_unique<KeymasterSigningKey>();
-    virtual ~KeymasterSigningKey(){};
-
-    // Allow the key to be moved around
-    KeymasterSigningKey& operator=(KeymasterSigningKey&& other) = default;
-    KeymasterSigningKey(KeymasterSigningKey&& other) = default;
-
-    static android::base::Result<SigningKey*> getInstance();
-
-    virtual android::base::Result<std::string> sign(const std::string& message) const;
-    virtual android::base::Result<std::vector<uint8_t>> getPublicKey() const;
-
-  private:
-    KeymasterSigningKey();
-
-    static android::base::Result<std::unique_ptr<KeymasterSigningKey>> createAndPersistNewKey();
-    static android::base::Result<std::unique_ptr<KeymasterSigningKey>>
-    loadFromBlobAndVerify(const std::string& path);
-
-    android::base::Result<void> createSigningKey();
-    android::base::Result<void> initializeFromKeyblob(const std::string& path);
-    android::base::Result<void> saveKeyblob(const std::string& path) const;
-
-    static android::base::Result<KeymasterSigningKey> createNewKey();
-
-    std::optional<Keymaster> mKeymaster;
-    std::vector<uint8_t> mVerifiedKeyBlob;
-
-    DISALLOW_COPY_AND_ASSIGN(KeymasterSigningKey);
-};
diff --git a/ondevice-signing/KeystoreHmacKey.cpp b/ondevice-signing/KeystoreHmacKey.cpp
new file mode 100644
index 0000000..db8d7d9
--- /dev/null
+++ b/ondevice-signing/KeystoreHmacKey.cpp
@@ -0,0 +1,267 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <string>
+
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <binder/IServiceManager.h>
+
+#include <fcntl.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#include "CertUtils.h"
+#include "KeyConstants.h"
+#include "KeystoreHmacKey.h"
+
+using android::sp;
+using android::String16;
+
+using android::hardware::security::keymint::Algorithm;
+using android::hardware::security::keymint::Digest;
+using android::hardware::security::keymint::KeyParameter;
+using android::hardware::security::keymint::KeyParameterValue;
+using android::hardware::security::keymint::KeyPurpose;
+using android::hardware::security::keymint::Tag;
+
+using android::system::keystore2::CreateOperationResponse;
+using android::system::keystore2::Domain;
+using android::system::keystore2::KeyDescriptor;
+using android::system::keystore2::KeyEntryResponse;
+using android::system::keystore2::KeyMetadata;
+
+using android::base::Error;
+using android::base::Result;
+
+using android::base::unique_fd;
+
+// Keystore boot level that the odsign key uses
+static const int kOdsignBootLevel = 30;
+
+static KeyDescriptor getHmacKeyDescriptor() {
+    // AIDL parcelable objects don't have constructor
+    static KeyDescriptor descriptor;
+    static std::once_flag flag;
+    std::call_once(flag, [&]() {
+        descriptor.domain = Domain::SELINUX;
+        descriptor.alias = String16("ondevice-signing-hmac");
+        descriptor.nspace = 101;  // odsign_key
+    });
+
+    return descriptor;
+}
+
+Result<void> KeystoreHmacKey::createKey() {
+    std::vector<KeyParameter> params;
+
+    KeyParameter algo;
+    algo.tag = Tag::ALGORITHM;
+    algo.value = KeyParameterValue::make<KeyParameterValue::algorithm>(Algorithm::HMAC);
+    params.push_back(algo);
+
+    KeyParameter key_size;
+    key_size.tag = Tag::KEY_SIZE;
+    key_size.value = KeyParameterValue::make<KeyParameterValue::integer>(kHmacKeySize);
+    params.push_back(key_size);
+
+    KeyParameter min_mac_length;
+    min_mac_length.tag = Tag::MIN_MAC_LENGTH;
+    min_mac_length.value = KeyParameterValue::make<KeyParameterValue::integer>(256);
+    params.push_back(min_mac_length);
+
+    KeyParameter digest;
+    digest.tag = Tag::DIGEST;
+    digest.value = KeyParameterValue::make<KeyParameterValue::digest>(Digest::SHA_2_256);
+    params.push_back(digest);
+
+    KeyParameter purposeSign;
+    purposeSign.tag = Tag::PURPOSE;
+    purposeSign.value = KeyParameterValue::make<KeyParameterValue::keyPurpose>(KeyPurpose::SIGN);
+    params.push_back(purposeSign);
+
+    KeyParameter purposeVerify;
+    purposeVerify.tag = Tag::PURPOSE;
+    purposeVerify.value =
+        KeyParameterValue::make<KeyParameterValue::keyPurpose>(KeyPurpose::VERIFY);
+    params.push_back(purposeVerify);
+
+    KeyParameter auth;
+    auth.tag = Tag::NO_AUTH_REQUIRED;
+    auth.value = KeyParameterValue::make<KeyParameterValue::boolValue>(true);
+    params.push_back(auth);
+
+    KeyParameter boot_level;
+    boot_level.tag = Tag::MAX_BOOT_LEVEL;
+    boot_level.value = KeyParameterValue::make<KeyParameterValue::integer>(kOdsignBootLevel);
+    params.push_back(boot_level);
+
+    KeyMetadata metadata;
+    auto status = mSecurityLevel->generateKey(mDescriptor, {}, params, 0, {}, &metadata);
+    if (!status.isOk()) {
+        return Error() << "Failed to create new HMAC key";
+    }
+
+    return {};
+}
+
+Result<void> KeystoreHmacKey::initialize(sp<IKeystoreService> service,
+                                         sp<IKeystoreSecurityLevel> securityLevel) {
+    mService = std::move(service);
+    mSecurityLevel = std::move(securityLevel);
+
+    // See if we can fetch an existing key
+    KeyEntryResponse keyEntryResponse;
+    LOG(INFO) << "Trying to retrieve existing HMAC key...";
+    auto status = mService->getKeyEntry(mDescriptor, &keyEntryResponse);
+    bool keyValid = false;
+
+    if (status.isOk()) {
+        // Make sure this is an early boot key
+        for (const auto& auth : keyEntryResponse.metadata.authorizations) {
+            if (auth.keyParameter.tag == Tag::MAX_BOOT_LEVEL) {
+                if (auth.keyParameter.value.get<KeyParameterValue::integer>() == kOdsignBootLevel) {
+                    keyValid = true;
+                    break;
+                }
+            }
+        }
+        if (!keyValid) {
+            LOG(WARNING) << "Found invalid HMAC key without MAX_BOOT_LEVEL tag";
+        }
+    }
+
+    if (!keyValid) {
+        LOG(INFO) << "Existing HMAC key not found or invalid, creating new key";
+        return createKey();
+    } else {
+        return {};
+    }
+}
+
+KeystoreHmacKey::KeystoreHmacKey() {
+    mDescriptor = getHmacKeyDescriptor();
+}
+
+static std::vector<KeyParameter> getVerifyOpParameters() {
+    std::vector<KeyParameter> opParameters;
+
+    KeyParameter algo;
+    algo.tag = Tag::ALGORITHM;
+    algo.value = KeyParameterValue::make<KeyParameterValue::algorithm>(Algorithm::HMAC);
+    opParameters.push_back(algo);
+
+    KeyParameter digest;
+    digest.tag = Tag::DIGEST;
+    digest.value = KeyParameterValue::make<KeyParameterValue::digest>(Digest::SHA_2_256);
+    opParameters.push_back(digest);
+
+    KeyParameter mac_length;
+    mac_length.tag = Tag::MAC_LENGTH;
+    mac_length.value = KeyParameterValue::make<KeyParameterValue::integer>(256);
+    opParameters.push_back(mac_length);
+
+    KeyParameter purpose;
+    purpose.tag = Tag::PURPOSE;
+    purpose.value = KeyParameterValue::make<KeyParameterValue::keyPurpose>(KeyPurpose::VERIFY);
+    opParameters.push_back(purpose);
+
+    return opParameters;
+}
+
+static std::vector<KeyParameter> getSignOpParameters() {
+    std::vector<KeyParameter> opParameters;
+
+    KeyParameter algo;
+    algo.tag = Tag::ALGORITHM;
+    algo.value = KeyParameterValue::make<KeyParameterValue::algorithm>(Algorithm::HMAC);
+    opParameters.push_back(algo);
+
+    KeyParameter mac_length;
+    mac_length.tag = Tag::MAC_LENGTH;
+    mac_length.value = KeyParameterValue::make<KeyParameterValue::integer>(256);
+    opParameters.push_back(mac_length);
+
+    KeyParameter digest;
+    digest.tag = Tag::DIGEST;
+    digest.value = KeyParameterValue::make<KeyParameterValue::digest>(Digest::SHA_2_256);
+    opParameters.push_back(digest);
+
+    KeyParameter purpose;
+    purpose.tag = Tag::PURPOSE;
+    purpose.value = KeyParameterValue::make<KeyParameterValue::keyPurpose>(KeyPurpose::SIGN);
+    opParameters.push_back(purpose);
+
+    return opParameters;
+}
+
+Result<std::string> KeystoreHmacKey::sign(const std::string& message) const {
+    CreateOperationResponse opResponse;
+    static auto params = getSignOpParameters();
+
+    auto status = mSecurityLevel->createOperation(mDescriptor, params, false, &opResponse);
+    if (!status.isOk()) {
+        return Error() << "Failed to create keystore signing operation: "
+                       << status.serviceSpecificErrorCode();
+    }
+    auto operation = opResponse.iOperation;
+
+    std::optional<std::vector<uint8_t>> out;
+    status = operation->update({message.begin(), message.end()}, &out);
+    if (!status.isOk()) {
+        return Error() << "Failed to call keystore update operation.";
+    }
+
+    std::optional<std::vector<uint8_t>> signature;
+    status = operation->finish({}, {}, &signature);
+    if (!status.isOk()) {
+        return Error() << "Failed to call keystore finish operation.";
+    }
+
+    if (!signature.has_value()) {
+        return Error() << "Didn't receive a signature from keystore finish operation.";
+    }
+
+    return std::string{signature.value().begin(), signature.value().end()};
+}
+
+Result<void> KeystoreHmacKey::verify(const std::string& message,
+                                     const std::string& signature) const {
+    CreateOperationResponse opResponse;
+    static auto params = getVerifyOpParameters();
+
+    auto status = mSecurityLevel->createOperation(mDescriptor, params, false, &opResponse);
+    if (!status.isOk()) {
+        return Error() << "Failed to create keystore verification operation: "
+                       << status.serviceSpecificErrorCode();
+    }
+    auto operation = opResponse.iOperation;
+
+    std::optional<std::vector<uint8_t>> out;
+    status = operation->update({message.begin(), message.end()}, &out);
+    if (!status.isOk()) {
+        return Error() << "Failed to call keystore update operation.";
+    }
+
+    std::optional<std::vector<uint8_t>> out_signature;
+    std::vector<uint8_t> in_signature{signature.begin(), signature.end()};
+    status = operation->finish({}, in_signature, &out_signature);
+    if (!status.isOk()) {
+        return Error() << "Failed to call keystore finish operation.";
+    }
+
+    return {};
+}
diff --git a/ondevice-signing/KeystoreHmacKey.h b/ondevice-signing/KeystoreHmacKey.h
new file mode 100644
index 0000000..fbad0fd
--- /dev/null
+++ b/ondevice-signing/KeystoreHmacKey.h
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <optional>
+
+#include <android-base/macros.h>
+#include <android-base/result.h>
+
+#include <utils/StrongPointer.h>
+
+#include <android/system/keystore2/IKeystoreService.h>
+
+class KeystoreHmacKey {
+    using IKeystoreService = ::android::system::keystore2::IKeystoreService;
+    using IKeystoreSecurityLevel = ::android::system::keystore2::IKeystoreSecurityLevel;
+    using KeyDescriptor = ::android::system::keystore2::KeyDescriptor;
+
+  public:
+    KeystoreHmacKey();
+    android::base::Result<void> initialize(android::sp<IKeystoreService> service,
+                                           android::sp<IKeystoreSecurityLevel> securityLevel);
+    android::base::Result<std::string> sign(const std::string& message) const;
+    android::base::Result<void> verify(const std::string& message,
+                                       const std::string& signature) const;
+
+  private:
+    android::base::Result<void> createKey();
+    KeyDescriptor mDescriptor;
+    android::sp<IKeystoreService> mService;
+    android::sp<IKeystoreSecurityLevel> mSecurityLevel;
+};
diff --git a/ondevice-signing/KeystoreKey.cpp b/ondevice-signing/KeystoreKey.cpp
index 4e59c58..0951d92 100644
--- a/ondevice-signing/KeystoreKey.cpp
+++ b/ondevice-signing/KeystoreKey.cpp
@@ -46,7 +46,6 @@
 using android::system::keystore2::Domain;
 using android::system::keystore2::KeyDescriptor;
 using android::system::keystore2::KeyEntryResponse;
-using android::system::keystore2::KeyMetadata;
 
 using android::base::Error;
 using android::base::Result;
@@ -54,6 +53,8 @@
 // Keystore boot level that the odsign key uses
 static const int kOdsignBootLevel = 30;
 
+const std::string kPublicKeySignature = "/data/misc/odsign/publickey.signature";
+
 static KeyDescriptor getKeyDescriptor() {
     // AIDL parcelable objects don't have constructor
     static KeyDescriptor descriptor;
@@ -67,9 +68,11 @@
     return descriptor;
 }
 
-KeystoreKey::KeystoreKey() {}
+KeystoreKey::KeystoreKey() {
+    mDescriptor = getKeyDescriptor();
+}
 
-Result<KeyMetadata> KeystoreKey::createNewKey(const KeyDescriptor& descriptor) {
+Result<std::vector<uint8_t>> KeystoreKey::createKey() {
     std::vector<KeyParameter> params;
 
     KeyParameter algo;
@@ -114,12 +117,31 @@
     params.push_back(boot_level);
 
     KeyMetadata metadata;
-    auto status = mSecurityLevel->generateKey(descriptor, {}, params, 0, {}, &metadata);
+    auto status = mSecurityLevel->generateKey(mDescriptor, {}, params, 0, {}, &metadata);
     if (!status.isOk()) {
         return Error() << "Failed to create new key";
     }
 
-    return metadata;
+    // Extract the public key from the certificate, HMAC it and store the signature
+    auto cert = metadata.certificate;
+    if (!cert) {
+        return Error() << "Key did not have a certificate.";
+    }
+    auto publicKey = extractPublicKeyFromX509(cert.value());
+    if (!publicKey.ok()) {
+        return publicKey.error();
+    }
+    std::string publicKeyString = {publicKey->begin(), publicKey->end()};
+    auto signature = mHmacKey.sign(publicKeyString);
+    if (!signature.ok()) {
+        return Error() << "Failed to sign public key.";
+    }
+
+    if (!android::base::WriteStringToFile(*signature, kPublicKeySignature)) {
+        return Error() << "Can't write public key signature.";
+    }
+
+    return *publicKey;
 }
 
 bool KeystoreKey::initialize() {
@@ -136,50 +158,100 @@
         return false;
     }
 
-    auto status = mService->getSecurityLevel(SecurityLevel::STRONGBOX, &mSecurityLevel);
+    auto status = mService->getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT, &mSecurityLevel);
     if (!status.isOk()) {
-        status = mService->getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT, &mSecurityLevel);
-        if (!status.isOk()) {
-            return false;
-        }
+        return false;
     }
 
-    auto descriptor = getKeyDescriptor();
+    // Initialize the HMAC key we use to sign/verify information about this key
+    auto hmacStatus = mHmacKey.initialize(mService, mSecurityLevel);
+    if (!hmacStatus.ok()) {
+        LOG(ERROR) << hmacStatus.error().message();
+        return false;
+    }
+
+    auto key = getOrCreateKey();
+    if (!key.ok()) {
+        LOG(ERROR) << key.error().message();
+        return false;
+    }
+    mPublicKey = *key;
+    LOG(ERROR) << "Initialized Keystore key.";
+    return true;
+}
+
+Result<std::vector<uint8_t>> KeystoreKey::verifyExistingKey() {
     // See if we can fetch an existing key
     KeyEntryResponse keyEntryResponse;
     LOG(INFO) << "Trying to retrieve existing keystore key...";
-    status = mService->getKeyEntry(descriptor, &keyEntryResponse);
-    bool keyValid = false;
+    auto status = mService->getKeyEntry(mDescriptor, &keyEntryResponse);
 
-    if (status.isOk()) {
-        // Make sure this is an early boot key
-        for (const auto& auth : keyEntryResponse.metadata.authorizations) {
-            if (auth.keyParameter.tag == Tag::MAX_BOOT_LEVEL) {
-                if (auth.keyParameter.value.get<KeyParameterValue::integer>() == kOdsignBootLevel) {
-                    keyValid = true;
-                    break;
-                }
+    if (!status.isOk()) {
+        return Error() << "Failed to find keystore key...";
+    }
+
+    // On some earlier builds, we created this key on the Strongbox security level;
+    // we now use TEE keys instead (mostly for speed). It shouldn't matter since
+    // verified boot is protected by the TEE anyway. If the key happens to be on
+    // the wrong security level, delete it (this should happen just once).
+    if (keyEntryResponse.metadata.keySecurityLevel != SecurityLevel::TRUSTED_ENVIRONMENT) {
+        return Error() << "Found invalid keystore key with security level: "
+                       << android::hardware::security::keymint::toString(
+                              keyEntryResponse.metadata.keySecurityLevel);
+    }
+
+    // Make sure this is an early boot key
+    bool foundBootLevel = false;
+    for (const auto& auth : keyEntryResponse.metadata.authorizations) {
+        if (auth.keyParameter.tag == Tag::MAX_BOOT_LEVEL) {
+            if (auth.keyParameter.value.get<KeyParameterValue::integer>() == kOdsignBootLevel) {
+                foundBootLevel = true;
+                break;
             }
         }
-        if (!keyValid) {
-            LOG(WARNING) << "Found invalid keystore key without MAX_BOOT_LEVEL tag";
-        }
+    }
+    if (!foundBootLevel) {
+        return Error() << "Found invalid keystore key without MAX_BOOT_LEVEL tag";
     }
 
-    if (!keyValid) {
+    // If the key is still considered valid at this point, extract the public
+    // key from the certificate. Note that we cannot trust this public key,
+    // because it is a part of the keystore2 database, which can be modified by
+    // an attacker.  So instead, when creating the key we HMAC the public key
+    // with a key of the same boot level, and verify the signature here.
+    auto cert = keyEntryResponse.metadata.certificate;
+    if (!cert) {
+        return Error() << "Key did not have a certificate.";
+    }
+    auto publicKey = extractPublicKeyFromX509(cert.value());
+    if (!publicKey.ok()) {
+        return publicKey.error();
+    }
+    std::string publicKeyString = {publicKey->begin(), publicKey->end()};
+
+    std::string signature;
+    if (!android::base::ReadFileToString(kPublicKeySignature, &signature)) {
+        return Error() << "Can't find signature for public key.";
+    }
+
+    auto signatureValid = mHmacKey.verify(publicKeyString, signature);
+    if (!signatureValid.ok()) {
+        return Error() << "Signature of public key did not match.";
+    }
+    LOG(INFO) << "Verified public key signature.";
+
+    return *publicKey;
+}
+
+Result<std::vector<uint8_t>> KeystoreKey::getOrCreateKey() {
+    auto existingKey = verifyExistingKey();
+    if (!existingKey.ok()) {
+        LOG(INFO) << existingKey.error().message();
         LOG(INFO) << "Existing keystore key not found or invalid, creating new key";
-        auto newKeyStatus = createNewKey(descriptor);
-        if (!newKeyStatus.ok()) {
-            LOG(ERROR) << "Failed to create new key";
-            return false;
-        }
-        mKeyMetadata = *newKeyStatus;
-    } else {
-        mKeyMetadata = keyEntryResponse.metadata;
+        return createKey();
     }
 
-    LOG(ERROR) << "Initialized Keystore key.";
-    return true;
+    return *existingKey;
 }
 
 Result<SigningKey*> KeystoreKey::getInstance() {
@@ -221,11 +293,9 @@
 
 Result<std::string> KeystoreKey::sign(const std::string& message) const {
     static auto opParameters = getSignOpParameters();
-
     CreateOperationResponse opResponse;
 
-    auto status =
-        mSecurityLevel->createOperation(getKeyDescriptor(), opParameters, false, &opResponse);
+    auto status = mSecurityLevel->createOperation(mDescriptor, opParameters, false, &opResponse);
     if (!status.isOk()) {
         return Error() << "Failed to create keystore signing operation: "
                        << status.serviceSpecificErrorCode();
@@ -248,16 +318,9 @@
         return Error() << "Didn't receive a signature from keystore finish operation.";
     }
 
-    std::string result{signature.value().begin(), signature.value().end()};
-
-    return result;
+    return std::string{signature.value().begin(), signature.value().end()};
 }
 
 Result<std::vector<uint8_t>> KeystoreKey::getPublicKey() const {
-    auto cert = mKeyMetadata.certificate;
-    if (cert) {
-        return extractPublicKeyFromX509(cert.value());
-    } else {
-        return Error() << "Key did not have a certificate";
-    }
+    return mPublicKey;
 }
diff --git a/ondevice-signing/KeystoreKey.h b/ondevice-signing/KeystoreKey.h
index 6b9cb57..1257cbb 100644
--- a/ondevice-signing/KeystoreKey.h
+++ b/ondevice-signing/KeystoreKey.h
@@ -26,6 +26,7 @@
 
 #include <android/system/keystore2/IKeystoreService.h>
 
+#include "KeystoreHmacKey.h"
 #include "SigningKey.h"
 
 class KeystoreKey : public SigningKey {
@@ -44,9 +45,13 @@
   private:
     KeystoreKey();
     bool initialize();
-    android::base::Result<KeyMetadata> createNewKey(const KeyDescriptor& descriptor);
+    android::base::Result<std::vector<uint8_t>> verifyExistingKey();
+    android::base::Result<std::vector<uint8_t>> createKey();
+    android::base::Result<std::vector<uint8_t>> getOrCreateKey();
 
+    KeyDescriptor mDescriptor;
+    KeystoreHmacKey mHmacKey;
     android::sp<IKeystoreService> mService;
     android::sp<IKeystoreSecurityLevel> mSecurityLevel;
-    KeyMetadata mKeyMetadata;
+    std::vector<uint8_t> mPublicKey;
 };
diff --git a/ondevice-signing/odsign_main.cpp b/ondevice-signing/odsign_main.cpp
index 6cab8b6..0991704 100644
--- a/ondevice-signing/odsign_main.cpp
+++ b/ondevice-signing/odsign_main.cpp
@@ -32,7 +32,6 @@
 #include <odrefresh/odrefresh.h>
 
 #include "CertUtils.h"
-#include "KeymasterSigningKey.h"
 #include "KeystoreKey.h"
 #include "VerityUtils.h"
 
@@ -57,7 +56,6 @@
 static const char* kFsVerityProcPath = "/proc/sys/fs/verity";
 
 static const bool kForceCompilation = false;
-static const bool kUseKeystore = true;
 
 static const char* kOdsignVerificationDoneProp = "odsign.verification.done";
 static const char* kOdsignKeyDoneProp = "odsign.key.done";
@@ -95,10 +93,8 @@
         return publicKey.error();
     }
 
-    auto keymasterSignFunction = [&](const std::string& to_be_signed) {
-        return key.sign(to_be_signed);
-    };
-    createSelfSignedCertificate(*publicKey, keymasterSignFunction, outPath);
+    auto keySignFunction = [&](const std::string& to_be_signed) { return key.sign(to_be_signed); };
+    createSelfSignedCertificate(*publicKey, keySignFunction, outPath);
     return {};
 }
 
@@ -302,23 +298,12 @@
         return 0;
     }
 
-    SigningKey* key;
-    if (kUseKeystore) {
-        auto keystoreResult = KeystoreKey::getInstance();
-        if (!keystoreResult.ok()) {
-            LOG(ERROR) << "Could not create keystore key: " << keystoreResult.error().message();
-            return -1;
-        }
-        key = keystoreResult.value();
-    } else {
-        // TODO - keymaster will go away
-        auto keymasterResult = KeymasterSigningKey::getInstance();
-        if (!keymasterResult.ok()) {
-            LOG(ERROR) << "Failed to create keymaster key: " << keymasterResult.error().message();
-            return -1;
-        }
-        key = keymasterResult.value();
+    auto keystoreResult = KeystoreKey::getInstance();
+    if (!keystoreResult.ok()) {
+        LOG(ERROR) << "Could not create keystore key: " << keystoreResult.error().message();
+        return -1;
     }
+    SigningKey* key = keystoreResult.value();
 
     bool supportsFsVerity = access(kFsVerityProcPath, F_OK) == 0;
     if (!supportsFsVerity) {