From 51a48ecc6a569291fa7192b1f488cc3ee95eaa3d Mon Sep 17 00:00:00 2001
From: Serhii Volovyk <SergeyVolovyk@gmail.com>
Date: Fri, 29 Mar 2024 14:16:18 +0200
Subject: [PATCH] Request counter (#531)

* request counter

* fmt

* clippy

* migration function added

* second migration function

* redundant field deleted from old states

* migration functions deleted

* clippy
---
 contract/src/lib.rs | 41 +++++++++++++++++++++++++----------------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/contract/src/lib.rs b/contract/src/lib.rs
index 092c0009d..3121b690f 100644
--- a/contract/src/lib.rs
+++ b/contract/src/lib.rs
@@ -1,7 +1,7 @@
 pub mod primitives;
 
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
-use near_sdk::collections::TreeMap;
+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};
@@ -49,7 +49,8 @@ pub enum ProtocolContractState {
 #[derive(BorshDeserialize, BorshSerialize, PanicOnDefault)]
 pub struct MpcContract {
     protocol_state: ProtocolContractState,
-    pending_requests: TreeMap<[u8; 32], Option<(String, String)>>,
+    pending_requests: LookupMap<[u8; 32], Option<(String, String)>>,
+    request_counter: u32,
 }
 
 #[near_bindgen]
@@ -68,7 +69,8 @@ impl MpcContract {
                 threshold,
                 pk_votes: PkVotes::new(),
             }),
-            pending_requests: TreeMap::new(b"m"),
+            pending_requests: LookupMap::new(b"m"),
+            request_counter: 0,
         }
     }
 
@@ -296,9 +298,6 @@ impl MpcContract {
     #[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 {
-        if self.pending_requests.len() > 8 {
-            env::panic_str("Too many pending requests. Please, try again later.");
-        }
         let latest_key_version: u32 = self.latest_key_version();
         assert!(
             key_version <= latest_key_version,
@@ -314,7 +313,7 @@ impl MpcContract {
         );
         match self.pending_requests.get(&payload) {
             None => {
-                self.pending_requests.insert(&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)
             }
@@ -336,12 +335,12 @@ impl MpcContract {
                         signature,
                         depth
                     );
-                    self.pending_requests.remove(&payload);
+                    self.remove_request(&payload);
                     PromiseOrValue::Value(signature)
                 }
                 None => {
                     if depth > 30 {
-                        self.pending_requests.remove(&payload);
+                        self.remove_request(&payload);
                         env::panic_str("Signature was not provided in time. Please, try again.");
                     }
                     log!(&format!(
@@ -388,16 +387,11 @@ impl MpcContract {
         }
         Self {
             protocol_state: ProtocolContractState::NotInitialized,
-            pending_requests: TreeMap::new(b"m"),
+            pending_requests: LookupMap::new(b"m"),
+            request_counter: 0,
         }
     }
 
-    #[private]
-    pub fn clean_requests(&mut self) {
-        log!("clean requests");
-        self.pending_requests.clear();
-    }
-
     /// 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 {
@@ -414,4 +408,19 @@ impl MpcContract {
     pub const fn latest_key_version(&self) -> u32 {
         0
     }
+
+    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);
+    }
+
+    fn remove_request(&mut self, payload: &[u8; 32]) {
+        self.pending_requests.remove(payload);
+        self.request_counter -= 1;
+    }
 }