feat: add public key to report data and enforcement on the contract

Author: barakeinav1

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

@@ -570,10 +570,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 +588,7 @@ impl MpcContract {
             NodeId {
                 account_id: account_id.clone(),
                 tls_public_key,
+                account_public_key: Some(account_key),
             },
             proposed_participant_attestation,
         );
@@ -1433,6 +1435,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 +2380,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)

crates/contract/src/tee/tee_state.rs

@@ -1,3 +1,4 @@
+
 use crate::{
     primitives::{key_state::AuthenticatedParticipantId, participants::Participants},
     storage_keys::StorageKey,
@@ -16,14 +17,35 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use mpc_primitives::hash::LauncherDockerComposeHash;
 use near_sdk::{env, near, store::IterableMap, AccountId, PublicKey};
 use std::{collections::HashSet, time::Duration};
+use std::hash::{Hash, Hasher};
 
 #[near(serializers=[borsh, json])]
-#[derive(Debug, Eq, Ord, PartialEq, PartialOrd, Clone, Hash)]
+#[derive(Debug, Ord, PartialOrd, Clone)]
 pub struct NodeId {
-    /// Operator account
+    /// Operator account (on-chain identity)
     pub account_id: AccountId,
     /// TLS public key
     pub tls_public_key: PublicKey,
+    /// Account signing public key (optional for backward compatibility)
+    pub account_public_key: Option<PublicKey>,
+}
+
+// Implement Eq + Hash ignoring account_public_key
+impl PartialEq for NodeId {
+    fn eq(&self, other: &Self) -> bool {
+        self.account_id == other.account_id
+            && self.tls_public_key == other.tls_public_key
+    }
+}
+
+impl Eq for NodeId {}
+
+impl Hash for NodeId {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.account_id.hash(state);
+        self.tls_public_key.hash(state);
+        // intentionally ignoring account_public_key
+    }
 }
 
 pub enum TeeValidationResult {
@@ -64,9 +86,13 @@ impl TeeState {
         &mut self,
         attestation: &Attestation,
         tls_public_key: PublicKey,
+        account_public_key: PublicKey,
         tee_upgrade_deadline_duration: Duration,
     ) -> TeeQuoteStatus {
-        let expected_report_data = ReportData::V1(ReportDataV1::new(tls_public_key));
+        // Recreate the exact same ReportData that the enclave produced
+        let expected_report_data =
+            ReportData::V1(ReportDataV1::new(tls_public_key, account_public_key));
+
         let is_valid = attestation.verify(
             expected_report_data,
             Self::current_time_seconds(),
@@ -96,8 +122,11 @@ impl TeeState {
             return TeeQuoteStatus::None;
         };
 
-        let expected_report_data =
-            ReportData::V1(ReportDataV1::new(node_id.tls_public_key.clone()));
+        let expected_report_data = ReportData ::new(
+            node_id.tls_public_key.clone(),
+            node_id.account_public_key.clone(),
+        );
+
         let time_stamp_seconds = Self::current_time_seconds();
 
         let quote_result = participant_attestation.verify(
@@ -132,12 +161,14 @@ impl TeeState {
             .iter()
             .filter(|(account_id, _, participant_info)| {
                 let tls_public_key = participant_info.sign_pk.clone();
+               
 
                 matches!(
                     self.verify_tee_participant(
                         &NodeId {
                             account_id: account_id.clone(),
-                            tls_public_key
+                            tls_public_key : tls_public_key,
+                            account_public_key: None,
                         },
                         tee_upgrade_deadline_duration
                     ),
@@ -222,6 +253,7 @@ impl TeeState {
             .map(|(account_id, _, p_info)| NodeId {
                 account_id: account_id.clone(),
                 tls_public_key: p_info.sign_pk.clone(),
+                account_public_key: None,
             })
             .collect();
 
@@ -244,6 +276,14 @@ impl TeeState {
     pub fn get_tee_accounts(&self) -> Vec<NodeId> {
         self.participants_attestations.keys().cloned().collect()
     }
+
+     /// Find a NodeId by its TLS public key.
+    pub fn find_node_id_by_tls_key(&self, tls_public_key: &PublicKey) -> Option<NodeId> {
+        self.participants_attestations
+            .keys()
+            .find(|node_id| &node_id.tls_public_key == tls_public_key)
+            .cloned()
+    }
 }
 
 #[cfg(test)]
@@ -269,6 +309,7 @@ mod tests {
             .map(|(account_id, _, p_info)| NodeId {
                 account_id: account_id.clone(),
                 tls_public_key: p_info.sign_pk.clone(),
+                account_public_key:  Some(bogus_ed25519_near_public_key()),//TODO check if this is ok.
             })
             .collect();
 
@@ -277,6 +318,7 @@ mod tests {
 
         let non_participant_uid = NodeId {
             account_id: non_participant.clone(),
+            account_public_key: Some(bogus_ed25519_near_public_key()), 
             tls_public_key: bogus_ed25519_near_public_key(),
         };
         for node_id in &participant_nodes {