Merge pull request #613 from near/serhii/contract-api

Use structs in sign request and response fn

.vscode/settings.json

@@ -9,5 +9,6 @@
     "integration-tests/chain-signatures/Cargo.toml",
     "integration-tests/fastauth/Cargo.toml",
     "mpc-recovery/Cargo.toml",
+    "load-tests/Cargo.toml",
   ],
 }

contract/src/lib.rs

@@ -3,10 +3,15 @@ pub mod primitives;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::collections::LookupMap;
 use near_sdk::serde::{Deserialize, Serialize};
+
 use near_sdk::{
     env, log, near_bindgen, AccountId, BorshStorageKey, Gas, Promise, PromiseOrValue, PublicKey,
 };
-use primitives::{CandidateInfo, Candidates, ParticipantInfo, Participants, PkVotes, Votes};
+
+use primitives::{
+    CandidateInfo, Candidates, ParticipantInfo, Participants, PkVotes, SignRequest, SignResult,
+    Votes,
+};
 use std::collections::{BTreeMap, HashSet};
 
 const GAS_FOR_SIGN_CALL: Gas = Gas::from_tgas(250);
@@ -68,35 +73,35 @@ impl Default for VersionedMpcContract {
 #[derive(BorshDeserialize, BorshSerialize, Debug)]
 pub struct MpcContract {
     protocol_state: ProtocolContractState,
-    pending_requests: LookupMap<[u8; 32], Option<(String, String)>>,
+    pending_requests: LookupMap<SignRequest, Option<SignResult>>,
     request_counter: u32,
 }
 
 impl MpcContract {
-    fn add_request(&mut self, payload: &[u8; 32], signature: &Option<(String, String)>) {
+    fn add_request(&mut self, request: &SignRequest, sign_result: &Option<SignResult>) {
         if self.request_counter > 8 {
             env::panic_str("Too many pending requests. Please, try again later.");
         }
-        if !self.pending_requests.contains_key(payload) {
+        if !self.pending_requests.contains_key(request) {
             self.request_counter += 1;
         }
-        self.pending_requests.insert(payload, signature);
+        self.pending_requests.insert(request, sign_result);
     }
 
-    fn remove_request(&mut self, payload: &[u8; 32]) {
-        self.pending_requests.remove(payload);
+    fn remove_request(&mut self, request: &SignRequest) {
+        self.pending_requests.remove(request);
         self.request_counter -= 1;
     }
 
-    fn add_signature(&mut self, payload: &[u8; 32], signature: (String, String)) {
-        if self.pending_requests.contains_key(payload) {
-            self.pending_requests.insert(payload, &Some(signature));
+    fn add_sig_result(&mut self, request: &SignRequest, sign_result: SignResult) {
+        if self.pending_requests.contains_key(request) {
+            self.pending_requests.insert(request, &Some(sign_result));
         }
     }
 
-    fn clean_payloads(&mut self, payloads: Vec<[u8; 32]>, counter: u32) {
+    fn clean_payloads(&mut self, requests: Vec<SignRequest>, counter: u32) {
         log!("clean_payloads");
-        for payload in payloads.iter() {
+        for payload in requests.iter() {
             self.pending_requests.remove(payload);
         }
         self.request_counter = counter;
@@ -120,10 +125,10 @@ impl MpcContract {
 impl VersionedMpcContract {
     #[allow(unused_variables)]
     /// `key_version` must be less than or equal to the value at `latest_key_version`
-    pub fn sign(&mut self, payload: [u8; 32], path: String, key_version: u32) -> Promise {
+    pub fn sign(&mut self, request: SignRequest) -> Promise {
         let latest_key_version: u32 = self.latest_key_version();
         assert!(
-            key_version <= latest_key_version,
+            request.key_version <= latest_key_version,
             "This version of the signer contract doesn't support versions greater than {}",
             latest_key_version,
         );
@@ -137,15 +142,15 @@ impl VersionedMpcContract {
         log!(
             "sign: signer={}, payload={:?}, path={:?}, key_version={}",
             env::signer_account_id(),
-            payload,
-            path,
-            key_version
+            request.payload,
+            request.path,
+            request.key_version
         );
-        match self.signature_per_payload(payload) {
+        match self.sign_result(&request) {
             None => {
-                self.add_request(&payload, &None);
+                self.add_sign_request(&request, &None);
                 log!(&serde_json::to_string(&near_sdk::env::random_seed_array()).unwrap());
-                Self::ext(env::current_account_id()).sign_helper(payload, 0)
+                Self::ext(env::current_account_id()).sign_helper(request, 0)
             }
             Some(_) => env::panic_str("Signature for this payload already requested"),
         }
@@ -177,19 +182,18 @@ impl VersionedMpcContract {
 // MPC Node API
 #[near_bindgen]
 impl VersionedMpcContract {
-    pub fn respond(&mut self, payload: [u8; 32], big_r: String, s: String) {
+    pub fn respond(&mut self, request: SignRequest, result: SignResult) {
         let protocol_state = self.mutable_state();
         if let ProtocolContractState::Running(state) = protocol_state {
             let signer = env::signer_account_id();
             if state.participants.contains_key(&signer) {
                 log!(
-                    "respond: signer={}, payload={:?} big_r={} s={}",
+                    "respond: signer={}, request={:?} result:{:?}",
                     signer,
-                    payload,
-                    big_r,
-                    s
+                    request,
+                    result
                 );
-                self.add_signature(&payload, (big_r, s));
+                self.add_sign_result(&request, result);
             } else {
                 env::panic_str("only participants can respond");
             }
@@ -470,18 +474,18 @@ impl VersionedMpcContract {
     #[private]
     pub fn sign_helper(
         &mut self,
-        payload: [u8; 32],
+        request: SignRequest,
         depth: usize,
-    ) -> PromiseOrValue<(String, String)> {
-        if let Some(signature) = self.signature_per_payload(payload) {
+    ) -> PromiseOrValue<SignResult> {
+        if let Some(signature) = self.sign_result(&request) {
             match signature {
                 Some(signature) => {
                     log!(
                         "sign_helper: signature ready: {:?}, depth: {:?}",
                         signature,
                         depth
                     );
-                    self.remove_request(&payload);
+                    self.remove_sign_request(&request);
                     PromiseOrValue::Value(signature)
                 }
                 None => {
@@ -490,7 +494,7 @@ impl VersionedMpcContract {
                     // We keep one call back so we can cleanup then call panic on the next call
                     // Start cleaning up if there's less than 25 teragas left regardless of how deep you are.
                     if depth > 30 || env::prepaid_gas() < Gas::from_tgas(25) {
-                        self.remove_request(&payload);
+                        self.remove_sign_request(&request);
                         let self_id = env::current_account_id();
                         PromiseOrValue::Promise(Self::ext(self_id).fail_helper(
                             "Signature was not provided in time. Please, try again.".to_string(),
@@ -502,7 +506,7 @@ impl VersionedMpcContract {
                         ));
                         let account_id = env::current_account_id();
                         PromiseOrValue::Promise(
-                            Self::ext(account_id).sign_helper(payload, depth + 1),
+                            Self::ext(account_id).sign_helper(request, depth + 1),
                         )
                     }
                 }
@@ -539,10 +543,10 @@ impl VersionedMpcContract {
     }
 
     #[private]
-    pub fn clean_payloads(&mut self, payloads: Vec<[u8; 32]>, counter: u32) {
+    pub fn clean_payloads(&mut self, requests: Vec<SignRequest>, counter: u32) {
         match self {
             Self::V0(mpc_contract) => {
-                mpc_contract.clean_payloads(payloads, counter);
+                mpc_contract.clean_payloads(requests, counter);
             }
         }
     }
@@ -562,26 +566,26 @@ impl VersionedMpcContract {
 // Helper functions
 #[near_bindgen]
 impl VersionedMpcContract {
-    fn remove_request(&mut self, payload: &[u8; 32]) {
+    fn remove_sign_request(&mut self, request: &SignRequest) {
         match self {
             Self::V0(mpc_contract) => {
-                mpc_contract.remove_request(payload);
+                mpc_contract.remove_request(request);
             }
         }
     }
 
-    fn add_request(&mut self, payload: &[u8; 32], signature: &Option<(String, String)>) {
+    fn add_sign_request(&mut self, request: &SignRequest, sign_result: &Option<SignResult>) {
         match self {
             Self::V0(mpc_contract) => {
-                mpc_contract.add_request(payload, signature);
+                mpc_contract.add_request(request, sign_result);
             }
         }
     }
 
-    fn add_signature(&mut self, payload: &[u8; 32], signature: (String, String)) {
+    fn add_sign_result(&mut self, request: &SignRequest, sign_result: SignResult) {
         match self {
             Self::V0(mpc_contract) => {
-                mpc_contract.add_signature(payload, signature);
+                mpc_contract.add_sig_result(request, sign_result);
             }
         }
     }
@@ -592,9 +596,9 @@ impl VersionedMpcContract {
         }
     }
 
-    fn signature_per_payload(&self, payload: [u8; 32]) -> Option<Option<(String, String)>> {
+    fn sign_result(&self, request: &SignRequest) -> Option<Option<SignResult>> {
         match self {
-            Self::V0(mpc_contract) => mpc_contract.pending_requests.get(&payload),
+            Self::V0(mpc_contract) => mpc_contract.pending_requests.get(request),
         }
     }
 }

contract/src/primitives.rs

@@ -207,3 +207,16 @@ impl PkVotes {
         self.votes.entry(public_key).or_default()
     }
 }
+
+#[derive(Serialize, Deserialize, BorshDeserialize, BorshSerialize, Debug)]
+pub struct SignRequest {
+    pub payload: [u8; 32],
+    pub path: String,
+    pub key_version: u32,
+}
+
+#[derive(Serialize, Deserialize, BorshDeserialize, BorshSerialize, Debug)]
+pub struct SignResult {
+    pub big_r: String,
+    pub s: String,
+}

integration-tests/chain-signatures/tests/actions/mod.rs

@@ -11,6 +11,7 @@ use k256::elliptic_curve::scalar::FromUintUnchecked;
 use k256::elliptic_curve::sec1::FromEncodedPoint;
 use k256::elliptic_curve::ProjectivePoint;
 use k256::{AffinePoint, EncodedPoint, Scalar, Secp256k1};
+use mpc_contract::primitives::SignRequest;
 use mpc_contract::RunningContractState;
 use mpc_recovery_node::kdf;
 use mpc_recovery_node::util::ScalarExt;
@@ -48,6 +49,12 @@ pub async fn request_sign(
         .rpc_client
         .fetch_nonce(&signer.account_id, &signer.public_key)
         .await?;
+
+    let request = SignRequest {
+        payload: payload_hashed,
+        path: "test".to_string(),
+        key_version: 0,
+    };
     let tx_hash = ctx
         .jsonrpc_client
         .call(&RpcBroadcastTxAsyncRequest {
@@ -60,9 +67,7 @@ pub async fn request_sign(
                 actions: vec![Action::FunctionCall(Box::new(FunctionCallAction {
                     method_name: "sign".to_string(),
                     args: serde_json::to_vec(&serde_json::json!({
-                        "payload": payload_hashed,
-                        "path": "test",
-                        "key_version": 0,
+                        "request": request,
                     }))?,
                     gas: 300_000_000_000_000,
                     deposit: 0,
@@ -118,6 +123,13 @@ pub async fn request_sign_non_random(
         .rpc_client
         .fetch_nonce(&signer.account_id, &signer.public_key)
         .await?;
+
+    let request = SignRequest {
+        payload: payload_hashed,
+        path: "test".to_string(),
+        key_version: 0,
+    };
+
     let tx_hash = ctx
         .jsonrpc_client
         .call(&RpcBroadcastTxAsyncRequest {
@@ -130,9 +142,7 @@ pub async fn request_sign_non_random(
                 actions: vec![Action::FunctionCall(Box::new(FunctionCallAction {
                     method_name: "sign".to_string(),
                     args: serde_json::to_vec(&serde_json::json!({
-                        "payload": payload_hashed,
-                        "path": "test",
-                        "key_version": 0,
+                        "request": request,
                     }))?,
                     gas: 300_000_000_000_000,
                     deposit: 0,

integration-tests/chain-signatures/tests/actions/wait_for.rs

@@ -16,6 +16,8 @@ use near_jsonrpc_client::methods::tx::TransactionInfo;
 use near_lake_primitives::CryptoHash;
 use near_primitives::views::FinalExecutionStatus;
 use near_workspaces::Account;
+use serde::Deserialize;
+use serde::Serialize;
 
 pub async fn running_mpc<'a>(
     ctx: &MultichainTestContext<'a>,
@@ -203,6 +205,13 @@ pub async fn has_at_least_mine_presignatures<'a>(
     Ok(state_views)
 }
 
+// TODO: use structure from contract when the internal types are the same
+#[derive(Serialize, Deserialize, Debug)]
+pub struct SignResult {
+    pub big_r: AffinePoint,
+    pub s: Scalar,
+}
+
 pub async fn signature_responded(
     ctx: &MultichainTestContext<'_>,
     tx_hash: CryptoHash,
@@ -224,14 +233,19 @@ pub async fn signature_responded(
         let Some(outcome) = outcome_view.final_execution_outcome else {
             anyhow::bail!("final execution outcome not available");
         };
+
         let outcome = outcome.into_outcome();
 
         let FinalExecutionStatus::SuccessValue(payload) = outcome.status else {
             anyhow::bail!("tx finished unsuccessfully: {:?}", outcome.status);
         };
 
-        let (big_r, s): (AffinePoint, Scalar) = serde_json::from_slice(&payload)?;
-        let signature = cait_sith::FullSignature::<Secp256k1> { big_r, s };
+        let result: SignResult = serde_json::from_slice(&payload)?;
+        let signature = cait_sith::FullSignature::<Secp256k1> {
+            big_r: result.big_r,
+            s: result.s,
+        };
+
         Ok(signature)
     };