fix: address comments

Author: gilcu3

crates/contract/src/api/node_migration.rs

@@ -169,6 +169,7 @@ impl MpcContract {
     /// Returns the following errors:
     /// - [`InvalidState::ProtocolStateNotRunning`]: if protocol is not in [`Running`](ProtocolContractState::Running) state
     /// - [`InvalidState::NotParticipant`]: if caller is not a current participant
+    /// - [`ConversionError::DataConversion`](crate::errors::ConversionError::DataConversion): if the provided keyset contains a malformed public key
     /// - [`NodeMigrationError::KeysetMismatch`](crate::errors::NodeMigrationError::KeysetMismatch): if provided keyset does not match the expected keyset
     /// - [`NodeMigrationError::MigrationNotFound`](crate::errors::NodeMigrationError::MigrationNotFound): if no migration record exists for the caller
     /// - [`NodeMigrationError::AccountPublicKeyMismatch`](crate::errors::NodeMigrationError::AccountPublicKeyMismatch): if caller’s public key does not match the expected destination node

crates/contract/src/dto_mapping.rs

@@ -18,7 +18,7 @@ use near_sdk::env::sha256_array;
 
 use crate::{
     config::Config,
-    crypto_shared::types::PublicKeyExtended,
+    crypto_shared::types::{PublicKeyExtended, serializable::SerializableEdwardsPoint},
     errors::{ConversionError, Error},
     primitives::{
         domain::{AddDomainsVotes, DomainRegistry},
@@ -682,35 +682,53 @@ impl IntoInterfaceType<dtos::DomainRegistry> for &DomainRegistry {
 impl TryIntoContractType<PublicKeyExtended> for dtos::PublicKeyExtended {
     type Error = Error;
     fn try_into_contract_type(self) -> Result<PublicKeyExtended, Self::Error> {
-        let public_key =
-            dtos::PublicKey::try_from(&self).map_err(|err| ConversionError::DataConversion {
-                reason: format!("Failed to parse public key: {err}"),
-            })?;
-        let extended: PublicKeyExtended =
-            public_key
-                .try_into()
-                .map_err(|err| ConversionError::DataConversion {
-                    reason: format!("Failed to extend public key: {err}"),
-                })?;
+        let parse_failed = |err| ConversionError::DataConversion {
+            reason: format!("Failed to parse public key: {err}"),
+        };
 
-        // The DTO carries the Edwards point alongside the compressed key; the contract type
-        // derives it instead, so reject a pair that disagrees rather than silently dropping it.
-        if let (
-            dtos::PublicKeyExtended::Ed25519 { edwards_point, .. },
-            PublicKeyExtended::Ed25519 {
-                edwards_point: derived,
-                ..
-            },
-        ) = (&self, &extended)
-            && derived.to_bytes() != *edwards_point
-        {
-            return Err(ConversionError::DataConversion {
-                reason: "The Edwards point does not match the compressed public key.".to_string(),
+        match self {
+            dtos::PublicKeyExtended::Secp256k1 { near_public_key } => {
+                Ok(PublicKeyExtended::Secp256k1 {
+                    near_public_key: near_public_key.parse().map_err(parse_failed)?,
+                })
             }
-            .into());
-        }
+            dtos::PublicKeyExtended::Ed25519 {
+                near_public_key_compressed,
+                edwards_point,
+            } => {
+                let near_public_key_compressed: dtos::Ed25519PublicKey =
+                    near_public_key_compressed.parse().map_err(parse_failed)?;
+                let derived = SerializableEdwardsPoint::from_bytes(&near_public_key_compressed)
+                    .into_option()
+                    .ok_or_else(|| ConversionError::DataConversion {
+                        reason: "The compressed key is not a valid Edwards point.".to_string(),
+                    })?;
+                // The DTO carries the Edwards point alongside the compressed key; the contract
+                // type derives it, so a pair that disagrees is rejected rather than dropped.
+                if derived.to_bytes() != edwards_point {
+                    return Err(ConversionError::DataConversion {
+                        reason: "The Edwards point does not match the compressed public key."
+                            .to_string(),
+                    }
+                    .into());
+                }
 
-        Ok(extended)
+                Ok(PublicKeyExtended::Ed25519 {
+                    near_public_key_compressed,
+                    edwards_point: derived,
+                })
+            }
+            dtos::PublicKeyExtended::Bls12381 { public_key } => {
+                let dtos::PublicKey::Bls12381(public_key) = public_key else {
+                    return Err(ConversionError::DataConversion {
+                        reason: "Expected a bls12381g2 public key.".to_string(),
+                    }
+                    .into());
+                };
+
+                Ok(PublicKeyExtended::Bls12381 { public_key })
+            }
+        }
     }
 }
 
@@ -1137,6 +1155,31 @@ mod tests {
         );
     }
 
+    /// The variant tag is not what decides the curve — the `<curve>:` prefix inside the key
+    /// string is. A pair that disagrees must be rejected, not silently reinterpreted.
+    #[rstest]
+    #[case::ed25519_tag_holding_a_secp256k1_key(dtos::PublicKeyExtended::Ed25519 {
+        near_public_key_compressed: String::from(&dtos::Secp256k1PublicKey([1u8; 64])),
+        edwards_point: [0u8; 32],
+    })]
+    #[case::bls12381_tag_holding_an_ed25519_key(dtos::PublicKeyExtended::Bls12381 {
+        public_key: dtos::PublicKey::Ed25519(bogus_ed25519_public_key()),
+    })]
+    fn public_key_extended__should_reject_a_variant_tag_that_disagrees_with_the_key(
+        #[case] dto: dtos::PublicKeyExtended,
+    ) {
+        // When
+        let result: Result<PublicKeyExtended, Error> = dto.try_into_contract_type();
+
+        // Then
+        assert_matches!(
+            result,
+            Err(Error::ConversionError(
+                ConversionError::DataConversion { .. }
+            ))
+        );
+    }
+
     /// A threshold below the relative (>= 60%) requirement must be rejected at the
     /// DTO boundary rather than deferred to a later validation step.
     #[test]

