Keystore 2.0: Add watch dog instrumentations in the calls into keystore.

Bug: 184006658
Test: Reduced the duration to 1ms and observed logs from watchdog.
Change-Id: I711b93881b164dfa7a4947ffe91886d4c8823d76
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/authorization.rs b/keystore2/src/authorization.rs
index cac75c0..d07dab5 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,
 };
@@ -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/maintenance.rs b/keystore2/src/maintenance.rs
index 13a5b43..a099d18 100644
--- a/keystore2/src/maintenance.rs
+++ b/keystore2/src/maintenance.rs
@@ -225,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)
     }
 
@@ -257,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 3ef3c1c..8d7ad0a 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -842,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"),
@@ -852,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"),
@@ -865,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"),
@@ -875,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/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index f8ee369..fc1a6ad 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -45,7 +45,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.
@@ -392,6 +392,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 +406,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 +429,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 +441,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 84d1963..d10aba0 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -951,6 +951,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(
@@ -961,6 +962,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());
@@ -974,6 +978,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());
@@ -987,6 +992,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);
@@ -997,9 +1003,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..3ce0550 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,
@@ -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/vpnprofilestore/lib.rs b/keystore2/vpnprofilestore/lib.rs
index d92e045..8b3bc2b 100644
--- a/keystore2/vpnprofilestore/lib.rs
+++ b/keystore2/vpnprofilestore/lib.rs
@@ -23,7 +23,7 @@
     ThreadState,
 };
 use anyhow::{Context, Result};
-use keystore2::{async_task::AsyncTask, legacy_blob::LegacyBlobLoader};
+use keystore2::{async_task::AsyncTask, legacy_blob::LegacyBlobLoader, utils::watchdog as wd};
 use rusqlite::{
     params, Connection, OptionalExtension, Transaction, TransactionBehavior, NO_PARAMS,
 };
@@ -366,15 +366,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)
     }
 }