Add, standardise or temporarily opt out of safety comments for keystore2.

These will soon be required by a lint.

Some functions were incorrectly marked as safe which were not actually
safe, so I've fixed those too.

Bug: 290018030
Test: m rust
Change-Id: I38df6a8162d430617f123ab1aace38b741458fce
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/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/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/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,