crates/contract/tests/snapshots/abi__abi_has_not_changed.snap

@@ -549,7 +549,7 @@ expression: abi
       },
       {
         "name": "conclude_node_migration",
-        "doc": " Finalizes a node migration for the calling account.\n\n This method can only be called while the protocol is in a [`Running`](ProtocolContractState::Running) state\n and by an existing participant. On success, the participant’s information is\n updated to the new destination node.\n\n # Errors\n Returns the following errors:\n - [`InvalidState::ProtocolStateNotRunning`]: if protocol is not in [`Running`](ProtocolContractState::Running) state\n - [`InvalidState::NotParticipant`]: if caller is not a current participant\n - [`NodeMigrationError::KeysetMismatch`](crate::errors::NodeMigrationError::KeysetMismatch): if provided keyset does not match the expected keyset\n - [`NodeMigrationError::MigrationNotFound`](crate::errors::NodeMigrationError::MigrationNotFound): if no migration record exists for the caller\n - [`NodeMigrationError::AccountPublicKeyMismatch`](crate::errors::NodeMigrationError::AccountPublicKeyMismatch): if caller’s public key does not match the expected destination node\n - [`InvalidParameters::InvalidTeeRemoteAttestation`]: if destination node’s TEE quote is invalid",
+        "doc": " Finalizes a node migration for the calling account.\n\n This method can only be called while the protocol is in a [`Running`](ProtocolContractState::Running) state\n and by an existing participant. On success, the participant’s information is\n updated to the new destination node.\n\n # Errors\n Returns the following errors:\n - [`InvalidState::ProtocolStateNotRunning`]: if protocol is not in [`Running`](ProtocolContractState::Running) state\n - [`InvalidState::NotParticipant`]: if caller is not a current participant\n - [`ConversionError::DataConversion`](crate::errors::ConversionError::DataConversion): if the provided keyset contains a malformed public key\n - [`NodeMigrationError::KeysetMismatch`](crate::errors::NodeMigrationError::KeysetMismatch): if provided keyset does not match the expected keyset\n - [`NodeMigrationError::MigrationNotFound`](crate::errors::NodeMigrationError::MigrationNotFound): if no migration record exists for the caller\n - [`NodeMigrationError::AccountPublicKeyMismatch`](crate::errors::NodeMigrationError::AccountPublicKeyMismatch): if caller’s public key does not match the expected destination node\n - [`InvalidParameters::InvalidTeeRemoteAttestation`]: if destination node’s TEE quote is invalid",
         "kind": "call",
         "params": {
           "serialization_type": "json",

crates/devnet/src/mpc.rs

@@ -495,11 +495,11 @@ impl MpcProposeUpdateContractCmd {
             .await
             .into_return_value()
             .expect("Failed to propose update");
-        let update_id: u64 = serde_json::from_slice(&result).expect(&format!(
+        let update_id: UpdateId = serde_json::from_slice(&result).expect(&format!(
             "Failed to deserialize result: {}",
             String::from_utf8_lossy(&result)
         ));
-        println!("Proposed update with ID {}", update_id);
+        println!("Proposed update with ID {}", update_id.0);
         println!("Run the following command to vote for the update:");
         let self_exe = std::env::current_exe()
             .expect("Failed to get current executable path")
@@ -508,7 +508,7 @@ impl MpcProposeUpdateContractCmd {
             .to_string();
         println!(
             "{} mpc {} vote-update --update-id={}",
-            self_exe, name, update_id
+            self_exe, name, update_id.0
         );
     }
 }

crates/near-mpc-contract-interface/src/types/updates.rs

@@ -25,7 +25,7 @@ type Sha256Digest = [u8; 32];
 )]
 #[cfg_attr(
     all(feature = "abi", not(target_arch = "wasm32")),
-    derive(schemars::JsonSchema, borsh::BorshSchema)
+    derive(schemars::JsonSchema)
 )]
 pub struct UpdateId(pub u64);