diff --git a/crates/contract/Cargo.toml b/crates/contract/Cargo.toml index c44270435c..3e931d116f 100644 --- a/crates/contract/Cargo.toml +++ b/crates/contract/Cargo.toml @@ -50,6 +50,7 @@ __abi-generate = ["near-sdk/__abi-generate"] [dependencies] anyhow = { workspace = true } +assert_matches = { workspace = true } borsh = { workspace = true } contract-interface = { workspace = true } curve25519-dalek = { workspace = true } @@ -85,7 +86,6 @@ utilities = { workspace = true } getrandom = { workspace = true, features = ["custom"] } [dev-dependencies] -assert_matches = { workspace = true } cargo-near-build = { workspace = true } contract-history = { workspace = true } digest = { workspace = true } diff --git a/crates/contract/src/lib.rs b/crates/contract/src/lib.rs index 54833f79b7..4d47b542b3 100644 --- a/crates/contract/src/lib.rs +++ b/crates/contract/src/lib.rs @@ -1424,12 +1424,15 @@ impl MpcContract { Self::assert_caller_is_signer(); - if !self + let attestation_check = self .tee_state - .is_caller_an_attested_participant(participants) - { - panic!("Caller must be an attested participant"); - } + .is_caller_an_attested_participant(participants); + + assert_matches::assert_matches!( + attestation_check, + Ok(()), + "Caller must be an attested participant" + ); } /// Ensures the current call originates from the signer account itself. diff --git a/crates/contract/src/tee/tee_state.rs b/crates/contract/src/tee/tee_state.rs index b3e99d2cf2..43f18129d9 100644 --- a/crates/contract/src/tee/tee_state.rs +++ b/crates/contract/src/tee/tee_state.rs @@ -72,7 +72,7 @@ pub struct TeeState { pub(crate) allowed_docker_image_hashes: AllowedDockerImageHashes, pub(crate) allowed_launcher_compose_hashes: Vec, pub(crate) votes: CodeHashesVotes, - /// Mapping of TLS public key of a participant to its NodeId and Attestation. + /// 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, @@ -342,40 +342,57 @@ impl TeeState { .map(|(node_id, _)| node_id.clone()) } - /// Returns true if the caller has at least one participant entry + /// Returns Ok(()) if the caller has at least one participant entry /// whose TLS key matches an attested node belonging to the caller account. - /// - /// Handles multiple participants per account and supports legacy mock nodes. - pub fn is_caller_an_attested_participant(&self, participants: &Participants) -> bool { + pub(crate) fn is_caller_an_attested_participant( + &self, + participants: &Participants, + ) -> Result<(), AttestationCheckError> { let signer_pk = env::signer_account_pk(); let signer_id = env::signer_account_id().as_v2_account_id(); - match participants.info(&signer_id) { - None => false, - Some(info) => { - match self.stored_attestations.get(&info.sign_pk) { - None => false, - Some((node_id, _attestation)) => { - node_id.account_id == signer_id - && node_id - .account_public_key - .as_ref() - .map(|pk| pk == &signer_pk) - .unwrap_or(true) // TODO (#823) Legacy fallback for mock nodes - } - } + let info = participants + .info(&signer_id) + .ok_or(AttestationCheckError::CallerNotParticipant)?; + + let attestation = self + .stored_attestations + .get(&info.sign_pk) + .ok_or(AttestationCheckError::AttestationNotFound)?; + + if attestation.0.account_id != signer_id { + return Err(AttestationCheckError::AttestationOwnerMismatch); + } + + if let Some(node_pk) = &attestation.0.account_public_key { + if node_pk != &signer_pk { + return Err(AttestationCheckError::AttestationKeyMismatch); } } + + Ok(()) } } +#[derive(Debug)] +pub(crate) enum AttestationCheckError { + CallerNotParticipant, + AttestationNotFound, + AttestationOwnerMismatch, + AttestationKeyMismatch, +} + #[cfg(test)] mod tests { - use crate::primitives::test_utils::bogus_ed25519_near_public_key; - use super::*; + use crate::primitives::test_utils::bogus_ed25519_near_public_key; + use crate::primitives::test_utils::gen_participants; + use assert_matches::assert_matches; use mpc_attestation::attestation::{Attestation, MockAttestation}; use near_account_id::AccountId; + use near_sdk::test_utils::VMContextBuilder; + use near_sdk::testing_env; + use utilities::AccountIdExtV2; #[test] fn test_clean_non_participants() { @@ -435,4 +452,147 @@ mod tests { .stored_attestations .contains_key(&non_participant_uid.tls_public_key)); } + + /// Helper to set up the testing environment with a specific signer + fn set_signer(account_id: &AccountId, public_key: &near_sdk::PublicKey) { + let mut builder = VMContextBuilder::new(); + builder + .signer_account_id(account_id.as_v1_account_id()) + .signer_account_pk(public_key.clone()); + testing_env!(builder.build()); + } + + #[test] + fn test_is_caller_attested_success() { + let mut tee_state = TeeState::default(); + // Generate 1 participant + let participants = gen_participants(1); + let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); + + // 1. Define the Signer's NEAR Public Key (Wallet Key) + let signer_pk = bogus_ed25519_near_public_key(); + + // 2. Set the environment so the caller is the participant + set_signer(account_id, &signer_pk); + + // 3. Register the attestation in TeeState + // The TLS key comes from participant_info, the Account Key must match the signer_pk + let node_id = NodeId { + account_id: account_id.clone(), + tls_public_key: participant_info.sign_pk.clone(), + account_public_key: Some(signer_pk), + }; + tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + + // 4. Verify check passes + let result = tee_state.is_caller_an_attested_participant(&participants); + assert!(result.is_ok()); + } + + #[test] + fn test_is_caller_attested_success_legacy_no_account_key() { + // Tests the case where account_public_key is None (legacy/mock nodes) + let mut tee_state = TeeState::default(); + let participants = gen_participants(1); + let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); + + let signer_pk = bogus_ed25519_near_public_key(); + set_signer(account_id, &signer_pk); + + // Register attestation with None for account_public_key + let node_id = NodeId { + account_id: account_id.clone(), + tls_public_key: participant_info.sign_pk.clone(), + account_public_key: None, + }; + tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + + let result = tee_state.is_caller_an_attested_participant(&participants); + assert_matches!(result, Ok(())); + } + + #[test] + fn test_err_caller_not_participant() { + let tee_state = TeeState::default(); + let participants = gen_participants(1); + + // Caller is NOT in the participants list + let random_account: AccountId = "random_guy.near".parse().unwrap(); + let random_pk = bogus_ed25519_near_public_key(); + set_signer(&random_account, &random_pk); + + let result = tee_state.is_caller_an_attested_participant(&participants); + + assert_matches!(result, Err(AttestationCheckError::CallerNotParticipant)); + } + + #[test] + fn test_err_attestation_not_found() { + let tee_state = TeeState::default(); + let participants = gen_participants(1); + let (account_id, _, _) = participants.participants().iter().next().unwrap(); + + let signer_pk = bogus_ed25519_near_public_key(); + set_signer(account_id, &signer_pk); + + // We do NOT add the participant to tee_state.stored_attestations + + let result = tee_state.is_caller_an_attested_participant(&participants); + + assert_matches!(result, Err(AttestationCheckError::AttestationNotFound)); + } + + #[test] + fn test_err_attestation_owner_mismatch() { + let mut tee_state = TeeState::default(); + let participants = gen_participants(1); + let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); + + let signer_pk = bogus_ed25519_near_public_key(); + set_signer(account_id, &signer_pk); + + // Create a data integrity issue: + // The TLS key in `participants` points to an entry in `tee_state`... + // ...but that entry claims it belongs to a different AccountId. + let other_account: AccountId = "imposter.near".parse().unwrap(); + + let node_id = NodeId { + account_id: other_account, // Mismatch here + tls_public_key: participant_info.sign_pk.clone(), + account_public_key: Some(signer_pk), + }; + tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + + let result = tee_state.is_caller_an_attested_participant(&participants); + + assert_matches!(result, Err(AttestationCheckError::AttestationOwnerMismatch)); + } + + #[test] + fn test_err_attestation_key_mismatch() { + let mut tee_state = TeeState::default(); + let participants = gen_participants(1); + let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); + + let signer_pk = bogus_ed25519_near_public_key(); + set_signer(account_id, &signer_pk); + + // Generate a different key for the stored attestation + // (e.g., The user rotated their wallet key, but hasn't updated the TEE registry) + let old_signer_pk: near_sdk::PublicKey = + "ed25519:3t4M1gXg2Qd5g6X8z1g2X3t4M1gXg2Qd5g6X8z1g2X3t" + .parse() + .unwrap(); + + let node_id = NodeId { + account_id: account_id.clone(), + tls_public_key: participant_info.sign_pk.clone(), + account_public_key: Some(old_signer_pk), // Mismatch here + }; + tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + + let result = tee_state.is_caller_an_attested_participant(&participants); + + assert_matches!(result, Err(AttestationCheckError::AttestationKeyMismatch)); + } }