chore: use UpdateId and Keyset from dtos

Author: gilcu3

crates/contract/src/api/attestation.rs

@@ -1028,7 +1028,7 @@ mod tests {
         let mut contract = MpcContract::init_running(
             domains.clone(),
             1,
-            keyset.clone(),
+            (&keyset).into_dto_type(),
             (&parameters).into_dto_type(),
             None,
         )
@@ -1149,9 +1149,14 @@ mod tests {
                 attempt: AttemptId::new(),
             }],
         );
-        let mut contract =
-            MpcContract::init_running(domains, 1, keyset, (&parameters).into_dto_type(), None)
-                .unwrap();
+        let mut contract = MpcContract::init_running(
+            domains,
+            1,
+            (&keyset).into_dto_type(),
+            (&parameters).into_dto_type(),
+            None,
+        )
+        .unwrap();
 
         // Expire the last participant's attestation so a kickout drops the set to 4.
         let participant_list: Vec<_> = participants.participants().to_vec();

crates/contract/src/api/foreign_chain_support.rs

@@ -946,7 +946,7 @@ mod tests {
         let mut env = Environment::new(None, Some(operator4.clone()), None);
         env.set_pk(new_signer_near_pk);
         contract
-            .conclude_node_migration(&keyset)
+            .conclude_node_migration((&keyset).into_dto_type())
             .expect("migration should succeed");
 
         // Then: all 4 chains available — no manual recompute needed.
@@ -1214,9 +1214,14 @@ mod tests {
         let parameters =
             GovernanceThresholdParameters::new(gen_participants(4), GovernanceThreshold::new(3))
                 .unwrap();
-        let mut contract =
-            MpcContract::init_running(domains, 1, keyset, (&parameters).into_dto_type(), None)
-                .unwrap();
+        let mut contract = MpcContract::init_running(
+            domains,
+            1,
+            (&keyset).into_dto_type(),
+            (&parameters).into_dto_type(),
+            None,
+        )
+        .unwrap();
         let participants = participant_account_ids(&contract);
         whitelist_chain(&mut contract, dtos::ForeignChain::Bitcoin);
 
@@ -1276,9 +1281,14 @@ mod tests {
         let parameters =
             GovernanceThresholdParameters::new(gen_participants(4), GovernanceThreshold::new(3))
                 .unwrap();
-        let mut contract =
-            MpcContract::init_running(domains, 2, keyset, (&parameters).into_dto_type(), None)
-                .unwrap();
+        let mut contract = MpcContract::init_running(
+            domains,
+            2,
+            (&keyset).into_dto_type(),
+            (&parameters).into_dto_type(),
+            None,
+        )
+        .unwrap();
         let participants = participant_account_ids(&contract);
         whitelist_chain(&mut contract, dtos::ForeignChain::Bitcoin);
 

crates/contract/src/api/governance.rs

@@ -292,9 +292,14 @@ mod tests {
                 attempt: AttemptId::new(),
             }],
         );
-        let contract =
-            MpcContract::init_running(domains, 1, keyset, (&parameters).into_dto_type(), None)
-                .unwrap();
+        let contract = MpcContract::init_running(
+            domains,
+            1,
+            (&keyset).into_dto_type(),
+            (&parameters).into_dto_type(),
+            None,
+        )
+        .unwrap();
         (contract, participants, first_participant_id, domain_id)
     }
 

crates/contract/src/api/lifecycle.rs

