Merge "Add safety comments." into main
diff --git a/OWNERS b/OWNERS
index 03e5769..6fdb550 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,5 +1,4 @@
 alanstokes@google.com
-cbrubaker@google.com
 drysdale@google.com
 eranm@google.com
 hasinitg@google.com
@@ -8,4 +7,5 @@
 kroot@google.com
 sethmo@google.com
 swillden@google.com
+trong@google.com
 zeuthen@google.com
diff --git a/diced/open_dice/Android.bp b/diced/open_dice/Android.bp
index 45cbb67..c59419b 100644
--- a/diced/open_dice/Android.bp
+++ b/diced/open_dice/Android.bp
@@ -7,14 +7,6 @@
     name: "libdiced_open_dice_defaults",
     crate_name: "diced_open_dice",
     srcs: ["src/lib.rs"],
-    static_libs: [
-        "libopen_dice_cbor",
-    ],
-    vendor_available: true,
-    apex_available: [
-        "//apex_available:platform",
-        "com.android.virt",
-    ],
 }
 
 rust_library_rlib {
@@ -37,6 +29,7 @@
 rust_library {
     name: "libdiced_open_dice",
     defaults: ["libdiced_open_dice_defaults"],
+    vendor_available: true,
     rustlibs: [
         "libopen_dice_bcc_bindgen",
         "libopen_dice_cbor_bindgen",
@@ -48,6 +41,9 @@
     shared_libs: [
         "libcrypto",
     ],
+    static_libs: [
+        "libopen_dice_cbor",
+    ],
     whole_static_libs: [
         "libopen_dice_bcc",
     ],
@@ -56,6 +52,10 @@
         "//packages/modules/Virtualization:__subpackages__",
         "//hardware/interfaces/security/dice/aidl:__subpackages__",
     ],
+    apex_available: [
+        "//apex_available:platform",
+        "com.android.virt",
+    ],
 }
 
 rust_defaults {
@@ -120,7 +120,6 @@
 
 rust_defaults {
     name: "libopen_dice_cbor_bindgen.rust_defaults",
-    defaults: ["libopen_dice.rust_defaults"],
     wrapper_src: "bindgen/dice.h",
     crate_name: "open_dice_cbor_bindgen",
     source_stem: "bindings",
@@ -156,7 +155,10 @@
 
 rust_bindgen {
     name: "libopen_dice_cbor_bindgen",
-    defaults: ["libopen_dice_cbor_bindgen.rust_defaults"],
+    defaults: [
+        "libopen_dice.rust_defaults",
+        "libopen_dice_cbor_bindgen.rust_defaults",
+    ],
     whole_static_libs: ["libopen_dice_cbor"],
 }
 
@@ -171,7 +173,6 @@
 
 rust_defaults {
     name: "libopen_dice_bcc_bindgen.rust_defaults",
-    defaults: ["libopen_dice.rust_defaults"],
     wrapper_src: "bindgen/android/bcc.h",
     crate_name: "open_dice_bcc_bindgen",
     source_stem: "bindings",
@@ -206,7 +207,10 @@
 
 rust_bindgen {
     name: "libopen_dice_bcc_bindgen",
-    defaults: ["libopen_dice_bcc_bindgen.rust_defaults"],
+    defaults: [
+        "libopen_dice.rust_defaults",
+        "libopen_dice_bcc_bindgen.rust_defaults",
+    ],
     rustlibs: [
         "libopen_dice_cbor_bindgen",
     ],
diff --git a/keystore/keystore_attestation_id.cpp b/keystore/keystore_attestation_id.cpp
index ccd3808..75c62dd 100644
--- a/keystore/keystore_attestation_id.cpp
+++ b/keystore/keystore_attestation_id.cpp
@@ -74,8 +74,8 @@
 }
 
 KeyAttestationApplicationIdProvider::KeyAttestationApplicationIdProvider()
-    : BpKeyAttestationApplicationIdProvider(
-          android::defaultServiceManager()->getService(String16("sec_key_att_app_id_provider"))) {}
+    : BpKeyAttestationApplicationIdProvider(android::defaultServiceManager()->waitForService(
+          String16("sec_key_att_app_id_provider"))) {}
 
 DECLARE_STACK_OF(ASN1_OCTET_STRING);
 
diff --git a/keystore2/selinux/src/lib.rs b/keystore2/selinux/src/lib.rs
index e5c3091..32fdb59 100644
--- a/keystore2/selinux/src/lib.rs
+++ b/keystore2/selinux/src/lib.rs
@@ -20,6 +20,9 @@
 //!  * selabel_lookup for the keystore2_key backend.
 //! And it provides an owning wrapper around context strings `Context`.
 
+// TODO(b/290018030): Remove this and add proper safety comments.
+#![allow(clippy::undocumented_unsafe_blocks)]
+
 use anyhow::Context as AnyhowContext;
 use anyhow::{anyhow, Result};
 use lazy_static::lazy_static;
@@ -160,8 +163,9 @@
     handle: *mut selinux::selabel_handle,
 }
 
-// KeystoreKeyBackend is Sync because selabel_lookup is thread safe.
+// SAFETY: KeystoreKeyBackend is Sync because selabel_lookup is thread safe.
 unsafe impl Sync for KeystoreKeyBackend {}
+// SAFETY: KeystoreKeyBackend is Send because selabel_lookup is thread safe.
 unsafe impl Send for KeystoreKeyBackend {}
 
 impl KeystoreKeyBackend {
diff --git a/keystore2/src/crypto/lib.rs b/keystore2/src/crypto/lib.rs
index 08b7589..f8fc574 100644
--- a/keystore2/src/crypto/lib.rs
+++ b/keystore2/src/crypto/lib.rs
@@ -49,8 +49,8 @@
 /// Generate an AES256 key, essentially 32 random bytes from the underlying
 /// boringssl library discretely stuffed into a ZVec.
 pub fn generate_aes256_key() -> Result<ZVec, Error> {
-    // Safety: key has the same length as the requested number of random bytes.
     let mut key = ZVec::new(AES_256_KEY_LENGTH)?;
+    // Safety: key has the same length as the requested number of random bytes.
     if unsafe { randomBytes(key.as_mut_ptr(), AES_256_KEY_LENGTH) } {
         Ok(key)
     } else {
@@ -65,8 +65,8 @@
 
 /// Generate random data of the given size.
 pub fn generate_random_data(size: usize) -> Result<Vec<u8>, Error> {
-    // Safety: data has the same length as the requested number of random bytes.
     let mut data = vec![0; size];
+    // Safety: data has the same length as the requested number of random bytes.
     if unsafe { randomBytes(data.as_mut_ptr(), size) } {
         Ok(data)
     } else {
@@ -209,6 +209,8 @@
         let pw = self.get_key();
         let mut result = ZVec::new(key_length)?;
 
+        // Safety: We checked that the salt is exactly 16 bytes long. The other pointers are valid,
+        // and have matching lengths.
         unsafe {
             generateKeyFromPassword(
                 result.as_mut_ptr(),
@@ -324,10 +326,10 @@
 /// Calls the boringssl ECDH_compute_key function.
 pub fn ecdh_compute_key(pub_key: &EC_POINT, priv_key: &ECKey) -> Result<ZVec, Error> {
     let mut buf = ZVec::new(EC_MAX_BYTES)?;
+    let result =
     // Safety: Our ECDHComputeKey wrapper passes EC_MAX_BYES to ECDH_compute_key, which
     // writes at most that many bytes to the output.
     // The two keys are valid objects.
-    let result =
         unsafe { ECDHComputeKey(buf.as_mut_ptr() as *mut std::ffi::c_void, pub_key, priv_key.0) };
     if result == -1 {
         return Err(Error::ECDHComputeKeyFailed);
@@ -490,6 +492,8 @@
         let key = vec![0; 16];
         let iv = vec![0; 12];
         let mut tag = vec![0; 16];
+        // SAFETY: The various pointers are obtained from references so they are valid, and
+        // `AES_gcm_encrypt` and `AES_gcm_decrypt` don't do anything with them after they return.
         unsafe {
             let res = AES_gcm_encrypt(
                 input.as_ptr(),
@@ -521,8 +525,10 @@
     fn test_create_key_id() {
         let blob = vec![0; 16];
         let mut out: u64 = 0;
+        // SAFETY: The pointers are obtained from references so they are valid, the length matches
+        // the length of the array, and `CreateKeyId` doesn't access them after it returns.
         unsafe {
-            let res = CreateKeyId(blob.as_ptr(), 16, &mut out);
+            let res = CreateKeyId(blob.as_ptr(), blob.len(), &mut out);
             assert!(res);
             assert_ne!(out, 0);
         }
@@ -533,8 +539,17 @@
         let mut key = vec![0; 16];
         let pw = vec![0; 16];
         let salt = vec![0; 16];
+        // SAFETY: The pointers are obtained from references so they are valid, the salt is the
+        // expected length, the other lengths match the lengths of the arrays, and
+        // `generateKeyFromPassword` doesn't access them after it returns.
         unsafe {
-            generateKeyFromPassword(key.as_mut_ptr(), 16, pw.as_ptr(), 16, salt.as_ptr());
+            generateKeyFromPassword(
+                key.as_mut_ptr(),
+                key.len(),
+                pw.as_ptr(),
+                pw.len(),
+                salt.as_ptr(),
+            );
         }
         assert_ne!(key, vec![0; 16]);
     }
diff --git a/keystore2/src/crypto/zvec.rs b/keystore2/src/crypto/zvec.rs
index 5a173c3..c917a89 100644
--- a/keystore2/src/crypto/zvec.rs
+++ b/keystore2/src/crypto/zvec.rs
@@ -45,6 +45,7 @@
         let v: Vec<u8> = vec![0; size];
         let b = v.into_boxed_slice();
         if size > 0 {
+            // SAFETY: The address range is part of our address space.
             unsafe { mlock(b.as_ptr() as *const std::ffi::c_void, b.len()) }?;
         }
         Ok(Self { elems: b, len: size })
@@ -71,11 +72,16 @@
 impl Drop for ZVec {
     fn drop(&mut self) {
         for i in 0..self.elems.len() {
-            unsafe { write_volatile(self.elems.as_mut_ptr().add(i), 0) };
+            // SAFETY: The pointer is valid and properly aligned because it came from a reference.
+            unsafe { write_volatile(&mut self.elems[i], 0) };
         }
         if !self.elems.is_empty() {
             if let Err(e) =
-                unsafe { munlock(self.elems.as_ptr() as *const std::ffi::c_void, self.elems.len()) }
+                // SAFETY: The address range is part of our address space, and was previously locked
+                // by `mlock` in `ZVec::new` or the `TryFrom<Vec<u8>>` implementation.
+                unsafe {
+                    munlock(self.elems.as_ptr() as *const std::ffi::c_void, self.elems.len())
+                }
             {
                 log::error!("In ZVec::drop: `munlock` failed: {:?}.", e);
             }
@@ -130,6 +136,7 @@
         v.resize(v.capacity(), 0);
         let b = v.into_boxed_slice();
         if !b.is_empty() {
+            // SAFETY: The address range is part of our address space.
             unsafe { mlock(b.as_ptr() as *const std::ffi::c_void, b.len()) }?;
         }
         Ok(Self { elems: b, len })
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 8d5e985..3bf582f 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -51,9 +51,9 @@
     /// An outstanding per operation authorization request.
     OpAuth,
     /// An outstanding request for per operation authorization and secure timestamp.
-    TimeStampedOpAuth(Receiver<Result<TimeStampToken, Error>>),
+    TimeStampedOpAuth(Mutex<Receiver<Result<TimeStampToken, Error>>>),
     /// An outstanding request for a timestamp token.
-    TimeStamp(Receiver<Result<TimeStampToken, Error>>),
+    TimeStamp(Mutex<Receiver<Result<TimeStampToken, Error>>>),
 }
 
 #[derive(Debug)]
@@ -64,8 +64,6 @@
     hat: Mutex<Option<HardwareAuthToken>>,
 }
 
-unsafe impl Sync for AuthRequest {}
-
 impl AuthRequest {
     fn op_auth() -> Arc<Self> {
         Arc::new(Self { state: AuthRequestState::OpAuth, hat: Mutex::new(None) })
@@ -73,7 +71,7 @@
 
     fn timestamped_op_auth(receiver: Receiver<Result<TimeStampToken, Error>>) -> Arc<Self> {
         Arc::new(Self {
-            state: AuthRequestState::TimeStampedOpAuth(receiver),
+            state: AuthRequestState::TimeStampedOpAuth(Mutex::new(receiver)),
             hat: Mutex::new(None),
         })
     }
@@ -82,7 +80,10 @@
         hat: HardwareAuthToken,
         receiver: Receiver<Result<TimeStampToken, Error>>,
     ) -> Arc<Self> {
-        Arc::new(Self { state: AuthRequestState::TimeStamp(receiver), hat: Mutex::new(Some(hat)) })
+        Arc::new(Self {
+            state: AuthRequestState::TimeStamp(Mutex::new(receiver)),
+            hat: Mutex::new(Some(hat)),
+        })
     }
 
     fn add_auth_token(&self, hat: HardwareAuthToken) {
@@ -100,7 +101,11 @@
 
         let tst = match &self.state {
             AuthRequestState::TimeStampedOpAuth(recv) | AuthRequestState::TimeStamp(recv) => {
-                let result = recv.recv().context("In get_auth_tokens: Sender disconnected.")?;
+                let result = recv
+                    .lock()
+                    .unwrap()
+                    .recv()
+                    .context("In get_auth_tokens: Sender disconnected.")?;
                 Some(result.context(ks_err!(
                     "Worker responded with error \
                     from generating timestamp token.",
diff --git a/keystore2/src/id_rotation.rs b/keystore2/src/id_rotation.rs
index 460caa7..5904047 100644
--- a/keystore2/src/id_rotation.rs
+++ b/keystore2/src/id_rotation.rs
@@ -26,7 +26,7 @@
 use std::fs;
 use std::io::ErrorKind;
 use std::path::{Path, PathBuf};
-use std::time::Duration;
+use std::time::{Duration, SystemTime};
 
 const ID_ROTATION_PERIOD: Duration = Duration::from_secs(30 * 24 * 60 * 60); // Thirty days.
 static TIMESTAMP_FILE_NAME: &str = "timestamp";
@@ -36,6 +36,8 @@
 /// and passed down to the users of the feature which can then query the timestamp on demand.
 #[derive(Debug, Clone)]
 pub struct IdRotationState {
+    /// We consider the time of last factory reset to be the point in time when this timestamp file
+    /// is created.
     timestamp_path: PathBuf,
 }
 
@@ -47,25 +49,41 @@
         Self { timestamp_path }
     }
 
-    /// Reads the metadata of or creates the timestamp file. It returns true if the timestamp
-    /// file is younger than `ID_ROTATION_PERIOD`, i.e., 30 days.
-    pub fn had_factory_reset_since_id_rotation(&self) -> Result<bool> {
+    /// Returns true iff a factory reset has occurred since the last ID rotation.
+    pub fn had_factory_reset_since_id_rotation(
+        &self,
+        creation_datetime: &SystemTime,
+    ) -> Result<bool> {
         match fs::metadata(&self.timestamp_path) {
             Ok(metadata) => {
-                let duration_since_factory_reset = metadata
-                    .modified()
-                    .context("File creation time not supported.")?
-                    .elapsed()
-                    .context("Failed to compute time elapsed since factory reset.")?;
-                Ok(duration_since_factory_reset < ID_ROTATION_PERIOD)
+                // For Tag::UNIQUE_ID, temporal counter value is defined as Tag::CREATION_DATETIME
+                // divided by 2592000000, dropping any remainder. Temporal counter value is
+                // effectively the index of the ID rotation period that we are currently in, with
+                // each ID rotation period being 30 days.
+                let temporal_counter_value = creation_datetime
+                    .duration_since(SystemTime::UNIX_EPOCH)
+                    .context(ks_err!("Failed to get epoch time"))?
+                    .as_millis()
+                    / ID_ROTATION_PERIOD.as_millis();
+
+                // Calculate the beginning of the current ID rotation period, which is also the
+                // last time ID was rotated.
+                let id_rotation_time: SystemTime = SystemTime::UNIX_EPOCH
+                    .checked_add(ID_ROTATION_PERIOD * temporal_counter_value.try_into()?)
+                    .context(ks_err!("Failed to get ID rotation time."))?;
+
+                let factory_reset_time =
+                    metadata.modified().context(ks_err!("File creation time not supported."))?;
+
+                Ok(id_rotation_time <= factory_reset_time)
             }
             Err(e) => match e.kind() {
                 ErrorKind::NotFound => {
                     fs::File::create(&self.timestamp_path)
-                        .context("Failed to create timestamp file.")?;
+                        .context(ks_err!("Failed to create timestamp file."))?;
                     Ok(true)
                 }
-                _ => Err(e).context("Failed to open timestamp file."),
+                _ => Err(e).context(ks_err!("Failed to open timestamp file.")),
             },
         }
         .context(ks_err!())
@@ -78,47 +96,75 @@
     use keystore2_test_utils::TempDir;
     use nix::sys::stat::utimes;
     use nix::sys::time::{TimeVal, TimeValLike};
-    use std::convert::TryInto;
-    use std::time::UNIX_EPOCH;
+    use std::thread::sleep;
 
-    #[test]
-    fn test_had_factory_reset_since_id_rotation() -> Result<()> {
-        let temp_dir = TempDir::new("test_had_factory_reset_since_id_rotation_")
-            .expect("Failed to create temp dir.");
+    static TEMP_DIR_NAME: &str = "test_had_factory_reset_since_id_rotation_";
+
+    fn set_up() -> (TempDir, PathBuf, IdRotationState) {
+        let temp_dir = TempDir::new(TEMP_DIR_NAME).expect("Failed to create temp dir.");
+        let mut timestamp_file_path = temp_dir.path().to_owned();
+        timestamp_file_path.push(TIMESTAMP_FILE_NAME);
         let id_rotation_state = IdRotationState::new(temp_dir.path());
 
-        let mut temp_file_path = temp_dir.path().to_owned();
-        temp_file_path.push(TIMESTAMP_FILE_NAME);
+        (temp_dir, timestamp_file_path, id_rotation_state)
+    }
+
+    #[test]
+    fn test_timestamp_creation() {
+        let (_temp_dir, timestamp_file_path, id_rotation_state) = set_up();
+        let creation_datetime = SystemTime::now();
 
         // The timestamp file should not exist.
-        assert!(!temp_file_path.exists());
+        assert!(!timestamp_file_path.exists());
 
-        // This should return true.
-        assert!(id_rotation_state.had_factory_reset_since_id_rotation()?);
+        // Trigger timestamp file creation one second later.
+        sleep(Duration::new(1, 0));
+        assert!(id_rotation_state.had_factory_reset_since_id_rotation(&creation_datetime).unwrap());
 
         // Now the timestamp file should exist.
-        assert!(temp_file_path.exists());
+        assert!(timestamp_file_path.exists());
 
-        // We should still return true because the timestamp file is young.
-        assert!(id_rotation_state.had_factory_reset_since_id_rotation()?);
+        let metadata = fs::metadata(&timestamp_file_path).unwrap();
+        assert!(metadata.modified().unwrap() > creation_datetime);
+    }
 
-        // Now let's age the timestamp file by backdating the modification time.
-        let metadata = fs::metadata(&temp_file_path)?;
-        let mtime = metadata.modified()?;
-        let mtime = mtime.duration_since(UNIX_EPOCH)?;
-        let mtime =
-            mtime.checked_sub(ID_ROTATION_PERIOD).expect("Failed to subtract id rotation period");
-        let mtime = TimeVal::seconds(mtime.as_secs().try_into().unwrap());
+    #[test]
+    fn test_existing_timestamp() {
+        let (_temp_dir, timestamp_file_path, id_rotation_state) = set_up();
 
-        let atime = metadata.accessed()?;
-        let atime = atime.duration_since(UNIX_EPOCH)?;
-        let atime = TimeVal::seconds(atime.as_secs().try_into().unwrap());
+        // Let's start with at a known point in time, so that it's easier to control which ID
+        // rotation period we're in.
+        let mut creation_datetime = SystemTime::UNIX_EPOCH;
 
-        utimes(&temp_file_path, &atime, &mtime)?;
+        // Create timestamp file and backdate it back to Unix epoch.
+        fs::File::create(&timestamp_file_path).unwrap();
+        let mtime = TimeVal::seconds(0);
+        let atime = TimeVal::seconds(0);
+        utimes(&timestamp_file_path, &atime, &mtime).unwrap();
 
-        // Now that the file has aged we should see false.
-        assert!(!id_rotation_state.had_factory_reset_since_id_rotation()?);
+        // Timestamp file was backdated to the very beginning of the current ID rotation period.
+        // So, this should return true.
+        assert!(id_rotation_state.had_factory_reset_since_id_rotation(&creation_datetime).unwrap());
 
-        Ok(())
+        // Move time forward, but stay in the same ID rotation period.
+        creation_datetime += Duration::from_millis(1);
+
+        // We should still return true because we're in the same ID rotation period.
+        assert!(id_rotation_state.had_factory_reset_since_id_rotation(&creation_datetime).unwrap());
+
+        // Move time to the next ID rotation period.
+        creation_datetime += ID_ROTATION_PERIOD;
+
+        // Now we should see false.
+        assert!(!id_rotation_state
+            .had_factory_reset_since_id_rotation(&creation_datetime)
+            .unwrap());
+
+        // Move timestamp to the future. This shouldn't ever happen, but even in this edge case ID
+        // must be rotated.
+        let mtime = TimeVal::seconds((ID_ROTATION_PERIOD.as_secs() * 10).try_into().unwrap());
+        let atime = TimeVal::seconds((ID_ROTATION_PERIOD.as_secs() * 10).try_into().unwrap());
+        utimes(&timestamp_file_path, &atime, &mtime).unwrap();
+        assert!(id_rotation_state.had_factory_reset_since_id_rotation(&creation_datetime).unwrap());
     }
 }
diff --git a/keystore2/src/keystore2_main.rs b/keystore2/src/keystore2_main.rs
index 31c1e29..059d59d 100644
--- a/keystore2/src/keystore2_main.rs
+++ b/keystore2/src/keystore2_main.rs
@@ -68,6 +68,8 @@
     fn sqlite_log_handler(err: c_int, message: &str) {
         log::error!("[SQLITE3] {}: {}", err, message);
     }
+    // SAFETY: There are no other threads yet, `sqlite_log_handler` is threadsafe, and it doesn't
+    // invoke any SQLite calls.
     unsafe { sqlite_trace::config_log(Some(sqlite_log_handler)) }
         .expect("Error setting sqlite log callback.");
 
diff --git a/keystore2/src/km_compat/lib.rs b/keystore2/src/km_compat/lib.rs
index 2632ec4..e61a13a 100644
--- a/keystore2/src/km_compat/lib.rs
+++ b/keystore2/src/km_compat/lib.rs
@@ -21,6 +21,7 @@
 
 #[allow(missing_docs)] // TODO remove this
 pub fn add_keymint_device_service() -> i32 {
+    // SAFETY: This is always safe to call.
     unsafe { addKeyMintDeviceService() }
 }
 
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 5eed37c..db44d4b 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -405,8 +405,7 @@
     ) -> Result<Vec<KeyParameter>> {
         let mut result = params.to_vec();
 
-        // Unconditionally add the CREATION_DATETIME tag and prevent callers from
-        // specifying it.
+        // Prevent callers from specifying the CREATION_DATETIME tag.
         if params.iter().any(|kp| kp.tag == Tag::CREATION_DATETIME) {
             return Err(Error::Rc(ResponseCode::INVALID_ARGUMENT)).context(ks_err!(
                 "KeystoreSecurityLevel::add_required_parameters: \
@@ -414,12 +413,16 @@
             ));
         }
 
+        // Use this variable to refer to notion of "now". This eliminates discrepancies from
+        // quering the clock multiple times.
+        let creation_datetime = SystemTime::now();
+
         // Add CREATION_DATETIME only if the backend version Keymint V1 (100) or newer.
         if self.hw_info.versionNumber >= 100 {
             result.push(KeyParameter {
                 tag: Tag::CREATION_DATETIME,
                 value: KeyParameterValue::DateTime(
-                    SystemTime::now()
+                    creation_datetime
                         .duration_since(SystemTime::UNIX_EPOCH)
                         .context(ks_err!(
                             "KeystoreSecurityLevel::add_required_parameters: \
@@ -462,7 +465,7 @@
             }
             if self
                 .id_rotation_state
-                .had_factory_reset_since_id_rotation()
+                .had_factory_reset_since_id_rotation(&creation_datetime)
                 .context(ks_err!("Call to had_factory_reset_since_id_rotation failed."))?
             {
                 result.push(KeyParameter {
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index acac7ee..584a51c 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -215,8 +215,8 @@
 /// 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.
+    // SAFETY: The pointer is valid because it comes from a reference, and clock_gettime doesn't
+    // retain it beyond the call.
     unsafe { libc::clock_gettime(libc::CLOCK_MONOTONIC_RAW, &mut current_time) };
     current_time.tv_sec as i64 * 1000 + (current_time.tv_nsec as i64 / 1_000_000)
 }
diff --git a/keystore2/test_utils/key_generations.rs b/keystore2/test_utils/key_generations.rs
index 0a1ffb1..0f9ecbe 100644
--- a/keystore2/test_utils/key_generations.rs
+++ b/keystore2/test_utils/key_generations.rs
@@ -46,6 +46,52 @@
 /// Vold context
 pub const TARGET_VOLD_CTX: &str = "u:r:vold:s0";
 
+/// Allowed tags in generated/imported key authorizations.
+/// See hardware/interfaces/security/keymint/aidl/android/hardware/security/keymint/Tag.aidl for the
+/// list feature tags.
+/// Note: This list need to be updated whenever a new Tag is introduced and is expected to be added
+/// in key authorizations.
+pub const ALLOWED_TAGS_IN_KEY_AUTHS: &[Tag] = &[
+    Tag::ACTIVE_DATETIME,
+    Tag::ALGORITHM,
+    Tag::ALLOW_WHILE_ON_BODY,
+    Tag::AUTH_TIMEOUT,
+    Tag::BLOCK_MODE,
+    Tag::BOOTLOADER_ONLY,
+    Tag::BOOT_PATCHLEVEL,
+    Tag::CALLER_NONCE,
+    Tag::CREATION_DATETIME,
+    Tag::DIGEST,
+    Tag::EARLY_BOOT_ONLY,
+    Tag::EC_CURVE,
+    Tag::IDENTITY_CREDENTIAL_KEY,
+    Tag::INCLUDE_UNIQUE_ID,
+    Tag::KEY_SIZE,
+    Tag::MAX_BOOT_LEVEL,
+    Tag::MAX_USES_PER_BOOT,
+    Tag::MIN_MAC_LENGTH,
+    Tag::NO_AUTH_REQUIRED,
+    Tag::ORIGIN,
+    Tag::ORIGINATION_EXPIRE_DATETIME,
+    Tag::OS_PATCHLEVEL,
+    Tag::OS_VERSION,
+    Tag::PADDING,
+    Tag::PURPOSE,
+    Tag::ROLLBACK_RESISTANCE,
+    Tag::RSA_OAEP_MGF_DIGEST,
+    Tag::RSA_PUBLIC_EXPONENT,
+    Tag::STORAGE_KEY,
+    Tag::TRUSTED_CONFIRMATION_REQUIRED,
+    Tag::TRUSTED_USER_PRESENCE_REQUIRED,
+    Tag::UNLOCKED_DEVICE_REQUIRED,
+    Tag::USAGE_COUNT_LIMIT,
+    Tag::USAGE_EXPIRE_DATETIME,
+    Tag::USER_AUTH_TYPE,
+    Tag::USER_ID,
+    Tag::USER_SECURE_ID,
+    Tag::VENDOR_PATCHLEVEL,
+];
+
 /// Key parameters to generate a key.
 pub struct KeyParams {
     /// Key Size.
@@ -337,6 +383,35 @@
     })
 }
 
+/// Verify that given key param is listed in given authorizations list.
+pub fn check_key_param(authorizations: &[Authorization], key_param: &KeyParameter) -> bool {
+    authorizations.iter().any(|auth| &auth.keyParameter == key_param)
+}
+
+fn check_key_authorizations(authorizations: &[Authorization], expected_params: &[KeyParameter]) {
+    // Make sure key authorizations contains only `ALLOWED_TAGS_IN_KEY_AUTHS`
+    authorizations.iter().all(|auth| {
+        assert!(
+            ALLOWED_TAGS_IN_KEY_AUTHS.contains(&auth.keyParameter.tag),
+            "key authorization is not allowed: {:#?}",
+            auth.keyParameter
+        );
+        true
+    });
+
+    //Check allowed-expected-key-parameters are present in given key authorizations list.
+    expected_params.iter().all(|key_param| {
+        if ALLOWED_TAGS_IN_KEY_AUTHS.contains(&key_param.tag) {
+            assert!(
+                check_key_param(authorizations, key_param),
+                "Key parameter not found: {:#?}",
+                key_param
+            );
+        }
+        true
+    });
+}
+
 /// Generate EC Key using given security level and domain with below key parameters and
 /// optionally allow the generated key to be attested with factory provisioned attest key using
 /// given challenge and application id -
@@ -380,6 +455,7 @@
                 assert!(key_metadata.key.blob.is_some());
             }
 
+            check_key_authorizations(&key_metadata.authorizations, &gen_params);
             Ok(key_metadata)
         }
         Err(e) => Err(e),
@@ -422,6 +498,7 @@
     } else {
         assert!(key_metadata.key.blob.is_none());
     }
+    check_key_authorizations(&key_metadata.authorizations, &gen_params);
     Ok(key_metadata)
 }
 
@@ -483,6 +560,19 @@
             || key_metadata.key.blob.is_none()
     );
 
+    check_key_authorizations(&key_metadata.authorizations, &gen_params);
+    // If `RSA_OAEP_MGF_DIGEST` tag is not mentioned explicitly while generating/importing a key,
+    // then make sure `RSA_OAEP_MGF_DIGEST` tag with default value (SHA1) must not be included in
+    // key authorization list.
+    if key_params.mgf_digest.is_none() {
+        assert!(!check_key_param(
+            &key_metadata.authorizations,
+            &KeyParameter {
+                tag: Tag::RSA_OAEP_MGF_DIGEST,
+                value: KeyParameterValue::Digest(Digest::SHA1)
+            }
+        ));
+    }
     Ok(key_metadata)
 }
 
@@ -527,6 +617,7 @@
 
     // Should not have an attestation record.
     assert!(key_metadata.certificateChain.is_none());
+    check_key_authorizations(&key_metadata.authorizations, &gen_params);
     Ok(key_metadata)
 }
 
@@ -566,6 +657,7 @@
     // Should not have an attestation record.
     assert!(key_metadata.certificateChain.is_none());
 
+    check_key_authorizations(&key_metadata.authorizations, &gen_params);
     Ok(key_metadata)
 }
 
@@ -650,6 +742,7 @@
     // Should have an attestation record.
     assert!(attestation_key_metadata.certificateChain.is_some());
 
+    check_key_authorizations(&attestation_key_metadata.authorizations, &gen_params);
     Ok(attestation_key_metadata)
 }
 
@@ -684,20 +777,10 @@
     // Shouldn't have an attestation record.
     assert!(ec_key_metadata.certificateChain.is_none());
 
+    check_key_authorizations(&ec_key_metadata.authorizations, &ec_gen_params);
     Ok(ec_key_metadata)
 }
 
-/// Verify that given key param is listed in given authorizations list.
-pub fn check_key_param(authorizations: &[Authorization], key_param: KeyParameter) -> bool {
-    for authrization in authorizations {
-        if authrization.keyParameter == key_param {
-            return true;
-        }
-    }
-
-    false
-}
-
 /// Imports above defined RSA key - `RSA_2048_KEY` and validates imported key parameters.
 pub fn import_rsa_2048_key(
     sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
@@ -719,24 +802,27 @@
     assert!(key_metadata.certificate.is_some());
     assert!(key_metadata.certificateChain.is_none());
 
+    check_key_authorizations(&key_metadata.authorizations, &import_params);
+
+    // Check below auths explicitly, they might not be addd in import parameters.
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::RSA) }
+        &KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::RSA) }
     ));
 
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(2048) }
+        &KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(2048) }
     ));
 
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::DIGEST, value: KeyParameterValue::Digest(Digest::SHA_2_256) }
+        &KeyParameter { tag: Tag::DIGEST, value: KeyParameterValue::Digest(Digest::SHA_2_256) }
     ));
 
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter {
+        &KeyParameter {
             tag: Tag::RSA_PUBLIC_EXPONENT,
             value: KeyParameterValue::LongInteger(65537)
         }
@@ -744,7 +830,7 @@
 
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter {
+        &KeyParameter {
             tag: Tag::PADDING,
             value: KeyParameterValue::PaddingMode(PaddingMode::RSA_PSS)
         }
@@ -752,7 +838,7 @@
 
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
+        &KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
     ));
 
     Ok(key_metadata)
@@ -779,23 +865,26 @@
     assert!(key_metadata.certificate.is_some());
     assert!(key_metadata.certificateChain.is_none());
 
+    check_key_authorizations(&key_metadata.authorizations, &import_params);
+
+    // Check below auths explicitly, they might not be addd in import parameters.
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::EC) }
+        &KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::EC) }
     ));
 
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::EC_CURVE, value: KeyParameterValue::EcCurve(EcCurve::P_256) }
+        &KeyParameter { tag: Tag::EC_CURVE, value: KeyParameterValue::EcCurve(EcCurve::P_256) }
     ));
 
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::DIGEST, value: KeyParameterValue::Digest(Digest::SHA_2_256) }
+        &KeyParameter { tag: Tag::DIGEST, value: KeyParameterValue::Digest(Digest::SHA_2_256) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
+        &KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
     ));
 
     Ok(key_metadata)
