fix: address post-merge code review comments for PR #1183 (#1189)

Author: pbeza

crates/attestation/src/report_data.rs

@@ -3,6 +3,7 @@ use derive_more::Constructor;
 use near_sdk::PublicKey;
 use serde::{Deserialize, Serialize};
 use sha3::{Digest, Sha3_384};
+use alloc::vec;
 
 /// Number of bytes for the report data.
 const REPORT_DATA_SIZE: usize = 64;
@@ -36,10 +37,11 @@ impl ReportDataVersion {
 #[derive(Debug, Clone, Constructor)]
 pub struct ReportDataV1 {
     tls_public_key: PublicKey,
+    account_public_key: PublicKey,
 }
 
 /// report_data_v1: [u8; 64] =
-///   [version(2 bytes big endian) || sha384(TLS pub key) || zero padding]
+///   [version(2 bytes big endian) || sha384(TLS pub key || account_pubkey ) || zero padding]
 impl ReportDataV1 {
     /// V1-specific format constants
     const PUBLIC_KEYS_OFFSET: usize = BINARY_VERSION_OFFSET + BINARY_VERSION_SIZE;
@@ -84,12 +86,18 @@ impl ReportDataV1 {
         hash
     }
 
-    /// Generates SHA3-384 hash of TLS public key only.
+    /// Generates SHA3-384 hash of TLS + NEAR account keys together.
     fn public_keys_hash(&self) -> [u8; Self::PUBLIC_KEYS_HASH_SIZE] {
         let mut hasher = Sha3_384::new();
-        // Skip first byte as it is used for identifier for the curve type.
-        let key_data = &self.tls_public_key.as_bytes()[1..];
-        hasher.update(key_data);
+
+        // Hash TLS key (skip first byte = curve type)
+        let tls_data = &self.tls_public_key.as_bytes()[1..];
+        hasher.update(tls_data);
+
+        // Hash NEAR account key (also skip first byte)
+        let account_data = &self.account_public_key.as_bytes()[1..];
+        hasher.update(account_data);
+
         hasher.finalize().into()
     }
 }
@@ -100,8 +108,13 @@ pub enum ReportData {
 }
 
 impl ReportData {
-    pub fn new(tls_public_key: PublicKey) -> Self {
-        ReportData::V1(ReportDataV1::new(tls_public_key))
+   pub fn new(tls_public_key: PublicKey, account_public_key: Option<PublicKey>) -> Self {
+        let account_pk = account_public_key.unwrap_or_else(|| {
+            // Construct a "zero" public key. will not be used in practice, only for backward compatibility. remove this code once that network enforces real attestation  
+            PublicKey::from_parts(near_sdk::CurveType::ED25519, vec![0u8; 32]).expect("valid zero PublicKey")
+        });
+
+        ReportData::V1(ReportDataV1::new(tls_public_key, account_pk))
     }
 
     pub fn version(&self) -> ReportDataVersion {
@@ -125,8 +138,15 @@ mod tests {
     use alloc::vec::Vec;
     use dcap_qvl::quote::Quote;
     use near_sdk::PublicKey;
+    use sha3::{Digest, Sha3_384};
     use test_utils::attestation::{p2p_tls_key, quote};
 
+    fn create_test_key() -> PublicKey {
+        "secp256k1:qMoRgcoXai4mBPsdbHi1wfyxF9TdbPCF4qSDQTRP3TfescSRoUdSx6nmeQoN3aiwGzwMyGXAb1gUjBTv5AY8DXj"
+            .parse()
+            .unwrap()
+    }
+
     #[test]
     fn test_from_str_valid() {
         let valid_quote: Vec<u8> =
@@ -136,14 +156,9 @@ mod tests {
         let td_report = quote.report.as_td10().expect("Should be a TD 1.0 report");
 
         let near_p2p_public_key: PublicKey = p2p_tls_key();
-        let report_data = ReportData::V1(ReportDataV1::new(near_p2p_public_key));
-        assert_eq!(report_data.to_bytes(), td_report.report_data,);
-    }
-
-    fn create_test_key() -> PublicKey {
-        "secp256k1:qMoRgcoXai4mBPsdbHi1wfyxF9TdbPCF4qSDQTRP3TfescSRoUdSx6nmeQoN3aiwGzwMyGXAb1gUjBTv5AY8DXj"
-            .parse()
-            .unwrap()
+        let account_key = create_test_key();
+        let report_data = ReportData::V1(ReportDataV1::new(near_p2p_public_key, account_key));
+        assert_eq!(report_data.to_bytes(), td_report.report_data);
     }
 
     #[test]
@@ -160,11 +175,13 @@ mod tests {
     #[test]
     fn test_report_data_enum_structure() {
         let tls_key = create_test_key();
-        let data = ReportData::V1(ReportDataV1::new(tls_key.clone()));
+        let account_key = create_test_key();
+        let data = ReportData::V1(ReportDataV1::new(tls_key.clone(), account_key.clone()));
 
         match &data {
             ReportData::V1(v1) => {
                 assert_eq!(&v1.tls_public_key, &tls_key);
+                assert_eq!(&v1.account_public_key, &account_key);
             }
         }
 
@@ -174,15 +191,18 @@ mod tests {
     #[test]
     fn test_report_data_v1_struct() {
         let tls_key = create_test_key();
+        let account_key = create_test_key();
 
-        let v1 = ReportDataV1::new(tls_key.clone());
+        let v1 = ReportDataV1::new(tls_key.clone(), account_key.clone());
         assert_eq!(v1.tls_public_key, tls_key);
+        assert_eq!(v1.account_public_key, account_key);
     }
 
     #[test]
     fn test_from_bytes() {
         let tls_key = create_test_key();
-        let report_data_v1 = ReportDataV1::new(tls_key);
+        let account_key = create_test_key();
+        let report_data_v1 = ReportDataV1::new(tls_key.clone(), account_key.clone());
         let bytes = report_data_v1.to_bytes();
 
         let hash = ReportDataV1::from_bytes(&bytes);
@@ -195,7 +215,8 @@ mod tests {
     #[test]
     fn test_binary_version_placement() {
         let tls_key = create_test_key();
-        let bytes = ReportDataV1::new(tls_key).to_bytes();
+        let account_key = create_test_key();
+        let bytes = ReportDataV1::new(tls_key, account_key).to_bytes();
 
         let version_bytes =
             &bytes[BINARY_VERSION_OFFSET..BINARY_VERSION_OFFSET + BINARY_VERSION_SIZE];
@@ -205,20 +226,21 @@ mod tests {
     #[test]
     fn test_public_key_hash_placement() {
         let tls_key = create_test_key();
-        let report_data_v1 = ReportDataV1::new(tls_key.clone());
+        let account_key = create_test_key();
+        let report_data_v1 = ReportDataV1::new(tls_key.clone(), account_key.clone());
         let bytes = report_data_v1.to_bytes();
 
-        let report_data = ReportData::V1(report_data_v1);
+        let report_data = ReportData::V1(report_data_v1.clone());
         assert_eq!(report_data.to_bytes(), bytes);
 
         let hash_bytes = &bytes[ReportDataV1::PUBLIC_KEYS_OFFSET
             ..ReportDataV1::PUBLIC_KEYS_OFFSET + ReportDataV1::PUBLIC_KEYS_HASH_SIZE];
         assert_ne!(hash_bytes, &[0u8; ReportDataV1::PUBLIC_KEYS_HASH_SIZE]);
 
+        // Expected hash = sha3_384(tls || account)
         let mut hasher = Sha3_384::new();
-        // Skip first byte as it is used for identifier for the curve type.
-        let key_data = &tls_key.as_bytes()[1..];
-        hasher.update(key_data);
+        hasher.update(&tls_key.as_bytes()[1..]);
+        hasher.update(&account_key.as_bytes()[1..]);
         let expected: [u8; ReportDataV1::PUBLIC_KEYS_HASH_SIZE] = hasher.finalize().into();
 
         assert_eq!(hash_bytes, &expected);
@@ -227,7 +249,8 @@ mod tests {
     #[test]
     fn test_zero_padding() {
         let tls_key = create_test_key();
-        let bytes = ReportDataV1::new(tls_key).to_bytes();
+        let account_key = create_test_key();
+        let bytes = ReportDataV1::new(tls_key, account_key).to_bytes();
 
         let padding =
             &bytes[ReportDataV1::PUBLIC_KEYS_OFFSET + ReportDataV1::PUBLIC_KEYS_HASH_SIZE..];
@@ -237,7 +260,8 @@ mod tests {
     #[test]
     fn test_report_data_size() {
         let tls_key = create_test_key();
-        let bytes = ReportDataV1::new(tls_key);
+        let account_key = create_test_key();
+        let bytes = ReportDataV1::new(tls_key, account_key);
         assert_eq!(bytes.to_bytes().len(), REPORT_DATA_SIZE);
     }
 }

crates/attestation/tests/test_attestation_verification.rs

@@ -6,7 +6,7 @@ use mpc_primitives::hash::{LauncherDockerComposeHash, MpcDockerImageHash};
 use near_sdk::PublicKey;
 use rstest::rstest;
 use test_utils::attestation::{
-    image_digest, launcher_compose_digest, mock_dstack_attestation, p2p_tls_key,
+    image_digest, launcher_compose_digest, mock_dstack_attestation, p2p_tls_key, account_key,
 };
 
 #[rstest]
@@ -20,7 +20,10 @@ fn test_mock_attestation_verify(
     let tls_key = "ed25519:DcA2MzgpJbrUATQLLceocVckhhAqrkingax4oJ9kZ847"
         .parse()
         .unwrap();
-    let report_data = ReportData::V1(ReportDataV1::new(tls_key));
+    let account_key = "ed25519:5v8Y8ZLoxZzCVtYpjh1cYdFrRh1p9EXAMPLEaQJ5sP4o"
+        .parse()
+        .unwrap();
+    let report_data = ReportData::V1(ReportDataV1::new(tls_key,account_key));
     let attestation = Attestation::Mock(local_attestation);
 
     assert_eq!(
@@ -33,8 +36,9 @@ fn test_mock_attestation_verify(
 fn test_verify_method_signature() {
     let attestation = mock_dstack_attestation();
     let tls_key: PublicKey = p2p_tls_key();
+    let account_key: PublicKey = account_key();
 
-    let report_data = ReportData::V1(ReportDataV1::new(tls_key));
+    let report_data = ReportData::V1(ReportDataV1::new(tls_key,account_key));
     let timestamp_s = 1755186041_u64;
 
     let allowed_mpc_image_digest: MpcDockerImageHash = image_digest();

crates/contract/src/lib.rs

@@ -440,6 +440,7 @@ impl MpcContract {
         let signer = env::signer_account_id();
         log!("respond: signer={}, request={:?}", &signer, &request);
 
+        self.tee_state.assert_caller_is_attested_node();
         if !self.protocol_state.is_running_or_resharing() {
             return Err(InvalidState::ProtocolStateNotRunning.into());
         }
@@ -519,6 +520,8 @@ impl MpcContract {
     pub fn respond_ckd(&mut self, request: CKDRequest, response: CKDResponse) -> Result<(), Error> {
         let signer = env::signer_account_id();
         log!("respond_ckd: signer={}, request={:?}", &signer, &request);
+        
+        self.tee_state.assert_caller_is_attested_node();
 
         if !self.protocol_state.is_running_or_resharing() {
             return Err(InvalidState::ProtocolStateNotRunning.into());
@@ -570,10 +573,11 @@ impl MpcContract {
         let tee_upgrade_deadline_duration =
             Duration::from_secs(self.config.tee_upgrade_deadline_duration_seconds);
 
-        // Verify the TEE quote and Docker image for the proposed participant
+        // Verify the TEE quote (including TLS and account keys) and Docker image for the proposed participant
         let status = self.tee_state.verify_proposed_participant_attestation(
             &proposed_participant_attestation,
             tls_public_key.clone(),
+            account_key.clone(),
             tee_upgrade_deadline_duration,
         );
 
@@ -587,6 +591,7 @@ impl MpcContract {
             NodeId {
                 account_id: account_id.clone(),
                 tls_public_key,
+                account_public_key: Some(account_key),
             },
             proposed_participant_attestation,
         );
@@ -701,6 +706,9 @@ impl MpcContract {
     #[handle_result]
     pub fn start_keygen_instance(&mut self, key_event_id: KeyEventId) -> Result<(), Error> {
         log!("start_keygen_instance: signer={}", env::signer_account_id(),);
+        
+        self.tee_state.assert_caller_is_attested_node();
+
         self.protocol_state
             .start_keygen_instance(key_event_id, self.config.key_event_timeout_blocks)
     }
@@ -732,6 +740,8 @@ impl MpcContract {
             key_event_id,
             public_key,
         );
+        
+        self.tee_state.assert_caller_is_attested_node();
 
         let extended_key =
             public_key
@@ -755,6 +765,8 @@ impl MpcContract {
             "start_reshare_instance: signer={}",
             env::signer_account_id()
         );
+        
+        self.tee_state.assert_caller_is_attested_node();
         self.protocol_state
             .start_reshare_instance(key_event_id, self.config.key_event_timeout_blocks)
     }
@@ -779,7 +791,8 @@ impl MpcContract {
             env::signer_account_id(),
             key_event_id,
         );
-
+        
+        self.tee_state.assert_caller_is_attested_node();
         let resharing_concluded =
             if let Some(new_state) = self.protocol_state.vote_reshared(key_event_id)? {
                 // Resharing has concluded, transition to running state
@@ -856,7 +869,8 @@ impl MpcContract {
             "vote_abort_key_event_instance: signer={}",
             env::signer_account_id()
         );
-
+        
+        self.tee_state.assert_caller_is_attested_node();
         self.protocol_state
             .vote_abort_key_event_instance(key_event_id)
     }
@@ -1004,6 +1018,8 @@ impl MpcContract {
     #[handle_result]
     pub fn verify_tee(&mut self) -> Result<bool, Error> {
         log!("verify_tee: signer={}", env::signer_account_id());
+        //caller must be a participant (node or operator)
+        let voter = self.voter_or_panic();
         let ProtocolContractState::Running(running_state) = &mut self.protocol_state else {
             return Err(InvalidState::ProtocolStateNotRunning.into());
         };
@@ -1391,8 +1407,8 @@ impl MpcContract {
     /// - `InvalidParameters::InvalidTeeRemoteAttestation`: if destination node’s TEE quote is invalid
     #[handle_result]
     pub fn conclude_node_migration(&mut self, keyset: &Keyset) -> Result<(), Error> {
-        let account_id = env::signer_account_id();
-        let signer_pk = env::signer_account_pk();
+        let account_id: AccountId = env::signer_account_id();
+        let signer_pk: PublicKey = env::signer_account_pk();
         log!(
             "conclude_node_migration: signer={:?}, signer_pk={:?} keyset={:?}",
             account_id,
@@ -1433,6 +1449,7 @@ impl MpcContract {
         // ensure that this node has a valid TEE quote
         let node_id = NodeId {
             account_id: account_id.clone(),
+            account_public_key: Some(expected_destination_node.signer_account_pk.clone()),
             tls_public_key: expected_destination_node
                 .destination_node_info
                 .sign_pk
@@ -2377,6 +2394,7 @@ mod tests {
                 NodeId {
                     account_id: self.signer_account_id.clone(),
                     tls_public_key: self.attestation_tls_key.clone(),
+                    account_public_key: Some(self.signer_account_pk.clone()),
                 },
                 valid_participant_attestation,
             );

crates/contract/src/primitives/participants.rs

@@ -210,6 +210,7 @@ impl Participants {
             .map(|(account_id, _, p_info)| NodeId {
                 account_id: account_id.clone(),
                 tls_public_key: p_info.sign_pk.clone(),
+                account_public_key: None,
             })
             .collect()
     }

crates/contract/src/state/key_event.rs

@@ -105,7 +105,7 @@ impl KeyEvent {
             .as_mut()
             .unwrap()
             .vote_success(candidate, public_key)?
-        {
+        {   
             VoteSuccessResult::Voted(count) => {
                 if count == self.parameters.participants().len() {
                     Ok(true)