Better test coverage for MPC context lifecycle (PR 3/7)

ai-docs/ARCHITECTURE.md

@@ -190,22 +190,24 @@ The primary service is `CoreServiceEndpoint`. Its RPCs group into:
   the key and logs a warning instead. When resharing legacy key material that
   has no dedicated OPRF secret-key share, the OPRF sub-protocol is skipped and
   the reshared private keyset keeps that field absent. A storage failure during
-  resharing rolls the new epoch back on the party that fails. 
-  That party attempts to delete the key shares, the CRS metadata and the epoch data of the new epoch. 
-  Observe that no public data is deleted as this is, and should be, unaffected by an epoch change. 
-  If cleanup succeeds, it forgets the epoch; otherwise, it keeps the epoch registered so that deletion can be retried. 
-  `DestroyMpcContext` carries
-  the context's epoch IDs and erases their secret shares (cascading to the
-  existing per-epoch deletion) before forgetting the context, so retiring a
-  party set leaves no usable key shares behind; the kms-connector is the source
-  of truth for which epochs belong to a context. In-memory lifecycle leases
-  serialize creation against destruction: `NewMpcEpoch` holds shared leases for
-  its target context and epoch through all PRSS, resharing and persistence work,
+  resharing rolls the new epoch back on the party that fails. That party attempts
+  to delete the key shares, the CRS metadata and the epoch data of the new epoch.
+  Public data remains because an epoch change does not affect it. If cleanup
+  succeeds, the party forgets the epoch. Otherwise, the party keeps the epoch
+  registered so that deletion can be retried. `DestroyMpcContext` takes a stable
+  snapshot of the context's registered epochs and erases their secret shares
+  before it forgets the context. This order leaves no usable key shares after the
+  party set retires. Its response lists the deleted epoch IDs. In-memory
+  lifecycle leases serialize creation against destruction: `NewMpcEpoch` holds
+  shared leases for its target context and epoch through all PRSS, resharing and
+  persistence work,
   while `DestroyMpcEpoch` and `DestroyMpcContext` require exclusive leases before
   taking snapshots or deleting data. A conflicting destruction is refused with
   `FailedPrecondition`, including while PRSS is still running and the new epoch
   has not yet been registered in the session maker; callers retry once creation
-  has settled.
+  has settled. MPC context updates serialize the existence check with storage and
+  cache or session updates. A failed deletion keeps the in-memory context if its
+  persistent entry remains, which permits a retry before or after restart.
 - **Session management** — creation, result retrieval, and cleanup for
   long-running threshold sessions.
 
@@ -282,6 +284,25 @@ end-to-end tests live at
 and
 [core/service/src/client/tests/threshold/custodian_backup_tests.rs](core/service/src/client/tests/threshold/custodian_backup_tests.rs).
 
+## Paired material writes
+
+Threshold calls to `CryptoMaterialStorage::write_all` use two public/private pairs:
+
+- `PublicKey` and `FheKeyInfo` share a key ID.
+- `CRS` and `CrsInfo` share a CRS ID.
+
+The public half has no epoch. The private half has an epoch and contains one party's material.
+Initial generation writes both halves through `CryptoMaterialStorage::write_all`. The method also
+accepts one-sided writes. Resharing writes only the private half for the new epoch and reuses the
+public half. A `ContextInfo` write stores one request-scoped private entry with no public half.
+
+Storage never overwrites an entry. If one requested half exists, storage keeps its bytes and writes
+the missing half. The caller must ensure that the two halves belong together. If either write
+fails, cleanup removes only entries created by that call. It retains each entry that existed
+before the call. A backend can apply a write and then return an error, so cleanup checks the
+earlier state. Callers must serialize writes to the same entries until cleanup finishes. A later
+backup failure does not purge the primary material.
+
 ## Boot-time storage verification
 
 Every node checks its storage during service construction, before it serves any request.

core/service/src/engine/context_manager.rs

@@ -19,7 +19,7 @@ use crate::util::meta_store::{
     lock_entry_in_meta_store, update_err_req_in_meta_store,
 };
 use crate::vault::keychain::KeychainProxy;
-use crate::vault::storage::crypto_material::{CryptoMaterialStorage, data_exists};
+use crate::vault::storage::crypto_material::{CryptoMaterialStorage, StorageError, data_exists};
 use crate::vault::storage::{
     StorageExt, delete_context_at_id, delete_custodian_context_at_id, store_context_at_id,
 };
@@ -60,6 +60,11 @@ struct SharedContextManager<
     base_kms: BaseKmsStruct,
     crypto_storage: CryptoMaterialStorage<PubS, PrivS>,
     custodian_meta_store: Arc<RwLock<CustodianMetaStore>>,
