test(contract): sandbox tests for the attestation flow against the real tee-verifier (#4165)

Author: pbeza

.github/workflows/ci.yml

@@ -295,11 +295,14 @@ jobs:
         if: github.ref == 'refs/heads/main'
         run: nix develop --command cargo near build reproducible-wasm --manifest-path crates/contract/Cargo.toml
 
+      - name: Build tee-verifier
+        run: nix develop --command cargo make build-tee-verifier-optimized
+
       - name: Build test-parallel-contract
         run: nix develop --command cargo make build-test-parallel-contract-optimized
 
       - name: Build backup-cli
-        run: nix develop --command cargo build -p backup-cli --release --locked
+        run: nix develop --command cargo make build-backup-cli
 
       - name: Run E2E tests
         run: nix develop --command cargo make e2e-tests-skip-build

Makefile.toml

@@ -165,7 +165,7 @@ args = ["scripts/check-sandbox-image-version.sh"]
 
 # These build tasks are the single source of truth for both local and CI builds.
 # CI's `mpc-e2e-tests` job invokes them via `cargo make`.
-# All three are skipped when `E2E_SKIP_BUILD` is set (used by `e2e-tests-skip-build`).
+# Each is skipped when `E2E_SKIP_BUILD` is set (used by `e2e-tests-skip-build`).
 
 [tasks.build-mpc-node-network-hardship-simulation]
 description = "Build the mpc-node binary used by the E2E tests"
@@ -202,7 +202,8 @@ args = [
 ]
 
 [tasks.build-tee-verifier-optimized]
-description = "Build the tee-verifier WASM for localnet"
+description = "Build the tee-verifier WASM for localnet and the E2E tests"
+condition = { env_not_set = ["E2E_SKIP_BUILD"] }
 command = "cargo"
 args = [
   "near",
@@ -252,6 +253,7 @@ private = true
 dependencies = [
   "build-mpc-node-network-hardship-simulation",
   "build-mpc-contract-optimized",
+  "build-tee-verifier-optimized",
   "build-test-parallel-contract-optimized",
   "build-backup-cli",
 ]
@@ -272,8 +274,8 @@ args = [
 [tasks._run-e2e-logic.env]
 MPC_CONTRACT_WASM = "${CARGO_MAKE_WORKING_DIRECTORY}/target/near/mpc_contract/mpc_contract.wasm"
 MPC_PARALLEL_CONTRACT_WASM = "${CARGO_MAKE_WORKING_DIRECTORY}/target/near/test_parallel_contract/test_parallel_contract.wasm"
+MPC_TEE_VERIFIER_WASM = "${CARGO_MAKE_WORKING_DIRECTORY}/target/near/tee_verifier/tee_verifier.wasm"
 
-# Build the mpc-node binary and both contract WASMs, then run the E2E tests.
 [tasks.e2e-tests]
 description = "Build required binaries and run the E2E tests"
 run_task = "_run-e2e-logic"

crates/contract/src/api/attestation.rs

@@ -593,12 +593,8 @@ mod tests {
             &mut contract.tee_state,
             image_digest(),
             launcher_image_hash(),
+            Some(launcher_compose_digest()),
         );
-        // The fixture's launcher compose carries the key-export service, so its hash is not derivable.
-        contract
-            .tee_state
-            .allowed_launcher_images
-            .allow_compose_hash(&launcher_image_hash(), launcher_compose_digest());
 
         // Storing a new entry consumes a grant, so stand in for the operator's prepayment.
         contract

crates/contract/src/tee/proposal.rs

@@ -491,7 +491,7 @@ impl AllowedLauncherImages {
     /// Test-only: allows one more compose hash for an already-allowed launcher. The attestation
     /// fixture is captured from a CVM whose launcher compose carries a key-export service, so
     /// [`get_docker_compose_hash`] cannot derive its hash.
-    #[cfg(test)]
+    #[cfg(any(test, feature = "test-utils"))]
     pub(crate) fn allow_compose_hash(
         &mut self,
         launcher_hash: &LauncherImageHash,

crates/contract/src/tee/tee_state.rs

@@ -1601,11 +1601,12 @@ mod tests {
         set_block_timestamp(VALID_ATTESTATION_TIMESTAMP * 1_000_000_000);
         let mut tee_state = TeeState::default();
         assert_eq!(tee_state.stored_attestations.len(), 0);
-        whitelist_dstack_measurements(&mut tee_state, image_digest(), launcher_image_hash());
-        // The fixture's launcher compose carries the key-export service, so its hash is not derivable.
-        tee_state
-            .allowed_launcher_images
-            .allow_compose_hash(&launcher_image_hash(), launcher_compose_digest());
+        whitelist_dstack_measurements(
+            &mut tee_state,
+            image_digest(),
+            launcher_image_hash(),
+            Some(launcher_compose_digest()),
+        );
         let node_id = NodeId {
             account_id: "alice.near".parse().unwrap(),
             tls_public_key: Ed25519PublicKey(p2p_tls_key()),

crates/contract/src/tee/test_utils.rs

@@ -3,11 +3,13 @@
 //! This module provides helper functions and types for testing TEE state,
 //! attestation behavior, and general contract state management.
 
+use crate::MpcContract;
 use crate::primitives::test_utils::{gen_account_id, gen_seed};
 use crate::tee::{measurements::ContractExpectedMeasurements, tee_state::TeeState};
 use mpc_attestation::attestation::default_measurements;
-use mpc_primitives::hash::{LauncherImageHash, NodeImageHash};
+use mpc_primitives::hash::{LauncherDockerComposeHash, LauncherImageHash, NodeImageHash};
 use near_account_id::AccountId;
+use near_sdk::borsh::{self, BorshDeserialize};
 use near_sdk::{BlockHeight, NearToken, PublicKey, test_utils::VMContextBuilder, testing_env};
 use rand::Rng;
 use std::time::Duration;
@@ -104,14 +106,41 @@ pub fn set_block_timestamp(timestamp_nanos: u64) {
     );
 }
 
+/// Fills the allowlists a Dstack submission is checked against: MPC image,
+/// launcher image, expected measurements, and optionally one compose hash.
+///
+/// The compose hash is injected directly because it is the one value no vote can
+/// authorize: [`get_docker_compose_hash`](crate::tee::proposal::get_docker_compose_hash)
+/// derives the accepted hashes from the voted launcher and MPC images through a
+/// compiled-in template, and the fixture's CVM ran a modified compose carrying a
+/// key-export service (the only way to get the signer key out of the CVM, which the
+/// quote's `report_data` binds and tests must sign with). Pass [`None`] to exercise
+/// the rejection path.
 pub fn whitelist_dstack_measurements(
     tee_state: &mut TeeState,
     image: NodeImageHash,
     launcher: LauncherImageHash,
+    compose_hash: Option<LauncherDockerComposeHash>,
 ) {
     tee_state.whitelist_tee_proposal(image, Duration::MAX);
     tee_state.add_launcher_image(launcher, Duration::MAX, Duration::MAX);
     for &measurements in default_measurements() {
         tee_state.add_measurement(ContractExpectedMeasurements::from(measurements));
     }
+    if let Some(compose_hash) = compose_hash {
+        tee_state
+            .allowed_launcher_images
+            .allow_compose_hash(&launcher, compose_hash);
+    }
+}
+
+pub fn whitelist_dstack_in_state(
+    state: &[u8],
+    image: NodeImageHash,
+    launcher: LauncherImageHash,
+    compose_hash: Option<LauncherDockerComposeHash>,
+) -> Vec<u8> {
+    let mut contract = MpcContract::try_from_slice(state).expect("STATE deserializes");
+    whitelist_dstack_measurements(&mut contract.tee_state, image, launcher, compose_hash);
+    borsh::to_vec(&contract).expect("STATE serializes")
 }

crates/contract/tests/sandbox/tee.rs

@@ -10,9 +10,9 @@ use crate::sandbox::{
         interface::IntoContractType,
         mpc_contract::{
             assert_running_return_participants, assert_running_return_threshold,
-            get_participant_attestation, get_state, get_tee_accounts,
-            prepay_and_submit_participant_info, prepay_attestation_grants, submit_participant_info,
-            vote_add_launcher_hash, vote_for_hash,
+            available_attestation_grants, get_config, get_participant_attestation, get_state,
+            get_tee_accounts, prepay_and_submit_participant_info, prepay_attestation_grants,
+            submit_participant_info, vote_add_launcher_hash, vote_for_hash,
         },
         resharing_utils::conclude_resharing,
         sign_utils::DomainResponseTest,
@@ -25,13 +25,22 @@ use mpc_primitives::hash::{LauncherDockerComposeHash, LauncherImageHash, NodeIma
 use near_mpc_contract_interface::deposits::STORAGE_BYTE_COST_YOCTONEAR;
 use near_mpc_contract_interface::method_names;
 use near_mpc_contract_interface::types::{
-    self as dtos, Attestation, Config, MockAttestation, Protocol,
+    self as dtos, Attestation, MockAttestation, Protocol, VerifiedAttestation,
 };
 use near_workspaces::types::{KeyType, NearToken, SecretKey};
 use near_workspaces::{AccessKey, Account, Contract};
 use rand::SeedableRng;
 use test_utils::attestation::{image_digest, p2p_tls_key};
 
+fn mock_expiring_at(expiry_timestamp_seconds: u64) -> MockAttestation {
+    MockAttestation::WithConstraints {
+        mpc_docker_image_hash: None,
+        launcher_docker_compose_hash: None,
+        expiry_timestamp_seconds: Some(expiry_timestamp_seconds),
+        expected_measurements: None,
+    }
+}
+
 /// Tests the basic code hash voting mechanism including threshold behavior and vote stability.
 /// Validates that votes below threshold don't allow hashes, reaching threshold allows them,
 /// and additional votes don't change the allowed state or latest hash.
@@ -457,12 +466,7 @@ async fn clean_invalid_attestations__should_remove_expired_entries() -> Result<(
     let block_info = worker.view_block().await?;
     let expiry_timestamp_seconds =
         block_info.timestamp() / 1_000_000_000 + ATTESTATION_EXPIRY_SECONDS;
-    let expiring_attestation = Attestation::Mock(MockAttestation::WithConstraints {
-        mpc_docker_image_hash: None,
-        launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(expiry_timestamp_seconds),
-        expected_measurements: None,
-    });
+    let expiring_attestation = Attestation::Mock(mock_expiring_at(expiry_timestamp_seconds));
     let submit_result = prepay_and_submit_participant_info(
         stale_account,
         &contract,
@@ -595,10 +599,9 @@ async fn get_attestation_returns_none_when_tls_key_is_not_associated_with_an_att
 
     assert!(validation_success);
 
-    let attestation_for_tls_key_2: Option<Attestation> =
-        get_participant_attestation(&contract, &tls_key_2)
-            .await
-            .unwrap();
+    let attestation_for_tls_key_2 = get_participant_attestation(&contract, &tls_key_2)
+        .await
+        .unwrap();
 
     assert_eq!(attestation_for_tls_key_2, None);
 }
@@ -630,19 +633,10 @@ async fn get_attestation_returns_some_when_tls_key_associated_with_an_attestatio
     // 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(now_seconds + 1_000),
-        expected_measurements: None,
-    });
+    let participant_1_attestation = Attestation::Mock(mock_expiring_at(now_seconds + 1_000));
 
-    let participant_2_attestation = Attestation::Mock(MockAttestation::WithConstraints {
-        mpc_docker_image_hash: None,
-        launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(now_seconds + 2_000),
-        expected_measurements: None,
-    });
+    let participant_2_mock = mock_expiring_at(now_seconds + 2_000);
+    let participant_2_attestation = Attestation::Mock(participant_2_mock.clone());
 
     assert_ne!(
         participant_1_attestation, participant_2_attestation,
@@ -671,12 +665,14 @@ async fn get_attestation_returns_some_when_tls_key_associated_with_an_attestatio
     .is_success();
     assert!(validation_success, "Submitting attestation failed.");
 
-    let attestation_for_tls_key_2: Option<Attestation> =
-        get_participant_attestation(&contract, &tls_key_2)
-            .await
-            .unwrap();
+    let attestation_for_tls_key_2 = get_participant_attestation(&contract, &tls_key_2)
+        .await
+        .unwrap();
 
-    assert_eq!(attestation_for_tls_key_2, Some(participant_2_attestation));
+    assert_eq!(
+        attestation_for_tls_key_2,
+        Some(VerifiedAttestation::Mock(participant_2_mock))
+    );
 }
 
 #[tokio::test]
@@ -698,19 +694,10 @@ async fn get_attestation_overwrites_when_same_tls_key_is_reused() {
     // 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(now_seconds + 1_000),
-        expected_measurements: None,
-    });
+    let first_attestation = Attestation::Mock(mock_expiring_at(now_seconds + 1_000));
 
-    let second_attestation = Attestation::Mock(MockAttestation::WithConstraints {
-        mpc_docker_image_hash: None,
-        launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(now_seconds + 2_000),
-        expected_measurements: None,
-    });
+    let second_mock = mock_expiring_at(now_seconds + 2_000);
+    let second_attestation = Attestation::Mock(second_mock.clone());
 
     assert_ne!(
         first_attestation, second_attestation,
@@ -742,14 +729,13 @@ async fn get_attestation_overwrites_when_same_tls_key_is_reused() {
     assert!(validation_success, "Second attestation submission failed");
 
     // Now the latest attestation should be returned
-    let attestation_for_tls_key: Option<Attestation> =
-        get_participant_attestation(&contract, &tls_key)
-            .await
-            .unwrap();
+    let attestation_for_tls_key = get_participant_attestation(&contract, &tls_key)
+        .await
+        .unwrap();
 
     assert_eq!(
         attestation_for_tls_key,
-        Some(second_attestation),
+        Some(VerifiedAttestation::Mock(second_mock)),
         "Expected the second attestation to overwrite the first for the same TLS key"
     );
 }
@@ -841,12 +827,7 @@ async fn test_verify_tee_expired_attestation_triggers_resharing() -> Result<()>
         .find(|node| node.account_id == *target_account.id())
         .expect("target participant not found");
 
-    let expiring_attestation = Attestation::Mock(MockAttestation::WithConstraints {
-        mpc_docker_image_hash: None,
-        launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(expiry_timestamp),
-        expected_measurements: None,
-    });
+    let expiring_attestation = Attestation::Mock(mock_expiring_at(expiry_timestamp));
 
     let submit_result = submit_participant_info(
         target_account,
@@ -963,12 +944,7 @@ async fn verify_tee__should_keep_participants_and_stop_signing_when_kickout_drop
     // Compute the expiry timestamp from the current block time.
     let block_info = worker.view_block().await?;
     let expiry_timestamp = block_info.timestamp() / 1_000_000_000 + ATTESTATION_EXPIRY_SECONDS;
-    let expiring_attestation = Attestation::Mock(MockAttestation::WithConstraints {
-        mpc_docker_image_hash: None,
-        launcher_docker_compose_hash: None,
-        expiry_timestamp_seconds: Some(expiry_timestamp),
-        expected_measurements: None,
-    });
+    let expiring_attestation = Attestation::Mock(mock_expiring_at(expiry_timestamp));
 
     // Submit an expiring attestation for every participant past the first `remaining_valid`.
     let internal_participants: Participants = (&initial_participants).into_contract_type();
@@ -1088,11 +1064,7 @@ async fn prepay_and_submit_a_constrained_mock__should_use_at_most_half_a_grant_f
     });
     let node = worker.dev_create_account().await?;
     let tls_key = bogus_ed25519_public_key();
-    let config: Config = contract
-        .view(method_names::CONFIG)
-        .args_json(serde_json::json!({}))
-        .await?
-        .json()?;
+    let config = get_config(&contract).await?;
     let before = contract.as_account().view_account().await?;
 
     // When
@@ -1102,11 +1074,7 @@ async fn prepay_and_submit_a_constrained_mock__should_use_at_most_half_a_grant_f
     assert!(submission.is_success(), "submission failed: {submission:?}");
 
     // Then
-    let remaining: u32 = contract
-        .view(method_names::AVAILABLE_ATTESTATION_GRANTS)
-        .args_json(serde_json::json!({ "account_id": node.id() }))
-        .await?
-        .json()?;
+    let remaining = available_attestation_grants(&contract, node.id()).await?;
     assert_eq!(remaining, GRANTS - 1, "the row must outlive the submission");
 
     let after = contract.as_account().view_account().await?;
@@ -1163,11 +1131,7 @@ async fn submit_participant_info__should_store_a_new_entry_against_a_prepaid_gra
     );
     let stored = get_participant_attestation(&contract, &fresh_tls_key).await?;
     assert!(stored.is_some(), "the entry should be stored on-chain");
-    let remaining: u32 = contract
-        .view(method_names::AVAILABLE_ATTESTATION_GRANTS)
-        .args_json(serde_json::json!({ "account_id": outsider.id() }))
-        .await?
-        .json()?;
+    let remaining = available_attestation_grants(&contract, outsider.id()).await?;
     assert_eq!(remaining, 0, "storing a new entry should consume the grant");
     Ok(())
 }