@@ -828,28 +917,31 @@
         AES_KEY,
     )?;
 
+    check_key_authorizations(&key_metadata.authorizations, &import_params);
+
+    // Check below auths explicitly, they might not be addd in import parameters.
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::AES) }
+        &KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::AES) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(128) }
+        &KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(128) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter {
+        &KeyParameter {
             tag: Tag::PADDING,
             value: KeyParameterValue::PaddingMode(PaddingMode::PKCS7)
         }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::BLOCK_MODE, value: KeyParameterValue::BlockMode(BlockMode::ECB) }
+        &KeyParameter { tag: Tag::BLOCK_MODE, value: KeyParameterValue::BlockMode(BlockMode::ECB) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
+        &KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
     ));
 
     Ok(key_metadata)
@@ -884,31 +976,34 @@
         TRIPLE_DES_KEY,
     )?;
 
+    check_key_authorizations(&key_metadata.authorizations, &import_params);
+
+    // Check below auths explicitly, they might not be addd in import parameters.
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter {
+        &KeyParameter {
             tag: Tag::ALGORITHM,
             value: KeyParameterValue::Algorithm(Algorithm::TRIPLE_DES)
         }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(168) }
+        &KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(168) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter {
+        &KeyParameter {
             tag: Tag::PADDING,
             value: KeyParameterValue::PaddingMode(PaddingMode::PKCS7)
         }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::BLOCK_MODE, value: KeyParameterValue::BlockMode(BlockMode::ECB) }