@@ -86,10 +86,11 @@ impl MpcContract {
     pub fn init_running(
         domains: Vec<DomainConfig>,
         next_domain_id: u64,
-        keyset: Keyset,
+        keyset: dtos::Keyset,
         parameters: dtos::GovernanceThresholdParameters,
         init_config: Option<dtos::InitConfig>,
     ) -> Result<Self, Error> {
+        let keyset: Keyset = keyset.try_into_contract_type()?;
         let parameters: GovernanceThresholdParameters = parameters.try_into_contract_type()?;
         // Log participant count and hash - full parameters exceed NEAR's 16KB log limit at ~100 participants
         let params_hash = env::sha256_array(borsh::to_vec(&parameters).unwrap());

crates/contract/src/api/node_migration.rs

@@ -2,7 +2,7 @@
 //! migration to new hardware, and updating a participant's URL.
 
 use crate::api::common::require_deposit;
-use crate::dto_mapping::IntoContractType;
+use crate::dto_mapping::{IntoContractType, TryIntoContractType};
 use crate::errors::{self, Error, InvalidParameters, InvalidState};
 use crate::primitives::key_state::Keyset;
 use crate::primitives::participants::ParticipantInfo;
@@ -174,7 +174,7 @@ impl MpcContract {
     /// - [`NodeMigrationError::AccountPublicKeyMismatch`](crate::errors::NodeMigrationError::AccountPublicKeyMismatch): if caller’s public key does not match the expected destination node
     /// - [`InvalidParameters::InvalidTeeRemoteAttestation`]: if destination node’s TEE quote is invalid
     #[handle_result]
-    pub fn conclude_node_migration(&mut self, keyset: &Keyset) -> Result<(), Error> {
+    pub fn conclude_node_migration(&mut self, keyset: dtos::Keyset) -> Result<(), Error> {
         let account_id = Self::assert_caller_is_signer();
         let signer_pk = env::signer_account_pk();
         log!(
@@ -194,10 +194,12 @@ impl MpcContract {
             .into());
         }
 
+        // Converted after the participant/state validation so those errors take precedence.
+        let keyset: Keyset = keyset.try_into_contract_type()?;
         let expected_keyset = &running_state.keyset;
-        if expected_keyset != keyset {
+        if *expected_keyset != keyset {
             return Err(errors::NodeMigrationError::KeysetMismatch {
-                found: keyset.clone(),
+                found: keyset,
                 expected: expected_keyset.clone(),
             }
             .into());
@@ -310,6 +312,7 @@ pub const MINIMUM_NODE_MANAGEMENT_DEPOSIT: NearToken =
 mod tests {
     use super::*;
     use crate::api::test_utils::NUM_DOMAINS;
+    use crate::dto_mapping::IntoInterfaceType;
     use crate::errors::NodeMigrationError;
     use crate::primitives::participants::{ParticipantId, Participants};
     use crate::primitives::test_utils::{
@@ -929,7 +932,7 @@ mod tests {
             test_env.set_signer(&self.signer_account_id);
             test_env.set_pk(self.signer_account_pk.clone());
 
-            let res = contract.conclude_node_migration(keyset);
+            let res = contract.conclude_node_migration(keyset.into_dto_type());
 
             if let Some(check) = &self.expected_error_check {
                 let err = res.unwrap_err();

crates/contract/src/api/test_utils.rs

@@ -139,8 +139,14 @@ pub(crate) fn basic_setup_with_protocol(
     let parameters =
         GovernanceThresholdParameters::new(gen_participants(4), GovernanceThreshold::new(3))
             .unwrap();
-    let contract =
-        MpcContract::init_running(domains, 1, keyset, (&parameters).into_dto_type(), None).unwrap();
+    let contract = MpcContract::init_running(
+        domains,
+        1,
+        (&keyset).into_dto_type(),
+        (&parameters).into_dto_type(),
+        None,
+    )
+    .unwrap();
     (context, contract, sk)
 }
 

crates/contract/src/api/update.rs

@@ -3,7 +3,7 @@
 
 use crate::api::common::refund_to;
 use crate::config::Config;
-use crate::dto_mapping::IntoInterfaceType;
+use crate::dto_mapping::{IntoContractType, IntoInterfaceType};
 use crate::errors::{Error, InvalidParameters, InvalidState};
 use crate::state::ProtocolContractState;
 use crate::update::{ProposedUpdates, Update, UpdateId};
@@ -19,7 +19,7 @@ impl MpcContract {
     pub fn propose_update(
         &mut self,
         #[serializer(borsh)] args: ProposeUpdateArgs,
-    ) -> Result<UpdateId, Error> {
+    ) -> Result<dtos::UpdateId, Error> {
         // Only voters can propose updates:
         let proposer = self.voter_or_panic();
         let payload_bytes =
@@ -56,22 +56,23 @@ impl MpcContract {
             refund_to(&proposer, diff);
         }
 
-        Ok(id)
+        Ok(id.into_dto_type())
     }
 
-    /// Vote for a proposed update given the [`UpdateId`] of the update.
+    /// Vote for a proposed update, given the id returned by [`Self::propose_update`].
     ///
     /// Returns `Ok(true)` if the amount of voters surpassed the threshold and the update was
     /// executed. Returns `Ok(false)` if the amount of voters did not surpass the threshold.
     /// Returns [`Error`] if the update was not found or if the voter is not a participant
     /// in the protocol.
     #[handle_result]
-    pub fn vote_update(&mut self, id: UpdateId) -> Result<bool, Error> {
+    pub fn vote_update(&mut self, id: dtos::UpdateId) -> Result<bool, Error> {
         log!(
             "vote_update: signer={}, id={:?}",
             env::signer_account_id(),
             id,
         );
+        let id: UpdateId = id.into_contract_type();
 
         let ProtocolContractState::Running(running_state) = &self.protocol_state else {
             env::panic_str("protocol must be in running state");
@@ -195,10 +196,10 @@ mod tests {
     fn propose_and_vote(
         contract: &mut MpcContract,
         update: Update,
-        expected_update_id: u64,
+        expected_update_id: UpdateId,
     ) -> Vec<dtos::AccountId> {
         let update_id = contract.proposed_updates.propose(update.clone());
-        assert_eq!(update_id.0, expected_update_id);
+        assert_eq!(update_id, expected_update_id);
         // generate two accounts for voting
         let account_id_0 = gen_account_id();
         let account_id_1 = gen_account_id();
@@ -220,14 +221,14 @@ mod tests {
     /// Used to convert BTreeMap-based [`ProposedUpdates`] into a sortable vector format for assertions.
     #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
     struct TestUpdate {
-        update_id: u64,
+        update_id: dtos::UpdateId,
         update_hash: dtos::UpdateHash,
         votes: Vec<dtos::AccountId>,
     }
 
     impl TestUpdate {
         fn from_proposed_updates(
-            update_id: u64,
+            update_id: dtos::UpdateId,
             update_hash: dtos::UpdateHash,
             proposed_updates: &dtos::ProposedUpdates,
         ) -> Self {
@@ -245,14 +246,17 @@ mod tests {
         }
     }
 
-    fn propose_and_vote_code(expected_update_id: u64, contract: &mut MpcContract) -> TestUpdate {
+    fn propose_and_vote_code(
+        expected_update_id: UpdateId,
+        contract: &mut MpcContract,
+    ) -> TestUpdate {
         let code: [u8; 1000] = std::array::from_fn(|_| rand::random());
         let hash = Sha256::digest(code);
         let update = Update::Contract(code.into());
         let expected_update_hash = dtos::UpdateHash::Code(hash.into());
         let expected_votes = propose_and_vote(contract, update, expected_update_id);
         TestUpdate {
-            update_id: expected_update_id,
+            update_id: expected_update_id.into_dto_type(),
             update_hash: expected_update_hash,
             votes: expected_votes,
         }
@@ -290,18 +294,18 @@ mod tests {
         assert_eq!(empty_result.updates, BTreeMap::new());
 
         // Propose and vote for code update
-        let code_update_id = 0;
+        let code_update_id = UpdateId(0);
         let mut code_update = propose_and_vote_code(code_update_id, &mut contract);
 
         // Propose and vote for config update
         let mut config_update = {
             let update_config = dummy_config(1);
             let config_hash = Sha256::digest(serde_json::to_vec(&update_config).unwrap());
             let config_update_obj = Update::Config(update_config.clone());
-            let config_update_id = 1;
+            let config_update_id = UpdateId(1);
             let config_votes = propose_and_vote(&mut contract, config_update_obj, config_update_id);
             TestUpdate {
-                update_id: config_update_id,
+                update_id: config_update_id.into_dto_type(),
                 update_hash: dtos::UpdateHash::Config(config_hash.into()),
                 votes: config_votes,
             }
@@ -364,9 +368,8 @@ mod tests {
         let mut contract = MpcContract::new_from_protocol_state(protocol_contract_state);
 
         // Propose and vote for code update
-        let update_id_u64 = 0;
-        let test_update = propose_and_vote_code(update_id_u64, &mut contract);
-        let update_id = UpdateId::from(update_id_u64);
+        let update_id = UpdateId(0);
+        let test_update = propose_and_vote_code(update_id, &mut contract);
 
         for (account_id, _, _) in participants.participants() {
             contract
@@ -376,7 +379,10 @@ mod tests {
             let proposed_updates = contract.proposed_updates();
             assert_eq!(proposed_updates.updates.len(), 1);
             assert_eq!(
-                *proposed_updates.updates.get(&update_id.0).unwrap(),
+                *proposed_updates
+                    .updates
+                    .get(&update_id.into_dto_type())
+                    .unwrap(),
                 test_update.update_hash
             );
 
@@ -386,7 +392,7 @@ mod tests {
             let actual_voters: Vec<_> = proposed_updates
                 .votes
                 .iter()
-                .filter(|&(_, &uid)| uid == update_id.0)
+                .filter(|&(_, &uid)| uid == update_id.into_dto_type())
                 .map(|(voter, _)| voter.clone())
                 .collect();
             assert_eq!(actual_voters.len(), expected_voters.len());
@@ -411,7 +417,7 @@ mod tests {
             let actual_voters: Vec<_> = res
                 .votes
                 .iter()
-                .filter(|&(_, &uid)| uid == update_id.0)
+                .filter(|&(_, &uid)| uid == update_id.into_dto_type())
                 .map(|(voter, _)| voter.clone())
                 .collect();
             assert_eq!(actual_voters.len(), test_update.votes.len());
@@ -429,7 +435,7 @@ mod tests {
         let mut contract = MpcContract::new_from_protocol_state(protocol_contract_state);
 
         // Propose and vote for code update
-        let update_id = 0;
+        let update_id = UpdateId(0);
         let test_update = propose_and_vote_code(update_id, &mut contract);
 
         let mut rng = rand::rngs::StdRng::seed_from_u64(42);
@@ -521,7 +527,7 @@ mod tests {
                 .build()
         );
         // then: threshold not met (need 2 valid votes, have only 1)
-        assert!(!contract.vote_update(update_id).unwrap());
+        assert!(!contract.vote_update(update_id.into_dto_type()).unwrap());
 
         // given: a 2nd participant vote is added
         contract
@@ -536,7 +542,7 @@ mod tests {
                 .build()
         );
         // then: threshold met (have 2 valid votes, need 2)
-        assert!(contract.vote_update(update_id).unwrap());
+        assert!(contract.vote_update(update_id.into_dto_type()).unwrap());
     }
 
     #[test]
@@ -608,9 +614,8 @@ mod tests {
             MpcContract::new_from_protocol_state(ProtocolContractState::Running(running_state));
 
         // propose_and_vote_code adds 2 non-participant votes.
-        let update_id_u64 = 0;
-        let _ = propose_and_vote_code(update_id_u64, &mut contract);
-        let update_id: UpdateId = update_id_u64.into();
+        let update_id = UpdateId(0);
+        let _ = propose_and_vote_code(update_id, &mut contract);
 
         // Add votes from 2 current participants.
         let participants = participants.participants();
@@ -654,9 +659,8 @@ mod tests {
         let mut contract =
             MpcContract::new_from_protocol_state(ProtocolContractState::Running(running_state));
 
-        let update_id_u64 = 0;
-        let test_update = propose_and_vote_code(update_id_u64, &mut contract);
-        let update_id: UpdateId = update_id_u64.into();
+        let update_id = UpdateId(0);
+        let test_update = propose_and_vote_code(update_id, &mut contract);
         let non_participants: HashSet<AccountId> = test_update.votes.iter().cloned().collect();
 
         let participants = participants.participants();