Merge pull request #781 from near/refund-sign

refund after sign

chain-signatures/contract/src/lib.rs

@@ -21,8 +21,8 @@ use near_sdk::{
     PromiseError, PublicKey,
 };
 use primitives::{
-    CandidateInfo, Candidates, Participants, PkVotes, SignRequest, SignaturePromiseError,
-    SignatureRequest, SignatureResult, StorageKey, Votes, YieldIndex,
+    CandidateInfo, Candidates, ContractSignatureRequest, Participants, PkVotes, SignRequest,
+    SignaturePromiseError, SignatureRequest, SignatureResult, StorageKey, Votes, YieldIndex,
 };
 use std::collections::{BTreeMap, HashSet};
 
@@ -140,7 +140,7 @@ impl VersionedMpcContract {
             return Err(InvalidParameters::InsufficientDeposit.message(format!(
                 "Attached {}, Required {}",
                 deposit.as_yoctonear(),
-                required_deposit
+                required_deposit,
             )));
         }
         // Make sure sign call will not run out of gas doing recursive calls because the payload will never be removed
@@ -166,7 +166,13 @@ impl VersionedMpcContract {
                 "sign: predecessor={predecessor}, payload={payload:?}, path={path:?}, key_version={key_version}",
             );
             env::log_str(&serde_json::to_string(&near_sdk::env::random_seed_array()).unwrap());
-            Ok(Self::ext(env::current_account_id()).sign_helper(request))
+            let contract_signature_request = ContractSignatureRequest {
+                request,
+                requester: predecessor,
+                deposit,
+                required_deposit: NearToken::from_yoctonear(required_deposit),
+            };
+            Ok(Self::ext(env::current_account_id()).sign_helper(contract_signature_request))
         } else {
             Err(SignError::PayloadCollision.into())
         }
@@ -677,12 +683,12 @@ impl VersionedMpcContract {
     }
 
     #[private]
