clear_payloads func (#551)

* clear_payloads func
* reset counter
* added add_signature in the contract
* gas check in the contract

contract/src/lib.rs

@@ -2,13 +2,15 @@ pub mod primitives;
 
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::collections::LookupMap;
-use near_sdk::log;
 use near_sdk::serde::{Deserialize, Serialize};
 use near_sdk::{env, near_bindgen, AccountId, PanicOnDefault, Promise, PromiseOrValue, PublicKey};
+use near_sdk::{log, Gas};
 use primitives::ParticipantInfo;
 use primitives::{CandidateInfo, Candidates, Participants, PkVotes, Votes};
 use std::collections::{BTreeMap, HashSet};
 
+const GAS_FOR_SIGN_CALL: Gas = Gas::from_gas(3 * 100_000_000_000_000);
+
 #[derive(BorshDeserialize, BorshSerialize, Serialize, Deserialize, Debug)]
 pub struct InitializingContractState {
     pub candidates: Candidates,
@@ -336,16 +338,24 @@ impl MpcContract {
             "This version of the signer contract doesn't support versions greater than {}",
             latest_key_version,
         );
+        // Make sure sign call will not run out of gas doing recursive calls because the payload will never be removed
+        assert!(
+            env::prepaid_gas() >= GAS_FOR_SIGN_CALL,
+            "Insufficient gas provided. Provided: {} Required: {}",
+            env::prepaid_gas(),
+            GAS_FOR_SIGN_CALL
+        );
         log!(
-            "sign: signer={}, payload={:?}, path={:?}, key_version={}",
+            "sign: signer={}, predecessor={}, payload={:?}, path={:?}, key_version={}",
             env::signer_account_id(),
+            env::predecessor_account_id(),
             payload,
             path,
             key_version
         );
         match self.pending_requests.get(&payload) {
             None => {
-                self.add_request(&payload, &None);
+                self.add_request(&payload, None);
                 log!(&serde_json::to_string(&near_sdk::env::random_seed_array()).unwrap());
                 Self::ext(env::current_account_id()).sign_helper(payload, 0)
             }
@@ -399,9 +409,7 @@ impl MpcContract {
                     big_r,
                     s
                 );
-                if self.pending_requests.contains_key(&payload) {
-                    self.pending_requests.insert(&payload, &Some((big_r, s)));
-                }
+                self.add_signature(&payload, (big_r, s));
             } else {
                 env::panic_str("only participants can respond");
             }
@@ -410,20 +418,6 @@ impl MpcContract {
         }
     }
 
-    #[private]
-    #[init(ignore_state)]
-    pub fn clean(keys: Vec<near_sdk::json_types::Base64VecU8>) -> Self {
-        log!("clean: keys={:?}", keys);
-        for key in keys.iter() {
-            env::storage_remove(&key.0);
-        }
-        Self {
-            protocol_state: ProtocolContractState::NotInitialized,
-            pending_requests: LookupMap::new(b"m"),
-            request_counter: 0,
-        }
-    }
-
     /// This is the root public key combined from all the public keys of the participants.
     pub fn public_key(&self) -> PublicKey {
         match &self.protocol_state {
@@ -445,18 +439,48 @@ impl MpcContract {
         0
     }
 
-    fn add_request(&mut self, payload: &[u8; 32], signature: &Option<(String, String)>) {
+    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_request(&mut self, payload: &[u8; 32], signature: Option<(String, String)>) {
         if self.request_counter > 8 {
             env::panic_str("Too many pending requests. Please, try again later.");
         }
         if !self.pending_requests.contains_key(payload) {
             self.request_counter += 1;
         }
-        self.pending_requests.insert(payload, signature);
+        self.pending_requests.insert(payload, &signature);
     }
 
     fn remove_request(&mut self, payload: &[u8; 32]) {
         self.pending_requests.remove(payload);
         self.request_counter -= 1;
     }
+
+    // Helper functions
+    #[private]
+    #[init(ignore_state)]
+    pub fn clean(keys: Vec<near_sdk::json_types::Base64VecU8>) -> Self {
+        log!("clean: keys={:?}", keys);
+        for key in keys.iter() {
+            env::storage_remove(&key.0);
+        }
+        Self {
+            protocol_state: ProtocolContractState::NotInitialized,
+            pending_requests: LookupMap::new(b"m"),
+            request_counter: 0,
+        }
+    }
+
+    #[private]
+    pub fn clean_payloads(&mut self, payloads: Vec<[u8; 32]>, counter: u32) {
+        log!("clean_payloads");
+        for payload in payloads.iter() {
+            self.pending_requests.remove(payload);
+        }
+        self.request_counter = counter;
+    }
 }