Bug 2062062 - Cache the encryption key in NSSKeyManager

CHANGELOG.md

@@ -18,6 +18,10 @@
 - The `CheckAuthorizationStatus` and `Disconnect` events are now valid from all states except `Uninitialized`.
   In the cases where the failed before, they're now no-ops.
 
+### Logins
+
+- `NSSKeyManager` now caches the encryption key instead of fetching it from NSS on every `encrypt()`/`decrypt()` call. Bulk operations such as `add_many_with_meta()` previously paid at least two NSS token round-trips per record while holding the store mutex, which could stall `shutdown()` past the async shutdown timeout. The cache is dropped whenever the token is found locked again, so primary password re-authentication is unaffected. ([Bug 2062062](https://bugzilla.mozilla.org/show_bug.cgi?id=2062062))
+
 ### Nimbus
 
 - `NimbusClient::get_available_firefox_labs()` now includes detailed debug level logging for each processed lab. ([#7482](https://github.com/mozilla/application-services/pull/7482))

components/logins/src/encryption.rs

@@ -57,6 +57,9 @@ use futures::executor::block_on;
 #[cfg(feature = "keydb")]
 use async_trait::async_trait;
 
+#[cfg(feature = "keydb")]
+use parking_lot::RwLock;
+
 #[cfg(feature = "keydb")]
 use nss_as::assert_initialized as assert_nss_initialized;
 #[cfg(feature = "keydb")]
@@ -200,6 +203,9 @@ pub trait PrimaryPasswordAuthenticator: Send + Sync {
 /// Make sure to initialize NSS using `ensure_initialized_with_profile_dir` before creating a
 /// NSSKeyManager.
 ///
+/// The key is cached after the first retrieval, since fetching it from NSS costs at least one
+/// token round-trip. The cache is dropped whenever the token turns out to be locked again.
+///
 /// # Examples
 /// ```no_run
 /// use async_trait::async_trait;
@@ -234,6 +240,7 @@ pub trait PrimaryPasswordAuthenticator: Send + Sync {
 #[derive(uniffi::Object)]
 pub struct NSSKeyManager {
     primary_password_authenticator: Arc<dyn PrimaryPasswordAuthenticator>,
+    cached_key: RwLock<Option<Vec<u8>>>,
 }
 
 #[cfg(feature = "keydb")]
@@ -247,6 +254,7 @@ impl NSSKeyManager {
         assert_nss_initialized();
         Self {
             primary_password_authenticator,
+            cached_key: RwLock::new(None),
         }
     }
 
@@ -283,6 +291,9 @@ fn api_authenticate_with_primary_password(primary_password: &str) -> ApiResult<b
 impl KeyManager for NSSKeyManager {
     fn get_key(&self) -> ApiResult<Vec<u8>> {
         if api_authentication_with_primary_password_is_needed()? {
+            // The token locked again since we cached the key, so the cached copy must go.
+            *self.cached_key.write() = None;
+
             let primary_password =
                 block_on(self.primary_password_authenticator.get_primary_password())?;
             let mut result = api_authenticate_with_primary_password(&primary_password)?;
@@ -310,13 +321,19 @@ impl KeyManager for NSSKeyManager {
             }
         }
 
+        let cached = self.cached_key.read().clone();
+        if let Some(bytes) = cached {
+            return Ok(bytes);
+        }
+
         let key = get_or_create_aes256_key(KEY_NAME).map_err(|_| LoginsApiError::MissingKey)?;
         let mut bytes: Vec<u8> = Vec::new();
         serde_json::to_writer(
             &mut bytes,
             &jwcrypto::Jwk::new_direct_from_bytes(None, &key),
         )
         .unwrap();
+        *self.cached_key.write() = Some(bytes.clone());
         Ok(bytes)
     }
 }
@@ -500,20 +517,18 @@ mod tests_keydb {
         let mock_primary_password_authenticator = MockPrimaryPasswordAuthenticator {
             password: "password".to_string(),
         };
-        let nss_key_manager = NSSKeyManager {
-            primary_password_authenticator: Arc::new(mock_primary_password_authenticator),
-        };
+        let nss_key_manager = NSSKeyManager::new(Arc::new(mock_primary_password_authenticator));
         // key from fixtures/profile/key4.db
-        assert_eq!(
-            nss_key_manager.get_key().unwrap(),
-            [
-                123, 34, 107, 116, 121, 34, 58, 34, 111, 99, 116, 34, 44, 34, 107, 34, 58, 34, 66,
-                74, 104, 84, 108, 103, 51, 118, 56, 49, 65, 66, 51, 118, 87, 50, 71, 122, 54, 104,
-                69, 54, 84, 116, 75, 83, 112, 85, 102, 84, 86, 75, 73, 83, 99, 74, 45, 77, 78, 83,
-                67, 117, 99, 34, 125
-            ]
-            .to_vec()
-        )
+        let expected = [
+            123, 34, 107, 116, 121, 34, 58, 34, 111, 99, 116, 34, 44, 34, 107, 34, 58, 34, 66, 74,
+            104, 84, 108, 103, 51, 118, 56, 49, 65, 66, 51, 118, 87, 50, 71, 122, 54, 104, 69, 54,
+            84, 116, 75, 83, 112, 85, 102, 84, 86, 75, 73, 83, 99, 74, 45, 77, 78, 83, 67, 117, 99,
+            34, 125,
+        ]
+        .to_vec();
+        assert_eq!(nss_key_manager.get_key().unwrap(), expected);
+        // served from the cache
+        assert_eq!(nss_key_manager.get_key().unwrap(), expected);
     }
 
     #[test]