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,44 @@ impl TeeState {
         removed
     }
 
+    /// One-time migration helper: stamps an expiry on stored mock attestations
+    /// that lack one. 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, while
+    /// leaving entries that already carry an explicit expiry unchanged.
+    ///
+    // 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 +755,88 @@ 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 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: without migration the 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 helper stamps an expiry (as of the migration block time)
+        // and the clock later advances past that stamped window.
+        tee_state.stamp_expiry_on_legacy_mocks(0);
+        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 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/mpc-attestation/src/attestation.rs

@@ -84,10 +84,19 @@ impl AcceptedAttestation {
         }
     }
 
-    /// Assembles the acceptance for a verified `Mock` attestation.
-    fn mock(mock_attestation: &MockAttestation) -> Self {
+    /// Assembles the acceptance for a verified `Mock` attestation. Stamps a
+    /// [`DEFAULT_EXPIRATION_DURATION_SECONDS`] expiry (via [`MockAttestation::with_expiry`]),
+    /// mirroring [`AcceptedAttestation::dstack`], so a `Valid` mock does not pass
+    /// re-verification forever and can be cleaned up.
+    fn mock(mock_attestation: &MockAttestation, current_timestamp_seconds: u64) -> Self {
+        let expiry_timestamp_seconds =
+            current_timestamp_seconds + DEFAULT_EXPIRATION_DURATION_SECONDS;
         Self {
-            attestation: VerifiedAttestation::Mock(mock_attestation.clone()),
+            attestation: VerifiedAttestation::Mock(
+                mock_attestation
+                    .clone()
+                    .with_expiry(expiry_timestamp_seconds),
+            ),
             advisory_ids: Vec::new(),
         }
     }
@@ -131,7 +140,36 @@ impl MockAttestation {
             allowed_launcher_docker_compose_hashes,
             accepted_measurements,
         )?;
-        Ok(AcceptedAttestation::mock(self))
+        Ok(AcceptedAttestation::mock(self, current_timestamp_seconds))
+    }
+
+    /// Returns a copy stamped with `expiry_timestamp_seconds`, unless the mock
+    /// already carries an explicit expiry (which is preserved). A bare
+    /// [`MockAttestation::Valid`] becomes an otherwise-unconstrained
+    /// [`MockAttestation::WithConstraints`] so that, once stored, it eventually
+    /// expires and becomes eligible for cleanup. [`MockAttestation::Invalid`] is
+    /// returned unchanged — it never reaches acceptance.
+    pub fn with_expiry(self, expiry_timestamp_seconds: u64) -> Self {
+        match self {
+            MockAttestation::Valid => MockAttestation::WithConstraints {
+                mpc_docker_image_hash: None,
+                launcher_docker_compose_hash: None,
+                expiry_timestamp_seconds: Some(expiry_timestamp_seconds),
+                expected_measurements: None,
+            },
+            MockAttestation::Invalid => MockAttestation::Invalid,
+            MockAttestation::WithConstraints {
+                mpc_docker_image_hash,
+                launcher_docker_compose_hash,
+                expiry_timestamp_seconds: existing_expiry,
+                expected_measurements,
+            } => MockAttestation::WithConstraints {
+                mpc_docker_image_hash,
+                launcher_docker_compose_hash,
+                expiry_timestamp_seconds: existing_expiry.or(Some(expiry_timestamp_seconds)),
+                expected_measurements,
+            },
+        }
     }
 
     /// Checks the mock's constraints, returning only pass/fail. Lets the
@@ -491,11 +529,86 @@ fn verify_measurements(
 }
 
 #[cfg(test)]
