fix(contract): make mock attestations cleanable via expiry

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
     }
 
@@ -473,6 +473,46 @@ impl TeeState {
         removed
     }
 
+    /// One-time migration helper: 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 [`MockAttestation::Valid`] entries pass re-verification forever and
+    /// can therefore never be evicted by [`TeeState::clean_invalid_attestations`]. [`MockAttestation::with_expiry`]
+    /// rewrites them as expiring [`MockAttestation::WithConstraints`] 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 — remove this and its migration call site once the
+    // pre-expiry state migration is retired.
+    pub(crate) fn stamp_expiry_on_legacy_mocks(&mut self, 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> = self
+            .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) = self.stored_attestations.get_mut(&tls_pk) else {
+                continue;
+            };
+            if let VerifiedAttestation::Mock(mock) = &node_attestation.verified_attestation {
+                let stamped = mock.clone().with_expiry(expiry_timestamp_seconds);
+                node_attestation.verified_attestation = VerifiedAttestation::Mock(stamped);
+            }
+        }
+    }
+
     /// Returns the list of accounts that currently have TEE attestations stored.
     /// Note: This may include accounts that are no longer active protocol participants.
     pub fn get_tee_accounts(&self) -> Vec<NodeId> {
@@ -717,6 +757,124 @@ 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 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.
+        tee_state.stamp_expiry_on_legacy_mocks(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)
+        );
+    }
+
     #[test]
     fn clean_invalid_attestations__should_honor_max_scan() {
         // Given: ten expired attestations stored.

crates/contract/src/v3_13_0_state.rs

@@ -109,6 +109,12 @@ impl From<MpcContract> for crate::MpcContract {
             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;
+        tee_state.stamp_expiry_on_legacy_mocks(TeeState::current_time_seconds());
+
         crate::MpcContract {
             protocol_state: old.protocol_state,
             pending_signature_requests: old.pending_signature_requests,
@@ -117,7 +123,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,

crates/contract/tests/sandbox/tee.rs

@@ -542,6 +542,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,
         ..
@@ -561,17 +562,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,
     });
 
@@ -613,6 +618,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,
         ..
@@ -624,17 +630,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,
     });