fix(contract): make mock attestations cleanable via expiry (#3785)

Author: barakeinav1

crates/contract/README.md

@@ -359,7 +359,7 @@ The MPC nodes will eventually run inside a Trusted Execution Environments (TEE).
 
 Participants that run their node inside a TEE will have to submit the following TEE related data to the contract:
 
-```rust
+```rust,ignore
 pub struct DstackAttestation {
     /// TEE Remote Attestation Quote that proves the participant's identity.
     pub quote: Quote,

crates/contract/src/tee/tee_state.rs

@@ -146,7 +146,7 @@ impl TeeState {
         tee_state
     }
 
-    fn current_time_seconds() -> u64 {
+    pub(crate) fn current_time_seconds() -> u64 {
         env::block_timestamp_ms() / 1_000
     }
 
@@ -717,6 +717,81 @@ mod tests {
         );
     }
 
+    #[test]
+    fn clean_invalid_attestations__should_remove_accepted_mock_valid_after_expiry() {
+        // Given: a `MockAttestation::Valid` accepted via the normal submission path.
+        // It must be stamped with an expiry so it can eventually be cleaned up.
+        testing_env!(VMContextBuilder::new().block_timestamp(0).build());
+
+        let mut tee_state = TeeState::default();
+        let node_id = NodeId {
+            account_id: "alice.near".parse().unwrap(),
+            tls_public_key: bogus_ed25519_public_key(),
+            account_public_key: bogus_ed25519_public_key(),
+        };
+        tee_state
+            .verify_and_store_mock(
+                node_id.clone(),
+                MockAttestation::Valid,
+                Duration::from_secs(0),
+            )
+            .unwrap();
+
+        // Before expiry: cleanup keeps the entry.
+        assert_eq!(
+            tee_state.clean_invalid_attestations(Duration::from_secs(0), 100),
+            0
+        );
+
+        // When: the clock advances past the stamped expiry window and cleanup runs.
+        set_block_timestamp((attestation::DEFAULT_EXPIRATION_DURATION_SECONDS + 1) * 1_000_000_000);
+        let removed = tee_state.clean_invalid_attestations(Duration::from_secs(0), 100);
+
+        // Then: the previously uncleanable mock entry is removed.
+        assert_eq!(removed, 1);
+        assert!(
+            !tee_state
+                .stored_attestations
+                .contains_key(&node_id.tls_public_key)
+        );
+    }
+
+    #[test]
+    fn verify_and_store_mock__should_cap_a_submitter_supplied_expiry_beyond_the_default() {
+        // Given: a mock submitted with an expiry far beyond the contract's default
+        // window. A submitter must not be able to extend its lifetime past the cap.
+        testing_env!(VMContextBuilder::new().block_timestamp(0).build());
+
+        let mut tee_state = TeeState::default();
+        let node_id = NodeId {
+            account_id: "alice.near".parse().unwrap(),
+            tls_public_key: bogus_ed25519_public_key(),
+            account_public_key: bogus_ed25519_public_key(),
+        };
+        let oversized = MockAttestation::WithConstraints {
+            mpc_docker_image_hash: None,
+            launcher_docker_compose_hash: None,
+            expiry_timestamp_seconds: Some(attestation::DEFAULT_EXPIRATION_DURATION_SECONDS * 10),
+            expected_measurements: None,
+        };
+        tee_state
+            .verify_and_store_mock(node_id.clone(), oversized, Duration::from_secs(0))
+            .unwrap();
+
+        // When: the clock advances just past the contract's default window — well
+        // before the submitter's requested expiry.
+        set_block_timestamp((attestation::DEFAULT_EXPIRATION_DURATION_SECONDS + 1) * 1_000_000_000);
+        let removed = tee_state.clean_invalid_attestations(Duration::from_secs(0), 100);
+
+        // Then: the entry is cleaned — the submitted expiry was capped at the default.
+        assert_eq!(removed, 1);
+        assert!(
+            !tee_state
+                .stored_attestations
+                .contains_key(&node_id.tls_public_key)
+        );
+    }
+
     #[test]
     fn clean_invalid_attestations__should_honor_max_scan() {
         // Given: ten expired attestations stored.

crates/contract/src/v3_13_0_state.rs

@@ -8,7 +8,10 @@
 //! A better approach: only copy the structures that have changed and import the rest from the existing codebase.
 
 use borsh::{BorshDeserialize, BorshSerialize};
-use near_mpc_contract_interface::types::{Metrics, VerifyForeignTransactionRequest};
+use mpc_attestation::attestation::{self, VerifiedAttestation};
+use near_mpc_contract_interface::types::{
+    Ed25519PublicKey, Metrics, VerifyForeignTransactionRequest,
+};
 use near_sdk::{
     AccountId, env,
     store::{Lazy, LookupMap},
@@ -103,12 +106,59 @@ pub struct MpcContract {
     tee_verifier_votes: TeeVerifierVotes,
 }
 
+/// Stamps an expiry on every stored mock attestation that lacks or exceeds one —
+/// both user-submitted mocks and the genesis sentinels written by
+/// [`TeeState::with_mocked_participant_attestations`]. Legacy
+/// [`mpc_attestation::attestation::MockAttestation::Valid`] entries pass
+/// re-verification forever and can therefore never be evicted by
+/// [`TeeState::clean_invalid_attestations`];
+/// [`mpc_attestation::attestation::MockAttestation::with_expiry_capped_at`] rewrites them as
+/// expiring mocks so the normal cleanup flow can remove stale entries once the
+/// window elapses. An entry whose expiry is longer than (or missing) the default
+/// window is capped at it; a shorter existing expiry is left as-is.
+///
+// TODO(#3978): transitional one-time upgrade step — removed together with this
+// module when the pre-expiry migration is retired.
+fn stamp_expiry_on_legacy_mocks(tee_state: &mut TeeState, current_timestamp_seconds: u64) {
+    let expiry_timestamp_seconds =
+        current_timestamp_seconds + attestation::DEFAULT_EXPIRATION_DURATION_SECONDS;
+
+    // Collect keys before mutating to avoid iterator invalidation.
+    let mock_tls_keys: Vec<Ed25519PublicKey> = tee_state
+        .stored_attestations
+        .iter()
+        .filter(|(_, node_attestation)| {
+            matches!(
+                node_attestation.verified_attestation,
+                VerifiedAttestation::Mock(_)
+            )
+        })
+        .map(|(tls_pk, _)| tls_pk.clone())
+        .collect();
+
+    for tls_pk in mock_tls_keys {
+        let Some(node_attestation) = tee_state.stored_attestations.get_mut(&tls_pk) else {
+            continue;
+        };
+        if let VerifiedAttestation::Mock(mock) = &node_attestation.verified_attestation {
+            let stamped = mock.clone().with_expiry_capped_at(expiry_timestamp_seconds);
+            node_attestation.verified_attestation = VerifiedAttestation::Mock(stamped);
+        }
+    }
+}
+
 impl From<MpcContract> for crate::MpcContract {
     fn from(old: MpcContract) -> Self {
         if !matches!(old.protocol_state, ProtocolContractState::Running(_)) {
             env::panic_str("Contract must be in running state when migrating.");
         }
 
+        // Legacy `MockAttestation::Valid` entries never expire and can never be
+        // cleaned up. Stamp an expiry on them so the standard cleanup flow can
+        // evict stale mock entries after the upgrade.
+        let mut tee_state = old.tee_state;
+        stamp_expiry_on_legacy_mocks(&mut tee_state, TeeState::current_time_seconds());
+
         crate::MpcContract {
             protocol_state: old.protocol_state,
             pending_signature_requests: old.pending_signature_requests,
@@ -117,7 +167,7 @@ impl From<MpcContract> for crate::MpcContract {
             proposed_updates: old.proposed_updates,
             node_foreign_chain_support: old.node_foreign_chain_support,
             config: old.config.into(),
-            tee_state: old.tee_state,
+            tee_state,
             accept_requests: old.accept_requests,
             node_migrations: old.node_migrations,
             metrics: old.metrics,
@@ -127,3 +177,59 @@ impl From<MpcContract> for crate::MpcContract {
         }
     }
 }
+
+#[cfg(test)]
+#[expect(non_snake_case)]
+mod tests {
+    use super::{TeeState, VerifiedAttestation, attestation, stamp_expiry_on_legacy_mocks};
+    use crate::primitives::test_utils::bogus_ed25519_public_key;
+    use crate::tee::tee_state::{NodeAttestation, NodeId};
+    use crate::tee::test_utils::set_block_timestamp;
+    use mpc_attestation::attestation::MockAttestation;
+    use near_sdk::test_utils::VMContextBuilder;
+    use near_sdk::testing_env;
+    use std::time::Duration;
+
+    #[test]
+    fn stamp_expiry_on_legacy_mocks__should_make_valid_mock_cleanable() {
+        // Given: a legacy `MockAttestation::Valid` entry stored with no expiry, as
+        // written by older contract versions. Such entries pass re-verification
+        // forever and cannot be cleaned up.
+        testing_env!(VMContextBuilder::new().block_timestamp(0).build());
+
+        let mut tee_state = TeeState::default();
+        let node_id = NodeId {
+            account_id: "legacy.near".parse().unwrap(),
+            tls_public_key: bogus_ed25519_public_key(),
+            account_public_key: bogus_ed25519_public_key(),
+        };
+        tee_state.stored_attestations.insert(
+            node_id.tls_public_key.clone(),
+            NodeAttestation {
+                node_id: node_id.clone(),
+                verified_attestation: VerifiedAttestation::Mock(MockAttestation::Valid),
+            },
+        );
+
+        // Sanity: past the default window but without migration, the un-stamped
+        // entry survives cleanup indefinitely.
+        set_block_timestamp((attestation::DEFAULT_EXPIRATION_DURATION_SECONDS + 1) * 1_000_000_000);
+        assert_eq!(
+            tee_state.clean_invalid_attestations(Duration::from_secs(0), 100),
+            0
+        );
+
+        // When: the migration stamps an expiry as of block time 0 (window ends at
+        // DEFAULT), which the clock (already at DEFAULT + 1) is past.
+        stamp_expiry_on_legacy_mocks(&mut tee_state, 0);
+        let removed = tee_state.clean_invalid_attestations(Duration::from_secs(0), 100);
+
+        // Then: the stale legacy mock entry is removed.
+        assert_eq!(removed, 1);
+        assert!(
+            !tee_state
+                .stored_attestations
+                .contains_key(&node_id.tls_public_key)
+        );
+    }
+}

crates/contract/tests/sandbox/tee.rs

@@ -588,6 +588,7 @@ async fn get_attestation_returns_none_when_tls_key_is_not_associated_with_an_att
 #[tokio::test]
 async fn get_attestation_returns_some_when_tls_key_associated_with_an_attestation() {
     let SandboxTestSetup {
+        worker,
         contract,
         mpc_signer_accounts,
         ..
@@ -607,17 +608,21 @@ async fn get_attestation_returns_some_when_tls_key_associated_with_an_attestatio
         "Sanity check failed. Participant tls keys can not be equal for this test."
     );
 
+    // Expiries within the default window, so the contract's expiry cap keeps them
+    // as-is (a stored mock's expiry is `min(submitted, now + default window)`) and
+    // the two attestations stay distinct.
+    let now_seconds = worker.view_block().await.unwrap().timestamp() / 1_000_000_000;
     let participant_1_attestation = Attestation::Mock(MockAttestation::WithConstraints {
         mpc_docker_image_hash: None,
         launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(u64::MAX),
+        expiry_timestamp_seconds: Some(now_seconds + 1_000),
         expected_measurements: None,
     });
 
     let participant_2_attestation = Attestation::Mock(MockAttestation::WithConstraints {
         mpc_docker_image_hash: None,
         launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(u64::MAX - 1),
+        expiry_timestamp_seconds: Some(now_seconds + 2_000),
         expected_measurements: None,
     });
 
@@ -659,6 +664,7 @@ async fn get_attestation_returns_some_when_tls_key_associated_with_an_attestatio
 #[tokio::test]
 async fn get_attestation_overwrites_when_same_tls_key_is_reused() {
     let SandboxTestSetup {
+        worker,
         contract,
         mpc_signer_accounts,
         ..
@@ -670,17 +676,21 @@ async fn get_attestation_overwrites_when_same_tls_key_is_reused() {
     let participant_account = &mpc_signer_accounts[0];
     let tls_key = bogus_ed25519_public_key();
 
+    // Expiries within the default window, so the contract's expiry cap keeps them
+    // as-is (a stored mock's expiry is `min(submitted, now + default window)`) and
+    // the two attestations stay distinct.
+    let now_seconds = worker.view_block().await.unwrap().timestamp() / 1_000_000_000;
     let first_attestation = Attestation::Mock(MockAttestation::WithConstraints {
         mpc_docker_image_hash: None,
         launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(u64::MAX),
+        expiry_timestamp_seconds: Some(now_seconds + 1_000),
         expected_measurements: None,
     });
 
     let second_attestation = Attestation::Mock(MockAttestation::WithConstraints {
         mpc_docker_image_hash: None,
         launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(u64::MAX - 1),
+        expiry_timestamp_seconds: Some(now_seconds + 2_000),
         expected_measurements: None,
     });