+    /// Serializes each MPC context existence check with its storage and in-memory updates.
+    ///
+    /// One global lock covers the update's storage I/O. Context updates are rare administrative
+    /// operations.
+    mpc_context_update_lock: Mutex<()>,
     /// Serializes whole custodian-context setups; see `inner_new_custodian_context`.
     custodian_setup_lock: Mutex<()>,
 }
@@ -678,6 +683,7 @@ where
                 base_kms,
                 crypto_storage,
                 custodian_meta_store,
+                mpc_context_update_lock: Mutex::new(()),
                 custodian_setup_lock: Mutex::new(()),
             },
             cache: Arc::new(RwLock::new(HashSet::new())),
@@ -729,6 +735,8 @@ where
             .map_err(|e| {
                 MetricedError::new(OP_NEW_MPC_CONTEXT, None, e, tonic::Code::InvalidArgument)
             })?;
+        let _update_guard = self.inner.mpc_context_update_lock.lock().await;
+
         // Check if the context already exists
         if self
             .inner
@@ -745,14 +753,19 @@ where
             ));
         }
 
-        // store the new context
-        let res = self
+        let storage_result = self
             .inner
             .crypto_storage
             .write_context_info(new_context.context_id(), &new_context, OP_NEW_MPC_CONTEXT)
             .await;
+        let context_is_stored = match &storage_result {
+            Ok(()) => true,
+            Err(error) => {
+                context_write_persisted(&self.inner.crypto_storage, &new_context, error).await
+            }
+        };
 
-        {
+        if context_is_stored {
             let mut write_guard = self.cache.write().await;
             let is_new_insert = (*write_guard).insert(*new_context.context_id());
             if !is_new_insert {
@@ -763,7 +776,7 @@ where
             }
         }
 
-        res.map_err(|e| {
+        storage_result.map_err(|e| {
             MetricedError::new(
                 OP_NEW_MPC_CONTEXT,
                 Some((*new_context.context_id()).into()),
@@ -791,6 +804,7 @@ where
                     tonic::Code::InvalidArgument,
                 )
             })?;
+        let _update_guard = self.inner.mpc_context_update_lock.lock().await;
 
         let storage_ref = self.inner.crypto_storage.private_storage.clone();
         let mut guarded_priv_storage = storage_ref.lock().await;
@@ -836,7 +850,7 @@ where
             ));
         }
 
-        delete_context_at_id(&mut *guarded_priv_storage, &context_id)
+        delete_context_from_storage(&mut *guarded_priv_storage, &context_id)
             .await
             .map_err(|e| {
                 MetricedError::new(
@@ -923,6 +937,7 @@ where
                 base_kms,
                 crypto_storage,
                 custodian_meta_store,
+                mpc_context_update_lock: Mutex::new(()),
                 custodian_setup_lock: Mutex::new(()),
             },
             session_maker,
@@ -975,8 +990,11 @@ where
     }
 }
 
-/// Atomically update both the storage and the session maker with the new context info.
-/// If any of the two operations fail, rollback to the original state.
+/// Updates storage and the session maker with the new context information.
+///
+/// A storage failure needs no extra cleanup because [`CryptoMaterialStorage::write_all`] removes
+/// entries created by its failed write. If the session update fails, this function removes the
+/// context that it stored.
 ///
 /// This function should only be used in the threshold setting since SessionMaker does not exist in centralized mode.
 async fn atomic_update_context<