+        &KeyParameter { tag: Tag::BLOCK_MODE, value: KeyParameterValue::BlockMode(BlockMode::ECB) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
+        &KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
     ));
 
     Ok(key_metadata)
@@ -941,21 +1036,24 @@
         HMAC_KEY,
     )?;
 
+    check_key_authorizations(&key_metadata.authorizations, &import_params);
+
+    // Check below auths explicitly, they might not be addd in import parameters.
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::HMAC) }
+        &KeyParameter { tag: Tag::ALGORITHM, value: KeyParameterValue::Algorithm(Algorithm::HMAC) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(128) }
+        &KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(128) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::DIGEST, value: KeyParameterValue::Digest(Digest::SHA_2_256) }
+        &KeyParameter { tag: Tag::DIGEST, value: KeyParameterValue::Digest(Digest::SHA_2_256) }
     ));
     assert!(check_key_param(
         &key_metadata.authorizations,
-        KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
+        &KeyParameter { tag: Tag::ORIGIN, value: KeyParameterValue::Origin(KeyOrigin::IMPORTED) }
     ));
 
     Ok(key_metadata)
@@ -1090,6 +1188,7 @@
                 assert!(key_metadata.key.blob.is_some());
             }
 
+            check_key_authorizations(&key_metadata.authorizations, &gen_params);
             Ok(key_metadata)
         }
         Err(e) => Err(e),
