Add keystore API for metrics re-routing.

Keystore2 atoms need to be routed to statsd via a proxy.
The proxy needs to call this API in order to pull metrics from
keystore.

Bug: 188590587
Test: Statsd Testdrive script
Merged-In: I28f8675fe5467b0760418c4d2d87808e45657be1

Change-Id: I28f8675fe5467b0760418c4d2d87808e45657be1
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 36c722a..fa54623 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -69,8 +69,9 @@
 use android_security_remoteprovisioning::aidl::android::security::remoteprovisioning::{
     AttestationPoolStatus::AttestationPoolStatus,
 };
-use statslog_rust::keystore2_storage_stats::{
-	Keystore2StorageStats, StorageType as StatsdStorageType,
+use android_security_metrics::aidl::android::security::metrics::{
+    StorageStats::StorageStats,
+    Storage::Storage as MetricsStorage,
 };
 
 use keystore2_crypto::ZVec;
@@ -1021,23 +1022,23 @@
 
     fn do_table_size_query(
         &mut self,
-        storage_type: StatsdStorageType,
+        storage_type: MetricsStorage,
         query: &str,
         params: &[&str],
-    ) -> Result<Keystore2StorageStats> {
+    ) -> Result<StorageStats> {
         let (total, unused) = self.with_transaction(TransactionBehavior::Deferred, |tx| {
             tx.query_row(query, params_from_iter(params), |row| Ok((row.get(0)?, row.get(1)?)))
                 .with_context(|| {
-                    format!("get_storage_stat: Error size of storage type {}", storage_type as i32)
+                    format!("get_storage_stat: Error size of storage type {}", storage_type.0)
                 })
                 .no_gc()
         })?;
-        Ok(Keystore2StorageStats { storage_type, size: total, unused_size: unused })
+        Ok(StorageStats { storage_type, size: total, unused_size: unused })
     }
 
-    fn get_total_size(&mut self) -> Result<Keystore2StorageStats> {
+    fn get_total_size(&mut self) -> Result<StorageStats> {
         self.do_table_size_query(
-            StatsdStorageType::Database,
+            MetricsStorage::DATABASE,
             "SELECT page_count * page_size, freelist_count * page_size
              FROM pragma_page_count('persistent'),
                   pragma_page_size('persistent'),
@@ -1048,10 +1049,10 @@
 
     fn get_table_size(
         &mut self,
-        storage_type: StatsdStorageType,
+        storage_type: MetricsStorage,
         schema: &str,
         table: &str,
-    ) -> Result<Keystore2StorageStats> {
+    ) -> Result<StorageStats> {
         self.do_table_size_query(
             storage_type,
             "SELECT pgsize,unused FROM dbstat(?1)
@@ -1063,63 +1064,57 @@
     /// Fetches a storage statisitics atom for a given storage type. For storage
     /// types that map to a table, information about the table's storage is
     /// returned. Requests for storage types that are not DB tables return None.
-    pub fn get_storage_stat(
-        &mut self,
-        storage_type: StatsdStorageType,
-    ) -> Result<Keystore2StorageStats> {
+    pub fn get_storage_stat(&mut self, storage_type: MetricsStorage) -> Result<StorageStats> {
         let _wp = wd::watch_millis("KeystoreDB::get_storage_stat", 500);
 
         match storage_type {
-            StatsdStorageType::Database => self.get_total_size(),
-            StatsdStorageType::KeyEntry => {
+            MetricsStorage::DATABASE => self.get_total_size(),
+            MetricsStorage::KEY_ENTRY => {
                 self.get_table_size(storage_type, "persistent", "keyentry")
             }
-            StatsdStorageType::KeyEntryIdIndex => {
+            MetricsStorage::KEY_ENTRY_ID_INDEX => {
                 self.get_table_size(storage_type, "persistent", "keyentry_id_index")
             }
-            StatsdStorageType::KeyEntryDomainNamespaceIndex => {
+            MetricsStorage::KEY_ENTRY_DOMAIN_NAMESPACE_INDEX => {
                 self.get_table_size(storage_type, "persistent", "keyentry_domain_namespace_index")
             }
-            StatsdStorageType::BlobEntry => {
+            MetricsStorage::BLOB_ENTRY => {
                 self.get_table_size(storage_type, "persistent", "blobentry")
             }
-            StatsdStorageType::BlobEntryKeyEntryIdIndex => {
+            MetricsStorage::BLOB_ENTRY_KEY_ENTRY_ID_INDEX => {
                 self.get_table_size(storage_type, "persistent", "blobentry_keyentryid_index")
             }
-            StatsdStorageType::KeyParameter => {
+            MetricsStorage::KEY_PARAMETER => {
                 self.get_table_size(storage_type, "persistent", "keyparameter")
             }
-            StatsdStorageType::KeyParameterKeyEntryIdIndex => {
+            MetricsStorage::KEY_PARAMETER_KEY_ENTRY_ID_INDEX => {
                 self.get_table_size(storage_type, "persistent", "keyparameter_keyentryid_index")
             }
-            StatsdStorageType::KeyMetadata => {
+            MetricsStorage::KEY_METADATA => {
                 self.get_table_size(storage_type, "persistent", "keymetadata")
             }
-            StatsdStorageType::KeyMetadataKeyEntryIdIndex => {
+            MetricsStorage::KEY_METADATA_KEY_ENTRY_ID_INDEX => {
                 self.get_table_size(storage_type, "persistent", "keymetadata_keyentryid_index")
             }
-            StatsdStorageType::Grant => self.get_table_size(storage_type, "persistent", "grant"),
-            StatsdStorageType::AuthToken => {
+            MetricsStorage::GRANT => self.get_table_size(storage_type, "persistent", "grant"),
+            MetricsStorage::AUTH_TOKEN => {
                 // Since the table is actually a BTreeMap now, unused_size is not meaningfully
                 // reportable
                 // Size provided is only an approximation
-                Ok(Keystore2StorageStats {
+                Ok(StorageStats {
                     storage_type,
                     size: (self.perboot.auth_tokens_len() * std::mem::size_of::<AuthTokenEntry>())
-                        as i64,
+                        as i32,
                     unused_size: 0,
                 })
             }
-            StatsdStorageType::BlobMetadata => {
+            MetricsStorage::BLOB_METADATA => {
                 self.get_table_size(storage_type, "persistent", "blobmetadata")
             }
-            StatsdStorageType::BlobMetadataBlobEntryIdIndex => {
+            MetricsStorage::BLOB_METADATA_BLOB_ENTRY_ID_INDEX => {
                 self.get_table_size(storage_type, "persistent", "blobmetadata_blobentryid_index")
             }
-            _ => Err(anyhow::Error::msg(format!(
-                "Unsupported storage type: {}",
-                storage_type as i32
-            ))),
+            _ => Err(anyhow::Error::msg(format!("Unsupported storage type: {}", storage_type.0))),
         }
     }
 
@@ -5497,21 +5492,21 @@
         Ok(())
     }
 
-    fn get_valid_statsd_storage_types() -> Vec<StatsdStorageType> {
+    fn get_valid_statsd_storage_types() -> Vec<MetricsStorage> {
         vec![
-            StatsdStorageType::KeyEntry,
-            StatsdStorageType::KeyEntryIdIndex,
-            StatsdStorageType::KeyEntryDomainNamespaceIndex,
-            StatsdStorageType::BlobEntry,
-            StatsdStorageType::BlobEntryKeyEntryIdIndex,
-            StatsdStorageType::KeyParameter,
-            StatsdStorageType::KeyParameterKeyEntryIdIndex,
-            StatsdStorageType::KeyMetadata,
-            StatsdStorageType::KeyMetadataKeyEntryIdIndex,
-            StatsdStorageType::Grant,
-            StatsdStorageType::AuthToken,
-            StatsdStorageType::BlobMetadata,
-            StatsdStorageType::BlobMetadataBlobEntryIdIndex,
+            MetricsStorage::KEY_ENTRY,
+            MetricsStorage::KEY_ENTRY_ID_INDEX,
+            MetricsStorage::KEY_ENTRY_DOMAIN_NAMESPACE_INDEX,
+            MetricsStorage::BLOB_ENTRY,
+            MetricsStorage::BLOB_ENTRY_KEY_ENTRY_ID_INDEX,
+            MetricsStorage::KEY_PARAMETER,
+            MetricsStorage::KEY_PARAMETER_KEY_ENTRY_ID_INDEX,
+            MetricsStorage::KEY_METADATA,
+            MetricsStorage::KEY_METADATA_KEY_ENTRY_ID_INDEX,
+            MetricsStorage::GRANT,
+            MetricsStorage::AUTH_TOKEN,
+            MetricsStorage::BLOB_METADATA,
+            MetricsStorage::BLOB_METADATA_BLOB_ENTRY_ID_INDEX,
         ]
     }
 
@@ -5519,7 +5514,7 @@
     /// that are supported by the DB. Check for reasonable values.
     #[test]
     fn test_query_all_valid_table_sizes() -> Result<()> {
-        const PAGE_SIZE: i64 = 4096;
+        const PAGE_SIZE: i32 = 4096;
 
         let mut db = new_test_db()?;
 
@@ -5527,7 +5522,7 @@
             let stat = db.get_storage_stat(t)?;
             // 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 {
+            if let MetricsStorage::AUTH_TOKEN = t {
             } else {
                 assert!(stat.size >= PAGE_SIZE);
             }
@@ -5537,35 +5532,35 @@
         Ok(())
     }
 
-    fn get_storage_stats_map(db: &mut KeystoreDB) -> BTreeMap<i32, Keystore2StorageStats> {
+    fn get_storage_stats_map(db: &mut KeystoreDB) -> BTreeMap<i32, StorageStats> {
         get_valid_statsd_storage_types()
             .into_iter()
-            .map(|t| (t as i32, db.get_storage_stat(t).unwrap()))
+            .map(|t| (t.0, db.get_storage_stat(t).unwrap()))
             .collect()
     }
 
     fn assert_storage_increased(
         db: &mut KeystoreDB,
-        increased_storage_types: Vec<StatsdStorageType>,
-        baseline: &mut BTreeMap<i32, Keystore2StorageStats>,
+        increased_storage_types: Vec<MetricsStorage>,
+        baseline: &mut BTreeMap<i32, StorageStats>,
     ) {
         for storage in increased_storage_types {
             // Verify the expected storage increased.
             let new = db.get_storage_stat(storage).unwrap();
-            let storage = storage as i32;
-            let old = &baseline[&storage];
-            assert!(new.size >= old.size, "{}: {} >= {}", storage, new.size, old.size);
+            let storage = storage;
+            let old = &baseline[&storage.0];
+            assert!(new.size >= old.size, "{}: {} >= {}", storage.0, new.size, old.size);
             assert!(
                 new.unused_size <= old.unused_size,
                 "{}: {} <= {}",
-                storage,
+                storage.0,
                 new.unused_size,
                 old.unused_size
             );
 
             // Update the baseline with the new value so that it succeeds in the
             // later comparison.
-            baseline.insert(storage, new);
+            baseline.insert(storage.0, new);
         }
 
         // Get an updated map of the storage and verify there were no unexpected changes.
@@ -5573,7 +5568,7 @@
         assert_eq!(updated_stats.len(), baseline.len());
 
         for &k in baseline.keys() {
-            let stringify = |map: &BTreeMap<i32, Keystore2StorageStats>| -> String {
+            let stringify = |map: &BTreeMap<i32, StorageStats>| -> String {
                 let mut s = String::new();
                 for &k in map.keys() {
                     writeln!(&mut s, "  {}: {}, {}", &k, map[&k].size, map[&k].unused_size)
@@ -5601,9 +5596,9 @@
         assert_storage_increased(
             &mut db,
             vec![
-                StatsdStorageType::KeyEntry,
-                StatsdStorageType::KeyEntryIdIndex,
-                StatsdStorageType::KeyEntryDomainNamespaceIndex,
+                MetricsStorage::KEY_ENTRY,
+                MetricsStorage::KEY_ENTRY_ID_INDEX,
+                MetricsStorage::KEY_ENTRY_DOMAIN_NAMESPACE_INDEX,
             ],
             &mut working_stats,
         );
@@ -5614,10 +5609,10 @@
         assert_storage_increased(
             &mut db,
             vec![
-                StatsdStorageType::BlobEntry,
-                StatsdStorageType::BlobEntryKeyEntryIdIndex,
-                StatsdStorageType::BlobMetadata,
-                StatsdStorageType::BlobMetadataBlobEntryIdIndex,
+                MetricsStorage::BLOB_ENTRY,
+                MetricsStorage::BLOB_ENTRY_KEY_ENTRY_ID_INDEX,
+                MetricsStorage::BLOB_METADATA,
+                MetricsStorage::BLOB_METADATA_BLOB_ENTRY_ID_INDEX,
             ],
             &mut working_stats,
         );
@@ -5626,7 +5621,7 @@
         db.insert_keyparameter(&key_id, &params)?;
         assert_storage_increased(
             &mut db,
-            vec![StatsdStorageType::KeyParameter, StatsdStorageType::KeyParameterKeyEntryIdIndex],
+            vec![MetricsStorage::KEY_PARAMETER, MetricsStorage::KEY_PARAMETER_KEY_ENTRY_ID_INDEX],
             &mut working_stats,
         );
 
@@ -5635,7 +5630,7 @@
         db.insert_key_metadata(&key_id, &metadata)?;
         assert_storage_increased(
             &mut db,
-            vec![StatsdStorageType::KeyMetadata, StatsdStorageType::KeyMetadataKeyEntryIdIndex],
+            vec![MetricsStorage::KEY_METADATA, MetricsStorage::KEY_METADATA_KEY_ENTRY_ID_INDEX],
             &mut working_stats,
         );
 
@@ -5643,7 +5638,7 @@
         for stat in working_stats.values() {
             sum += stat.size;
         }
-        let total = db.get_storage_stat(StatsdStorageType::Database)?.size;
+        let total = db.get_storage_stat(MetricsStorage::DATABASE)?.size;
         assert!(sum <= total, "Expected sum <= total. sum: {}, total: {}", sum, total);
 
         Ok(())
@@ -5661,7 +5656,7 @@
             timestamp: Timestamp { milliSeconds: 10 },
             mac: b"mac".to_vec(),
         });
-        assert_storage_increased(&mut db, vec![StatsdStorageType::AuthToken], &mut working_stats);
+        assert_storage_increased(&mut db, vec![MetricsStorage::AUTH_TOKEN], &mut working_stats);
         Ok(())
     }
 
@@ -5685,7 +5680,7 @@
             |_, _| Ok(()),
         )?;
 
-        assert_storage_increased(&mut db, vec![StatsdStorageType::Grant], &mut working_stats);
+        assert_storage_increased(&mut db, vec![MetricsStorage::GRANT], &mut working_stats);
 
         Ok(())
     }
diff --git a/keystore2/src/keystore2_main.rs b/keystore2/src/keystore2_main.rs
index a6a25b2..2c7d4a0 100644
--- a/keystore2/src/keystore2_main.rs
+++ b/keystore2/src/keystore2_main.rs
@@ -17,7 +17,7 @@
 use keystore2::entropy;
 use keystore2::globals::ENFORCEMENTS;
 use keystore2::maintenance::Maintenance;
-use keystore2::metrics;
+use keystore2::metrics::Metrics;
 use keystore2::remote_provisioning::RemoteProvisioningService;
 use keystore2::service::KeystoreService;
 use keystore2::{apc::ApcManager, shared_secret_negotiation};
@@ -29,6 +29,7 @@
 static KS2_SERVICE_NAME: &str = "android.system.keystore2.IKeystoreService/default";
 static APC_SERVICE_NAME: &str = "android.security.apc";
 static AUTHORIZATION_SERVICE_NAME: &str = "android.security.authorization";
+static METRICS_SERVICE_NAME: &str = "android.security.metrics";
 static REMOTE_PROVISIONING_SERVICE_NAME: &str = "android.security.remoteprovisioning";
 static USER_MANAGER_SERVICE_NAME: &str = "android.security.maintenance";
 static LEGACY_KEYSTORE_SERVICE_NAME: &str = "android.security.legacykeystore";
@@ -109,6 +110,13 @@
         },
     );
 
+    let metrics_service = Metrics::new_native_binder().unwrap_or_else(|e| {
+        panic!("Failed to create service {} because of {:?}.", METRICS_SERVICE_NAME, e);
+    });
+    binder::add_service(METRICS_SERVICE_NAME, metrics_service.as_binder()).unwrap_or_else(|e| {
+        panic!("Failed to register service {} because of {:?}.", METRICS_SERVICE_NAME, e);
+    });
+
     // Devices with KS2 and KM 1.0 may not have any IRemotelyProvisionedComponent HALs at all. Do
     // not panic if new_native_binder returns failure because it could not find the TEE HAL.
     if let Ok(remote_provisioning_service) = RemoteProvisioningService::new_native_binder() {
@@ -133,8 +141,6 @@
         },
     );
 
-    metrics::register_pull_metrics_callbacks();
-
     info!("Successfully registered Keystore 2.0 service.");
 
     info!("Joining thread pool now.");
diff --git a/keystore2/src/lib.rs b/keystore2/src/lib.rs
index 51316d7..8b629b1 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -32,6 +32,7 @@
 pub mod legacy_migrator;
 pub mod maintenance;
 pub mod metrics;
+pub mod metrics_store;
 pub mod operation;
 pub mod permission;
 pub mod raw_device;
diff --git a/keystore2/src/metrics.rs b/keystore2/src/metrics.rs
index 10a465c..42295b7 100644
--- a/keystore2/src/metrics.rs
+++ b/keystore2/src/metrics.rs
@@ -12,533 +12,45 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! This module provides convenience functions for keystore2 logging.
-use crate::error::get_error_code;
-use crate::globals::{DB, LOGS_HANDLER};
-use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
-use crate::operation::Outcome;
-use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    Algorithm::Algorithm, BlockMode::BlockMode, Digest::Digest, EcCurve::EcCurve,
-    HardwareAuthenticatorType::HardwareAuthenticatorType, KeyOrigin::KeyOrigin,
-    KeyParameter::KeyParameter, KeyPurpose::KeyPurpose, PaddingMode::PaddingMode,
-    SecurityLevel::SecurityLevel,
+//! This module implements the IKeystoreMetrics AIDL interface, which exposes the API method for the
+//! proxy in the system server to pull the aggregated metrics in keystore.
+use crate::error::map_or_log_err;
+use crate::metrics_store::METRICS_STORE;
+use crate::permission::KeystorePerm;
+use crate::utils::{check_keystore_permission, watchdog as wd};
+use android_security_metrics::aidl::android::security::metrics::{
+    AtomID::AtomID,
+    IKeystoreMetrics::{BnKeystoreMetrics, IKeystoreMetrics},
+    KeystoreAtom::KeystoreAtom,
 };
-use anyhow::{anyhow, Result};
-use keystore2_system_property::PropertyWatcher;
-use statslog_rust::{
-    keystore2_key_creation_event_reported::{
-        Algorithm as StatsdAlgorithm, EcCurve as StatsdEcCurve, KeyOrigin as StatsdKeyOrigin,
-        Keystore2KeyCreationEventReported, SecurityLevel as StatsdKeyCreationSecurityLevel,
-        UserAuthType as StatsdUserAuthType,
-    },
-    keystore2_key_operation_event_reported::{
-        Keystore2KeyOperationEventReported, Outcome as StatsdOutcome, Purpose as StatsdKeyPurpose,
-        SecurityLevel as StatsdKeyOperationSecurityLevel,
-    },
-    keystore2_storage_stats::StorageType as StatsdStorageType,
-};
-use statslog_rust_header::Atoms;
-use statspull_rust::{set_pull_atom_callback, StatsPullResult};
+use android_security_metrics::binder::{BinderFeatures, Interface, Result as BinderResult, Strong};
+use anyhow::{Context, Result};
 
-// Waits and returns Ok if boot is completed.
-fn wait_for_boot_completed() -> Result<()> {
-    let watcher = PropertyWatcher::new("sys.boot_completed");
-    match watcher {
-        Ok(mut watcher) => {
-            loop {
-                let wait_result = watcher.wait();
-                match wait_result {
-                    Ok(_) => {
-                        let value_result =
-                            watcher.read(|_name, value| Ok(value.trim().to_string()));
-                        match value_result {
-                            Ok(value) => {
-                                if value == "1" {
-                                    break;
-                                }
-                            }
-                            Err(e) => {
-                                log::error!(
-                                    "In wait_for_boot_completed: Failed while reading property. {}",
-                                    e
-                                );
-                                return Err(anyhow!("Error in waiting for boot completion."));
-                            }
-                        }
-                    }
-                    Err(e) => {
-                        log::error!("In wait_for_boot_completed: Failed while waiting. {}", e);
-                        return Err(anyhow!("Error in waiting for boot completion."));
-                    }
-                }
-            }
-            Ok(())
-        }
-        Err(e) => {
-            log::error!("In wait_for_boot_completed: Failed to create PropertyWatcher. {}", e);
-            Err(anyhow!("Error in waiting for boot completion."))
-        }
+/// This struct is defined to implement IKeystoreMetrics AIDL interface.
+pub struct Metrics;
+
+impl Metrics {
+    /// Create a new instance of Keystore Metrics service.
+    pub fn new_native_binder() -> Result<Strong<dyn IKeystoreMetrics>> {
+        Ok(BnKeystoreMetrics::new_binder(
+            Self,
+            BinderFeatures { set_requesting_sid: true, ..BinderFeatures::default() },
+        ))
+    }
+
+    fn pull_metrics(&self, atom_id: AtomID) -> Result<Vec<KeystoreAtom>> {
+        // Check permission. Function should return if this failed. Therefore having '?' at the end
+        // is very important.
+        check_keystore_permission(KeystorePerm::pull_metrics()).context("In pull_metrics.")?;
+        METRICS_STORE.get_atoms(atom_id)
     }
 }
 
-fn create_default_key_creation_atom() -> Keystore2KeyCreationEventReported {
-    // If a value is not present, fields represented by bitmaps and i32 fields
-    // will take 0, except error_code which defaults to 1 indicating NO_ERROR and key_size,
-    // and auth_time_out which default to -1.
-    // The boolean fields are set to false by default.
-    // Some keymint enums do have 0 as an enum variant value. In such cases, the corresponding
-    // enum variant value in atoms.proto is incremented by 1, in order to have 0 as the reserved
-    // value for unspecified fields.
-    Keystore2KeyCreationEventReported {
-        algorithm: StatsdAlgorithm::AlgorithmUnspecified,
-        key_size: -1,
-        key_origin: StatsdKeyOrigin::OriginUnspecified,
-        user_auth_type: StatsdUserAuthType::AuthTypeUnspecified,
-        user_auth_key_timeout_seconds: -1,
-        padding_mode_bitmap: 0,
-        digest_bitmap: 0,
-        block_mode_bitmap: 0,
-        purpose_bitmap: 0,
-        ec_curve: StatsdEcCurve::EcCurveUnspecified,
-        // as per keystore2/ResponseCode.aidl, 1 is reserved for NO_ERROR
-        error_code: 1,
-        attestation_requested: false,
-        security_level: StatsdKeyCreationSecurityLevel::SecurityLevelUnspecified,
+impl Interface for Metrics {}
+
+impl IKeystoreMetrics for Metrics {
+    fn pullMetrics(&self, atom_id: AtomID) -> BinderResult<Vec<KeystoreAtom>> {
+        let _wp = wd::watch_millis("IKeystoreMetrics::pullMetrics", 500);
+        map_or_log_err(self.pull_metrics(atom_id), Ok)
     }
 }
-
-fn create_default_key_operation_atom() -> Keystore2KeyOperationEventReported {
-    Keystore2KeyOperationEventReported {
-        purpose: StatsdKeyPurpose::KeyPurposeUnspecified,
-        padding_mode_bitmap: 0,
-        digest_bitmap: 0,
-        block_mode_bitmap: 0,
-        outcome: StatsdOutcome::OutcomeUnspecified,
-        error_code: 1,
-        key_upgraded: false,
-        security_level: StatsdKeyOperationSecurityLevel::SecurityLevelUnspecified,
-    }
-}
-
-/// Log key creation events via statsd API.
-pub fn log_key_creation_event_stats<U>(
-    sec_level: SecurityLevel,
-    key_params: &[KeyParameter],
-    result: &Result<U>,
-) {
-    let key_creation_event_stats =
-        construct_key_creation_event_stats(sec_level, key_params, result);
-
-    LOGS_HANDLER.queue_lo(move |_| {
-        if let Ok(()) = wait_for_boot_completed() {
-            if let Err(e) = key_creation_event_stats.stats_write() {
-                log::error!("Error in logging key creation event in the async task. {:?}", e);
-            }
-        }
-    });
-}
-
-/// Log key operation events via statsd API.
-pub fn log_key_operation_event_stats(
-    sec_level: SecurityLevel,
-    key_purpose: KeyPurpose,
-    op_params: &[KeyParameter],
-    op_outcome: &Outcome,
-    key_upgraded: bool,
-) {
-    let key_operation_event_stats = construct_key_operation_event_stats(
-        sec_level,
-        key_purpose,
-        op_params,
-        op_outcome,
-        key_upgraded,
-    );
-
-    LOGS_HANDLER.queue_lo(move |_| {
-        if let Ok(()) = wait_for_boot_completed() {
-            if let Err(e) = key_operation_event_stats.stats_write() {
-                log::error!("Error in logging key operation event in the async task. {:?}", e);
-            }
-        }
-    });
-}
-
-fn construct_key_creation_event_stats<U>(
-    sec_level: SecurityLevel,
-    key_params: &[KeyParameter],
-    result: &Result<U>,
-) -> Keystore2KeyCreationEventReported {
-    let mut key_creation_event_atom = create_default_key_creation_atom();
-
-    if let Err(ref e) = result {
-        key_creation_event_atom.error_code = get_error_code(e);
-    }
-
-    key_creation_event_atom.security_level = match sec_level {
-        SecurityLevel::SOFTWARE => StatsdKeyCreationSecurityLevel::SecurityLevelSoftware,
-        SecurityLevel::TRUSTED_ENVIRONMENT => {
-            StatsdKeyCreationSecurityLevel::SecurityLevelTrustedEnvironment
-        }
-        SecurityLevel::STRONGBOX => StatsdKeyCreationSecurityLevel::SecurityLevelStrongbox,
-        //KEYSTORE is not a valid variant here
-        _ => StatsdKeyCreationSecurityLevel::SecurityLevelUnspecified,
-    };
-
-    for key_param in key_params.iter().map(KsKeyParamValue::from) {
-        match key_param {
-            KsKeyParamValue::Algorithm(a) => {
-                key_creation_event_atom.algorithm = match a {
-                    Algorithm::RSA => StatsdAlgorithm::Rsa,
-                    Algorithm::EC => StatsdAlgorithm::Ec,
-                    Algorithm::AES => StatsdAlgorithm::Aes,
-                    Algorithm::TRIPLE_DES => StatsdAlgorithm::TripleDes,
-                    Algorithm::HMAC => StatsdAlgorithm::Hmac,
-                    _ => StatsdAlgorithm::AlgorithmUnspecified,
-                }
-            }
-            KsKeyParamValue::KeySize(s) => {
-                key_creation_event_atom.key_size = s;
-            }
-            KsKeyParamValue::KeyOrigin(o) => {
-                key_creation_event_atom.key_origin = match o {
-                    KeyOrigin::GENERATED => StatsdKeyOrigin::Generated,
-                    KeyOrigin::DERIVED => StatsdKeyOrigin::Derived,
-                    KeyOrigin::IMPORTED => StatsdKeyOrigin::Imported,
-                    KeyOrigin::RESERVED => StatsdKeyOrigin::Reserved,
-                    KeyOrigin::SECURELY_IMPORTED => StatsdKeyOrigin::SecurelyImported,
-                    _ => StatsdKeyOrigin::OriginUnspecified,
-                }
-            }
-            KsKeyParamValue::HardwareAuthenticatorType(a) => {
-                key_creation_event_atom.user_auth_type = match a {
-                    HardwareAuthenticatorType::NONE => StatsdUserAuthType::None,
-                    HardwareAuthenticatorType::PASSWORD => StatsdUserAuthType::Password,
-                    HardwareAuthenticatorType::FINGERPRINT => StatsdUserAuthType::Fingerprint,
-                    HardwareAuthenticatorType::ANY => StatsdUserAuthType::Any,
-                    _ => StatsdUserAuthType::AuthTypeUnspecified,
-                }
-            }
-            KsKeyParamValue::AuthTimeout(t) => {
-                key_creation_event_atom.user_auth_key_timeout_seconds = t;
-            }
-            KsKeyParamValue::PaddingMode(p) => {
-                key_creation_event_atom.padding_mode_bitmap =
-                    compute_padding_mode_bitmap(&key_creation_event_atom.padding_mode_bitmap, p);
-            }
-            KsKeyParamValue::Digest(d) => {
-                key_creation_event_atom.digest_bitmap =
-                    compute_digest_bitmap(&key_creation_event_atom.digest_bitmap, d);
-            }
-            KsKeyParamValue::BlockMode(b) => {
-                key_creation_event_atom.block_mode_bitmap =
-                    compute_block_mode_bitmap(&key_creation_event_atom.block_mode_bitmap, b);
-            }
-            KsKeyParamValue::KeyPurpose(k) => {
-                key_creation_event_atom.purpose_bitmap =
-                    compute_purpose_bitmap(&key_creation_event_atom.purpose_bitmap, k);
-            }
-            KsKeyParamValue::EcCurve(e) => {
-                key_creation_event_atom.ec_curve = match e {
-                    EcCurve::P_224 => StatsdEcCurve::P224,
-                    EcCurve::P_256 => StatsdEcCurve::P256,
-                    EcCurve::P_384 => StatsdEcCurve::P384,
-                    EcCurve::P_521 => StatsdEcCurve::P521,
-                    _ => StatsdEcCurve::EcCurveUnspecified,
-                }
-            }
-            KsKeyParamValue::AttestationChallenge(_) => {
-                key_creation_event_atom.attestation_requested = true;
-            }
-            _ => {}
-        }
-    }
-    key_creation_event_atom
-}
-
-fn construct_key_operation_event_stats(
-    sec_level: SecurityLevel,
-    key_purpose: KeyPurpose,
-    op_params: &[KeyParameter],
-    op_outcome: &Outcome,
-    key_upgraded: bool,
-) -> Keystore2KeyOperationEventReported {
-    let mut key_operation_event_atom = create_default_key_operation_atom();
-
-    key_operation_event_atom.security_level = match sec_level {
-        SecurityLevel::SOFTWARE => StatsdKeyOperationSecurityLevel::SecurityLevelSoftware,
-        SecurityLevel::TRUSTED_ENVIRONMENT => {
-            StatsdKeyOperationSecurityLevel::SecurityLevelTrustedEnvironment
-        }
-        SecurityLevel::STRONGBOX => StatsdKeyOperationSecurityLevel::SecurityLevelStrongbox,
-        //KEYSTORE is not a valid variant here
-        _ => StatsdKeyOperationSecurityLevel::SecurityLevelUnspecified,
-    };
-
-    key_operation_event_atom.key_upgraded = key_upgraded;
-
-    key_operation_event_atom.purpose = match key_purpose {
-        KeyPurpose::ENCRYPT => StatsdKeyPurpose::Encrypt,
-        KeyPurpose::DECRYPT => StatsdKeyPurpose::Decrypt,
-        KeyPurpose::SIGN => StatsdKeyPurpose::Sign,
-        KeyPurpose::VERIFY => StatsdKeyPurpose::Verify,
-        KeyPurpose::WRAP_KEY => StatsdKeyPurpose::WrapKey,
-        KeyPurpose::AGREE_KEY => StatsdKeyPurpose::AgreeKey,
-        KeyPurpose::ATTEST_KEY => StatsdKeyPurpose::AttestKey,
-        _ => StatsdKeyPurpose::KeyPurposeUnspecified,
-    };
-
-    key_operation_event_atom.outcome = match op_outcome {
-        Outcome::Unknown | Outcome::Dropped => StatsdOutcome::Dropped,
-        Outcome::Success => StatsdOutcome::Success,
-        Outcome::Abort => StatsdOutcome::Abort,
-        Outcome::Pruned => StatsdOutcome::Pruned,
-        Outcome::ErrorCode(e) => {
-            key_operation_event_atom.error_code = e.0;
-            StatsdOutcome::Error
-        }
-    };
-
-    for key_param in op_params.iter().map(KsKeyParamValue::from) {
-        match key_param {
-            KsKeyParamValue::PaddingMode(p) => {
-                key_operation_event_atom.padding_mode_bitmap =
-                    compute_padding_mode_bitmap(&key_operation_event_atom.padding_mode_bitmap, p);
-            }
-            KsKeyParamValue::Digest(d) => {
-                key_operation_event_atom.digest_bitmap =
-                    compute_digest_bitmap(&key_operation_event_atom.digest_bitmap, d);
-            }
-            KsKeyParamValue::BlockMode(b) => {
-                key_operation_event_atom.block_mode_bitmap =
-                    compute_block_mode_bitmap(&key_operation_event_atom.block_mode_bitmap, b);
-            }
-            _ => {}
-        }
-    }
-
-    key_operation_event_atom
-}
-
-fn compute_purpose_bitmap(purpose_bitmap: &i32, purpose: KeyPurpose) -> i32 {
-    let mut bitmap = *purpose_bitmap;
-    match purpose {
-        KeyPurpose::ENCRYPT => {
-            bitmap |= 1 << KeyPurposeBitPosition::ENCRYPT_BIT_POS as i32;
-        }
-        KeyPurpose::DECRYPT => {
-            bitmap |= 1 << KeyPurposeBitPosition::DECRYPT_BIT_POS as i32;
-        }
-        KeyPurpose::SIGN => {
-            bitmap |= 1 << KeyPurposeBitPosition::SIGN_BIT_POS as i32;
-        }
-        KeyPurpose::VERIFY => {
-            bitmap |= 1 << KeyPurposeBitPosition::VERIFY_BIT_POS as i32;
-        }
-        KeyPurpose::WRAP_KEY => {
-            bitmap |= 1 << KeyPurposeBitPosition::WRAP_KEY_BIT_POS as i32;
-        }
-        KeyPurpose::AGREE_KEY => {
-            bitmap |= 1 << KeyPurposeBitPosition::AGREE_KEY_BIT_POS as i32;
-        }
-        KeyPurpose::ATTEST_KEY => {
-            bitmap |= 1 << KeyPurposeBitPosition::ATTEST_KEY_BIT_POS as i32;
-        }
-        _ => {}
-    }
-    bitmap
-}
-
-fn compute_padding_mode_bitmap(padding_mode_bitmap: &i32, padding_mode: PaddingMode) -> i32 {
-    let mut bitmap = *padding_mode_bitmap;
-    match padding_mode {
-        PaddingMode::NONE => {
-            bitmap |= 1 << PaddingModeBitPosition::NONE_BIT_POSITION as i32;
-        }
-        PaddingMode::RSA_OAEP => {
-            bitmap |= 1 << PaddingModeBitPosition::RSA_OAEP_BIT_POS as i32;
-        }
-        PaddingMode::RSA_PSS => {
-            bitmap |= 1 << PaddingModeBitPosition::RSA_PSS_BIT_POS as i32;
-        }
-        PaddingMode::RSA_PKCS1_1_5_ENCRYPT => {
-            bitmap |= 1 << PaddingModeBitPosition::RSA_PKCS1_1_5_ENCRYPT_BIT_POS as i32;
-        }
-        PaddingMode::RSA_PKCS1_1_5_SIGN => {
-            bitmap |= 1 << PaddingModeBitPosition::RSA_PKCS1_1_5_SIGN_BIT_POS as i32;
-        }
-        PaddingMode::PKCS7 => {
-            bitmap |= 1 << PaddingModeBitPosition::PKCS7_BIT_POS as i32;
-        }
-        _ => {}
-    }
-    bitmap
-}
-
-fn compute_digest_bitmap(digest_bitmap: &i32, digest: Digest) -> i32 {
-    let mut bitmap = *digest_bitmap;
-    match digest {
-        Digest::NONE => {
-            bitmap |= 1 << DigestBitPosition::NONE_BIT_POSITION as i32;
-        }
-        Digest::MD5 => {
-            bitmap |= 1 << DigestBitPosition::MD5_BIT_POS as i32;
-        }
-        Digest::SHA1 => {
-            bitmap |= 1 << DigestBitPosition::SHA_1_BIT_POS as i32;
-        }
-        Digest::SHA_2_224 => {
-            bitmap |= 1 << DigestBitPosition::SHA_2_224_BIT_POS as i32;
-        }
-        Digest::SHA_2_256 => {
-            bitmap |= 1 << DigestBitPosition::SHA_2_256_BIT_POS as i32;
-        }
-        Digest::SHA_2_384 => {
-            bitmap |= 1 << DigestBitPosition::SHA_2_384_BIT_POS as i32;
-        }
-        Digest::SHA_2_512 => {
-            bitmap |= 1 << DigestBitPosition::SHA_2_512_BIT_POS as i32;
-        }
-        _ => {}
-    }
-    bitmap
-}
-
-fn compute_block_mode_bitmap(block_mode_bitmap: &i32, block_mode: BlockMode) -> i32 {
-    let mut bitmap = *block_mode_bitmap;
-    match block_mode {
-        BlockMode::ECB => {
-            bitmap |= 1 << BlockModeBitPosition::ECB_BIT_POS as i32;
-        }
-        BlockMode::CBC => {
-            bitmap |= 1 << BlockModeBitPosition::CBC_BIT_POS as i32;
-        }
-        BlockMode::CTR => {
-            bitmap |= 1 << BlockModeBitPosition::CTR_BIT_POS as i32;
-        }
-        BlockMode::GCM => {
-            bitmap |= 1 << BlockModeBitPosition::GCM_BIT_POS as i32;
-        }
-        _ => {}
-    }
-    bitmap
-}
-
-/// Registers pull metrics callbacks
-pub fn register_pull_metrics_callbacks() {
-    // Before registering the callbacks with statsd, we have to wait for the system to finish
-    // booting up. This avoids possible races that may occur at startup. For example, statsd
-    // depends on a companion service, and if registration happens too soon it will fail since
-    // the companion service isn't up yet.
-    LOGS_HANDLER.queue_lo(move |_| {
-        if let Ok(()) = wait_for_boot_completed() {
-            set_pull_atom_callback(Atoms::Keystore2StorageStats, None, pull_metrics_callback);
-            log::info!("Pull metrics callbacks successfully registered.")
-        }
-    });
-}
-
-fn pull_metrics_callback() -> StatsPullResult {
-    let mut result = StatsPullResult::new();
-    let mut append = |stat| {
-        match stat {
-            Ok(s) => result.push(Box::new(s)),
-            Err(error) => {
-                log::error!("pull_metrics_callback: Error getting storage stat: {}", error)
-            }
-        };
-    };
-    DB.with(|db| {
-        let mut db = db.borrow_mut();
-        append(db.get_storage_stat(StatsdStorageType::Database));
-        append(db.get_storage_stat(StatsdStorageType::KeyEntry));
-        append(db.get_storage_stat(StatsdStorageType::KeyEntryIdIndex));
-        append(db.get_storage_stat(StatsdStorageType::KeyEntryDomainNamespaceIndex));
-        append(db.get_storage_stat(StatsdStorageType::BlobEntry));
-        append(db.get_storage_stat(StatsdStorageType::BlobEntryKeyEntryIdIndex));
-        append(db.get_storage_stat(StatsdStorageType::KeyParameter));
-        append(db.get_storage_stat(StatsdStorageType::KeyParameterKeyEntryIdIndex));
-        append(db.get_storage_stat(StatsdStorageType::KeyMetadata));
-        append(db.get_storage_stat(StatsdStorageType::KeyMetadataKeyEntryIdIndex));
-        append(db.get_storage_stat(StatsdStorageType::Grant));
-        append(db.get_storage_stat(StatsdStorageType::AuthToken));
-        append(db.get_storage_stat(StatsdStorageType::BlobMetadata));
-        append(db.get_storage_stat(StatsdStorageType::BlobMetadataBlobEntryIdIndex));
-    });
-    result
-}
-
-/// Enum defining the bit position for each padding mode. Since padding mode can be repeatable, it
-/// is represented using a bitmap.
-#[allow(non_camel_case_types)]
-#[repr(i32)]
-pub enum PaddingModeBitPosition {
-    ///Bit position in the PaddingMode bitmap for NONE.
-    NONE_BIT_POSITION = 0,
-    ///Bit position in the PaddingMode bitmap for RSA_OAEP.
-    RSA_OAEP_BIT_POS = 1,
-    ///Bit position in the PaddingMode bitmap for RSA_PSS.
-    RSA_PSS_BIT_POS = 2,
-    ///Bit position in the PaddingMode bitmap for RSA_PKCS1_1_5_ENCRYPT.
-    RSA_PKCS1_1_5_ENCRYPT_BIT_POS = 3,
-    ///Bit position in the PaddingMode bitmap for RSA_PKCS1_1_5_SIGN.
-    RSA_PKCS1_1_5_SIGN_BIT_POS = 4,
-    ///Bit position in the PaddingMode bitmap for RSA_PKCS7.
-    PKCS7_BIT_POS = 5,
-}
-
-/// Enum defining the bit position for each digest type. Since digest can be repeatable in
-/// key parameters, it is represented using a bitmap.
-#[allow(non_camel_case_types)]
-#[repr(i32)]
-pub enum DigestBitPosition {
-    ///Bit position in the Digest bitmap for NONE.
-    NONE_BIT_POSITION = 0,
-    ///Bit position in the Digest bitmap for MD5.
-    MD5_BIT_POS = 1,
-    ///Bit position in the Digest bitmap for SHA1.
-    SHA_1_BIT_POS = 2,
-    ///Bit position in the Digest bitmap for SHA_2_224.
-    SHA_2_224_BIT_POS = 3,
-    ///Bit position in the Digest bitmap for SHA_2_256.
-    SHA_2_256_BIT_POS = 4,
-    ///Bit position in the Digest bitmap for SHA_2_384.
-    SHA_2_384_BIT_POS = 5,
-    ///Bit position in the Digest bitmap for SHA_2_512.
-    SHA_2_512_BIT_POS = 6,
-}
-
-/// Enum defining the bit position for each block mode type. Since block mode can be repeatable in
-/// key parameters, it is represented using a bitmap.
-#[allow(non_camel_case_types)]
-#[repr(i32)]
-enum BlockModeBitPosition {
-    ///Bit position in the BlockMode bitmap for ECB.
-    ECB_BIT_POS = 1,
-    ///Bit position in the BlockMode bitmap for CBC.
-    CBC_BIT_POS = 2,
-    ///Bit position in the BlockMode bitmap for CTR.
-    CTR_BIT_POS = 3,
-    ///Bit position in the BlockMode bitmap for GCM.
-    GCM_BIT_POS = 4,
-}
-
-/// Enum defining the bit position for each key purpose. Since key purpose can be repeatable in
-/// key parameters, it is represented using a bitmap.
-#[allow(non_camel_case_types)]
-#[repr(i32)]
-enum KeyPurposeBitPosition {
-    ///Bit position in the KeyPurpose bitmap for Encrypt.
-    ENCRYPT_BIT_POS = 1,
-    ///Bit position in the KeyPurpose bitmap for Decrypt.
-    DECRYPT_BIT_POS = 2,
-    ///Bit position in the KeyPurpose bitmap for Sign.
-    SIGN_BIT_POS = 3,
-    ///Bit position in the KeyPurpose bitmap for Verify.
-    VERIFY_BIT_POS = 4,
-    ///Bit position in the KeyPurpose bitmap for Wrap Key.
-    WRAP_KEY_BIT_POS = 5,
-    ///Bit position in the KeyPurpose bitmap for Agree Key.
-    AGREE_KEY_BIT_POS = 6,
-    ///Bit position in the KeyPurpose bitmap for Attest Key.
-    ATTEST_KEY_BIT_POS = 7,
-}
diff --git a/keystore2/src/metrics_store.rs b/keystore2/src/metrics_store.rs
new file mode 100644
index 0000000..e3de035
--- /dev/null
+++ b/keystore2/src/metrics_store.rs
@@ -0,0 +1,611 @@
+// 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 is the metrics store module of keystore. It does the following tasks:
+//! 1. Processes the data about keystore events asynchronously, and
+//!    stores them in an in-memory store.
+//! 2. Returns the collected metrics when requested by the statsd proxy.
+
+use crate::error::get_error_code;
+use crate::globals::DB;
+use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
+use crate::operation::Outcome;
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    Algorithm::Algorithm, BlockMode::BlockMode, Digest::Digest, EcCurve::EcCurve,
+    HardwareAuthenticatorType::HardwareAuthenticatorType, KeyOrigin::KeyOrigin,
+    KeyParameter::KeyParameter, KeyPurpose::KeyPurpose, PaddingMode::PaddingMode,
+    SecurityLevel::SecurityLevel,
+};
+use android_security_metrics::aidl::android::security::metrics::{
+    Algorithm::Algorithm as MetricsAlgorithm, AtomID::AtomID, EcCurve::EcCurve as MetricsEcCurve,
+    HardwareAuthenticatorType::HardwareAuthenticatorType as MetricsHardwareAuthenticatorType,
+    KeyCreationWithAuthInfo::KeyCreationWithAuthInfo,
+    KeyCreationWithGeneralInfo::KeyCreationWithGeneralInfo,
+    KeyCreationWithPurposeAndModesInfo::KeyCreationWithPurposeAndModesInfo,
+    KeyOperationWithGeneralInfo::KeyOperationWithGeneralInfo,
+    KeyOperationWithPurposeAndModesInfo::KeyOperationWithPurposeAndModesInfo,
+    KeyOrigin::KeyOrigin as MetricsKeyOrigin, Keystore2AtomWithOverflow::Keystore2AtomWithOverflow,
+    KeystoreAtom::KeystoreAtom, KeystoreAtomPayload::KeystoreAtomPayload,
+    Outcome::Outcome as MetricsOutcome, Purpose::Purpose as MetricsPurpose,
+    SecurityLevel::SecurityLevel as MetricsSecurityLevel, Storage::Storage as MetricsStorage,
+};
+use anyhow::Result;
+use lazy_static::lazy_static;
+use std::collections::HashMap;
+use std::sync::Mutex;
+
+lazy_static! {
+    /// Singleton for MetricsStore.
+    pub static ref METRICS_STORE: MetricsStore = Default::default();
+}
+
+/// MetricsStore stores the <atom object, count> as <key, value> in the inner hash map,
+/// indexed by the atom id, in the outer hash map.
+/// There can be different atom objects with the same atom id based on the values assigned to the
+/// fields of the atom objects. When an atom object with a particular combination of field values is
+/// inserted, we first check if that atom object is in the inner hash map. If one exists, count
+/// is inceremented. Otherwise, the atom object is inserted with count = 1. Note that count field
+/// of the atom object itself is set to 0 while the object is stored in the hash map. When the atom
+/// objects are queried by the atom id, the corresponding atom objects are retrieved, cloned, and
+/// the count field of the cloned objects is set to the corresponding value field in the inner hash
+/// map before the query result is returned.
+#[derive(Default)]
+pub struct MetricsStore {
+    metrics_store: Mutex<HashMap<AtomID, HashMap<KeystoreAtomPayload, i32>>>,
+}
+
+impl MetricsStore {
+    /// There are some atoms whose maximum cardinality exceeds the cardinality limits tolerated
+    /// by statsd. Statsd tolerates cardinality between 200-300. Therefore, the in-memory storage
+    /// limit for a single atom is set to 250. If the number of atom objects created for a
+    /// particular atom exceeds this limit, an overflow atom object is created to track the ID of
+    /// such atoms.
+    const SINGLE_ATOM_STORE_MAX_SIZE: usize = 250;
+
+    /// Return a vector of atom objects with the given atom ID, if one exists in the metrics_store.
+    /// If any atom object does not exist in the metrics_store for the given atom ID, return an
+    /// empty vector.
+    pub fn get_atoms(&self, atom_id: AtomID) -> Result<Vec<KeystoreAtom>> {
+        // StorageStats is an original pulled atom (i.e. not a pushed atom converted to a
+        // pulledd atom). Therefore, it is handled separately.
+        if AtomID::STORAGE_STATS == atom_id {
+            return pull_storage_stats();
+        }
+
+        // TODO (b/184301651): process and return RKP pool stats.
+
+        // It is safe to call unwrap here since the lock can not be poisoned based on its usage
+        // in this module and the lock is not acquired in the same thread before.
+        let metrics_store_guard = self.metrics_store.lock().unwrap();
+        metrics_store_guard.get(&atom_id).map_or(Ok(Vec::<KeystoreAtom>::new()), |atom_count_map| {
+            Ok(atom_count_map
+                .iter()
+                .map(|(atom, count)| KeystoreAtom { payload: atom.clone(), count: *count })
+                .collect())
+        })
+    }
+
+    /// Insert an atom object to the metrics_store indexed by the atom ID.
+    fn insert_atom(&self, atom_id: AtomID, atom: KeystoreAtomPayload) {
+        // It is ok to unwrap here since the mutex cannot be poisoned according to the way it is
+        // used in this module. And the lock is not acquired by this thread before.
+        let mut metrics_store_guard = self.metrics_store.lock().unwrap();
+        let atom_count_map = metrics_store_guard.entry(atom_id).or_insert_with(HashMap::new);
+        if atom_count_map.len() < MetricsStore::SINGLE_ATOM_STORE_MAX_SIZE {
+            let atom_count = atom_count_map.entry(atom).or_insert(0);
+            *atom_count += 1;
+        } else {
+            // Insert an overflow atom
+            let overflow_atom_count_map = metrics_store_guard
+                .entry(AtomID::KEYSTORE2_ATOM_WITH_OVERFLOW)
+                .or_insert_with(HashMap::new);
+
+            if overflow_atom_count_map.len() < MetricsStore::SINGLE_ATOM_STORE_MAX_SIZE {
+                let overflow_atom = Keystore2AtomWithOverflow { atom_id };
+                let atom_count = overflow_atom_count_map
+                    .entry(KeystoreAtomPayload::Keystore2AtomWithOverflow(overflow_atom))
+                    .or_insert(0);
+                *atom_count += 1;
+            } else {
+                // This is a rare case, if at all.
+                log::error!("In insert_atom: Maximum storage limit reached for overflow atom.")
+            }
+        }
+    }
+}
+
+/// Log key creation events to be sent to statsd.
+pub fn log_key_creation_event_stats<U>(
+    sec_level: SecurityLevel,
+    key_params: &[KeyParameter],
+    result: &Result<U>,
+) {
+    let (
+        key_creation_with_general_info,
+        key_creation_with_auth_info,
+        key_creation_with_purpose_and_modes_info,
+    ) = process_key_creation_event_stats(sec_level, key_params, result);
+
+    METRICS_STORE
+        .insert_atom(AtomID::KEY_CREATION_WITH_GENERAL_INFO, key_creation_with_general_info);
+    METRICS_STORE.insert_atom(AtomID::KEY_CREATION_WITH_AUTH_INFO, key_creation_with_auth_info);
+    METRICS_STORE.insert_atom(
+        AtomID::KEY_CREATION_WITH_PURPOSE_AND_MODES_INFO,
+        key_creation_with_purpose_and_modes_info,
+    );
+}
+
+// Process the statistics related to key creations and return the three atom objects related to key
+// creations: i) KeyCreationWithGeneralInfo ii) KeyCreationWithAuthInfo
+// iii) KeyCreationWithPurposeAndModesInfo
+fn process_key_creation_event_stats<U>(
+    sec_level: SecurityLevel,
+    key_params: &[KeyParameter],
+    result: &Result<U>,
+) -> (KeystoreAtomPayload, KeystoreAtomPayload, KeystoreAtomPayload) {
+    // In the default atom objects, fields represented by bitmaps and i32 fields
+    // will take 0, except error_code which defaults to 1 indicating NO_ERROR and key_size,
+    // and auth_time_out which defaults to -1.
+    // The boolean fields are set to false by default.
+    // Some keymint enums do have 0 as an enum variant value. In such cases, the corresponding
+    // enum variant value in atoms.proto is incremented by 1, in order to have 0 as the reserved
+    // value for unspecified fields.
+    let mut key_creation_with_general_info = KeyCreationWithGeneralInfo {
+        algorithm: MetricsAlgorithm::ALGORITHM_UNSPECIFIED,
+        key_size: -1,
+        ec_curve: MetricsEcCurve::EC_CURVE_UNSPECIFIED,
+        key_origin: MetricsKeyOrigin::ORIGIN_UNSPECIFIED,
+        error_code: 1,
+        // Default for bool is false (for attestation_requested field).
+        ..Default::default()
+    };
+
+    let mut key_creation_with_auth_info = KeyCreationWithAuthInfo {
+        user_auth_type: MetricsHardwareAuthenticatorType::AUTH_TYPE_UNSPECIFIED,
+        log10_auth_key_timeout_seconds: -1,
+        security_level: MetricsSecurityLevel::SECURITY_LEVEL_UNSPECIFIED,
+    };
+
+    let mut key_creation_with_purpose_and_modes_info = KeyCreationWithPurposeAndModesInfo {
+        algorithm: MetricsAlgorithm::ALGORITHM_UNSPECIFIED,
+        // Default for i32 is 0 (for the remaining bitmap fields).
+        ..Default::default()
+    };
+
+    if let Err(ref e) = result {
+        key_creation_with_general_info.error_code = get_error_code(e);
+    }
+
+    key_creation_with_auth_info.security_level = process_security_level(sec_level);
+
+    for key_param in key_params.iter().map(KsKeyParamValue::from) {
+        match key_param {
+            KsKeyParamValue::Algorithm(a) => {
+                let algorithm = match a {
+                    Algorithm::RSA => MetricsAlgorithm::RSA,
+                    Algorithm::EC => MetricsAlgorithm::EC,
+                    Algorithm::AES => MetricsAlgorithm::AES,
+                    Algorithm::TRIPLE_DES => MetricsAlgorithm::TRIPLE_DES,
+                    Algorithm::HMAC => MetricsAlgorithm::HMAC,
+                    _ => MetricsAlgorithm::ALGORITHM_UNSPECIFIED,
+                };
+                key_creation_with_general_info.algorithm = algorithm;
+                key_creation_with_purpose_and_modes_info.algorithm = algorithm;
+            }
+            KsKeyParamValue::KeySize(s) => {
+                key_creation_with_general_info.key_size = s;
+            }
+            KsKeyParamValue::KeyOrigin(o) => {
+                key_creation_with_general_info.key_origin = match o {
+                    KeyOrigin::GENERATED => MetricsKeyOrigin::GENERATED,
+                    KeyOrigin::DERIVED => MetricsKeyOrigin::DERIVED,
+                    KeyOrigin::IMPORTED => MetricsKeyOrigin::IMPORTED,
+                    KeyOrigin::RESERVED => MetricsKeyOrigin::RESERVED,
+                    KeyOrigin::SECURELY_IMPORTED => MetricsKeyOrigin::SECURELY_IMPORTED,
+                    _ => MetricsKeyOrigin::ORIGIN_UNSPECIFIED,
+                }
+            }
+            KsKeyParamValue::HardwareAuthenticatorType(a) => {
+                key_creation_with_auth_info.user_auth_type = match a {
+                    HardwareAuthenticatorType::NONE => MetricsHardwareAuthenticatorType::NONE,
+                    HardwareAuthenticatorType::PASSWORD => {
+                        MetricsHardwareAuthenticatorType::PASSWORD
+                    }
+                    HardwareAuthenticatorType::FINGERPRINT => {
+                        MetricsHardwareAuthenticatorType::FINGERPRINT
+                    }
+                    HardwareAuthenticatorType::ANY => MetricsHardwareAuthenticatorType::ANY,
+                    _ => MetricsHardwareAuthenticatorType::AUTH_TYPE_UNSPECIFIED,
+                }
+            }
+            KsKeyParamValue::AuthTimeout(t) => {
+                key_creation_with_auth_info.log10_auth_key_timeout_seconds =
+                    f32::log10(t as f32) as i32;
+            }
+            KsKeyParamValue::PaddingMode(p) => {
+                compute_padding_mode_bitmap(
+                    &mut key_creation_with_purpose_and_modes_info.padding_mode_bitmap,
+                    p,
+                );
+            }
+            KsKeyParamValue::Digest(d) => {
+                // key_creation_with_purpose_and_modes_info.digest_bitmap =
+                compute_digest_bitmap(
+                    &mut key_creation_with_purpose_and_modes_info.digest_bitmap,
+                    d,
+                );
+            }
+            KsKeyParamValue::BlockMode(b) => {
+                compute_block_mode_bitmap(
+                    &mut key_creation_with_purpose_and_modes_info.block_mode_bitmap,
+                    b,
+                );
+            }
+            KsKeyParamValue::KeyPurpose(k) => {
+                compute_purpose_bitmap(
+                    &mut key_creation_with_purpose_and_modes_info.purpose_bitmap,
+                    k,
+                );
+            }
+            KsKeyParamValue::EcCurve(e) => {
+                key_creation_with_general_info.ec_curve = match e {
+                    EcCurve::P_224 => MetricsEcCurve::P_224,
+                    EcCurve::P_256 => MetricsEcCurve::P_256,
+                    EcCurve::P_384 => MetricsEcCurve::P_384,
+                    EcCurve::P_521 => MetricsEcCurve::P_521,
+                    _ => MetricsEcCurve::EC_CURVE_UNSPECIFIED,
+                }
+            }
+            KsKeyParamValue::AttestationChallenge(_) => {
+                key_creation_with_general_info.attestation_requested = true;
+            }
+            _ => {}
+        }
+    }
+    if key_creation_with_general_info.algorithm == MetricsAlgorithm::EC {
+        // Do not record key sizes if Algorithm = EC, in order to reduce cardinality.
+        key_creation_with_general_info.key_size = -1;
+    }
+
+    (
+        KeystoreAtomPayload::KeyCreationWithGeneralInfo(key_creation_with_general_info),
+        KeystoreAtomPayload::KeyCreationWithAuthInfo(key_creation_with_auth_info),
+        KeystoreAtomPayload::KeyCreationWithPurposeAndModesInfo(
+            key_creation_with_purpose_and_modes_info,
+        ),
+    )
+}
+
+/// Log key operation events to be sent to statsd.
+pub fn log_key_operation_event_stats(
+    sec_level: SecurityLevel,
+    key_purpose: KeyPurpose,
+    op_params: &[KeyParameter],
+    op_outcome: &Outcome,
+    key_upgraded: bool,
+) {
+    let (key_operation_with_general_info, key_operation_with_purpose_and_modes_info) =
+        process_key_operation_event_stats(
+            sec_level,
+            key_purpose,
+            op_params,
+            op_outcome,
+            key_upgraded,
+        );
+    METRICS_STORE
+        .insert_atom(AtomID::KEY_OPERATION_WITH_GENERAL_INFO, key_operation_with_general_info);
+    METRICS_STORE.insert_atom(
+        AtomID::KEY_OPERATION_WITH_PURPOSE_AND_MODES_INFO,
+        key_operation_with_purpose_and_modes_info,
+    );
+}
+
+// Process the statistics related to key operations and return the two atom objects related to key
+// operations: i) KeyOperationWithGeneralInfo ii) KeyOperationWithPurposeAndModesInfo
+fn process_key_operation_event_stats(
+    sec_level: SecurityLevel,
+    key_purpose: KeyPurpose,
+    op_params: &[KeyParameter],
+    op_outcome: &Outcome,
+    key_upgraded: bool,
+) -> (KeystoreAtomPayload, KeystoreAtomPayload) {
+    let mut key_operation_with_general_info = KeyOperationWithGeneralInfo {
+        outcome: MetricsOutcome::OUTCOME_UNSPECIFIED,
+        error_code: 1,
+        security_level: MetricsSecurityLevel::SECURITY_LEVEL_UNSPECIFIED,
+        // Default for bool is false (for key_upgraded field).
+        ..Default::default()
+    };
+
+    let mut key_operation_with_purpose_and_modes_info = KeyOperationWithPurposeAndModesInfo {
+        purpose: MetricsPurpose::KEY_PURPOSE_UNSPECIFIED,
+        // Default for i32 is 0 (for the remaining bitmap fields).
+        ..Default::default()
+    };
+
+    key_operation_with_general_info.security_level = process_security_level(sec_level);
+
+    key_operation_with_general_info.key_upgraded = key_upgraded;
+
+    key_operation_with_purpose_and_modes_info.purpose = match key_purpose {
+        KeyPurpose::ENCRYPT => MetricsPurpose::ENCRYPT,
+        KeyPurpose::DECRYPT => MetricsPurpose::DECRYPT,
+        KeyPurpose::SIGN => MetricsPurpose::SIGN,
+        KeyPurpose::VERIFY => MetricsPurpose::VERIFY,
+        KeyPurpose::WRAP_KEY => MetricsPurpose::WRAP_KEY,
+        KeyPurpose::AGREE_KEY => MetricsPurpose::AGREE_KEY,
+        KeyPurpose::ATTEST_KEY => MetricsPurpose::ATTEST_KEY,
+        _ => MetricsPurpose::KEY_PURPOSE_UNSPECIFIED,
+    };
+
+    key_operation_with_general_info.outcome = match op_outcome {
+        Outcome::Unknown | Outcome::Dropped => MetricsOutcome::DROPPED,
+        Outcome::Success => MetricsOutcome::SUCCESS,
+        Outcome::Abort => MetricsOutcome::ABORT,
+        Outcome::Pruned => MetricsOutcome::PRUNED,
+        Outcome::ErrorCode(e) => {
+            key_operation_with_general_info.error_code = e.0;
+            MetricsOutcome::ERROR
+        }
+    };
+
+    for key_param in op_params.iter().map(KsKeyParamValue::from) {
+        match key_param {
+            KsKeyParamValue::PaddingMode(p) => {
+                compute_padding_mode_bitmap(
+                    &mut key_operation_with_purpose_and_modes_info.padding_mode_bitmap,
+                    p,
+                );
+            }
+            KsKeyParamValue::Digest(d) => {
+                compute_digest_bitmap(
+                    &mut key_operation_with_purpose_and_modes_info.digest_bitmap,
+                    d,
+                );
+            }
+            KsKeyParamValue::BlockMode(b) => {
+                compute_block_mode_bitmap(
+                    &mut key_operation_with_purpose_and_modes_info.block_mode_bitmap,
+                    b,
+                );
+            }
+            _ => {}
+        }
+    }
+
+    (
+        KeystoreAtomPayload::KeyOperationWithGeneralInfo(key_operation_with_general_info),
+        KeystoreAtomPayload::KeyOperationWithPurposeAndModesInfo(
+            key_operation_with_purpose_and_modes_info,
+        ),
+    )
+}
+
+fn process_security_level(sec_level: SecurityLevel) -> MetricsSecurityLevel {
+    match sec_level {
+        SecurityLevel::SOFTWARE => MetricsSecurityLevel::SECURITY_LEVEL_SOFTWARE,
+        SecurityLevel::TRUSTED_ENVIRONMENT => {
+            MetricsSecurityLevel::SECURITY_LEVEL_TRUSTED_ENVIRONMENT
+        }
+        SecurityLevel::STRONGBOX => MetricsSecurityLevel::SECURITY_LEVEL_STRONGBOX,
+        SecurityLevel::KEYSTORE => MetricsSecurityLevel::SECURITY_LEVEL_KEYSTORE,
+        _ => MetricsSecurityLevel::SECURITY_LEVEL_UNSPECIFIED,
+    }
+}
+
+fn compute_padding_mode_bitmap(padding_mode_bitmap: &mut i32, padding_mode: PaddingMode) {
+    match padding_mode {
+        PaddingMode::NONE => {
+            *padding_mode_bitmap |= 1 << PaddingModeBitPosition::NONE_BIT_POSITION as i32;
+        }
+        PaddingMode::RSA_OAEP => {
+            *padding_mode_bitmap |= 1 << PaddingModeBitPosition::RSA_OAEP_BIT_POS as i32;
+        }
+        PaddingMode::RSA_PSS => {
+            *padding_mode_bitmap |= 1 << PaddingModeBitPosition::RSA_PSS_BIT_POS as i32;
+        }
+        PaddingMode::RSA_PKCS1_1_5_ENCRYPT => {
+            *padding_mode_bitmap |=
+                1 << PaddingModeBitPosition::RSA_PKCS1_1_5_ENCRYPT_BIT_POS as i32;
+        }
+        PaddingMode::RSA_PKCS1_1_5_SIGN => {
+            *padding_mode_bitmap |= 1 << PaddingModeBitPosition::RSA_PKCS1_1_5_SIGN_BIT_POS as i32;
+        }
+        PaddingMode::PKCS7 => {
+            *padding_mode_bitmap |= 1 << PaddingModeBitPosition::PKCS7_BIT_POS as i32;
+        }
+        _ => {}
+    }
+}
+
+fn compute_digest_bitmap(digest_bitmap: &mut i32, digest: Digest) {
+    match digest {
+        Digest::NONE => {
+            *digest_bitmap |= 1 << DigestBitPosition::NONE_BIT_POSITION as i32;
+        }
+        Digest::MD5 => {
+            *digest_bitmap |= 1 << DigestBitPosition::MD5_BIT_POS as i32;
+        }
+        Digest::SHA1 => {
+            *digest_bitmap |= 1 << DigestBitPosition::SHA_1_BIT_POS as i32;
+        }
+        Digest::SHA_2_224 => {
+            *digest_bitmap |= 1 << DigestBitPosition::SHA_2_224_BIT_POS as i32;
+        }
+        Digest::SHA_2_256 => {
+            *digest_bitmap |= 1 << DigestBitPosition::SHA_2_256_BIT_POS as i32;
+        }
+        Digest::SHA_2_384 => {
+            *digest_bitmap |= 1 << DigestBitPosition::SHA_2_384_BIT_POS as i32;
+        }
+        Digest::SHA_2_512 => {
+            *digest_bitmap |= 1 << DigestBitPosition::SHA_2_512_BIT_POS as i32;
+        }
+        _ => {}
+    }
+}
+
+fn compute_block_mode_bitmap(block_mode_bitmap: &mut i32, block_mode: BlockMode) {
+    match block_mode {
+        BlockMode::ECB => {
+            *block_mode_bitmap |= 1 << BlockModeBitPosition::ECB_BIT_POS as i32;
+        }
+        BlockMode::CBC => {
+            *block_mode_bitmap |= 1 << BlockModeBitPosition::CBC_BIT_POS as i32;
+        }
+        BlockMode::CTR => {
+            *block_mode_bitmap |= 1 << BlockModeBitPosition::CTR_BIT_POS as i32;
+        }
+        BlockMode::GCM => {
+            *block_mode_bitmap |= 1 << BlockModeBitPosition::GCM_BIT_POS as i32;
+        }
+        _ => {}
+    }
+}
+
+fn compute_purpose_bitmap(purpose_bitmap: &mut i32, purpose: KeyPurpose) {
+    match purpose {
+        KeyPurpose::ENCRYPT => {
+            *purpose_bitmap |= 1 << KeyPurposeBitPosition::ENCRYPT_BIT_POS as i32;
+        }
+        KeyPurpose::DECRYPT => {
+            *purpose_bitmap |= 1 << KeyPurposeBitPosition::DECRYPT_BIT_POS as i32;
+        }
+        KeyPurpose::SIGN => {
+            *purpose_bitmap |= 1 << KeyPurposeBitPosition::SIGN_BIT_POS as i32;
+        }
+        KeyPurpose::VERIFY => {
+            *purpose_bitmap |= 1 << KeyPurposeBitPosition::VERIFY_BIT_POS as i32;
+        }
+        KeyPurpose::WRAP_KEY => {
+            *purpose_bitmap |= 1 << KeyPurposeBitPosition::WRAP_KEY_BIT_POS as i32;
+        }
+        KeyPurpose::AGREE_KEY => {
+            *purpose_bitmap |= 1 << KeyPurposeBitPosition::AGREE_KEY_BIT_POS as i32;
+        }
+        KeyPurpose::ATTEST_KEY => {
+            *purpose_bitmap |= 1 << KeyPurposeBitPosition::ATTEST_KEY_BIT_POS as i32;
+        }
+        _ => {}
+    }
+}
+
+fn pull_storage_stats() -> Result<Vec<KeystoreAtom>> {
+    let mut atom_vec: Vec<KeystoreAtom> = Vec::new();
+    let mut append = |stat| {
+        match stat {
+            Ok(s) => atom_vec.push(KeystoreAtom {
+                payload: KeystoreAtomPayload::StorageStats(s),
+                ..Default::default()
+            }),
+            Err(error) => {
+                log::error!("pull_metrics_callback: Error getting storage stat: {}", error)
+            }
+        };
+    };
+    DB.with(|db| {
+        let mut db = db.borrow_mut();
+        append(db.get_storage_stat(MetricsStorage::DATABASE));
+        append(db.get_storage_stat(MetricsStorage::KEY_ENTRY));
+        append(db.get_storage_stat(MetricsStorage::KEY_ENTRY_ID_INDEX));
+        append(db.get_storage_stat(MetricsStorage::KEY_ENTRY_DOMAIN_NAMESPACE_INDEX));
+        append(db.get_storage_stat(MetricsStorage::BLOB_ENTRY));
+        append(db.get_storage_stat(MetricsStorage::BLOB_ENTRY_KEY_ENTRY_ID_INDEX));
+        append(db.get_storage_stat(MetricsStorage::KEY_PARAMETER));
+        append(db.get_storage_stat(MetricsStorage::KEY_PARAMETER_KEY_ENTRY_ID_INDEX));
+        append(db.get_storage_stat(MetricsStorage::KEY_METADATA));
+        append(db.get_storage_stat(MetricsStorage::KEY_METADATA_KEY_ENTRY_ID_INDEX));
+        append(db.get_storage_stat(MetricsStorage::GRANT));
+        append(db.get_storage_stat(MetricsStorage::AUTH_TOKEN));
+        append(db.get_storage_stat(MetricsStorage::BLOB_METADATA));
+        append(db.get_storage_stat(MetricsStorage::BLOB_METADATA_BLOB_ENTRY_ID_INDEX));
+    });
+    Ok(atom_vec)
+}
+
+/// Enum defining the bit position for each padding mode. Since padding mode can be repeatable, it
+/// is represented using a bitmap.
+#[allow(non_camel_case_types)]
+#[repr(i32)]
+enum PaddingModeBitPosition {
+    ///Bit position in the PaddingMode bitmap for NONE.
+    NONE_BIT_POSITION = 0,
+    ///Bit position in the PaddingMode bitmap for RSA_OAEP.
+    RSA_OAEP_BIT_POS = 1,
+    ///Bit position in the PaddingMode bitmap for RSA_PSS.
+    RSA_PSS_BIT_POS = 2,
+    ///Bit position in the PaddingMode bitmap for RSA_PKCS1_1_5_ENCRYPT.
+    RSA_PKCS1_1_5_ENCRYPT_BIT_POS = 3,
+    ///Bit position in the PaddingMode bitmap for RSA_PKCS1_1_5_SIGN.
+    RSA_PKCS1_1_5_SIGN_BIT_POS = 4,
+    ///Bit position in the PaddingMode bitmap for RSA_PKCS7.
+    PKCS7_BIT_POS = 5,
+}
+
+/// Enum defining the bit position for each digest type. Since digest can be repeatable in
+/// key parameters, it is represented using a bitmap.
+#[allow(non_camel_case_types)]
+#[repr(i32)]
+enum DigestBitPosition {
+    ///Bit position in the Digest bitmap for NONE.
+    NONE_BIT_POSITION = 0,
+    ///Bit position in the Digest bitmap for MD5.
+    MD5_BIT_POS = 1,
+    ///Bit position in the Digest bitmap for SHA1.
+    SHA_1_BIT_POS = 2,
+    ///Bit position in the Digest bitmap for SHA_2_224.
+    SHA_2_224_BIT_POS = 3,
+    ///Bit position in the Digest bitmap for SHA_2_256.
+    SHA_2_256_BIT_POS = 4,
+    ///Bit position in the Digest bitmap for SHA_2_384.
+    SHA_2_384_BIT_POS = 5,
+    ///Bit position in the Digest bitmap for SHA_2_512.
+    SHA_2_512_BIT_POS = 6,
+}
+
+/// Enum defining the bit position for each block mode type. Since block mode can be repeatable in
+/// key parameters, it is represented using a bitmap.
+#[allow(non_camel_case_types)]
+#[repr(i32)]
+enum BlockModeBitPosition {
+    ///Bit position in the BlockMode bitmap for ECB.
+    ECB_BIT_POS = 1,
+    ///Bit position in the BlockMode bitmap for CBC.
+    CBC_BIT_POS = 2,
+    ///Bit position in the BlockMode bitmap for CTR.
+    CTR_BIT_POS = 3,
+    ///Bit position in the BlockMode bitmap for GCM.
+    GCM_BIT_POS = 4,
+}
+
+/// Enum defining the bit position for each key purpose. Since key purpose can be repeatable in
+/// key parameters, it is represented using a bitmap.
+#[allow(non_camel_case_types)]
+#[repr(i32)]
+enum KeyPurposeBitPosition {
+    ///Bit position in the KeyPurpose bitmap for Encrypt.
+    ENCRYPT_BIT_POS = 1,
+    ///Bit position in the KeyPurpose bitmap for Decrypt.
+    DECRYPT_BIT_POS = 2,
+    ///Bit position in the KeyPurpose bitmap for Sign.
+    SIGN_BIT_POS = 3,
+    ///Bit position in the KeyPurpose bitmap for Verify.
+    VERIFY_BIT_POS = 4,
+    ///Bit position in the KeyPurpose bitmap for Wrap Key.
+    WRAP_KEY_BIT_POS = 5,
+    ///Bit position in the KeyPurpose bitmap for Agree Key.
+    AGREE_KEY_BIT_POS = 6,
+    ///Bit position in the KeyPurpose bitmap for Attest Key.
+    ATTEST_KEY_BIT_POS = 7,
+}
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 1d595b3..7e08f4e 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -127,7 +127,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::metrics_store::log_key_operation_event_stats;
 use crate::utils::watchdog as wd;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     IKeyMintOperation::IKeyMintOperation, KeyParameter::KeyParameter, KeyPurpose::KeyPurpose,
diff --git a/keystore2/src/permission.rs b/keystore2/src/permission.rs
index e7999bc..8343a29 100644
--- a/keystore2/src/permission.rs
+++ b/keystore2/src/permission.rs
@@ -315,6 +315,8 @@
         EarlyBootEnded = 0x800,   selinux name: early_boot_ended;
         /// Checked when IKeystoreMaintenance::onDeviceOffBody is called.
         ReportOffBody = 0x1000, selinux name: report_off_body;
+        /// Checked when IkeystoreMetrics::pullMetris is called.
+        PullMetrics = 0x2000, selinux name: pull_metrics;
     }
 );
 
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 2fddc18..b66c778 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -23,7 +23,7 @@
 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::metrics_store::log_key_creation_event_stats;
 use crate::remote_provisioning::RemProvState;
 use crate::super_key::{KeyBlob, SuperKeyManager};
 use crate::utils::{