+#[expect(non_snake_case)]
 mod tests {
     use alloc::vec;
 
     use super::*;
 
+    #[test]
+    fn with_expiry__should_convert_valid_to_expiring_constraints() {
+        // Given / When
+        let stamped = MockAttestation::Valid.with_expiry(42);
+
+        // Then
+        assert_matches::assert_matches!(
+            stamped,
+            MockAttestation::WithConstraints {
+                mpc_docker_image_hash: None,
+                launcher_docker_compose_hash: None,
+                expiry_timestamp_seconds: Some(42),
+                expected_measurements: None,
+            }
+        );
+    }
+
+    #[test]
+    fn with_expiry__should_fill_missing_expiry_and_keep_other_constraints() {
+        // Given: a constrained mock with an image-hash constraint but no expiry.
+        let image_hash = NodeImageHash::from([7; 32]);
+        let mock = MockAttestation::WithConstraints {
+            mpc_docker_image_hash: Some(image_hash),
+            launcher_docker_compose_hash: None,
+            expiry_timestamp_seconds: None,
+            expected_measurements: None,
+        };
+
+        // When
+        let stamped = mock.with_expiry(42);
+
+        // Then: the expiry is filled in and the other constraints are preserved.
+        assert_matches::assert_matches!(
+            stamped,
+            MockAttestation::WithConstraints {
+                mpc_docker_image_hash: Some(hash),
+                expiry_timestamp_seconds: Some(42),
+                ..
+            } if hash == image_hash
+        );
+    }
+
+    #[test]
+    fn with_expiry__should_preserve_existing_expiry() {
+        // Given: a mock that already carries an explicit expiry.
+        let mock = MockAttestation::WithConstraints {
+            mpc_docker_image_hash: None,
+            launcher_docker_compose_hash: None,
+            expiry_timestamp_seconds: Some(100),
+            expected_measurements: None,
+        };
+
+        // When
+        let stamped = mock.with_expiry(42);
+
+        // Then: the caller-set expiry wins over the default.
+        assert_matches::assert_matches!(
+            stamped,
+            MockAttestation::WithConstraints {
+                expiry_timestamp_seconds: Some(100),
+                ..
+            }
+        );
+    }
+
+    #[test]
+    fn with_expiry__should_leave_invalid_unchanged() {
+        // Given / When
+        let stamped = MockAttestation::Invalid.with_expiry(42);
+
+        // Then
+        assert_matches::assert_matches!(stamped, MockAttestation::Invalid);
+    }
+
     #[test]
     fn mock_constrained_verification_passes_if_hash_in_allowed_list() {
         let allowed_hash = NodeImageHash::from([42; 32]);

crates/mpc-attestation/tests/test_attestation_verification.rs

@@ -24,12 +24,19 @@ fn valid_mock_attestation_succeeds_verification() {
     let account_key = account_key();
     let report_data = ReportData::V1(ReportDataV1::new(tls_key, account_key));
 
+    // A `Valid` mock is accepted as an expiring `WithConstraints` mock so the
+    // stored attestation can later be cleaned up (#3293).
     assert_matches!(
         valid_attestation.verify_locally(report_data.into(), timestamp_s, &[], &[], &[]),
         Ok(AcceptedAttestation {
-            attestation: VerifiedAttestation::Mock(MockAttestation::Valid),
+            attestation: VerifiedAttestation::Mock(MockAttestation::WithConstraints {
+                mpc_docker_image_hash: None,
+                launcher_docker_compose_hash: None,
+                expiry_timestamp_seconds: Some(expiry),
+                expected_measurements: None,
+            }),
             advisory_ids,
-        }) if advisory_ids.is_empty()
+        }) if advisory_ids.is_empty() && expiry == timestamp_s + DEFAULT_EXPIRATION_DURATION_SECONDS
     );
 }
 

crates/near-mpc-contract-interface/src/types/attestation.rs

@@ -75,6 +75,18 @@ pub enum VerifiedAttestation {
     Mock(MockAttestation),
 }
 
+impl VerifiedAttestation {
+    /// The stored expiry timestamp, if the attestation carries one. `Dstack`
+    /// entries always do; a `Mock` entry does only when it was stamped with an
+    /// expiry — contracts predating #3293 store `Mock::Valid` without one.
+    pub fn expiry_timestamp_seconds(&self) -> Option<u64> {
+        match self {
+            VerifiedAttestation::Dstack(attestation) => Some(attestation.expiry_timestamp_seconds),
+            VerifiedAttestation::Mock(attestation) => attestation.expiry_timestamp_seconds(),
+        }
+    }
+}
+
 #[derive(
     Clone,
     Debug,
@@ -184,6 +196,20 @@ pub enum MockAttestation {
     },
 }
 
+impl MockAttestation {
+    /// The configured expiry timestamp, if any. `Valid` and `Invalid` never
+    /// carry one.
+    pub fn expiry_timestamp_seconds(&self) -> Option<u64> {
+        match self {
+            MockAttestation::WithConstraints {
+                expiry_timestamp_seconds,
+                ..
+            } => *expiry_timestamp_seconds,
+            MockAttestation::Valid | MockAttestation::Invalid => None,
+        }
+    }
+}
+
 // TODO(#3494): superseded by `tee_verifier_interface::Collateral`; remove
 // this serde-carrying copy once `mpc-contract` consumes the Borsh mirrors.
 #[derive(