-    pub fn sign_helper(&mut self, request: SignatureRequest) {
+    pub fn sign_helper(&mut self, contract_signature_request: ContractSignatureRequest) {
         match self {
             Self::V0(mpc_contract) => {
                 let yield_promise = env::promise_yield_create(
                     "clear_state_on_finish",
-                    &serde_json::to_vec(&(&request,)).unwrap(),
+                    &serde_json::to_vec(&(&contract_signature_request,)).unwrap(),
                     CLEAR_STATE_ON_FINISH_CALL_GAS,
                     GasWeight(0),
                     DATA_ID_REGISTER,
@@ -694,7 +700,7 @@ impl VersionedMpcContract {
                     .try_into()
                     .expect("conversion to CryptoHash failed");
 
-                mpc_contract.add_request(&request, data_id);
+                mpc_contract.add_request(&contract_signature_request.request, data_id);
 
                 // NOTE: there's another promise after the clear_state_on_finish to avoid any errors
                 // that would rollback the state.
@@ -730,20 +736,54 @@ impl VersionedMpcContract {
         }
     }
 
+    fn refund_on_fail(request: &ContractSignatureRequest) {
+        let amount = request.deposit;
+        let to = request.requester.clone();
+        log!("refund {amount} to {to} due to fail");
+        Promise::new(to).transfer(amount);
+    }
+
+    fn refund_on_success(request: &ContractSignatureRequest) {
+        let deposit = request.deposit;
+        let required = request.required_deposit;
+        if let Some(diff) = deposit.checked_sub(required) {
+            if diff > NearToken::from_yoctonear(0) {
+                let to = request.requester.clone();
+                log!("refund more than required deposit {diff} to {to}");
+                Promise::new(to).transfer(diff);
+            }
+        }
+    }
+
     #[private]
     #[handle_result]
     pub fn clear_state_on_finish(
         &mut self,
-        request: SignatureRequest,
+        contract_signature_request: ContractSignatureRequest,
         #[callback_result] signature: Result<SignatureResponse, PromiseError>,
     ) -> Result<SignatureResult<SignatureResponse, SignaturePromiseError>, Error> {
         match self {
             Self::V0(mpc_contract) => {
                 // Clean up the local state
-                mpc_contract.remove_request(request)?;
+                let result =
+                    mpc_contract.remove_request(contract_signature_request.request.clone());
+                if result.is_err() {
+                    // refund must happen in clear_state_on_finish, because regardless of this success or fail
+                    // the promise created by clear_state_on_finish is executed, because of callback_unwrap and
+                    // promise_then. but if return_signature_on_finish fail (returns error), the promise created
+                    // by it won't execute.
+                    Self::refund_on_fail(&contract_signature_request);
+                    result?;
+                }
                 match signature {
-                    Ok(signature) => Ok(SignatureResult::Ok(signature)),
-                    Err(_) => Ok(SignatureResult::Err(SignaturePromiseError::Failed)),
+                    Ok(signature) => {
+                        Self::refund_on_success(&contract_signature_request);
+                        Ok(SignatureResult::Ok(signature))
+                    }
+                    Err(_) => {
+                        Self::refund_on_fail(&contract_signature_request);
+                        Ok(SignatureResult::Err(SignaturePromiseError::Failed))
+                    }
                 }
             }
         }

chain-signatures/contract/src/primitives.rs

@@ -2,7 +2,7 @@ use crypto_shared::{derive_epsilon, SerializableScalar};
 use k256::Scalar;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::serde::{Deserialize, Serialize};
-use near_sdk::{AccountId, BorshStorageKey, CryptoHash, PublicKey};
+use near_sdk::{AccountId, BorshStorageKey, CryptoHash, NearToken, PublicKey};
 use std::collections::{BTreeMap, HashMap, HashSet};
 
 pub mod hpke {
@@ -31,6 +31,15 @@ pub struct SignatureRequest {
     pub payload_hash: SerializableScalar,
 }
 
+#[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Debug, Clone)]
+#[borsh(crate = "near_sdk::borsh")]
+pub struct ContractSignatureRequest {
+    pub request: SignatureRequest,
+    pub requester: AccountId,
+    pub deposit: NearToken,
+    pub required_deposit: NearToken,
+}
+
 impl SignatureRequest {
     pub fn new(payload_hash: Scalar, predecessor_id: &AccountId, path: &str) -> Self {
         let epsilon = derive_epsilon(predecessor_id, path);

chain-signatures/contract/tests/sign.rs

@@ -5,6 +5,8 @@ use mpc_contract::errors;
 use mpc_contract::primitives::{CandidateInfo, SignRequest};
 use near_workspaces::types::AccountId;
 
+use crypto_shared::SignatureResponse;
+use near_sdk::NearToken;
 use std::collections::HashMap;
 
 #[tokio::test]
@@ -57,6 +59,142 @@ async fn test_contract_sign_request() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_contract_sign_success_refund() -> anyhow::Result<()> {
+    let (worker, contract, _, sk) = init_env().await;
+    let alice = worker.dev_create_account().await?;
+    let balance = alice.view_account().await?.balance;
+    let contract_balance = contract.view_account().await?.balance;
+    let path = "test";
+
+    let msg = "hello world!";
+    println!("submitting: {msg}");
+    let (payload_hash, respond_req, respond_resp) =
+        create_response(alice.id(), msg, path, &sk).await;
+    let request = SignRequest {
+        payload: payload_hash,
+        path: path.into(),
+        key_version: 0,
+    };
+
+    let status = alice
+        .call(contract.id(), "sign")
+        .args_json(serde_json::json!({
+            "request": request,
+        }))
+        .deposit(NearToken::from_near(1))
+        .max_gas()
+        .transact_async()
+        .await?;
+    dbg!(&status);
+    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
+
+    // Call `respond` as if we are the MPC network itself.
+    let respond = contract
+        .call("respond")
+        .args_json(serde_json::json!({
+            "request": respond_req,
+            "response": respond_resp
+        }))
+        .max_gas()
+        .transact()
+        .await?;
+    dbg!(&respond);
+
+    let execution = status.await?;
+    dbg!(&execution);
+
+    let execution = execution.into_result()?;
+
+    // Finally wait the result:
+    let returned_resp: SignatureResponse = execution.json()?;
+    assert_eq!(
+        returned_resp, respond_resp,
+        "Returned signature request does not match"
+    );
+
+    let new_balance = alice.view_account().await?.balance;
+    let new_contract_balance = contract.view_account().await?.balance;
+    assert!(
+        balance.as_millinear() - new_balance.as_millinear() < 10,
+        "refund should happen"
+    );
+    println!(
+        "{} {} {} {}",
+        balance.as_millinear(),
+        new_balance.as_millinear(),
+        contract_balance.as_millinear(),
+        new_contract_balance.as_millinear(),
+    );
+    assert!(
+        contract_balance.as_millinear() - new_contract_balance.as_millinear() < 20,
+        "respond should take less than 0.02 NEAR"
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_contract_sign_fail_refund() -> anyhow::Result<()> {
+    let (worker, contract, _, sk) = init_env().await;
+    let alice = worker.dev_create_account().await?;
+    let balance = alice.view_account().await?.balance;
+    let contract_balance = contract.view_account().await?.balance;
+    let path = "test";
+
+    let msg = "hello world!";
+    println!("submitting: {msg}");
+    let (payload_hash, _, _) = create_response(alice.id(), msg, path, &sk).await;
+    let request = SignRequest {
+        payload: payload_hash,
+        path: path.into(),
+        key_version: 0,
+    };
+
+    let status = alice
+        .call(contract.id(), "sign")
+        .args_json(serde_json::json!({
+            "request": request,
+        }))
+        .deposit(NearToken::from_near(1))
+        .max_gas()
+        .transact_async()
+        .await?;
+    dbg!(&status);
+    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
+
+    // we do not respond, sign will fail due to timeout
+    let execution = status.await;
+    dbg!(&execution);
+    let err = execution
+        .unwrap()
+        .into_result()
+        .expect_err("should have failed with timeout");
+    assert!(err
+        .to_string()
+        .contains(&errors::SignError::Timeout.to_string()));
+
+    let new_balance = alice.view_account().await?.balance;
+    let new_contract_balance = contract.view_account().await?.balance;
+    println!(
+        "{} {} {} {}",
+        balance.as_millinear(),
+        new_balance.as_millinear(),
+        contract_balance.as_yoctonear(),
+        new_contract_balance.as_yoctonear(),
+    );
+    assert!(
+        balance.as_millinear() - new_balance.as_millinear() < 10,
+        "refund should happen"
+    );
+    assert!(
+        contract_balance.as_millinear() - new_contract_balance.as_millinear() <= 1,
+        "refund transfer should take less than 0.001 NEAR"
+    );
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_contract_sign_request_deposits() -> anyhow::Result<()> {
     let (_, contract, _, sk) = init_env().await;

chain-signatures/contract/tests/updates.rs

@@ -34,7 +34,7 @@ pub fn invalid_contract() -> ProposeUpdateArgs {
 /// This is the current deposit required for a contract deploy. This is subject to change but make
 /// sure that it's not larger than 2mb. We can go up to 4mb technically but our contract should
 /// not be getting that big.
-const CURRENT_CONTRACT_DEPLOY_DEPOSIT: NearToken = NearToken::from_millinear(8600);
+const CURRENT_CONTRACT_DEPLOY_DEPOSIT: NearToken = NearToken::from_millinear(8700);
 
 #[tokio::test]
 async fn test_propose_contract_max_size_upload() {