diff --git a/Cargo.lock b/Cargo.lock index d3da2b405..4574ad9fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4793,6 +4793,7 @@ checksum = "7843ec2de400bcbc6a6328c958dc38e5359da6e93e72e37bc5246bf1ae776389" name = "mpc-attestation" version = "3.2.0" dependencies = [ + "assert_matches", "attestation", "borsh", "dcap-qvl", @@ -4800,7 +4801,6 @@ dependencies = [ "hex", "include-measurements", "mpc-primitives", - "rstest", "serde", "serde_json", "sha2", diff --git a/crates/contract-interface/src/lib.rs b/crates/contract-interface/src/lib.rs index cc0b4fea3..8ca071114 100644 --- a/crates/contract-interface/src/lib.rs +++ b/crates/contract-interface/src/lib.rs @@ -3,6 +3,7 @@ pub mod types { pub use attestation::{ AppCompose, Attestation, Collateral, DstackAttestation, EventLog, MockAttestation, TcbInfo, + VerifiedAttestation, VerifiedDstackAttestation, }; pub use config::{Config, InitConfig}; pub use crypto::{ diff --git a/crates/contract-interface/src/types/attestation.rs b/crates/contract-interface/src/types/attestation.rs index a5bba3802..ebb88d4e2 100644 --- a/crates/contract-interface/src/types/attestation.rs +++ b/crates/contract-interface/src/types/attestation.rs @@ -29,6 +29,54 @@ pub enum Attestation { Mock(MockAttestation), } +#[derive( + Clone, + Debug, + Eq, + PartialEq, + Ord, + PartialOrd, + Hash, + Serialize, + Deserialize, + BorshSerialize, + BorshDeserialize, +)] +#[cfg_attr( + all(feature = "abi", not(target_arch = "wasm32")), + derive(schemars::JsonSchema) +)] +pub enum VerifiedAttestation { + Dtack(VerifiedDstackAttestation), + Mock(MockAttestation), +} + +#[derive( + Clone, + Debug, + Eq, + PartialEq, + Ord, + PartialOrd, + Hash, + Serialize, + Deserialize, + BorshSerialize, + BorshDeserialize, +)] +#[cfg_attr( + all(feature = "abi", not(target_arch = "wasm32")), + derive(schemars::JsonSchema) +)] +pub struct VerifiedDstackAttestation { + /// The digest of the MPC image running. + pub mpc_image_hash: Sha256Digest, + /// The digest of the launcher compose file running. + pub launcher_compose_hash: Sha256Digest, + /// Unix time stamp for when this attestation expires. + pub expiry_timestamp_seconds: u64, +} + #[derive( Clone, Eq, @@ -78,8 +126,8 @@ pub enum MockAttestation { WithConstraints { mpc_docker_image_hash: Option, launcher_docker_compose_hash: Option, - /// Unix time stamp for when this attestation expires. - expiry_time_stamp_seconds: Option, + /// Unix time stamp for when this attestation expires. + expiry_timestamp_seconds: Option, }, } diff --git a/crates/contract/src/dto_mapping.rs b/crates/contract/src/dto_mapping.rs index 6ce6e4eb7..e966ef091 100644 --- a/crates/contract/src/dto_mapping.rs +++ b/crates/contract/src/dto_mapping.rs @@ -6,7 +6,7 @@ use contract_interface::types as dtos; use mpc_attestation::{ - attestation::{Attestation, DstackAttestation, MockAttestation}, + attestation::{Attestation, DstackAttestation, MockAttestation, VerifiedAttestation}, collateral::{Collateral, QuoteCollateralV3}, tcb_info::{EventLog, HexBytes, TcbInfo}, }; @@ -71,11 +71,11 @@ impl IntoContractType for dtos::MockAttestation { dtos::MockAttestation::WithConstraints { mpc_docker_image_hash, launcher_docker_compose_hash, - expiry_time_stamp_seconds, + expiry_timestamp_seconds, } => MockAttestation::WithConstraints { mpc_docker_image_hash: mpc_docker_image_hash.map(Into::into), launcher_docker_compose_hash: launcher_docker_compose_hash.map(Into::into), - expiry_time_stamp_seconds, + expiry_timestamp_seconds, }, } } @@ -197,14 +197,20 @@ impl TryIntoContractType for dtos::EventLog { } } -impl IntoInterfaceType for Attestation { - fn into_dto_type(self) -> dtos::Attestation { +impl IntoInterfaceType for VerifiedAttestation { + fn into_dto_type(self) -> dtos::VerifiedAttestation { match self { - Attestation::Dstack(dstack_attestation) => { - dtos::Attestation::Dstack(dstack_attestation.into_dto_type()) + VerifiedAttestation::Mock(mock_attestation) => { + dtos::VerifiedAttestation::Mock(mock_attestation.into_dto_type()) } - Attestation::Mock(mock_attestation) => { - dtos::Attestation::Mock(mock_attestation.into_dto_type()) + VerifiedAttestation::Dstack(validated_dstack_attestation) => { + dtos::VerifiedAttestation::Dtack(dtos::VerifiedDstackAttestation { + mpc_image_hash: validated_dstack_attestation.mpc_image_hash.into(), + launcher_compose_hash: validated_dstack_attestation + .launcher_compose_hash + .into(), + expiry_timestamp_seconds: validated_dstack_attestation.expiry_timestamp_seconds, + }) } } } @@ -218,11 +224,11 @@ impl IntoInterfaceType for MockAttestation { MockAttestation::WithConstraints { mpc_docker_image_hash, launcher_docker_compose_hash, - expiry_time_stamp_seconds, + expiry_timestamp_seconds, } => dtos::MockAttestation::WithConstraints { mpc_docker_image_hash: mpc_docker_image_hash.map(Into::into), launcher_docker_compose_hash: launcher_docker_compose_hash.map(Into::into), - expiry_time_stamp_seconds, + expiry_timestamp_seconds, }, } } diff --git a/crates/contract/src/lib.rs b/crates/contract/src/lib.rs index f9bab19e8..fbe59bd84 100644 --- a/crates/contract/src/lib.rs +++ b/crates/contract/src/lib.rs @@ -19,6 +19,7 @@ pub mod update; #[cfg(feature = "dev-utils")] pub mod utils; pub mod v3_0_2_state; +pub mod v3_2_0_state; mod dto_mapping; @@ -47,13 +48,14 @@ use errors::{ }; use k256::elliptic_curve::PrimeField; +use mpc_attestation::attestation::Attestation; use mpc_primitives::hash::LauncherDockerComposeHash; use near_account_id::AccountId; use near_sdk::{ env::{self, ed25519_verify}, log, near_bindgen, state::ContractState, - store::LookupMap, + store::{IterableMap, LookupMap}, CryptoHash, Gas, GasWeight, NearToken, Promise, PromiseError, PromiseOrValue, }; use node_migrations::{BackupServiceInfo, DestinationNodeInfo, NodeMigrations}; @@ -67,7 +69,7 @@ use primitives::{ use state::{running::RunningContractState, ProtocolContractState}; use tee::{ proposal::MpcDockerImageHash, - tee_state::{NodeId, TeeValidationResult}, + tee_state::{NodeId, ParticipantInsertion, TeeValidationResult}, }; use utilities::{AccountIdExtV1, AccountIdExtV2}; @@ -99,6 +101,27 @@ pub struct MpcContract { tee_state: TeeState, accept_requests: bool, node_migrations: NodeMigrations, + stale_data: StaleData, +} + +/// A container for "orphaned" state that persists across contract migrations. +/// +/// ### Why this exists +/// On the NEAR blockchain, the `migrate` function is limited by the maximum transaction gas +/// (300 Tgas). Large data structures, specifically `IterableMap` or `LookupMap` +/// often cannot be cleared in a single block without hitting this limit. +/// +/// ### The Pattern +/// 1. During `migrate()`, expensive-to-delete fields are moved from the main state into [`StaleData`]. +/// 2. The main contract state becomes usable immediately. +/// 3. "Lazy cleanup" methods (like `post_upgrade_cleanup`) are then called in subsequent, +/// separate transactions to gradually deallocate this storage. +#[derive(Debug, Default, BorshSerialize, BorshDeserialize)] +struct StaleData { + /// Holds the TEE attestations from the previous contract version. + /// This is stored as an `Option` so it can be `.take()`n during the cleanup process, + /// ensuring the `IterableMap` handle is properly dropped. + participant_attestations: Option>, } impl MpcContract { @@ -579,31 +602,28 @@ impl MpcContract { let tee_upgrade_deadline_duration = Duration::from_secs(self.config.tee_upgrade_deadline_duration_seconds); - // Verify the TEE quote (including TLS and account keys) and Docker image for the proposed participant - let account_key_dto = account_key.clone().try_into_dto_type()?; - let status = self.tee_state.verify_proposed_participant_attestation( - &proposed_participant_attestation, - tls_public_key.clone(), - account_key_dto, - tee_upgrade_deadline_duration, - ); - - if let TeeQuoteStatus::Invalid(reason) = status { - return Err(InvalidParameters::InvalidTeeRemoteAttestation - .message(format!("TeeQuoteStatus is invalid: {reason}"))); - } - // Add the participant information to the contract state - let is_new_attestation = self.tee_state.add_participant( - NodeId { - account_id: account_id.clone(), - tls_public_key: tls_public_key.into_contract_type(), - account_public_key: Some(account_key), - }, - proposed_participant_attestation, - ); + let attestation_insertion_result = self + .tee_state + .add_participant( + NodeId { + account_id: account_id.clone(), + tls_public_key: tls_public_key.into_contract_type(), + account_public_key: Some(account_key), + }, + proposed_participant_attestation, + tee_upgrade_deadline_duration, + ) + .map_err(|err| { + InvalidParameters::InvalidTeeRemoteAttestation + .message(format!("TeeQuoteStatus is invalid: {err}")) + })?; let caller_is_not_participant = self.voter_account().is_err(); + let is_new_attestation = matches!( + attestation_insertion_result, + ParticipantInsertion::NewlyInsertedParticipant + ); let attestation_storage_must_be_paid_by_caller = is_new_attestation || caller_is_not_participant; @@ -638,15 +658,19 @@ impl MpcContract { pub fn get_attestation( &self, tls_public_key: dtos::Ed25519PublicKey, - ) -> Result, Error> { + ) -> Result, Error> { let tls_public_key = tls_public_key.into_contract_type(); Ok(self .tee_state .stored_attestations - .iter() - .find(|(stored_tls_pk, _)| **stored_tls_pk == tls_public_key) - .map(|(_, (_, attestation))| attestation.clone().into_dto_type())) + .get(&tls_public_key) + .map(|node_attestation| { + node_attestation + .verified_attestation + .clone() + .into_dto_type() + })) } /// Propose a new set of parameters (participants and threshold) for the MPC network. @@ -671,9 +695,10 @@ impl MpcContract { let tee_upgrade_deadline_duration = Duration::from_secs(self.config.tee_upgrade_deadline_duration_seconds); - let validation_result = self - .tee_state - .validate_tee(proposal.participants(), tee_upgrade_deadline_duration); + let validation_result = self.tee_state.reverify_and_cleanup_participants( + proposal.participants(), + tee_upgrade_deadline_duration, + ); let proposed_participants = proposal.participants(); match validation_result { @@ -1076,10 +1101,10 @@ impl MpcContract { let tee_upgrade_deadline_duration = Duration::from_secs(self.config.tee_upgrade_deadline_duration_seconds); - match self - .tee_state - .validate_tee(current_params.participants(), tee_upgrade_deadline_duration) - { + match self.tee_state.reverify_and_cleanup_participants( + current_params.participants(), + tee_upgrade_deadline_duration, + ) { TeeValidationResult::Full => { self.accept_requests = true; log!("All participants have an accepted Tee status"); @@ -1201,6 +1226,7 @@ impl MpcContract { tee_state, accept_requests: true, node_migrations: NodeMigrations::default(), + stale_data: Default::default(), }) } @@ -1251,6 +1277,7 @@ impl MpcContract { tee_state, accept_requests: true, node_migrations: NodeMigrations::default(), + stale_data: Default::default(), }) } @@ -1269,7 +1296,17 @@ impl MpcContract { match try_state_read::() { Ok(Some(state)) => return Ok(state.into()), Ok(None) => return Err(InvalidState::ContractStateIsMissing.into()), - Err(_) => (), // Try read as "Self" instead + Err(err) => { + log!("failed to deserialize state into 3_0_2 state: {:?}", err); + } + }; + + match try_state_read::() { + Ok(Some(state)) => return Ok(state.into()), + Ok(None) => return Err(InvalidState::ContractStateIsMissing.into()), + Err(err) => { + log!("failed to deserialize state into 3_2_0 state: {:?}", err); + } }; match try_state_read::() { @@ -1279,6 +1316,19 @@ impl MpcContract { } } + /// Removes stale data from the contract to be removed after a contract upgrade. Some + /// containers are expensive to run destructors on, thus we don't include it in the contract upgrade itself, + /// as it can run out of gas. Thus we create methods to run these destructors manually post upgrade. + pub fn post_upgrade_cleanup(&mut self) { + let Some(mut attestations) = self.stale_data.participant_attestations.take() else { + panic!("stale participant_attestations data has already been cleared"); + }; + + attestations.clear(); + + log!("Successfully cleared stale TEE attestations."); + } + pub fn state(&self) -> &ProtocolContractState { &self.protocol_state } @@ -1596,9 +1646,9 @@ impl MpcContract { }; if !(matches!( - self.tee_state.verify_tee_participant( + self.tee_state.reverify_participants( &node_id, - Duration::from_secs(self.config.tee_upgrade_deadline_duration_seconds) + Duration::from_secs(self.config.tee_upgrade_deadline_duration_seconds), ), TeeQuoteStatus::Valid )) { @@ -1834,7 +1884,7 @@ mod tests { .find(|(public_key, _)| active_participant_pks.contains(public_key)) .expect("No attested participants in tee_state") .1 - .0 + .node_id .clone(); // Build a new simulated environment with this node as caller @@ -2346,6 +2396,7 @@ mod tests { .predecessor_account_id(outsider_id.clone().as_v1_account_id()) .attached_deposit(NearToken::from_near(1)) .build()); + contract .submit_participant_info(Attestation::Mock(MockAttestation::Valid), dto_public_key) .unwrap(); @@ -2410,11 +2461,12 @@ mod tests { protocol_state, pending_signature_requests: LookupMap::new(StorageKey::PendingSignatureRequestsV2), pending_ckd_requests: LookupMap::new(StorageKey::PendingCKDRequests), - proposed_updates: ProposedUpdates::default(), - config: Config::default(), - tee_state: TeeState::default(), accept_requests: true, - node_migrations: NodeMigrations::default(), + proposed_updates: Default::default(), + config: Default::default(), + tee_state: Default::default(), + node_migrations: Default::default(), + stale_data: Default::default(), } } } @@ -2844,13 +2896,22 @@ mod tests { let valid_participant_attestation = mpc_attestation::attestation::Attestation::Mock( mpc_attestation::attestation::MockAttestation::Valid, ); - contract.tee_state.add_participant( + + let tee_upgrade_duration = + Duration::from_secs(contract.config.tee_upgrade_deadline_duration_seconds); + + let insertion_result = contract.tee_state.add_participant( NodeId { account_id: self.signer_account_id.clone(), tls_public_key: self.attestation_tls_key.clone().into_contract_type(), account_public_key: Some(self.signer_account_pk.clone()), }, valid_participant_attestation, + tee_upgrade_duration, + ); + assert_matches::assert_matches!( + insertion_result, + Ok(ParticipantInsertion::NewlyInsertedParticipant) ); } @@ -3389,6 +3450,7 @@ mod tests { fn test_verify_tee_triggers_resharing_and_kickout_on_expired_attestation() { const PARTICIPANT_COUNT: usize = 3; const ATTESTATION_EXPIRY_SECONDS: u64 = 5; + const TEE_UPGRADE_DURATION: Duration = Duration::MAX; let participants = gen_participants(PARTICIPANT_COUNT); let parameters = ThresholdParameters::new(participants.clone(), Threshold::new(2)).unwrap(); @@ -3429,11 +3491,12 @@ mod tests { let expiring_attestation = MpcAttestation::Mock(MpcMockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(ATTESTATION_EXPIRY_SECONDS), + expiry_timestamp_seconds: Some(ATTESTATION_EXPIRY_SECONDS), }); contract .tee_state - .add_participant(node_id, expiring_attestation); + .add_participant(node_id, expiring_attestation, TEE_UPGRADE_DURATION) + .expect("mock attestation is not yet expired and valid"); // Capture the running state before verify_tee for comparison let ProtocolContractState::Running(running_state_before) = &contract.protocol_state else { @@ -3489,4 +3552,53 @@ mod tests { assert_eq!(*resharing_state, expected_resharing_state); } + + #[test] + fn test_post_upgrade_cleanup_success() { + // given + let mut contract = MpcContract::init( + ThresholdParameters::new(gen_participants(3), Threshold::new(2)).unwrap(), + None, + ) + .unwrap(); + + let mut mock_stale_map = IterableMap::new(StorageKey::_DeprecatedTeeParticipantAttestation); + let node_pk = bogus_ed25519_near_public_key(); + let node_id = NodeId { + account_id: gen_account_id(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + let attestation = mpc_attestation::attestation::Attestation::Mock( + mpc_attestation::attestation::MockAttestation::Valid, + ); + + mock_stale_map.insert(node_pk.clone(), (node_id, attestation)); + + contract.stale_data.participant_attestations = Some(mock_stale_map); + + // when + contract.post_upgrade_cleanup(); + + // then + assert_matches::assert_matches!(contract.stale_data.participant_attestations, None); + } + + #[test] + #[should_panic(expected = "stale participant_attestations data has already been cleared")] + fn test_post_upgrade_cleanup_panics_if_already_cleared() { + // given + let mut contract = MpcContract::init( + ThresholdParameters::new(gen_participants(3), Threshold::new(2)).unwrap(), + None, + ) + .unwrap(); + + contract.stale_data.participant_attestations = None; + + // when + contract.post_upgrade_cleanup(); + + // then panic + } } diff --git a/crates/contract/src/primitives/time.rs b/crates/contract/src/primitives/time.rs index a849a5a6b..b2441e9b8 100644 --- a/crates/contract/src/primitives/time.rs +++ b/crates/contract/src/primitives/time.rs @@ -1,7 +1,5 @@ -use near_sdk::near; use std::time::Duration; -#[near(serializers=[json])] #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] pub(crate) struct Timestamp { duration_since_unix_epoch: Duration, diff --git a/crates/contract/src/storage_keys.rs b/crates/contract/src/storage_keys.rs index e8cabe6d9..17897946a 100644 --- a/crates/contract/src/storage_keys.rs +++ b/crates/contract/src/storage_keys.rs @@ -13,7 +13,7 @@ pub enum StorageKey { PendingSignatureRequestsV2, ProposedUpdatesEntriesV2, ProposedUpdatesVotesV2, - TeeParticipantAttestation, + _DeprecatedTeeParticipantAttestation, PendingCKDRequests, BackupServicesInfo, NodeMigrations, diff --git a/crates/contract/src/tee/proposal.rs b/crates/contract/src/tee/proposal.rs index 43a46b47b..d814e0cd7 100644 --- a/crates/contract/src/tee/proposal.rs +++ b/crates/contract/src/tee/proposal.rs @@ -54,7 +54,6 @@ impl CodeHashesVotes { /// An allowed Docker image configuration entry containing both the MPC image hash and its /// corresponding launcher compose hash, along with when it was added to the allowlist. -#[near(serializers=[json])] #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct AllowedMpcDockerImage { pub(crate) image_hash: MpcDockerImageHash, @@ -63,7 +62,6 @@ pub struct AllowedMpcDockerImage { } /// Collection of whitelisted Docker code hashes that are the only ones MPC nodes are allowed to /// run. -#[near(serializers=[json])] #[derive(Clone, Default, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub(crate) struct AllowedDockerImageHashes { /// Whitelisted code hashes, sorted by when they were added (oldest first). Expired entries are diff --git a/crates/contract/src/tee/tee_state.rs b/crates/contract/src/tee/tee_state.rs index f58204468..c40575bcc 100644 --- a/crates/contract/src/tee/tee_state.rs +++ b/crates/contract/src/tee/tee_state.rs @@ -1,21 +1,22 @@ use crate::{ primitives::{key_state::AuthenticatedParticipantId, participants::Participants}, - storage_keys::StorageKey, tee::proposal::{ AllowedDockerImageHashes, AllowedMpcDockerImage, CodeHashesVotes, MpcDockerImageHash, }, TryIntoInterfaceType, }; use borsh::{BorshDeserialize, BorshSerialize}; -use contract_interface::types::Ed25519PublicKey; use mpc_attestation::{ - attestation::{Attestation, MockAttestation}, + attestation::{self, Attestation, VerifiedAttestation}, report_data::{ReportData, ReportDataV1}, }; use mpc_primitives::hash::LauncherDockerComposeHash; use near_account_id::AccountId; -use near_sdk::{env, near, store::IterableMap}; -use std::hash::{Hash, Hasher}; +use near_sdk::{env, near}; +use std::{ + collections::BTreeMap, + hash::{Hash, Hasher}, +}; use std::{collections::HashSet, time::Duration}; use utilities::AccountIdExtV1; @@ -57,6 +58,21 @@ pub enum TeeQuoteStatus { /// The participant should not be trusted for TEE-dependent operations. Invalid(String), } + +#[derive(Debug, Clone, thiserror::Error)] +pub(crate) enum AttestationSubmissionError { + #[error("the submitted attestation failed verification, reason: {:?}", .0)] + InvalidAttestation(#[from] attestation::VerificationError), + #[error("the submitted attestation's TLS key is not a valid ED25519 key")] + InvalidTlsKey, +} + +#[derive(Debug)] +pub(crate) enum ParticipantInsertion { + NewlyInsertedParticipant, + UpdatedExistingParticipant, +} + #[derive(Debug)] pub enum TeeValidationResult { /// All participants are valid @@ -68,31 +84,26 @@ pub enum TeeValidationResult { } #[derive(Debug, BorshSerialize, BorshDeserialize)] +pub(crate) struct NodeAttestation { + pub(crate) node_id: NodeId, + pub(crate) verified_attestation: VerifiedAttestation, +} + +#[derive(Default, Debug, BorshSerialize, BorshDeserialize)] 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 [`NodeAttestation`]. /// 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 { - fn default() -> Self { - Self { - allowed_docker_image_hashes: AllowedDockerImageHashes::default(), - allowed_launcher_compose_hashes: vec![], - votes: CodeHashesVotes::default(), - stored_attestations: IterableMap::new(StorageKey::TeeParticipantAttestation), - } - } + pub(crate) stored_attestations: BTreeMap, } impl TeeState { /// Creates a [`TeeState`] with an initial set of participants that will receive a valid mocked attestation. pub(crate) fn with_mocked_participant_attestations(participants: &Participants) -> Self { - let mut participants_attestations = IterableMap::new(StorageKey::TeeParticipantAttestation); + let mut participants_attestations = BTreeMap::new(); participants .participants() @@ -106,7 +117,12 @@ impl TeeState { participants_attestations.insert( participant_info.sign_pk.clone(), - (node_id, Attestation::Mock(MockAttestation::Valid)), + NodeAttestation { + node_id, + verified_attestation: VerifiedAttestation::Mock( + attestation::MockAttestation::Valid, + ), + }, ); }); @@ -115,56 +131,25 @@ impl TeeState { ..Default::default() } } + fn current_time_seconds() -> u64 { let current_time_milliseconds = env::block_timestamp_ms(); current_time_milliseconds / 1_000 } - pub(crate) fn verify_proposed_participant_attestation( + /// Adds a participant attestation for the given node iff the attestation succeeds verification. + pub(crate) fn add_participant( &mut self, - attestation: &Attestation, - tls_public_key: Ed25519PublicKey, - account_public_key: Ed25519PublicKey, + node_id: NodeId, + attestation: Attestation, tee_upgrade_deadline_duration: Duration, - ) -> TeeQuoteStatus { - let expected_report_data: ReportData = - ReportDataV1::new(*tls_public_key.as_bytes(), *account_public_key.as_bytes()).into(); - - match attestation.verify( - expected_report_data.into(), - Self::current_time_seconds(), - &self.get_allowed_mpc_docker_image_hashes(tee_upgrade_deadline_duration), - &self.allowed_launcher_compose_hashes, - ) { - Ok(()) => TeeQuoteStatus::Valid, - Err(err) => TeeQuoteStatus::Invalid(err.to_string()), - } - } - - /// Verifies the TEE quote and Docker image - pub(crate) fn verify_tee_participant( - &mut self, - node_id: &NodeId, - tee_upgrade_deadline_duration: Duration, - ) -> TeeQuoteStatus { - let allowed_mpc_docker_image_hashes = - 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.stored_attestations.get(&node_id.tls_public_key); - let Some(participant_attestation) = participant_attestation else { - return TeeQuoteStatus::Invalid("participant has no attestation".to_string()); - }; - + ) -> Result { // Convert TLS public key - let tls_public_key = match node_id.tls_public_key.clone().try_into_dto_type() { - Ok(value) => value, - Err(err) => { - return TeeQuoteStatus::Invalid(format!( - "could not convert TLS pub key to DTO type: {err}" - )) - } - }; + let tls_public_key = node_id + .tls_public_key + .clone() + .try_into_dto_type() + .map_err(|_| AttestationSubmissionError::InvalidTlsKey)?; // Convert account public key if available // @@ -187,10 +172,47 @@ impl TeeState { let expected_report_data: ReportData = ReportDataV1::new(*tls_public_key.as_bytes(), account_key_bytes).into(); + let verified_attestation = attestation.verify( + expected_report_data.into(), + Self::current_time_seconds(), + &self.get_allowed_mpc_docker_image_hashes(tee_upgrade_deadline_duration), + &self.allowed_launcher_compose_hashes, + )?; + + let tls_pk = node_id.tls_public_key.clone(); + + let insertion = self.stored_attestations.insert( + tls_pk, + NodeAttestation { + node_id, + verified_attestation, + }, + ); + + Ok(match insertion { + Some(_previous_attestation) => ParticipantInsertion::UpdatedExistingParticipant, + None => ParticipantInsertion::NewlyInsertedParticipant, + }) + } + + /// reverifies stored participant attestations. + pub(crate) fn reverify_participants( + &self, + node_id: &NodeId, + tee_upgrade_deadline_duration: Duration, + ) -> TeeQuoteStatus { + let allowed_mpc_docker_image_hashes = + 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.stored_attestations.get(&node_id.tls_public_key); + let Some(participant_attestation) = participant_attestation else { + return TeeQuoteStatus::Invalid("participant has no attestation".to_string()); + }; + // Verify the attestation quote let time_stamp_seconds = Self::current_time_seconds(); - match participant_attestation.1.verify( - expected_report_data.into(), + match participant_attestation.verified_attestation.re_verify( time_stamp_seconds, &allowed_mpc_docker_image_hashes, allowed_launcher_compose_hashes, @@ -200,7 +222,11 @@ impl TeeState { } } - pub fn validate_tee( + /// reverifies stored participant attestations and removes any participant attestation + /// from the internal state that fails reverifications. Reverification can fail, for example, + /// the MPC image hash the attestation was tied to is no longer allowed, or due to certificate + /// expiries. + pub fn reverify_and_cleanup_participants( &mut self, participants: &Participants, tee_upgrade_deadline_duration: Duration, @@ -227,7 +253,7 @@ impl TeeState { }; let tee_status = - self.verify_tee_participant(&node_id, tee_upgrade_deadline_duration); + self.reverify_participants(&node_id, tee_upgrade_deadline_duration); matches!(tee_status, TeeQuoteStatus::Valid) }) @@ -246,23 +272,6 @@ impl TeeState { } } - /// Adds a participant attestation for the given node. - /// - /// Returns: - /// - `true` if this is the first attestation for the node (i.e., a new participant was added). - /// - `false` if the node already had an attestation (the existing one was replaced). - 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.stored_attestations.contains_key(&tls_pk); - - // Must pass owned values, not references - self.stored_attestations - .insert(tls_pk, (node_id, attestation)); - - is_new - } - pub fn vote( &mut self, code_hash: MpcDockerImageHash, @@ -331,7 +340,7 @@ impl TeeState { pub fn get_tee_accounts(&self) -> Vec { self.stored_attestations .values() - .map(|(node_id, _)| node_id.clone()) + .map(|node_attestation| node_attestation.node_id.clone()) .collect() } @@ -339,11 +348,13 @@ impl TeeState { pub fn find_node_id_by_tls_key(&self, tls_public_key: &near_sdk::PublicKey) -> Option { self.stored_attestations .get(tls_public_key) - .map(|(node_id, _)| node_id.clone()) + .map(|node_attestation| node_attestation.node_id.clone()) } /// 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(crate) fn is_caller_an_attested_participant( &self, participants: &Participants, @@ -360,11 +371,11 @@ impl TeeState { .get(&info.sign_pk) .ok_or(AttestationCheckError::AttestationNotFound)?; - if attestation.0.account_id != signer_id { + if attestation.node_id.account_id != signer_id { return Err(AttestationCheckError::AttestationOwnerMismatch); } - if let Some(node_pk) = &attestation.0.account_public_key { + if let Some(node_pk) = &attestation.node_id.account_public_key { if node_pk != &signer_pk { return Err(AttestationCheckError::AttestationKeyMismatch); } @@ -393,10 +404,22 @@ mod tests { use near_account_id::AccountId; use near_sdk::test_utils::VMContextBuilder; use near_sdk::testing_env; + use std::time::Duration; use utilities::AccountIdExtV2; + /// 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_clean_non_participants() { + const TEE_UPGRADE_DURATION: Duration = Duration::from_secs(10000); + let mut tee_state = TeeState::default(); // Create some test participants using test utils @@ -424,9 +447,26 @@ mod tests { }; for node_id in &participant_nodes { - tee_state.add_participant(node_id.clone(), local_attestation.clone()); + let insertion_result = tee_state.add_participant( + node_id.clone(), + local_attestation.clone(), + TEE_UPGRADE_DURATION, + ); + + assert_matches!( + insertion_result, + Ok(ParticipantInsertion::NewlyInsertedParticipant) + ); } - tee_state.add_participant(non_participant_uid.clone(), local_attestation.clone()); + let insertion_result = tee_state.add_participant( + non_participant_uid.clone(), + local_attestation.clone(), + TEE_UPGRADE_DURATION, + ); + assert_matches!( + insertion_result, + Ok(ParticipantInsertion::NewlyInsertedParticipant) + ); // Verify all 4 accounts have TEE info initially assert_eq!(tee_state.stored_attestations.len(), 4); @@ -454,18 +494,292 @@ mod tests { .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 updating_existing_participant_returns_existing_participant() { + // given + const TEE_UPGRADE_DURATION: Duration = Duration::from_secs(10000); + let mut tee_state = TeeState::default(); + + let participant: AccountId = "dave.near".parse().unwrap(); + let local_attestation = Attestation::Mock(MockAttestation::Valid); + + let participant_id = NodeId { + account_id: participant.clone(), + account_public_key: Some(bogus_ed25519_near_public_key()), + tls_public_key: bogus_ed25519_near_public_key(), + }; + + let insertion_result = tee_state.add_participant( + participant_id.clone(), + local_attestation.clone(), + TEE_UPGRADE_DURATION, + ); + assert_matches!( + insertion_result, + Ok(ParticipantInsertion::NewlyInsertedParticipant) + ); + + // when + let re_insertion_result = tee_state.add_participant( + participant_id.clone(), + local_attestation.clone(), + TEE_UPGRADE_DURATION, + ); + + // then + assert_matches!( + re_insertion_result, + Ok(ParticipantInsertion::UpdatedExistingParticipant) + ); + } + + #[test] + fn add_participant_increases_storage_size() { + // given + let mut tee_state = TeeState::default(); + let node_id = NodeId { + account_id: "alice.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + let attestation = Attestation::Mock(MockAttestation::Valid); + + // when + tee_state + .add_participant(node_id, attestation, Duration::from_secs(0)) + .unwrap(); + + // then + assert_eq!( + tee_state.stored_attestations.len(), + 1, + "Internal storage count should increase by exactly one" + ); + } + + #[test] + fn add_participant_indexes_by_tls_key() { + // given + let mut tee_state = TeeState::default(); + let node_id = NodeId { + account_id: "alice.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + let attestation = Attestation::Mock(MockAttestation::Valid); + + // when + tee_state + .add_participant(node_id.clone(), attestation, Duration::from_secs(0)) + .unwrap(); + + // then + assert!( + tee_state + .stored_attestations + .contains_key(&node_id.tls_public_key), + "Entry should be strictly retrievable using the TLS public key" + ); + } + + #[test] + fn add_participant_preserves_node_id_integrity() { + // given + let mut tee_state = TeeState::default(); + let node_id = NodeId { + account_id: "alice.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + let attestation = Attestation::Mock(MockAttestation::Valid); + + // when + tee_state + .add_participant(node_id.clone(), attestation, Duration::from_secs(0)) + .unwrap(); + + // then + let stored_entry = tee_state + .stored_attestations + .get(&node_id.tls_public_key) + .unwrap(); + + assert_eq!( + stored_entry.node_id, node_id, + "The stored NodeId struct must exactly match the inserted one" + ); + } + + #[test] + fn internal_storage_distinguishes_participants_by_tls_key() { + // given + let mut tee_state = TeeState::default(); + + let node_1 = NodeId { + account_id: "alice.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + + let node_2 = NodeId { + account_id: "bob.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + + // when + tee_state + .add_participant( + node_1.clone(), + Attestation::Mock(MockAttestation::Valid), + Duration::from_secs(0), + ) + .unwrap(); + tee_state + .add_participant( + node_2.clone(), + Attestation::Mock(MockAttestation::Valid), + Duration::from_secs(0), + ) + .unwrap(); + + // then + assert_eq!(tee_state.stored_attestations.len(), 2); + assert!(tee_state + .stored_attestations + .contains_key(&node_1.tls_public_key)); + assert!(tee_state + .stored_attestations + .contains_key(&node_2.tls_public_key)); + } + + #[test] + fn re_verify_validates_fresh_attestation() { + // given + let mut tee_state = TeeState::default(); + let node_id = NodeId { + account_id: "fresh.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + + const NOW_SECONDS: u64 = 1000; + + testing_env!(VMContextBuilder::new().block_timestamp(NOW_SECONDS).build()); + + let attestation = Attestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: None, + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: Some(NOW_SECONDS), + }); + + tee_state + .add_participant(node_id.clone(), attestation, Duration::from_secs(0)) + .unwrap(); + + // when + let status = tee_state.reverify_participants(&node_id, Duration::from_secs(0)); + + // then + assert_eq!(status, TeeQuoteStatus::Valid); + } + + #[test] + fn test_re_verify_rejects_expired_attestation() { + // given + let mut tee_state = TeeState::default(); + let node_id = NodeId { + account_id: "about_to_be_expired.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + + const EXPIRY_TIMESTAMP_SECONDS: u64 = 1000; + const ELAPSED_SECONDS: u64 = 200; + + testing_env!(VMContextBuilder::new().block_timestamp(0).build()); + + let attestation = Attestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: None, + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: Some(EXPIRY_TIMESTAMP_SECONDS), + }); + + tee_state + .add_participant(node_id.clone(), attestation, Duration::from_secs(0)) + .unwrap(); + + // when + testing_env!(VMContextBuilder::new() + .block_timestamp( + Duration::from_secs(EXPIRY_TIMESTAMP_SECONDS + ELAPSED_SECONDS).as_nanos() as u64 + ) + .build()); + + let status = tee_state.reverify_participants(&node_id, Duration::from_secs(0)); + + // then + assert_matches!(status, TeeQuoteStatus::Invalid(_)); + } + + #[test] + fn re_verify_succeeds_within_expiry_time() { + // given + let mut tee_state = TeeState::default(); + let node_id = NodeId { + account_id: "valid_check.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + + const EXPIRY_TIMESTAMP_SECONDS: u64 = 1000; + + testing_env!(VMContextBuilder::new() + .block_timestamp(Duration::from_secs(0).as_nanos() as u64) + .build()); + + let attestation = Attestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: None, + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: Some(EXPIRY_TIMESTAMP_SECONDS), + }); + + tee_state + .add_participant(node_id.clone(), attestation, Duration::from_secs(0)) + .unwrap(); + + // when + testing_env!(VMContextBuilder::new() + .block_timestamp(Duration::from_secs(EXPIRY_TIMESTAMP_SECONDS - 1).as_nanos() as u64) + .build()); + + let status = tee_state.reverify_participants(&node_id, Duration::from_secs(0)); + + // then + assert_eq!(status, TeeQuoteStatus::Valid); + } + + #[test] + fn test_re_verify_returns_invalid_for_missing_node() { + // given + let tee_state = TeeState::default(); + let node_id = NodeId { + account_id: "ghost.near".parse().unwrap(), + tls_public_key: bogus_ed25519_near_public_key(), + account_public_key: Some(bogus_ed25519_near_public_key()), + }; + + // when + let status = tee_state.reverify_participants(&node_id, Duration::from_secs(0)); + + // then + assert_matches!(status, TeeQuoteStatus::Invalid(msg) if msg.contains("participant has no attestation")); } #[test] fn test_is_caller_attested_success() { let mut tee_state = TeeState::default(); + let tee_upgrade_duration = Duration::MAX; // Generate 1 participant let participants = gen_participants(1); let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); @@ -483,7 +797,13 @@ mod tests { tls_public_key: participant_info.sign_pk.clone(), account_public_key: Some(signer_pk), }; - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("Attestation is valid on insertion"); // 4. Verify check passes let result = tee_state.is_caller_an_attested_participant(&participants); @@ -493,6 +813,7 @@ mod tests { #[test] fn test_is_caller_attested_success_legacy_no_account_key() { // Tests the case where account_public_key is None (legacy/mock nodes) + let tee_upgrade_duration = Duration::MAX; let mut tee_state = TeeState::default(); let participants = gen_participants(1); let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); @@ -506,7 +827,13 @@ mod tests { tls_public_key: participant_info.sign_pk.clone(), account_public_key: None, }; - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("Attestation is valid on insertion"); let result = tee_state.is_caller_an_attested_participant(&participants); assert_matches!(result, Ok(())); @@ -548,6 +875,7 @@ mod tests { let mut tee_state = TeeState::default(); let participants = gen_participants(1); let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); + let tee_upgrade_duration = Duration::MAX; let signer_pk = bogus_ed25519_near_public_key(); set_signer(account_id, &signer_pk); @@ -562,7 +890,13 @@ mod tests { tls_public_key: participant_info.sign_pk.clone(), account_public_key: Some(signer_pk), }; - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("Attestation is valid on insertion"); let result = tee_state.is_caller_an_attested_participant(&participants); @@ -571,9 +905,11 @@ mod tests { #[test] fn test_err_attestation_key_mismatch() { + // given let mut tee_state = TeeState::default(); let participants = gen_participants(1); let (account_id, _, participant_info) = participants.participants().iter().next().unwrap(); + let tee_upgrade_duration = Duration::MAX; let signer_pk = bogus_ed25519_near_public_key(); set_signer(account_id, &signer_pk); @@ -590,10 +926,18 @@ mod tests { 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)); - + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("Attestation is valid on insertion"); + + // when let result = tee_state.is_caller_an_attested_participant(&participants); + // then assert_matches!(result, Err(AttestationCheckError::AttestationKeyMismatch)); } @@ -624,14 +968,22 @@ mod tests { fn validate_tee_returns_full_when_all_participants_have_valid_attestations() { let mut tee_state = TeeState::default(); let participants = gen_participants(3); + let tee_upgrade_duration = Duration::MAX; // Add valid attestations for all participants for (account_id, _, participant_info) in participants.participants().iter() { let node_id = create_node_id(account_id, &participant_info.sign_pk); - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("mock attestation is valid"); } - let validation_result = tee_state.validate_tee(&participants, TEST_GRACE_PERIOD); + let validation_result = + tee_state.reverify_and_cleanup_participants(&participants, TEST_GRACE_PERIOD); assert_matches!(validation_result, TeeValidationResult::Full); } @@ -641,15 +993,23 @@ mod tests { let mut tee_state = TeeState::default(); let participants = gen_participants(3); let participant_list: Vec<_> = participants.participants().to_vec(); + let tee_upgrade_duration = Duration::MAX; // Add valid attestations for only first 2 participants for (account_id, _, participant_info) in participant_list.iter().take(2) { let node_id = create_node_id(account_id, &participant_info.sign_pk); - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("mock attestation is valid"); } // Third participant has no attestation - let validation_result = tee_state.validate_tee(&participants, TEST_GRACE_PERIOD); + let validation_result = + tee_state.reverify_and_cleanup_participants(&participants, TEST_GRACE_PERIOD); let expected_valid_account_ids = account_ids(&participants)[..2].to_vec(); assert_matches!( @@ -663,6 +1023,7 @@ mod tests { fn validate_tee_returns_partial_when_attestation_is_expired() { let current_time_secs = env::block_timestamp() / 1_000_000_000; let expiry_time_secs = current_time_secs + TEST_GRACE_PERIOD.as_secs(); + let tee_upgrade_duration = Duration::MAX; let mut tee_state = TeeState::default(); let participants = gen_participants(3); @@ -671,7 +1032,13 @@ mod tests { // Add valid attestations for first 2 participants for (account_id, _, participant_info) in participant_list.iter().take(2) { let node_id = create_node_id(account_id, &participant_info.sign_pk); - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("mock attestation is valid"); } // Add expiring attestation for third participant @@ -680,14 +1047,17 @@ mod tests { let expiring_attestation = Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(expiry_time_secs), + expiry_timestamp_seconds: Some(expiry_time_secs), }); - tee_state.add_participant(node_id, expiring_attestation); + tee_state + .add_participant(node_id, expiring_attestation, tee_upgrade_duration) + .expect("mock attestation is valid"); // Advance time to exact expiry boundary set_block_timestamp(expiry_time_secs * 1_000_000_000); - let validation_result = tee_state.validate_tee(&participants, TEST_GRACE_PERIOD); + let validation_result = + tee_state.reverify_and_cleanup_participants(&participants, TEST_GRACE_PERIOD); let expected_valid_account_ids = account_ids(&participants)[..2].to_vec(); assert_matches!( @@ -702,6 +1072,7 @@ mod tests { let current_time_secs = env::block_timestamp() / 1_000_000_000; let expiry_time_secs = current_time_secs + 2 * TEST_GRACE_PERIOD.as_secs(); let before_expiry_time_secs = current_time_secs + TEST_GRACE_PERIOD.as_secs(); + let tee_upgrade_duration = Duration::MAX; let mut tee_state = TeeState::default(); let participants = gen_participants(3); @@ -715,18 +1086,21 @@ mod tests { Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(expiry_time_secs), + expiry_timestamp_seconds: Some(expiry_time_secs), }) } else { Attestation::Mock(MockAttestation::Valid) }; - tee_state.add_participant(node_id, attestation); + tee_state + .add_participant(node_id, attestation, tee_upgrade_duration) + .expect("mock attestation is valid"); } // Advance time, but still before expiry set_block_timestamp(before_expiry_time_secs * 1_000_000_000); - let validation_result = tee_state.validate_tee(&participants, TEST_GRACE_PERIOD); + let validation_result = + tee_state.reverify_and_cleanup_participants(&participants, TEST_GRACE_PERIOD); assert_matches!( validation_result, @@ -736,29 +1110,36 @@ mod tests { } #[test] - fn validate_tee_returns_partial_when_attestation_is_invalid() { + fn add_participant_rejects_invalid_attesations() { let mut tee_state = TeeState::default(); let participants = gen_participants(3); let participant_list: Vec<_> = participants.participants().to_vec(); + let tee_upgrade_duration = Duration::MAX; // Add valid attestations for first 2 participants for (account_id, _, participant_info) in participant_list.iter().take(2) { let node_id = create_node_id(account_id, &participant_info.sign_pk); - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Valid)); + tee_state + .add_participant( + node_id, + Attestation::Mock(MockAttestation::Valid), + tee_upgrade_duration, + ) + .expect("mock attestation is valid"); } // Add invalid attestation for third participant let (account_id, _, participant_info) = &participant_list[2]; let node_id = create_node_id(account_id, &participant_info.sign_pk); - tee_state.add_participant(node_id, Attestation::Mock(MockAttestation::Invalid)); - - let validation_result = tee_state.validate_tee(&participants, TEST_GRACE_PERIOD); + let add_participant_result = tee_state.add_participant( + node_id, + Attestation::Mock(MockAttestation::Invalid), + tee_upgrade_duration, + ); - let expected_valid_account_ids = account_ids(&participants)[..2].to_vec(); assert_matches!( - validation_result, - TeeValidationResult::Partial { participants_with_valid_attestation } - if account_ids(&participants_with_valid_attestation) == expected_valid_account_ids - ); + add_participant_result, + Err(AttestationSubmissionError::InvalidAttestation(_)) + ) } } diff --git a/crates/contract/src/v3_0_2_state.rs b/crates/contract/src/v3_0_2_state.rs index 4ac0bcc32..1a7634f68 100644 --- a/crates/contract/src/v3_0_2_state.rs +++ b/crates/contract/src/v3_0_2_state.rs @@ -8,8 +8,13 @@ //! A better approach: only copy the structures that have changed and import the rest from the existing codebase. use borsh::{BorshDeserialize, BorshSerialize}; +use mpc_attestation::attestation::Attestation; +use mpc_primitives::hash::LauncherDockerComposeHash; use near_account_id::AccountId; -use near_sdk::store::{IterableMap, LookupMap}; +use near_sdk::{ + env, + store::{IterableMap, LookupMap}, +}; use std::collections::HashSet; use crate::{ @@ -19,7 +24,10 @@ use crate::{ signature::{SignatureRequest, YieldIndex}, }, state::ProtocolContractState, - tee::tee_state::TeeState, + tee::{ + proposal::{AllowedDockerImageHashes, CodeHashesVotes}, + tee_state::NodeId, + }, update::{Update, UpdateId}, }; @@ -87,7 +95,15 @@ impl From for crate::update::ProposedUpdates { } } -#[derive(Debug, BorshSerialize, BorshDeserialize)] +#[derive(Debug, BorshDeserialize)] +struct TeeState { + _allowed_docker_image_hashes: AllowedDockerImageHashes, + _allowed_launcher_compose_hashes: Vec, + _votes: CodeHashesVotes, + participants_attestations: IterableMap, +} + +#[derive(Debug, BorshDeserialize)] pub struct MpcContract { protocol_state: ProtocolContractState, pending_signature_requests: LookupMap, @@ -101,15 +117,33 @@ pub struct MpcContract { impl From for crate::MpcContract { fn from(value: MpcContract) -> Self { + let protocol_state = value.protocol_state; + + let crate::ProtocolContractState::Running(running_state) = &protocol_state else { + env::panic_str("Contract must be in running state when migrating."); + }; + + // For the soft release we give every participant a mocked attestation. + // Since this upgrade has a non-backwards compatible change, instead of manually mapping the attestations + // we give everyone a new mock attestation again instead. + // clear previous attestations from the storage trie + let stale_participant_attestations = value.tee_state.participants_attestations; + + let threshold_parameters = &running_state.parameters.participants(); + let tee_state = crate::TeeState::with_mocked_participant_attestations(threshold_parameters); + Self { - protocol_state: value.protocol_state, + protocol_state, pending_signature_requests: value.pending_signature_requests, pending_ckd_requests: value.pending_ckd_requests, proposed_updates: value.proposed_updates.into(), config: value.config.into(), - tee_state: value.tee_state, + tee_state, accept_requests: value.accept_requests, node_migrations: value.node_migrations, + stale_data: crate::StaleData { + participant_attestations: Some(stale_participant_attestations), + }, } } } diff --git a/crates/contract/src/v3_2_0_state.rs b/crates/contract/src/v3_2_0_state.rs new file mode 100644 index 000000000..d307eab33 --- /dev/null +++ b/crates/contract/src/v3_2_0_state.rs @@ -0,0 +1,83 @@ +//! ## Overview +//! This module stores the previous contract state—the one you want to migrate from. +//! The goal is to describe the data layout _exactly_ as it existed before. +//! +//! ## Guideline +//! In theory, you could copy-paste every struct from the specific commit you're migrating from. +//! However, this approach (a) requires manual effort from a developer and (b) increases the binary size. +//! A better approach: only copy the structures that have changed and import the rest from the existing codebase. + +use borsh::BorshDeserialize; +use mpc_attestation::attestation::Attestation; +use mpc_primitives::hash::LauncherDockerComposeHash; +use near_sdk::{ + env, + store::{IterableMap, LookupMap}, +}; + +use crate::{ + node_migrations::NodeMigrations, + primitives::{ + ckd::CKDRequest, + signature::{SignatureRequest, YieldIndex}, + }, + state::ProtocolContractState, + tee::{ + proposal::{AllowedDockerImageHashes, CodeHashesVotes}, + tee_state::NodeId, + }, + update::ProposedUpdates, + Config, +}; + +#[derive(Debug, BorshDeserialize)] +struct TeeState { + _allowed_docker_image_hashes: AllowedDockerImageHashes, + _allowed_launcher_compose_hashes: Vec, + _votes: CodeHashesVotes, + participants_attestations: IterableMap, +} + +#[derive(Debug, BorshDeserialize)] +pub struct MpcContract { + protocol_state: ProtocolContractState, + pending_signature_requests: LookupMap, + pending_ckd_requests: LookupMap, + proposed_updates: ProposedUpdates, + config: Config, + tee_state: TeeState, + accept_requests: bool, + node_migrations: NodeMigrations, +} + +impl From for crate::MpcContract { + fn from(value: MpcContract) -> Self { + let protocol_state = value.protocol_state; + + let crate::ProtocolContractState::Running(running_state) = &protocol_state else { + env::panic_str("Contract must be in running state when migrating."); + }; + + // For the soft release we give every participant a mocked attestation. + // Since this upgrade has a non-backwards compatible change, instead of manually mapping the attestations + // we give everyone a new mock attestation again instead. + // clear previous attestations from the storage trie + let stale_participant_attestations = value.tee_state.participants_attestations; + let threshold_parameters = &running_state.parameters.participants(); + let tee_state = crate::TeeState::with_mocked_participant_attestations(threshold_parameters); + + Self { + protocol_state, + pending_signature_requests: value.pending_signature_requests, + pending_ckd_requests: value.pending_ckd_requests, + proposed_updates: value.proposed_updates, + config: value.config, + tee_state, + accept_requests: value.accept_requests, + node_migrations: value.node_migrations, + stale_data: crate::StaleData { + participant_attestations: Some(stale_participant_attestations), + }, + } + } +} diff --git a/crates/contract/tests/inprocess/attestation_submission.rs b/crates/contract/tests/inprocess/attestation_submission.rs index d21e2fbd7..81f173a40 100644 --- a/crates/contract/tests/inprocess/attestation_submission.rs +++ b/crates/contract/tests/inprocess/attestation_submission.rs @@ -249,7 +249,7 @@ impl TestSetup { Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: Some(hash), launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: None, + expiry_timestamp_seconds: None, }) } } diff --git a/crates/contract/tests/sandbox/common.rs b/crates/contract/tests/sandbox/common.rs index 8dac75179..48343e623 100644 --- a/crates/contract/tests/sandbox/common.rs +++ b/crates/contract/tests/sandbox/common.rs @@ -24,7 +24,11 @@ use mpc_contract::{ update::{ProposeUpdateArgs, UpdateId}, }; use near_account_id::AccountId; -use near_workspaces::{network::Sandbox, Contract}; +use near_workspaces::{ + network::Sandbox, + result::{ExecutionFailure, ExecutionSuccess}, + Contract, +}; use near_workspaces::{result::Execution, Account, Worker}; use rand_core::CryptoRngCore; use serde_json::json; @@ -489,3 +493,16 @@ pub async fn generate_participant_and_submit_attestation( .expect("Attestation submission for new account must succeed."); (new_account, account_id, new_participant) } + +pub async fn cleanup_post_migrate( + contract: &Contract, + account: &Account, +) -> Result { + account + .call(contract.id(), "post_upgrade_cleanup") + .max_gas() + .transact() + .await + .unwrap() + .into_result() +} diff --git a/crates/contract/tests/sandbox/tee.rs b/crates/contract/tests/sandbox/tee.rs index 0d8468241..0d6f2fc9b 100644 --- a/crates/contract/tests/sandbox/tee.rs +++ b/crates/contract/tests/sandbox/tee.rs @@ -436,7 +436,7 @@ async fn new_hash_and_previous_hashes_under_grace_period_pass_attestation_verifi let mock_attestation = MockAttestation::WithConstraints { mpc_docker_image_hash: Some(*approved_hash), launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: None, + expiry_timestamp_seconds: None, }; let attestation = Attestation::Mock(mock_attestation); @@ -514,13 +514,13 @@ async fn get_attestation_returns_some_when_tls_key_associated_with_an_attestatio let participant_1_attestation = Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(u64::MAX), + expiry_timestamp_seconds: Some(u64::MAX), }); let participant_2_attestation = Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(u64::MAX - 1), + expiry_timestamp_seconds: Some(u64::MAX - 1), }); assert_ne!( @@ -570,13 +570,13 @@ async fn get_attestation_overwrites_when_same_tls_key_is_reused() { let first_attestation = Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(u64::MAX), + expiry_timestamp_seconds: Some(u64::MAX), }); let second_attestation = Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(u64::MAX - 1), + expiry_timestamp_seconds: Some(u64::MAX - 1), }); assert_ne!( @@ -685,7 +685,7 @@ async fn test_verify_tee_expired_attestation_triggers_resharing() -> Result<()> let expiring_attestation = Attestation::Mock(MockAttestation::WithConstraints { mpc_docker_image_hash: None, launcher_docker_compose_hash: None, - expiry_time_stamp_seconds: Some(expiry_timestamp), + expiry_timestamp_seconds: Some(expiry_timestamp), }); let submit_result = submit_participant_info( diff --git a/crates/contract/tests/sandbox/upgrade_to_current_contract.rs b/crates/contract/tests/sandbox/upgrade_to_current_contract.rs index 580d5e194..f5194acb6 100644 --- a/crates/contract/tests/sandbox/upgrade_to_current_contract.rs +++ b/crates/contract/tests/sandbox/upgrade_to_current_contract.rs @@ -1,7 +1,8 @@ use crate::sandbox::{ common::{ - call_contract_key_generation, execute_key_generation_and_add_random_state, gen_accounts, - init, propose_and_vote_contract_binary, + call_contract_key_generation, cleanup_post_migrate, + execute_key_generation_and_add_random_state, gen_accounts, init, + propose_and_vote_contract_binary, }, utils::{ consts::PARTICIPANT_LEN, @@ -178,6 +179,10 @@ async fn propose_upgrade_from_production_to_current_binary( state_pre_upgrade, state_post_upgrade, "State of the contract should remain the same post upgrade." ); + + cleanup_post_migrate(&contract, &accounts[0]) + .await + .expect("post migration cleanup works"); } //// Verifies that upgrading the contract preserves state and functionality. diff --git a/crates/contract/tests/sandbox/utils/consts.rs b/crates/contract/tests/sandbox/utils/consts.rs index 3955797ca..f2ba226b9 100644 --- a/crates/contract/tests/sandbox/utils/consts.rs +++ b/crates/contract/tests/sandbox/utils/consts.rs @@ -34,7 +34,7 @@ pub const GAS_FOR_VOTE_UPDATE: Gas = Gas::from_tgas(232); /// Gas required for votes cast before the threshold is reached (votes 1 through N-1). /// These votes are cheap because they only record the vote without triggering the actual /// contract update deployment and migration. -pub const GAS_FOR_VOTE_BEFORE_THRESHOLD: Gas = Gas::from_tgas(4); +pub const GAS_FOR_VOTE_BEFORE_THRESHOLD: Gas = Gas::from_tgas(5); /// Maximum gas expected for the threshold vote that triggers the contract update. /// This vote is more expensive because it deploys the new contract code and executes /// the migration function. diff --git a/crates/contract/tests/snapshots/abi__abi_has_not_changed.snap b/crates/contract/tests/snapshots/abi__abi_has_not_changed.snap index 23fa58e8c..ee1586a72 100644 --- a/crates/contract/tests/snapshots/abi__abi_has_not_changed.snap +++ b/crates/contract/tests/snapshots/abi__abi_has_not_changed.snap @@ -190,7 +190,7 @@ expression: abi "type_schema": { "anyOf": [ { - "$ref": "#/definitions/Attestation" + "$ref": "#/definitions/VerifiedAttestation" }, { "type": "null" @@ -432,6 +432,11 @@ expression: abi } } }, + { + "name": "post_upgrade_cleanup", + "doc": " Removes stale data from the contract to be removed after a contract upgrade. Some\n containers are expensive to run destructors on, thus we don't include it in the contract upgrade itself,\n as it can run out of gas. Thus we create methods to run these destructors manually post upgrade.", + "kind": "call" + }, { "name": "propose_update", "doc": " Propose update to either code or config, but not both of them at the same time.", @@ -1961,7 +1966,7 @@ expression: abi "WithConstraints": { "type": "object", "properties": { - "expiry_time_stamp_seconds": { + "expiry_timestamp_seconds": { "description": "Unix time stamp for when this attestation expires.", "type": [ "integer", @@ -2798,6 +2803,72 @@ expression: abi "format": "uint64", "minimum": 0.0 }, + "VerifiedAttestation": { + "oneOf": [ + { + "type": "object", + "required": [ + "Dtack" + ], + "properties": { + "Dtack": { + "$ref": "#/definitions/VerifiedDstackAttestation" + } + }, + "additionalProperties": false + }, + { + "type": "object", + "required": [ + "Mock" + ], + "properties": { + "Mock": { + "$ref": "#/definitions/MockAttestation" + } + }, + "additionalProperties": false + } + ] + }, + "VerifiedDstackAttestation": { + "type": "object", + "required": [ + "expiry_timestamp_seconds", + "launcher_compose_hash", + "mpc_image_hash" + ], + "properties": { + "expiry_timestamp_seconds": { + "description": "Unix time stamp for when this attestation expires.", + "type": "integer", + "format": "uint64", + "minimum": 0.0 + }, + "launcher_compose_hash": { + "description": "The digest of the launcher compose file running.", + "type": "array", + "items": { + "type": "integer", + "format": "uint8", + "minimum": 0.0 + }, + "maxItems": 32, + "minItems": 32 + }, + "mpc_image_hash": { + "description": "The digest of the MPC image running.", + "type": "array", + "items": { + "type": "integer", + "format": "uint8", + "minimum": 0.0 + }, + "maxItems": 32, + "minItems": 32 + } + } + }, "YieldIndex": { "description": "The index into calling the YieldResume feature of NEAR. This will allow to resume a yield call after the contract has been called back via this index.", "type": "object", diff --git a/crates/mpc-attestation/Cargo.toml b/crates/mpc-attestation/Cargo.toml index d1dcccccd..9dab92cda 100644 --- a/crates/mpc-attestation/Cargo.toml +++ b/crates/mpc-attestation/Cargo.toml @@ -17,6 +17,6 @@ sha2 = { workspace = true } sha3 = { workspace = true } [dev-dependencies] +assert_matches = { workspace = true } dcap-qvl = { workspace = true } -rstest = { workspace = true } test-utils = { workspace = true } diff --git a/crates/mpc-attestation/src/attestation.rs b/crates/mpc-attestation/src/attestation.rs index fa0e852a6..ea2ec0641 100644 --- a/crates/mpc-attestation/src/attestation.rs +++ b/crates/mpc-attestation/src/attestation.rs @@ -1,20 +1,18 @@ +use alloc::vec::Vec; use attestation::{ app_compose::AppCompose, attestation::{GetSingleEvent as _, OrErr as _}, measurements::ExpectedMeasurements, measurements::Measurements, report_data::ReportData, - tcb_info::TcbInfo, }; use include_measurements::include_measurements; pub use attestation::attestation::{DstackAttestation, VerificationError}; - use mpc_primitives::hash::{LauncherDockerComposeHash, MpcDockerImageHash}; use borsh::{BorshDeserialize, BorshSerialize}; -use core::ops::Deref as _; use serde::{Deserialize, Serialize}; use sha2::{Digest as _, Sha256}; @@ -23,126 +21,217 @@ use crate::alloc::string::ToString; const MPC_IMAGE_HASH_EVENT: &str = "mpc-image-digest"; +// TODO(#1639): extract timestamp from certificate itself +pub const DEFAULT_EXPIRATION_DURATION_SECONDS: u64 = 60 * 60 * 24 * 7; // 7 days + #[allow(clippy::large_enum_variant)] -#[derive(Clone, Debug, Serialize, Deserialize, BorshDeserialize, BorshSerialize)] +#[derive(Clone, Debug, Serialize, Deserialize, BorshSerialize, BorshDeserialize)] pub enum Attestation { Dstack(DstackAttestation), Mock(MockAttestation), } +#[allow(clippy::large_enum_variant)] +#[derive(Clone, Debug, Serialize, Deserialize, BorshDeserialize, BorshSerialize)] +pub enum VerifiedAttestation { + Dstack(ValidatedDstackAttestation), + Mock(MockAttestation), +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize, BorshDeserialize, BorshSerialize)] +pub enum MockAttestation { + #[default] + /// Always pass validation + Valid, + /// Always fails validation + Invalid, + /// Pass validation depending on the set constraints + WithConstraints { + mpc_docker_image_hash: Option, + launcher_docker_compose_hash: Option, + /// Unix time stamp for when this attestation expires. + expiry_timestamp_seconds: Option, + }, +} + +#[derive(Clone, Debug, Serialize, Deserialize, BorshDeserialize, BorshSerialize)] +pub struct ValidatedDstackAttestation { + pub mpc_image_hash: MpcDockerImageHash, + pub launcher_compose_hash: LauncherDockerComposeHash, + // TODO(#1639): This timestamp can not come from the contract, + // but should be extracted from the certificate itself. + pub expiry_timestamp_seconds: u64, +} + +impl VerifiedAttestation { + pub fn re_verify( + &self, + timestamp_seconds: u64, + allowed_mpc_docker_image_hashes: &[MpcDockerImageHash], + allowed_launcher_docker_compose_hashes: &[LauncherDockerComposeHash], + ) -> Result<(), VerificationError> { + match self { + Self::Dstack(ValidatedDstackAttestation { + mpc_image_hash, + launcher_compose_hash, + expiry_timestamp_seconds: expiration_timestamp_seconds, + }) => { + let attestation_has_expired = *expiration_timestamp_seconds < timestamp_seconds; + + if attestation_has_expired { + return Err(VerificationError::Custom(format!( + "The attestation expired at t = {:?}, time_now = {:?}", + expiration_timestamp_seconds, timestamp_seconds + ))); + } + + let () = verify_mpc_hash(mpc_image_hash, allowed_mpc_docker_image_hashes)?; + let () = verify_launcher_compose_hash( + launcher_compose_hash, + allowed_launcher_docker_compose_hashes, + )?; + + Ok(()) + } + Self::Mock(mock_attestation) => verify_mock_attestation( + mock_attestation, + allowed_mpc_docker_image_hashes, + allowed_launcher_docker_compose_hashes, + timestamp_seconds, + ), + } + } +} + impl Attestation { pub fn verify( &self, expected_report_data: ReportData, - timestamp_seconds: u64, + current_timestamp_seconds: u64, allowed_mpc_docker_image_hashes: &[MpcDockerImageHash], allowed_launcher_docker_compose_hashes: &[LauncherDockerComposeHash], - ) -> Result<(), VerificationError> { - let attestation = match self { + ) -> Result { + match self { Self::Dstack(dstack_attestation) => { // Makes MPC related attestation verification first - if allowed_mpc_docker_image_hashes.is_empty() { - return Err(VerificationError::Custom( - "the allowed mpc image hashes list is empty".to_string(), - )); - } - if allowed_launcher_docker_compose_hashes.is_empty() { - return Err(VerificationError::Custom( - "the allowed mpc laucher compose hashes list is empty".to_string(), - )); - } - self.verify_mpc_hash( - &dstack_attestation.tcb_info, - allowed_mpc_docker_image_hashes, - )?; - self.verify_launcher_compose_hash( - &dstack_attestation.tcb_info, + let mpc_image_hash: MpcDockerImageHash = { + let mpc_image_hash_payload = &dstack_attestation + .tcb_info + .get_single_event(MPC_IMAGE_HASH_EVENT)? + .event_payload; + + let mpc_image_hash_bytes: Vec = hex::decode(mpc_image_hash_payload) + .map_err(|err| { + VerificationError::Custom(format!( + "provided mpc image is not hex encoded: {:?}", + err + )) + })?; + let mpc_image_hash_bytes: [u8; 32] = + mpc_image_hash_bytes.try_into().map_err(|_| { + VerificationError::Custom( + "The provided MPC image hash is not 32 bytes".to_string(), + ) + })?; + MpcDockerImageHash::from(mpc_image_hash_bytes) + }; + + let () = verify_mpc_hash(&mpc_image_hash, allowed_mpc_docker_image_hashes)?; + + let launcher_compose_hash: LauncherDockerComposeHash = { + let app_compose: AppCompose = + serde_json::from_str(&dstack_attestation.tcb_info.app_compose) + .map_err(|e| VerificationError::AppComposeParsing(e.to_string()))?; + + let launcher_compose_hash_bytes: [u8; 32] = + Sha256::digest(app_compose.docker_compose_file.as_bytes()).into(); + + LauncherDockerComposeHash::from(launcher_compose_hash_bytes) + }; + + let () = verify_launcher_compose_hash( + &launcher_compose_hash, allowed_launcher_docker_compose_hashes, )?; - dstack_attestation + let accepted_measurements = [ + include_measurements!("assets/tcb_info.json"), + // TODO(#1433): Security - remove dev measurements from production builds after testing is complete + include_measurements!("assets/tcb_info_dev.json"), + ]; + + dstack_attestation.verify( + expected_report_data, + current_timestamp_seconds, + &accepted_measurements, + )?; + + // TODO(#1639): extract timestamp from certificate itself + let expiration_timestamp_seconds = + current_timestamp_seconds + DEFAULT_EXPIRATION_DURATION_SECONDS; + Ok(VerifiedAttestation::Dstack(ValidatedDstackAttestation { + mpc_image_hash, + launcher_compose_hash, + expiry_timestamp_seconds: expiration_timestamp_seconds, + })) } Self::Mock(mock_attestation) => { // Override attestation verification for this case - return verify_mock_attestation( + let () = verify_mock_attestation( mock_attestation, allowed_mpc_docker_image_hashes, allowed_launcher_docker_compose_hashes, - timestamp_seconds, - ); + current_timestamp_seconds, + )?; + + Ok(VerifiedAttestation::Mock(mock_attestation.clone())) } - }; - - let accepted_measurements = [ - include_measurements!("assets/tcb_info.json"), - // TODO(#1433): Security - remove dev measurements from production builds after testing is complete - include_measurements!("assets/tcb_info_dev.json"), - ]; - - attestation.verify( - expected_report_data, - timestamp_seconds, - &accepted_measurements, - ) + } } +} - /// Verifies MPC node image hash is in allowed list. - fn verify_mpc_hash( - &self, - tcb_info: &TcbInfo, - allowed_hashes: &[MpcDockerImageHash], - ) -> Result<(), VerificationError> { - let event = tcb_info.get_single_event(MPC_IMAGE_HASH_EVENT)?; - - allowed_hashes - .iter() - .any(|hash| hash.as_hex() == *event.event_payload) - .or_err(|| { - VerificationError::Custom(format!( - "MPC image hash {} is not in the allowed hashes list", - event.event_payload.clone() - )) - }) +/// Verifies MPC node image hash is in allowed list. +fn verify_mpc_hash( + image_hash: &MpcDockerImageHash, + allowed_hashes: &[MpcDockerImageHash], +) -> Result<(), VerificationError> { + if allowed_hashes.is_empty() { + return Err(VerificationError::Custom( + "the allowed mpc image hashes list is empty".to_string(), + )); } - fn verify_launcher_compose_hash( - &self, - tcb_info: &TcbInfo, - allowed_hashes: &[LauncherDockerComposeHash], - ) -> Result<(), VerificationError> { - let app_compose: AppCompose = serde_json::from_str(&tcb_info.app_compose) - .map_err(|e| VerificationError::AppComposeParsing(e.to_string()))?; - - let launcher_bytes: [u8; 32] = - Sha256::digest(app_compose.docker_compose_file.as_bytes()).into(); - - allowed_hashes - .iter() - .any(|hash| hash.deref() == &launcher_bytes) - .or_err(|| { - VerificationError::Custom(format!( - "launcher compose hash {} is not in the allowed hashes list", - hex::encode(launcher_bytes.as_ref(),) - )) - }) + let image_hash_is_allowed = allowed_hashes.contains(image_hash); + if !image_hash_is_allowed { + return Err(VerificationError::Custom(format!( + "MPC image hash {:?} is not in the allowed hashes list", + image_hash + ))); } + + Ok(()) } -#[derive(Debug, Default, Clone, Serialize, Deserialize, BorshDeserialize, BorshSerialize)] -pub enum MockAttestation { - #[default] - /// Always pass validation - Valid, - /// Always fails validation - Invalid, - /// Pass validation depending on the set constraints - WithConstraints { - mpc_docker_image_hash: Option, - launcher_docker_compose_hash: Option, +fn verify_launcher_compose_hash( + launcher_compose_hash: &LauncherDockerComposeHash, + allowed_hashes: &[LauncherDockerComposeHash], +) -> Result<(), VerificationError> { + if allowed_hashes.is_empty() { + return Err(VerificationError::Custom( + "the allowed mpc launcher compose hashes list is empty".to_string(), + )); + } - /// Unix time stamp for when this attestation expires. - expiry_time_stamp_seconds: Option, - }, + let launcher_compose_hash_is_allowed = allowed_hashes.contains(launcher_compose_hash); + + if !launcher_compose_hash_is_allowed { + return Err(VerificationError::Custom(format!( + "MPC launcher compose hash {:?} is not in the allowed hashes list", + launcher_compose_hash + ))); + } + + Ok(()) } pub(crate) fn verify_mock_attestation( @@ -157,7 +246,7 @@ pub(crate) fn verify_mock_attestation( MockAttestation::WithConstraints { mpc_docker_image_hash, launcher_docker_compose_hash, - expiry_time_stamp_seconds: expiry_timestamp_seconds, + expiry_timestamp_seconds, } => { if let Some(hash) = mpc_docker_image_hash { if allowed_mpc_docker_image_hashes.is_empty() { @@ -176,7 +265,7 @@ pub(crate) fn verify_mock_attestation( if let Some(hash) = launcher_docker_compose_hash { if allowed_launcher_docker_compose_hashes.is_empty() { return Err(VerificationError::Custom( - "the allowed mpc laucher compose hashes list is empty".to_string(), + "the allowed mpc launcher compose hashes list is empty".to_string(), )); } allowed_launcher_docker_compose_hashes @@ -201,3 +290,164 @@ pub(crate) fn verify_mock_attestation( } } } + +#[cfg(test)] +mod tests { + use alloc::vec; + + use super::*; + + #[test] + fn mock_constrained_verification_passes_if_hash_in_allowed_list() { + let allowed_hash = MpcDockerImageHash::from([42; 32]); + + let hash_constrained_attestation = + VerifiedAttestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: Some(allowed_hash.clone()), + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: None, + }); + + let other_hash = MpcDockerImageHash::from([1; 32]); + let allowed_mpc_hashes: Vec = vec![other_hash, allowed_hash]; + + hash_constrained_attestation + .re_verify(0, &allowed_mpc_hashes, &[]) + .expect("constrained mpc image hash is allowed and should therefore pass validation"); + } + + #[test] + fn mock_constrained_verification_fails_if_hash_not_in_allowed_list() { + let restricted_hash = MpcDockerImageHash::from([42; 32]); + + let hash_constrained_attestation = + VerifiedAttestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: Some(restricted_hash), + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: None, + }); + + let other_hash = MpcDockerImageHash::from([1; 32]); + let allowed_mpc_hashes: Vec = vec![other_hash]; + + let result = hash_constrained_attestation.re_verify(0, &allowed_mpc_hashes, &[]); + + match result { + Err(VerificationError::Custom(msg)) => { + assert!( + msg.contains("MPC image hash"), + "Expected error message regarding MPC image hash, got: {}", + msg + ); + } + _ => panic!("Expected Custom VerificationError, got: {:?}", result), + } + } + + #[test] + fn mock_constrained_verification_fails_if_allowed_list_is_empty() { + let restricted_hash = MpcDockerImageHash::from([42; 32]); + + let hash_constrained_attestation = + VerifiedAttestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: Some(restricted_hash), + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: None, + }); + + let allowed_mpc_hashes: Vec = vec![]; + + let result = hash_constrained_attestation.re_verify(0, &allowed_mpc_hashes, &[]); + + match result { + Err(VerificationError::Custom(msg)) => { + assert!(msg.contains("list is empty")); + } + _ => panic!("Expected empty list error, got: {:?}", result), + } + } + + #[test] + fn launcher_constraint_passes_if_hash_in_allowed_list() { + let allowed_hash = LauncherDockerComposeHash::from([99; 32]); + + let hash_constrained_attestation = + VerifiedAttestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: None, + launcher_docker_compose_hash: Some(allowed_hash.clone()), + expiry_timestamp_seconds: None, + }); + + let other_hash = LauncherDockerComposeHash::from([1; 32]); + let allowed_launcher_hashes: Vec = + vec![other_hash, allowed_hash]; + + hash_constrained_attestation + .re_verify(0, &[], &allowed_launcher_hashes) + .expect("constrained launcher hash is allowed and should therefore pass validation"); + } + + #[test] + fn launcher_constraint_fails_if_hash_not_in_allowed_list() { + let restricted_hash = LauncherDockerComposeHash::from([99; 32]); + + let hash_constrained_attestation = + VerifiedAttestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: None, + launcher_docker_compose_hash: Some(restricted_hash), + expiry_timestamp_seconds: None, + }); + + let other_hash = LauncherDockerComposeHash::from([1; 32]); + let allowed_launcher_hashes: Vec = vec![other_hash]; + + let result = hash_constrained_attestation.re_verify(0, &[], &allowed_launcher_hashes); + + match result { + Err(VerificationError::Custom(msg)) => { + assert!(msg.contains("launcher compose hash")); + } + _ => panic!("Expected Custom VerificationError, got: {:?}", result), + } + } + + #[test] + fn mock_time_constraint_passes_if_time_is_within_expiry_window() { + let expiry_timestamp_seconds = 101; + let time_now = 100; + + let time_constrained_attestation = + VerifiedAttestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: None, + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: Some(expiry_timestamp_seconds), + }); + + time_constrained_attestation + .re_verify(time_now, &[], &[]) + .expect("Attestation is within valid time window and should pass"); + } + + #[test] + fn time_constraint_fails_if_time_is_past_expiry_window() { + let expiry_timestamp_seconds = 100; + let time_now = 101; + + let time_constrained_attestation = + VerifiedAttestation::Mock(MockAttestation::WithConstraints { + mpc_docker_image_hash: None, + launcher_docker_compose_hash: None, + expiry_timestamp_seconds: Some(expiry_timestamp_seconds), + }); + + let verification_result = time_constrained_attestation.re_verify(time_now, &[], &[]); + + assert_matches::assert_matches!( + verification_result, + Err(VerificationError::ExpiredCertificate { + attestation_time, + expiry_time, + }) if attestation_time == time_now && expiry_time == expiry_timestamp_seconds + ); + } +} diff --git a/crates/mpc-attestation/tests/test_attestation_verification.rs b/crates/mpc-attestation/tests/test_attestation_verification.rs index 2cfa43b22..cfbc4302d 100644 --- a/crates/mpc-attestation/tests/test_attestation_verification.rs +++ b/crates/mpc-attestation/tests/test_attestation_verification.rs @@ -1,53 +1,156 @@ +use assert_matches::assert_matches; use attestation::attestation::VerificationError; -use mpc_primitives::hash::{LauncherDockerComposeHash, MpcDockerImageHash}; -use rstest::rstest; +use mpc_attestation::attestation::{ + Attestation, DEFAULT_EXPIRATION_DURATION_SECONDS, MockAttestation, VerifiedAttestation, +}; +use mpc_attestation::report_data::{ReportData, ReportDataV1}; use test_utils::attestation::{ account_key, image_digest, launcher_compose_digest, mock_dstack_attestation, p2p_tls_key, }; -use mpc_attestation::attestation::{Attestation, MockAttestation}; -use mpc_attestation::report_data::{ReportData, ReportDataV1}; +#[test] +fn valid_mock_attestation_succeeds_verification() { + let valid_attestation = Attestation::Mock(MockAttestation::Valid); -#[rstest] -#[case(MockAttestation::Valid, Ok(()))] -#[case( - MockAttestation::Invalid, - Err(VerificationError::InvalidMockAttestation) -)] -fn test_mock_attestation_verify( - #[case] local_attestation: MockAttestation, - #[case] expected_quote_verification_result: Result<(), VerificationError>, -) { let timestamp_s = 0u64; let tls_key = p2p_tls_key(); let account_key = account_key(); let report_data = ReportData::V1(ReportDataV1::new(tls_key, account_key)); - let attestation = Attestation::Mock(local_attestation); + assert_matches!( + valid_attestation.verify(report_data.into(), timestamp_s, &[], &[]), + Ok(VerifiedAttestation::Mock(MockAttestation::Valid)) + ); +} + +#[test] +fn invalid_mock_attestation_fails_verification() { + let valid_attestation = Attestation::Mock(MockAttestation::Invalid); + + let timestamp_s = 0u64; + let tls_key = p2p_tls_key(); + let account_key = account_key(); + let report_data = ReportData::V1(ReportDataV1::new(tls_key, account_key)); - assert_eq!( - attestation.verify(report_data.into(), timestamp_s, &[], &[]), - expected_quote_verification_result + assert_matches!( + valid_attestation.verify(report_data.into(), timestamp_s, &[], &[]), + Err(VerificationError::InvalidMockAttestation) ); } #[test] -fn test_verify_method_signature() { +fn validated_dstack_attestation_can_be_reverified() { + // given let attestation = mock_dstack_attestation(); let tls_key = p2p_tls_key(); let account_key = account_key(); + let report_data = ReportData::V1(ReportDataV1::new(tls_key, account_key)); + let timestamp_s = 1763626832_u64; + let allowed_mpc_hashes = [image_digest()]; + let allowed_launcher_hashes = [launcher_compose_digest()]; + + let validated = attestation + .verify( + report_data.into(), + timestamp_s, + &allowed_mpc_hashes, + &allowed_launcher_hashes, + ) + .expect("Initial verification failed"); + + // when + let re_verification_result = validated.re_verify( + timestamp_s + DEFAULT_EXPIRATION_DURATION_SECONDS, + &allowed_mpc_hashes, + &allowed_launcher_hashes, + ); + + // then + assert_matches!(re_verification_result, Ok(())); +} +#[test] +fn validated_dstack_attestation_fails_reverification_when_expired() { + // given + let attestation = mock_dstack_attestation(); + let tls_key = p2p_tls_key(); + let account_key = account_key(); + let report_data = ReportData::V1(ReportDataV1::new(tls_key, account_key)); + let timestamp_s = 1763626832_u64; + let allowed_mpc_hashes = [image_digest()]; + let allowed_launcher_hashes = [launcher_compose_digest()]; + + let validated = attestation + .verify( + report_data.into(), + timestamp_s, + &allowed_mpc_hashes, + &allowed_launcher_hashes, + ) + .expect("Initial verification failed"); + + // when + let re_verification_result = validated.re_verify( + timestamp_s + DEFAULT_EXPIRATION_DURATION_SECONDS + 1, + &allowed_mpc_hashes, + &allowed_launcher_hashes, + ); + + // then + assert_matches!( + re_verification_result, + Err(VerificationError::Custom(msg)) if msg.contains("The attestation expired") + ); +} + +#[test] +fn validated_mock_attestation_passes_reverification() { + let valid_attestation = Attestation::Mock(MockAttestation::Valid); + let tls_key = p2p_tls_key(); + let account_key = account_key(); let report_data: ReportData = ReportDataV1::new(tls_key, account_key).into(); - let timestamp_s = 1763626832_u64; //Thursday, 20 November 2025 08:20:32 - let allowed_mpc_image_digest: MpcDockerImageHash = image_digest(); - let allowed_launcher_compose_digest: LauncherDockerComposeHash = launcher_compose_digest(); + let validated = valid_attestation + .verify(report_data.into(), 0, &[], &[]) + .expect("Initial verification failed"); + + // Mock should generally pass re-verify + assert_matches!(validated.re_verify(100, &[], &[]), Ok(())); +} + +#[test] +fn validated_dstack_attestation_fails_reverification_with_rotated_hashes() { + let attestation = mock_dstack_attestation(); + let tls_key = p2p_tls_key(); + let account_key = account_key(); + let report_data: ReportData = ReportDataV1::new(tls_key, account_key).into(); + let creation_time = 1763626832_u64; + + let allowed_mpc_hashes = [image_digest()]; + let allowed_launcher_hashes = [launcher_compose_digest()]; + + // 1. Initial verify succeeds with the "old" allowed list + let validated = attestation + .verify( + report_data.into(), + creation_time, + &allowed_mpc_hashes, + &allowed_launcher_hashes, + ) + .expect("Initial verification should succeed"); + + let new_allowed_mpc_docker_image_hashes = [[42; 32].into()]; + + // 2. Re-verify fails if we remove the allowed hash (e.g. strict rotation) + let result = validated.re_verify( + creation_time, + &new_allowed_mpc_docker_image_hashes, + &allowed_launcher_hashes, + ); - let verification_result = attestation.verify( - report_data.into(), - timestamp_s, - &[allowed_mpc_image_digest], - &[allowed_launcher_compose_digest], + assert_matches!( + result, + Err(VerificationError::Custom(msg)) + if msg.contains("not in the allowed hashes list") ); - assert!(verification_result.is_ok()); } diff --git a/crates/node/src/indexer.rs b/crates/node/src/indexer.rs index 598400635..7bef9c5ab 100644 --- a/crates/node/src/indexer.rs +++ b/crates/node/src/indexer.rs @@ -183,7 +183,7 @@ impl IndexerViewClient { &self, mpc_contract_id: &AccountId, participant_tls_public_key: &contract_interface::types::Ed25519PublicKey, - ) -> anyhow::Result> { + ) -> anyhow::Result> { let get_attestation_args: Vec = serde_json::to_string(&GetAttestationArgs { tls_public_key: participant_tls_public_key, }) @@ -210,7 +210,7 @@ impl IndexerViewClient { match query_response.kind { QueryResponseKind::CallResult(call_result) => serde_json::from_slice::< - Option, + Option, >(&call_result.result) .context("failed to deserialize pending request response"), _ => { diff --git a/crates/node/src/indexer/tx_sender.rs b/crates/node/src/indexer/tx_sender.rs index 709c7b6e7..583e2a9a4 100644 --- a/crates/node/src/indexer/tx_sender.rs +++ b/crates/node/src/indexer/tx_sender.rs @@ -4,17 +4,20 @@ use super::IndexerState; use crate::config::RespondConfig; use crate::metrics; use anyhow::Context; +use contract_interface::types::{Attestation, VerifiedAttestation}; use ed25519_dalek::SigningKey; +use mpc_attestation::attestation::DEFAULT_EXPIRATION_DURATION_SECONDS; use near_account_id::AccountId; use near_indexer_primitives::types::Gas; use std::future::Future; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::{mpsc, oneshot}; use tokio::time; const TRANSACTION_PROCESSOR_CHANNEL_SIZE: usize = 10000; const TRANSACTION_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_ATTESTATION_AGE: Duration = Duration::from_secs(60 * 2); pub trait TransactionSender: Clone + Send + Sync { fn send( @@ -207,11 +210,61 @@ async fn observe_tx_result( .get_participant_attestation(&indexer_state.mpc_contract_id, tls_public_key) .await?; + let Some(stored_attestation) = attestation_stored_on_contract else { + return Ok(TransactionStatus::NotExecuted); + }; + + let submitted_attestation = + &submit_participant_info_args.proposed_participant_attestation; + let submitted_attestation_is_on_chain = - attestation_stored_on_contract.is_some_and(|stored_attestation| { - stored_attestation - == submit_participant_info_args.proposed_participant_attestation - }); + match (stored_attestation, submitted_attestation) { + ( + VerifiedAttestation::Dtack(verified_dstack_attestation), + Attestation::Dstack(_), + ) => { + // Check if the attestation stored on chain is fresh by verifying its age + // is less than `MAX_ATTESTATION_AGE` + // + // TODO(#1637): extract expiration timestamp from the certificate itself, + // instead of using heuristics. + let expiry_timestamp_seconds = + verified_dstack_attestation.expiry_timestamp_seconds; + + let Some(attestation_duration_since_unix_epoch) = expiry_timestamp_seconds + .checked_sub(DEFAULT_EXPIRATION_DURATION_SECONDS) + .map(Duration::from_secs) + else { + tracing::error!( + ?expiry_timestamp_seconds, + "could not calculate attestation storage time" + ); + + return Ok(TransactionStatus::NotExecuted); + }; + + let timestamp_seconds_now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("could not calculate system time")?; + + let attestation_age = + attestation_duration_since_unix_epoch.abs_diff(timestamp_seconds_now); + let attestation_is_fresh = attestation_age < MAX_ATTESTATION_AGE; + + tracing::info!( + ?attestation_age, + ?attestation_is_fresh, + "node found dstack attestation on chain" + ); + + attestation_is_fresh + } + ( + VerifiedAttestation::Mock(stored_mock_attestation), + Attestation::Mock(submitted_mock_attestation), + ) => stored_mock_attestation == *submitted_mock_attestation, + _ => false, + }; if submitted_attestation_is_on_chain { Ok(TransactionStatus::Executed) diff --git a/crates/node/src/tee/remote_attestation.rs b/crates/node/src/tee/remote_attestation.rs index 32921f69a..5d14c1547 100644 --- a/crates/node/src/tee/remote_attestation.rs +++ b/crates/node/src/tee/remote_attestation.rs @@ -104,12 +104,14 @@ fn validate_remote_attestation( .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_secs(); - attestation.verify( - expected_report_data.into(), - now, - allowed_docker_image_hashes, - allowed_launcher_compose_hashes, - ) + attestation + .verify( + expected_report_data.into(), + now, + allowed_docker_image_hashes, + allowed_launcher_compose_hashes, + ) + .map(|_| ()) } pub async fn validate_and_submit_remote_attestation( diff --git a/crates/node/src/trait_extensions/convert_to_contract_dto.rs b/crates/node/src/trait_extensions/convert_to_contract_dto.rs index 018db5ba8..91793ab23 100644 --- a/crates/node/src/trait_extensions/convert_to_contract_dto.rs +++ b/crates/node/src/trait_extensions/convert_to_contract_dto.rs @@ -67,11 +67,11 @@ impl IntoContractInterfaceType for M MockAttestation::WithConstraints { mpc_docker_image_hash, launcher_docker_compose_hash, - expiry_time_stamp_seconds, + expiry_timestamp_seconds, } => contract_interface::types::MockAttestation::WithConstraints { mpc_docker_image_hash: mpc_docker_image_hash.map(Into::into), launcher_docker_compose_hash: launcher_docker_compose_hash.map(Into::into), - expiry_time_stamp_seconds, + expiry_timestamp_seconds, }, } } diff --git a/libs/nearcore b/libs/nearcore index d2a99fe98..3caaf4463 160000 --- a/libs/nearcore +++ b/libs/nearcore @@ -1 +1 @@ -Subproject commit d2a99fe9896b3eb5321b3237fbac2ca1ab511a25 +Subproject commit 3caaf44631859fb8a962b70c883f699bcc10f318