diff --git a/crates/contract/src/lib.rs b/crates/contract/src/lib.rs index f29bde4127..54833f79b7 100644 --- a/crates/contract/src/lib.rs +++ b/crates/contract/src/lib.rs @@ -643,7 +643,7 @@ impl MpcContract { Ok(self .tee_state - .participants_attestations + .stored_attestations .iter() .find(|(stored_tls_pk, _)| **stored_tls_pk == tls_public_key) .map(|(_, (_, attestation))| attestation.clone().into_dto_type())) @@ -1810,13 +1810,24 @@ mod tests { /// Temporarily sets the testing environment so that calls appear /// to come from an attested MPC node registered in the contract's `tee_state`. /// Returns the `AccountId` of the node used. - pub fn with_attested_context(contract: &MpcContract) -> AccountId { - let (_account_id, (node_id, _)) = contract + pub fn with_active_participant_and_attested_context(contract: &MpcContract) -> AccountId { + let active_participant_pks: Vec<_> = contract + .protocol_state + .active_participants() + .participants() + .iter() + .map(|(_, _, participant_info)| participant_info.sign_pk.clone()) + .collect(); + + let node_id = contract .tee_state - .participants_attestations + .stored_attestations .iter() - .next() - .expect("No attested participants in tee_state"); + .find(|(public_key, _)| active_participant_pks.contains(public_key)) + .expect("No attested participants in tee_state") + .1 + .0 + .clone(); // Build a new simulated environment with this node as caller let mut ctx_builder = VMContextBuilder::new(); @@ -1894,7 +1905,7 @@ mod tests { )) }; - with_attested_context(&contract); + with_active_participant_and_attested_context(&contract); match contract.respond(signature_request.clone(), signature_response.clone()) { Ok(_) => { @@ -1980,7 +1991,7 @@ mod tests { big_c: dtos::Bls12381G1PublicKey([2u8; 48]), }; - with_attested_context(&contract); + with_active_participant_and_attested_context(&contract); match contract.respond_ckd(ckd_request.clone(), response.clone()) { Ok(_) => { @@ -2332,7 +2343,7 @@ mod tests { .unwrap(); // --- Step 4: Verify that a participant can still respond successfully --- - with_attested_context(&contract); // sets env to a real attested participant + with_active_participant_and_attested_context(&contract); // sets env to a real attested participant let valid_response = CKDResponse { big_y: dtos::Bls12381G1PublicKey([1u8; 48]), diff --git a/crates/contract/src/tee/tee_state.rs b/crates/contract/src/tee/tee_state.rs index dd0e01e651..b3e99d2cf2 100644 --- a/crates/contract/src/tee/tee_state.rs +++ b/crates/contract/src/tee/tee_state.rs @@ -72,7 +72,10 @@ pub struct TeeState { pub(crate) allowed_docker_image_hashes: AllowedDockerImageHashes, pub(crate) allowed_launcher_compose_hashes: Vec, pub(crate) votes: CodeHashesVotes, - pub participants_attestations: IterableMap, + /// Mapping of TLS public key of a participant to its NodeId and Attestation. + /// Attestations are stored for any valid participant that has submitted one, not + /// just for the currently active participants. + pub(crate) stored_attestations: IterableMap, } impl Default for TeeState { @@ -81,7 +84,7 @@ impl Default for TeeState { allowed_docker_image_hashes: AllowedDockerImageHashes::default(), allowed_launcher_compose_hashes: vec![], votes: CodeHashesVotes::default(), - participants_attestations: IterableMap::new(StorageKey::TeeParticipantAttestation), + stored_attestations: IterableMap::new(StorageKey::TeeParticipantAttestation), } } } @@ -108,7 +111,7 @@ impl TeeState { }); Self { - participants_attestations, + stored_attestations: participants_attestations, ..Default::default() } } @@ -148,7 +151,7 @@ impl TeeState { self.get_allowed_mpc_docker_image_hashes(tee_upgrade_deadline_duration); let allowed_launcher_compose_hashes = &self.allowed_launcher_compose_hashes; - let participant_attestation = self.participants_attestations.get(&node_id.tls_public_key); + let participant_attestation = self.stored_attestations.get(&node_id.tls_public_key); let Some(participant_attestation) = participant_attestation else { return TeeQuoteStatus::Invalid("participant has no attestation".to_string()); }; @@ -251,10 +254,10 @@ impl TeeState { pub fn add_participant(&mut self, node_id: NodeId, attestation: Attestation) -> bool { let tls_pk = node_id.tls_public_key.clone(); - let is_new = !self.participants_attestations.contains_key(&tls_pk); + let is_new = !self.stored_attestations.contains_key(&tls_pk); // Must pass owned values, not references - self.participants_attestations + self.stored_attestations .insert(tls_pk, (node_id, attestation)); is_new @@ -311,7 +314,7 @@ impl TeeState { // Collect TLS keys that are *not* in the active participants list let stale_keys: Vec = self - .participants_attestations + .stored_attestations .keys() .filter(|tls_pk| !active_tls_keys.contains(*tls_pk)) .cloned() @@ -319,14 +322,14 @@ impl TeeState { // Remove all stale TEE entries for tls_pk in stale_keys { - self.participants_attestations.remove(&tls_pk); + self.stored_attestations.remove(&tls_pk); } } /// Returns the list of accounts that currently have TEE attestations stored. /// Note: This may include accounts that are no longer active protocol participants. pub fn get_tee_accounts(&self) -> Vec { - self.participants_attestations + self.stored_attestations .values() .map(|(node_id, _)| node_id.clone()) .collect() @@ -334,7 +337,7 @@ impl TeeState { /// Find a NodeId by its TLS public key. pub fn find_node_id_by_tls_key(&self, tls_public_key: &near_sdk::PublicKey) -> Option { - self.participants_attestations + self.stored_attestations .get(tls_public_key) .map(|(node_id, _)| node_id.clone()) } @@ -350,7 +353,7 @@ impl TeeState { match participants.info(&signer_id) { None => false, Some(info) => { - match self.participants_attestations.get(&info.sign_pk) { + match self.stored_attestations.get(&info.sign_pk) { None => false, Some((node_id, _attestation)) => { node_id.account_id == signer_id @@ -408,28 +411,28 @@ mod tests { tee_state.add_participant(non_participant_uid.clone(), local_attestation.clone()); // Verify all 4 accounts have TEE info initially - assert_eq!(tee_state.participants_attestations.len(), 4); + assert_eq!(tee_state.stored_attestations.len(), 4); for node_id in &participant_nodes { assert!(tee_state - .participants_attestations + .stored_attestations .contains_key(&node_id.tls_public_key)); } assert!(tee_state - .participants_attestations + .stored_attestations .contains_key(&non_participant_uid.tls_public_key)); // Clean non-participants tee_state.clean_non_participants(&participants); // Verify only participants remain - assert_eq!(tee_state.participants_attestations.len(), 3); + assert_eq!(tee_state.stored_attestations.len(), 3); for node_id in &participant_nodes { assert!(tee_state - .participants_attestations + .stored_attestations .contains_key(&node_id.tls_public_key)); } assert!(!tee_state - .participants_attestations + .stored_attestations .contains_key(&non_participant_uid.tls_public_key)); } }