refactor(mock-attestation): address review

- Move stamp_expiry_on_legacy_mocks (+ its test) into the v3_13_0_state
  migration module, since it is migration-only (netrome).
- Simplify submitted_attestation_landed docstring: a changed stored expiry
  is enough to conclude the submit landed; drop the confusing 'identity'
  wording (the legacy equality fallback is covered by TODO(#3786)).

Author: barakeinav1

crates/contract/src/tee/tee_state.rs

@@ -473,46 +473,6 @@ 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> {
@@ -832,49 +792,6 @@ mod tests {
         );
     }
 
-    #[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

@@ -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,6 +106,47 @@ 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`] 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.
+///
+/// This is a one-time upgrade step; it is 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(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(_)) {
@@ -113,7 +157,7 @@ impl From<MpcContract> for crate::MpcContract {
         // 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());
+        stamp_expiry_on_legacy_mocks(&mut tee_state, TeeState::current_time_seconds());
 
         crate::MpcContract {
             protocol_state: old.protocol_state,
@@ -133,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/node/src/indexer/tx_sender.rs

@@ -192,13 +192,11 @@ fn attestation_expiry_changed(pre_submit_expiry: Option<u64>, stored_expiry: u64
 
 /// Whether the attestation we submitted is now the one stored on chain.
 ///
-/// An entry that carries an expiry keeps no stored per-submission identity, so it is confirmed via
-/// [`attestation_expiry_changed`]: an accepted submit re-stamps the entry's expiry (to the submit
-/// block time plus
+/// An accepted submit re-stamps the entry's expiry (to the submit block time plus
 /// [`DEFAULT_EXPIRATION_DURATION_SECONDS`](mpc_attestation::attestation::DEFAULT_EXPIRATION_DURATION_SECONDS)),
-/// and only the owning account may rewrite it, so a changed expiry means our submit landed. This
-/// covers every Dstack entry, and mock entries stored by contracts that stamp expiries.
-/// A mock stored by an older contract carries no expiry and is matched by identity instead.
+/// and only the owning account may rewrite it, so observing a **changed stored expiry** is enough
+/// to conclude our submit landed — for every Dstack entry and for mocks stored with an expiry. A
+/// legacy mock with no stored expiry falls back to an equality check (see the `TODO(#3786)` below).
 // TODO(#1639): match a certificate-derived identity instead of this expiry heuristic.
 fn submitted_attestation_landed(
     pre_submit_expiry: Option<u64>,