@@ -988,29 +1006,91 @@ async fn atomic_update_context<
     my_role: Option<Role>,
     new_context: &ContextInfo,
 ) -> anyhow::Result<()> {
-    let context_id = new_context.context_id();
-    let res1 = crypto_storage
+    let storage_result = crypto_storage
         .write_context_info(new_context.context_id(), new_context, OP_NEW_MPC_CONTEXT)
         .await;
+    let context_is_stored = match &storage_result {
+        Ok(()) => true,
+        Err(error) => context_write_persisted(crypto_storage, new_context, error).await,
+    };
+    if !context_is_stored {
+        return Err(anyhow::anyhow!(
+            "Failed to store context: {}",
+            storage_result.unwrap_err()
+        ));
+    }
 
-    let res2 = session_maker.add_context_info(my_role, new_context).await;
+    if let Err(session_error) = session_maker.add_context_info(my_role, new_context).await {
+        let context_id = new_context.context_id();
+        let storage_ref = crypto_storage.private_storage.clone();
+        let mut guarded_priv_storage = storage_ref.lock().await;
+        let cleanup_result =
+            delete_context_from_storage(&mut *guarded_priv_storage, context_id).await;
+        // Ensure no session state remains if `add_context_info` applied only part of the update.
+        session_maker.remove_context(context_id).await;
+
+        return match cleanup_result {
+            Ok(()) => Err(anyhow::anyhow!(
+                "Failed to add context to the session maker: {session_error}"
+            )),
+            Err(cleanup_error) => Err(anyhow::anyhow!(
+                "Failed to add context to the session maker: {session_error}; failed to remove the stored context: {cleanup_error}"
+            )),
+        };
+    }
 
-    match (res1, res2) {
-        (Ok(_), Ok(_)) => (),
-        _ => {
-            // Rollback if any operation failed
-            // first delete the context from storage
-            let storage_ref = crypto_storage.private_storage.clone();
-            let mut guarded_priv_storage = storage_ref.lock().await;
-            _ = delete_context_at_id(&mut *guarded_priv_storage, context_id).await;
+    storage_result.map_err(|error| anyhow::anyhow!("Failed to store context: {error}"))
+}
 
-            // next delete the context from session maker
-            session_maker.remove_context(context_id).await;
-            return Err(anyhow::anyhow!("Failed to atomically update context"));
-        }
+/// Returns whether a failed write left this call's context in primary storage.
+async fn context_write_persisted<
+    PubS: Storage + Sync + Send + 'static,
+    PrivS: StorageExt + Sync + Send + 'static,
+>(
+    crypto_storage: &CryptoMaterialStorage<PubS, PrivS>,
+    context: &ContextInfo,
+    error: &StorageError,
+) -> bool {
+    match error {
+        StorageError::Backup => true,
+        StorageError::Purging => crypto_storage
+            .read_context_info(context.context_id())
+            .await
+            .is_ok_and(|stored| stored == *context),
+        _ => false,
     }
+}
 
-    Ok(())
+/// Deletes a context and resolves an error that arrives after the backend applied the deletion.
+///
+/// The caller must serialize updates for `context_id` until the existence check completes.
+async fn delete_context_from_storage<PrivS: StorageExt + Sync + Send + 'static>(
+    storage: &mut PrivS,
+    context_id: &ContextId,
+) -> anyhow::Result<()> {
+    let delete_error = match delete_context_at_id(storage, context_id).await {
+        Ok(()) => return Ok(()),
+        Err(error) => error,
+    };
+
+    match storage
+        .data_exists(
+            &(*context_id).into(),
+            &PrivDataType::ContextInfo.to_string(),
+        )
+        .await
+    {
+        Ok(false) => {
+            tracing::warn!(
+                "Context {context_id} was removed although the storage backend reported an error: {delete_error}"
+            );
+            Ok(())
+        }
+        Ok(true) => Err(delete_error),
+        Err(check_error) => Err(anyhow::anyhow!(
+            "Failed to delete context: {delete_error}; failed to check the result: {check_error}"
+        )),
+    }
 }
 
 #[tonic::async_trait]
@@ -1030,6 +1110,7 @@ where
             .map_err(|e| {
                 MetricedError::new(OP_NEW_MPC_CONTEXT, None, e, tonic::Code::InvalidArgument)
             })?;
+        let _update_guard = self.inner.mpc_context_update_lock.lock().await;
 
         // First check if the context already exists
         if self
@@ -1085,6 +1166,8 @@ where
                     tonic::Code::InvalidArgument,
                 )
             })?;
+        let _update_guard = self.inner.mpc_context_update_lock.lock().await;
+
         if !self.session_maker.context_exists(&context_id).await {
             return Err(MetricedError::new(
                 OP_DESTROY_MPC_CONTEXT,
@@ -1107,11 +1190,7 @@ where
 
         let storage_ref = self.inner.crypto_storage.private_storage.clone();
         let mut guarded_priv_storage = storage_ref.lock().await;
-        self.session_maker.remove_context(&context_id).await;
-
-        // There is nothing we can do if deletion fails here.
-        // Note that it cannot fail if the context does not exist.
-        delete_context_at_id(&mut *guarded_priv_storage, &context_id)
+        delete_context_from_storage(&mut *guarded_priv_storage, &context_id)
             .await
             .map_err(|e| {
                 MetricedError::new(
@@ -1121,6 +1200,7 @@ where
                     tonic::Code::Internal,
                 )
             })?;
+        self.session_maker.remove_context(&context_id).await;
         Ok(())
     }
 
@@ -1207,6 +1287,8 @@ async fn gen_recovery_validation(
 
 #[cfg(test)]
 mod tests {
+    mod lifecycle_side_effects;
+
     use super::*;
     use crate::{
         backup::{

core/service/src/engine/context_manager/tests/lifecycle_side_effects.rs

@@ -0,0 +1,7 @@
+//! Tests MPC context lifecycle side effects in persistent and in-memory state.
+//!
+//! Cases cover storage failures, backup failures, session rollback, duplicate creation and
+//! serialized updates. The fixture also stores a keeper context to detect broad cleanup.
+
+mod cases;
+mod support;