Skip to content

Commit

Permalink
Request counter (#531)
Browse files Browse the repository at this point in the history
* request counter

* fmt

* clippy

* migration function added

* second migration function

* redundant field deleted from old states

* migration functions deleted

* clippy
  • Loading branch information
volovyks authored Mar 29, 2024
1 parent bcbb5d4 commit 51a48ec
Showing 1 changed file with 25 additions and 16 deletions.
41 changes: 25 additions & 16 deletions contract/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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]
Expand All @@ -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,
}
}

Expand Down Expand Up @@ -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,
Expand All @@ -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)
}
Expand All @@ -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!(
Expand Down Expand Up @@ -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 {
Expand All @@ -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;
}
}

0 comments on commit 51a48ec

Please sign in to comment.