Added propose_update refund diff and migrate

chain-signatures/contract/src/errors.rs

@@ -50,6 +50,8 @@ pub enum PublicKeyError {
 pub enum InitError {
     #[error("Threshold cannot be greater than the number of candidates")]
     ThresholdTooHigh,
+    #[error("Cannot load in contract due to missing state")]
+    ContractStateIsMissing,
 }
 
 #[derive(Debug, thiserror::Error)]

chain-signatures/contract/src/lib.rs

@@ -13,8 +13,8 @@ use k256::Scalar;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::collections::LookupMap;
 use near_sdk::{
-    env, log, near_bindgen, AccountId, CryptoHash, Gas, GasWeight, NearToken, PromiseError,
-    PublicKey,
+    env, log, near_bindgen, AccountId, CryptoHash, Gas, GasWeight, NearToken, Promise,
+    PromiseError, PublicKey,
 };
 use primitives::{
     CandidateInfo, Candidates, Participants, PkVotes, SignRequest, SignaturePromiseError,
@@ -524,7 +524,7 @@ impl VersionedMpcContract {
         config: Option<Config>,
     ) -> Result<UpdateId, MpcContractError> {
         // Only voters can propose updates:
-        self.voter()?;
+        let proposer = self.voter()?;
 
         let attached = env::attached_deposit();
         let required = ProposedUpdates::required_deposit(&code, &config);
@@ -541,6 +541,13 @@ impl VersionedMpcContract {
             )));
         };
 
+        // Refund the difference if the propser attached more than required.
+        if let Some(diff) = attached.checked_sub(required) {
+            if diff > NearToken::from_yoctonear(0) {
+                Promise::new(proposer).transfer(diff);
+            }
+        }
+
         Ok(id)
     }
 
@@ -562,10 +569,7 @@ impl VersionedMpcContract {
             return Ok(false);
         }
 
-        let Some(_promise) =
-            self.proposed_updates()
-                .do_update(&id, "update_config", UPDATE_CONFIG_GAS)
-        else {
+        let Some(_promise) = self.proposed_updates().do_update(&id, UPDATE_CONFIG_GAS) else {
             return Err(MpcContractError::from(VoteError::UpdateNotFound));
         };
 
@@ -636,6 +640,22 @@ impl VersionedMpcContract {
         }))
     }
 
+    /// This will be called internally by the contract to migrate the state when a new contract
+    /// is deployed. This function should be changed every time state is changed to do the proper
+    /// migrate flow.
+    ///
+    /// If nothing is changed, then this function will just return the current state. If it fails
+    /// to read the state, then it will return an error.
+    #[private]
+    #[init(ignore_state)]
+    #[handle_result]
+    pub fn migrate() -> Result<Self, MpcContractError> {
+        let old: MpcContract = env::state_read().ok_or(MpcContractError::InitError(
+            InitError::ContractStateIsMissing,
+        ))?;
+        Ok(VersionedMpcContract::V0(old))
+    }
+
     pub fn state(&self) -> &ProtocolContractState {
         match self {
             Self::V0(mpc_contract) => &mpc_contract.protocol_state,

chain-signatures/contract/src/update.rs

@@ -102,21 +102,29 @@ impl ProposedUpdates {
         self.entries.remove(id)
     }
 
-    pub fn do_update(&mut self, id: &UpdateId, config_callback: &str, gas: Gas) -> Option<Promise> {
+    pub fn do_update(&mut self, id: &UpdateId, gas: Gas) -> Option<Promise> {
         let entry = self.remove(id)?;
 
         let mut promise = Promise::new(env::current_account_id());
         for update in entry.updates {
             match update {
                 Update::Config(config) => {
                     promise = promise.function_call(
-                        config_callback.into(),
+                        "update_config".into(),
                         serde_json::to_vec(&(&config,)).unwrap(),
                         NearToken::from_near(0),
                         gas,
                     );
                 }
-                Update::Contract(code) => promise = promise.deploy_contract(code),
+                Update::Contract(code) => {
+                    // deploy contract then do a `migrate` call to migrate state.
+                    promise = promise.deploy_contract(code).function_call(
+                        "migrate".into(),
+                        Vec::new(),
+                        NearToken::from_near(0),
+                        gas,
+                    );
+                }
             }
         }
         Some(promise)