diff --git a/keystore2/test_utils/run_as.rs b/keystore2/test_utils/run_as.rs
index 2485ab5..be643b6 100644
--- a/keystore2/test_utils/run_as.rs
+++ b/keystore2/test_utils/run_as.rs
@@ -255,7 +255,9 @@
     let (response_reader, mut response_writer) =
         pipe_channel().expect("Failed to create cmd pipe.");
 
-    match fork() {
+    // SAFETY: Our caller guarantees that the process only has a single thread, so calling
+    // non-async-signal-safe functions in the child is in fact safe.
+    match unsafe { fork() } {
         Ok(ForkResult::Parent { child, .. }) => {
             drop(response_writer);
             drop(cmd_reader);
@@ -314,7 +316,9 @@
         selinux::Context::new(se_context).expect("Unable to construct selinux::Context.");
     let (mut reader, mut writer) = pipe_channel::<R>().expect("Failed to create pipe.");
 
-    match fork() {
+    // SAFETY: Our caller guarantees that the process only has a single thread, so calling
+    // non-async-signal-safe functions in the child is in fact safe.
+    match unsafe { fork() } {
         Ok(ForkResult::Parent { child, .. }) => {
             drop(writer);
             let status = waitpid(child, None).expect("Failed while waiting for child.");
diff --git a/keystore2/tests/keystore2_client_ec_key_tests.rs b/keystore2/tests/keystore2_client_ec_key_tests.rs
index c2034de..8267140 100644
--- a/keystore2/tests/keystore2_client_ec_key_tests.rs
+++ b/keystore2/tests/keystore2_client_ec_key_tests.rs
@@ -432,15 +432,18 @@
     // Client#1: Generate a key and create an operation using generated key.
     // Wait until the parent notifies to continue. Once the parent notifies, this operation
     // is expected to be completed successfully.
-    let mut child_handle = execute_op_run_as_child(
-        TARGET_CTX,
-        Domain::APP,
-        -1,
-        Some(alias.to_string()),
-        Uid::from_raw(uid1),
-        Gid::from_raw(gid1),
-        ForcedOp(false),
-    );
+    // SAFETY: The test is run in a separate process with no other threads.
+    let mut child_handle = unsafe {
+        execute_op_run_as_child(
+            TARGET_CTX,
+            Domain::APP,
+            -1,
+            Some(alias.to_string()),
+            Uid::from_raw(uid1),
+            Gid::from_raw(gid1),
+            ForcedOp(false),
+        )
+    };
 
     // Wait until (client#1) child process notifies us to continue, so that there will be a key
     // generated by client#1.
@@ -450,6 +453,7 @@
     const APPLICATION_ID_2: u32 = 10602;
     let uid2 = USER_ID * AID_USER_OFFSET + APPLICATION_ID_2;
     let gid2 = USER_ID * AID_USER_OFFSET + APPLICATION_ID_2;
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(TARGET_CTX, Uid::from_raw(uid2), Gid::from_raw(gid2), move || {
             let keystore2_inst = get_keystore_service();
diff --git a/keystore2/tests/keystore2_client_grant_key_tests.rs b/keystore2/tests/keystore2_client_grant_key_tests.rs
index bde872d..516869a 100644
--- a/keystore2/tests/keystore2_client_grant_key_tests.rs
+++ b/keystore2/tests/keystore2_client_grant_key_tests.rs
@@ -114,6 +114,7 @@
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     static GRANTEE_GID: u32 = GRANTEE_UID;
 
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(TARGET_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let empty_access_vector = KeyPermission::NONE.0;
@@ -132,6 +133,7 @@
 
     // In grantee context try to load the key, it should fail to load the granted key as it is
     // granted with empty access vector.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -169,6 +171,7 @@
     static GRANTEE_GID: u32 = GRANTEE_UID;
 
     // Generate a key and grant it to a user with GET_INFO|USE key permissions.
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(TARGET_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let access_vector = KeyPermission::GET_INFO.0 | KeyPermission::USE.0;
@@ -185,6 +188,7 @@
     };
 
     // In grantee context load the key and try to perform crypto operation.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -251,6 +255,7 @@
     static ALIAS: &str = "ks_grant_key_delete_success";
 
     // Generate a key and grant it to a user with DELETE permission.
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -270,6 +275,7 @@
     };
 
     // Grantee context, delete the key.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -290,6 +296,7 @@
     };
 
     // Verify whether key got deleted in grantor's context.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), move || {
             let keystore2_inst = get_keystore_service();
@@ -325,6 +332,7 @@
     static SEC_GRANTEE_GID: u32 = SEC_GRANTEE_UID;
 
     // Generate a key and grant it to a user with GET_INFO permission.
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -345,6 +353,7 @@
     };
 
     // Grantee context, load the granted key and try to grant it to `SEC_GRANTEE_UID` grantee.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -375,6 +384,7 @@
     };
 
     // Make sure second grantee shouldn't have access to the above granted key.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -457,6 +467,7 @@
     static GRANTEE_GID: u32 = GRANTEE_UID;
 
     // Generate a key and grant it to a user with GET_INFO permission.
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -492,6 +503,7 @@
     };
 
     // Grantee context, try to load the ungranted key.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -527,6 +539,7 @@
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     static GRANTEE_GID: u32 = GRANTEE_UID;
 
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -576,6 +589,7 @@
     // Make sure grant did not persist, try to access the earlier granted key in grantee context.
     // Grantee context should fail to load the granted key as its associated key is deleted in
     // grantor context.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -614,6 +628,7 @@
     static GRANTEE_2_GID: u32 = GRANTEE_2_UID;
 
     // Generate a key and grant it to multiple users with GET_INFO|USE permissions.
+    // SAFETY: The test is run in a separate process with no other threads.
     let mut grant_keys = unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -636,6 +651,7 @@
         &[(GRANTEE_1_UID, GRANTEE_1_GID), (GRANTEE_2_UID, GRANTEE_2_GID)]
     {
         let grant_key_nspace = grant_keys.remove(0);
+        // SAFETY: The test is run in a separate process with no other threads.
         unsafe {
             run_as::run_as(
                 GRANTEE_CTX,
@@ -678,6 +694,7 @@
     static GRANTEE_2_GID: u32 = GRANTEE_2_UID;
 
     // Generate a key and grant it to multiple users with GET_INFO permission.
+    // SAFETY: The test is run in a separate process with no other threads.
     let mut grant_keys = unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -699,6 +716,7 @@
 
     // Grantee #1 context
     let grant_key1_nspace = grant_keys.remove(0);
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -733,6 +751,7 @@
 
     // Grantee #2 context
     let grant_key2_nspace = grant_keys.remove(0);
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
diff --git a/keystore2/tests/keystore2_client_keystore_engine_tests.rs b/keystore2/tests/keystore2_client_keystore_engine_tests.rs
index 1aed8e6..339eb60 100644
--- a/keystore2/tests/keystore2_client_keystore_engine_tests.rs
+++ b/keystore2/tests/keystore2_client_keystore_engine_tests.rs
@@ -167,6 +167,7 @@
     static GRANTEE_GID: u32 = GRANTEE_UID;
 
     // Generate a key and grant it to a user with GET_INFO|USE|DELETE key permissions.
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(TARGET_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -184,6 +185,7 @@
     };
 
     // In grantee context load the key and try to perform crypto operation.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -208,6 +210,7 @@
     static GRANTEE_GID: u32 = GRANTEE_UID;
 
     // Generate a key and grant it to a user with GET_INFO|USE|DELETE key permissions.
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(TARGET_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -225,6 +228,7 @@
     };
 
     // In grantee context load the key and try to perform crypto operation.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -250,6 +254,7 @@
 
     // Generate a key and re-encode it's certificate as PEM and update it and
     // grant it to a user with GET_INFO|USE|DELETE key permissions.
+    // SAFETY: The test is run in a separate process with no other threads.
     let grant_key_nspace = unsafe {
         run_as::run_as(TARGET_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -285,6 +290,7 @@
     };
 
     // In grantee context load the key and try to perform crypto operation.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
diff --git a/keystore2/tests/keystore2_client_list_entries_tests.rs b/keystore2/tests/keystore2_client_list_entries_tests.rs
index 809c01f..1c0c496 100644
--- a/keystore2/tests/keystore2_client_list_entries_tests.rs
+++ b/keystore2/tests/keystore2_client_list_entries_tests.rs
@@ -60,6 +60,7 @@
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     static GRANTEE_GID: u32 = GRANTEE_UID;
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -113,6 +114,7 @@
     };
 
     // In user context validate list of key entries associated with it.
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -161,6 +163,7 @@
     let agid = 91 * AID_USER_OFFSET + 10001;
     static TARGET_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(TARGET_CTX, Uid::from_raw(auid), Gid::from_raw(agid), move || {
             let keystore2 = get_keystore_service();
@@ -198,6 +201,7 @@
     static CLIENT_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     static CLIENT_GID: u32 = CLIENT_UID;
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(CLIENT_CTX, Uid::from_raw(CLIENT_UID), Gid::from_raw(CLIENT_GID), || {
             let keystore2 = get_keystore_service();
@@ -264,6 +268,7 @@
     static CLIENT_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     static CLIENT_GID: u32 = CLIENT_UID;
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(CLIENT_CTX, Uid::from_raw(CLIENT_UID), Gid::from_raw(CLIENT_GID), || {
             let keystore2 = get_keystore_service();
@@ -331,6 +336,7 @@
     static CLIENT_GID: u32 = CLIENT_UID;
     static ALIAS_PREFIX: &str = "key_test_batch_list";
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(CLIENT_CTX, Uid::from_raw(CLIENT_UID), Gid::from_raw(CLIENT_GID), || {
             let keystore2 = get_keystore_service();
@@ -361,6 +367,7 @@
         })
     };
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(CLIENT_CTX, Uid::from_raw(CLIENT_UID), Gid::from_raw(CLIENT_GID), || {
             let keystore2 = get_keystore_service();
@@ -428,6 +435,7 @@
     static CLIENT_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     static CLIENT_GID: u32 = CLIENT_UID;
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(CLIENT_CTX, Uid::from_raw(CLIENT_UID), Gid::from_raw(CLIENT_GID), || {
             let keystore2 = get_keystore_service();
@@ -511,6 +519,7 @@
     static CLIENT_GID: u32 = CLIENT_UID;
     static ALIAS_PREFIX: &str = "key_test_batch_list";
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(CLIENT_CTX, Uid::from_raw(CLIENT_UID), Gid::from_raw(CLIENT_GID), || {
             let keystore2 = get_keystore_service();
@@ -647,6 +656,7 @@
     let agid = 91 * AID_USER_OFFSET + 10001;
     static TARGET_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(TARGET_CTX, Uid::from_raw(auid), Gid::from_raw(agid), move || {
             let keystore2 = get_keystore_service();
@@ -686,6 +696,7 @@
     let agid = 91 * AID_USER_OFFSET + 10001;
     static TARGET_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(TARGET_CTX, Uid::from_raw(auid), Gid::from_raw(agid), move || {
             let keystore2 = get_keystore_service();
diff --git a/keystore2/tests/keystore2_client_operation_tests.rs b/keystore2/tests/keystore2_client_operation_tests.rs
index 19175dd..89b5a31 100644
--- a/keystore2/tests/keystore2_client_operation_tests.rs
+++ b/keystore2/tests/keystore2_client_operation_tests.rs
@@ -36,7 +36,11 @@
 
 /// Create `max_ops` number child processes with the given context and perform an operation under each
 /// child process.
-pub fn create_operations(
+///
+/// # Safety
+///
+/// Must be called from a process with no other threads.
+pub unsafe fn create_operations(
     target_ctx: &'static str,
     forced_op: ForcedOp,
     max_ops: i32,
@@ -45,7 +49,8 @@
     let base_gid = 99 * AID_USER_OFFSET + 10001;
     let base_uid = 99 * AID_USER_OFFSET + 10001;
     (0..max_ops)
-        .map(|i| {
+        // SAFETY: The caller guarantees that there are no other threads.
+        .map(|i| unsafe {
             execute_op_run_as_child(
                 target_ctx,
                 Domain::APP,
@@ -87,7 +92,8 @@
     const MAX_OPS: i32 = 100;
     static TARGET_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
 
-    let mut child_handles = create_operations(TARGET_CTX, ForcedOp(false), MAX_OPS);
+    // SAFETY: The test is run in a separate process with no other threads.
+    let mut child_handles = unsafe { create_operations(TARGET_CTX, ForcedOp(false), MAX_OPS) };
 
     // Wait until all child procs notifies us to continue,
     // so that there are definitely enough operations outstanding to trigger a BACKEND_BUSY.
@@ -120,7 +126,8 @@
     static TARGET_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
 
     // Create regular operations.
-    let mut child_handles = create_operations(TARGET_CTX, ForcedOp(false), MAX_OPS);
+    // SAFETY: The test is run in a separate process with no other threads.
+    let mut child_handles = unsafe { create_operations(TARGET_CTX, ForcedOp(false), MAX_OPS) };
 
     // Wait until all child procs notifies us to continue, so that there are enough
     // operations outstanding to trigger a BACKEND_BUSY.
@@ -131,6 +138,7 @@
     // Create a forced operation.
     let auid = 99 * AID_USER_OFFSET + 10604;
     let agid = 99 * AID_USER_OFFSET + 10604;
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             key_generations::TARGET_VOLD_CTX,
@@ -203,15 +211,18 @@
     // Create initial forced operation in a child process
     // and wait for the parent to notify to perform operation.
     let alias = format!("ks_forced_op_key_{}", getuid());
-    let mut first_op_handle = execute_op_run_as_child(
-        key_generations::TARGET_SU_CTX,
-        Domain::SELINUX,
-        key_generations::SELINUX_SHELL_NAMESPACE,
-        Some(alias),
-        Uid::from_raw(auid),
-        Gid::from_raw(agid),
-        ForcedOp(true),
-    );
+    // SAFETY: The test is run in a separate process with no other threads.
+    let mut first_op_handle = unsafe {
+        execute_op_run_as_child(
+            key_generations::TARGET_SU_CTX,
+            Domain::SELINUX,
+            key_generations::SELINUX_SHELL_NAMESPACE,
+            Some(alias),
+            Uid::from_raw(auid),
+            Gid::from_raw(agid),
+            ForcedOp(true),
+        )
+    };
 
     // Wait until above child proc notifies us to continue, so that there is definitely a forced
     // operation outstanding to perform a operation.
@@ -219,7 +230,8 @@
 
     // Create MAX_OPS number of forced operations.
     let mut child_handles =
-        create_operations(key_generations::TARGET_SU_CTX, ForcedOp(true), MAX_OPS);
+    // SAFETY: The test is run in a separate process with no other threads.
+        unsafe { create_operations(key_generations::TARGET_SU_CTX, ForcedOp(true), MAX_OPS) };
 
     // Wait until all child procs notifies us to continue, so that  there are enough operations
     // outstanding to trigger a BACKEND_BUSY.
@@ -282,15 +294,18 @@
     // Create an operation in an untrusted_app context. Wait until the parent notifies to continue.
     // Once the parent notifies, this operation is expected to be completed successfully.
     let alias = format!("ks_reg_op_key_{}", getuid());
-    let mut child_handle = execute_op_run_as_child(
-        TARGET_CTX,
-        Domain::APP,
-        -1,
-        Some(alias),
-        Uid::from_raw(uid),
-        Gid::from_raw(gid),
-        ForcedOp(false),
-    );
+    // SAFETY: The test is run in a separate process with no other threads.
+    let mut child_handle = unsafe {
+        execute_op_run_as_child(
+            TARGET_CTX,
+            Domain::APP,
+            -1,
+            Some(alias),
+            Uid::from_raw(uid),
+            Gid::from_raw(gid),
+            ForcedOp(false),
+        )
+    };
 
     // Wait until child process notifies us to continue, so that an operation from child process is
     // outstanding to complete the operation.
@@ -377,6 +392,7 @@
     let gid = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
 
     for context in TARGET_CTXS.iter() {
+        // SAFETY: The test is run in a separate process with no other threads.
         unsafe {
             run_as::run_as(context, Uid::from_raw(uid), Gid::from_raw(gid), move || {
                 let alias = format!("ks_app_forced_op_test_key_{}", getuid());
@@ -406,6 +422,7 @@
     let uid = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     let gid = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
 
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(TARGET_CTX, Uid::from_raw(uid), Gid::from_raw(gid), move || {
             let alias = format!("ks_vold_forced_op_key_{}", getuid());
diff --git a/keystore2/tests/keystore2_client_test_utils.rs b/keystore2/tests/keystore2_client_test_utils.rs
index 3354798..f7e7985 100644
--- a/keystore2/tests/keystore2_client_test_utils.rs
+++ b/keystore2/tests/keystore2_client_test_utils.rs
@@ -266,7 +266,11 @@
 }
 
 /// Create new operation on child proc and perform simple operation after parent notification.
-pub fn execute_op_run_as_child(
+///
+/// # Safety
+///
+/// Must only be called from a single-threaded process.
+pub unsafe fn execute_op_run_as_child(
     target_ctx: &'static str,
     domain: Domain,
     nspace: i64,
@@ -275,6 +279,7 @@
     agid: Gid,
     forced_op: ForcedOp,
 ) -> run_as::ChildHandle<TestOutcome, BarrierReached> {
+    // SAFETY: The caller guarantees that there are no other threads.
     unsafe {
         run_as::run_as_child(target_ctx, auid, agid, move |reader, writer| {
             let result = key_generations::map_ks_error(create_signing_operation(
diff --git a/keystore2/tests/keystore2_client_update_subcomponent_tests.rs b/keystore2/tests/keystore2_client_update_subcomponent_tests.rs
index 0be092f..d9576a8 100644
--- a/keystore2/tests/keystore2_client_update_subcomponent_tests.rs
+++ b/keystore2/tests/keystore2_client_update_subcomponent_tests.rs
@@ -167,6 +167,7 @@
     static GRANTEE_2_GID: u32 = GRANTEE_2_UID;
 
     // Generate a key and grant it to multiple users with different access permissions.
+    // SAFETY: The test is run in a separate process with no other threads.
     let mut granted_keys = unsafe {
         run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
             let keystore2 = get_keystore_service();
@@ -205,6 +206,7 @@
 
     // Grantee context, try to update the key public certs, permission denied error is expected.
     let granted_key1_nspace = granted_keys.remove(0);
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
@@ -234,6 +236,7 @@
 
     // Grantee context, update granted key public certs. Update should happen successfully.
     let granted_key2_nspace = granted_keys.remove(0);
+    // SAFETY: The test is run in a separate process with no other threads.
     unsafe {
         run_as::run_as(
             GRANTEE_CTX,
diff --git a/prng_seeder/src/cutils_socket.rs b/prng_seeder/src/cutils_socket.rs
index ab2c869..b408be6 100644
--- a/prng_seeder/src/cutils_socket.rs
+++ b/prng_seeder/src/cutils_socket.rs
@@ -19,7 +19,11 @@
 
 pub fn android_get_control_socket(name: &str) -> Result<UnixListener> {
     let name = CString::new(name)?;
+    // SAFETY: name is a valid C string, and android_get_control_socket doesn't retain it after it
+    // returns.
     let fd = unsafe { cutils_socket_bindgen::android_get_control_socket(name.as_ptr()) };
     ensure!(fd >= 0, "android_get_control_socket failed");
+    // SAFETY: android_get_control_socket either returns a valid and open FD or -1, and we checked
+    // that it's not -1.
     Ok(unsafe { UnixListener::from_raw_fd(fd) })
 }
diff --git a/prng_seeder/src/drbg.rs b/prng_seeder/src/drbg.rs
index 89c5a88..808ea18 100644
--- a/prng_seeder/src/drbg.rs
+++ b/prng_seeder/src/drbg.rs
@@ -23,6 +23,9 @@
 
 impl Drbg {
     pub fn new(entropy: &Entropy) -> Result<Drbg> {
+        // SAFETY: entropy must be a valid pointer because it comes from a reference, and a null
+        // pointer is allowed for personalization. CTR_DRBG_new doesn't retain the entropy pointer
+        // for use after it returns.
         let p = unsafe { bssl_sys::CTR_DRBG_new(entropy.as_ptr(), std::ptr::null(), 0) };
         ensure!(!p.is_null(), "CTR_DRBG_new failed");
         Ok(Drbg(p))
@@ -30,6 +33,9 @@
 
     pub fn reseed(&mut self, entropy: &Entropy) -> Result<()> {
         ensure!(
+            // SAFETY: We know that self.0 is valid because it was initialised from CTR_DRBG_new in
+            // Drbg::new above. The entropy pointer must be valid because it comes from a reference,
+            // and CTR_DRBG_reseed doesn't retain it after it returns.
             unsafe { bssl_sys::CTR_DRBG_reseed(self.0, entropy.as_ptr(), std::ptr::null(), 0) }
                 == 1,
             "CTR_DRBG_reseed failed"
@@ -39,6 +45,10 @@
 
     pub fn generate(&mut self, buf: &mut [u8]) -> Result<()> {
         ensure!(
+            // SAFETY: We know that self.0 is valid because it was initialised from CTR_DRBG_new in
+            // Drbg::new above. The out pointer and length must be valid and unaliased because they
+            // come from a mutable slice reference, and CTR_DRBG_generate doesn't retain them after
+            // it returns.
             unsafe {
                 bssl_sys::CTR_DRBG_generate(
                     self.0,
@@ -56,10 +66,13 @@
 
 impl Drop for Drbg {
     fn drop(&mut self) {
+        // SAFETY: We know that self.0 is valid because it was initialised from CTR_DRBG_new in
+        // Drbg::new above, and this is the only place that frees it.
         unsafe {
             bssl_sys::CTR_DRBG_free(self.0);
         }
     }
 }
 
+// SAFETY: CTR_DRBG functions can be called from any thread.
 unsafe impl Send for Drbg {}
diff --git a/prng_seeder/src/main.rs b/prng_seeder/src/main.rs
index 924481a..f8b0c63 100644
--- a/prng_seeder/src/main.rs
+++ b/prng_seeder/src/main.rs
@@ -70,6 +70,7 @@
 fn setup() -> Result<(ConditionerBuilder, UnixListener)> {
     configure_logging()?;
     let cli = Cli::try_parse()?;
+    // SAFETY: Nothing else sets the signal handler, so either it was set here or it is the default.
     unsafe { signal::signal(signal::Signal::SIGPIPE, signal::SigHandler::SigIgn) }
         .context("In setup, setting SIGPIPE to SIG_IGN")?;