change generalized cache to request limit

crates/threshold-signature-server/src/helpers/tests.rs

@@ -48,15 +48,7 @@ use entropy_protocol::PartyId;
 use entropy_shared::EncodedVerifyingKey;
 use entropy_shared::NETWORK_PARENT_KEY;
 use sp_keyring::AccountKeyring;
-use std::{
-    collections::HashMap,
-    fmt,
-    net::SocketAddr,
-    path::PathBuf,
-    str,
-    sync::{Arc, RwLock},
-    time::Duration,
-};
+use std::{fmt, net::SocketAddr, path::PathBuf, str, time::Duration};
 use subxt::{
     backend::legacy::LegacyRpcMethods, ext::sp_core::sr25519, tx::PairSigner,
     utils::AccountId32 as SubxtAccountId32, Config, OnlineClient,
@@ -85,15 +77,8 @@ pub async fn setup_client() -> KvManager {
         setup_kv_store(&Some(ValidatorName::Alice), Some(storage_path.clone())).await.unwrap();
 
     let _ = setup_latest_block_number(&kv_store).await;
-    let cache: HashMap<String, Vec<u8>> = HashMap::new();
 
-    let app_state = AppState::new(
-        configuration,
-        kv_store.clone(),
-        sr25519_pair,
-        x25519_secret,
-        Arc::new(RwLock::new(cache)),
-    );
+    let app_state = AppState::new(configuration, kv_store.clone(), sr25519_pair, x25519_secret);
 
     // Mock making the pre-requisite checks by setting the application state to ready
     app_state.make_ready().unwrap();
@@ -125,15 +110,8 @@ pub async fn create_clients(
         setup_kv_store(validator_name, Some(path.into())).await.unwrap();
 
     let _ = setup_latest_block_number(&kv_store).await;
-    let cache: HashMap<String, Vec<u8>> = HashMap::new();
-
-    let app_state = AppState::new(
-        configuration,
-        kv_store.clone(),
-        sr25519_pair,
-        x25519_secret,
-        Arc::new(RwLock::new(cache)),
-    );
+
+    let app_state = AppState::new(configuration, kv_store.clone(), sr25519_pair, x25519_secret);
 
     let _ = setup_latest_block_number(&kv_store).await;
 

crates/threshold-signature-server/src/lib.rs

@@ -191,7 +191,9 @@ use crate::{
     health::api::healthz,
     launch::Configuration,
     node_info::api::{hashes, info, version as get_version},
-    r#unsafe::api::{delete, put, read_from_cache, remove_keys, unsafe_get, write_to_cache},
+    r#unsafe::api::{
+        delete, put, read_from_request_limit, remove_keys, unsafe_get, write_to_request_limit,
+    },
     signing_client::{api::*, ListenerState},
     user::api::*,
     validator::api::{new_reshare, rotate_network_key},
@@ -248,8 +250,8 @@ pub struct AppState {
     pub configuration: Configuration,
     /// Key-value store
     pub kv_store: KvManager,
-    /// Cache for TSS
-    pub cache: Cache,
+    /// Storage for request limit
+    pub request_limit: Arc<RwLock<HashMap<String, RequestLimitStorage>>>,
     /// Storage for encryption key backups for other TSS nodes
     /// Maps TSS account id to encryption key
     pub encryption_key_backup_provider: Arc<RwLock<HashMap<AccountId32, [u8; 32]>>>,
@@ -258,17 +260,13 @@ pub struct AppState {
     pub attestation_nonces: Arc<RwLock<HashMap<X25519PublicKey, [u8; 32]>>>,
 }
 
-/// A global cache type for the TSS
-pub type Cache = Arc<RwLock<HashMap<String, Vec<u8>>>>;
-
 impl AppState {
     /// Setup AppState with given secret keys
     pub fn new(
         configuration: Configuration,
         kv_store: KvManager,
         pair: sr25519::Pair,
         x25519_secret: StaticSecret,
-        cache: Cache,
     ) -> Self {
         Self {
             tss_state: Arc::new(RwLock::new(TssState::new())),
@@ -279,7 +277,7 @@ impl AppState {
             kv_store,
             encryption_key_backup_provider: Default::default(),
             attestation_nonces: Default::default(),
-            cache,
+            request_limit: Default::default(),
         }
     }
 
@@ -352,45 +350,61 @@ impl AppState {
         ))
     }
 
-    pub fn write_to_cache(&self, key: String, value: Vec<u8>) -> anyhow::Result<()> {
-        self.clear_poisioned_cache();
-        let mut cache =
-            self.cache.write().map_err(|_| anyhow!("Error getting write write_to_cache lock"))?;
-        cache.insert(key, value);
+    /// Write to request limit
+    pub fn write_to_request_limit(
+        &self,
+        key: String,
+        value: RequestLimitStorage,
+    ) -> anyhow::Result<()> {
+        self.clear_poisioned_request_limit();
+        let mut request_limit = self
+            .request_limit
+            .write()
+            .map_err(|_| anyhow!("Error getting write write_to_request_limit lock"))?;
+        request_limit.insert(key, value);
         Ok(())
     }
 
-    pub fn exists_in_cache(&self, key: &String) -> anyhow::Result<bool> {
-        self.clear_poisioned_cache();
-        let cache =
-            self.cache.read().map_err(|_| anyhow!("Error getting read exists_in_cache lock"))?;
-        Ok(cache.contains_key(key))
+    /// Check if key exists in request limit
+    pub fn exists_in_request_limit(&self, key: &String) -> anyhow::Result<bool> {
+        self.clear_poisioned_request_limit();
+        let request_limit = self
+            .request_limit
+            .read()
+            .map_err(|_| anyhow!("Error getting read exists_in_request_limit lock"))?;
+        Ok(request_limit.contains_key(key))
     }
 
-    pub fn remove_from_cache(&self, key: &String) -> anyhow::Result<()> {
-        self.clear_poisioned_cache();
-        let mut cache = self
-            .cache
+    /// Remove key from request limt
+    pub fn remove_from_request_limit(&self, key: &String) -> anyhow::Result<()> {
+        self.clear_poisioned_request_limit();
+        let mut request_limit = self
+            .request_limit
             .write()
-            .map_err(|_| anyhow!("Error getting write remove_from_cache lock"))?;
-        cache.remove(key);
+            .map_err(|_| anyhow!("Error getting write remove_from_request_limit lock"))?;
+        request_limit.remove(key);
         Ok(())
     }
 
-    /// Reads from cache will error if no value, call exists_in_cache to check
-    pub fn read_from_cache(&self, key: &String) -> anyhow::Result<Option<Vec<u8>>> {
-        self.clear_poisioned_cache();
-        let cache =
-            self.cache.read().map_err(|_| anyhow!("Error getting read read_from_cache lock"))?;
-        Ok(cache.get(key).clone().cloned())
+    /// Reads from request_limit will error if no value, call exists_in_request_limit to check
+    pub fn read_from_request_limit(
+        &self,
+        key: &String,
+    ) -> anyhow::Result<Option<RequestLimitStorage>> {
+        self.clear_poisioned_request_limit();
+        let request_limit = self
+            .request_limit
+            .read()
+            .map_err(|_| anyhow!("Error getting read read_from_request_limit lock"))?;
+        Ok(request_limit.get(key).cloned())
     }
 
-    pub fn clear_poisioned_cache(&self) {
-        if self.cache.is_poisoned() {
-            self.cache.clear_poison()
+    /// Clears a poisioned lock from request limit
+    pub fn clear_poisioned_request_limit(&self) {
+        if self.request_limit.is_poisoned() {
+            self.request_limit.clear_poison()
         }
     }
-    // TODO delete from cache
 
     /// Gets the list of peers who haven't yet subscribed to us for this particular session.
     pub fn unsubscribed_peers(
@@ -434,8 +448,8 @@ pub fn app(app_state: AppState) -> Router {
         tracing::warn!("Server started in unsafe mode - do not use in production!");
         routes = routes
             .route("/unsafe/put", post(put))
-            .route("/unsafe/write_to_cache", post(write_to_cache))
-            .route("/unsafe/read_from_cache", post(read_from_cache))
+            .route("/unsafe/write_to_request_limit", post(write_to_request_limit))
+            .route("/unsafe/read_from_request_limit", post(read_from_request_limit))
             .route("/unsafe/get", post(unsafe_get))
             .route("/unsafe/delete", post(delete))
             .route("/unsafe/remove_keys", get(remove_keys));

crates/threshold-signature-server/src/main.rs

@@ -13,13 +13,7 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-use std::{
-    collections::HashMap,
-    net::SocketAddr,
-    process,
-    str::FromStr,
-    sync::{Arc, RwLock},
-};
+use std::{net::SocketAddr, process, str::FromStr};
 
 use anyhow::{anyhow, ensure};
 use clap::Parser;
@@ -66,15 +60,9 @@ async fn main() -> anyhow::Result<()> {
 
     let (kv_store, sr25519_pair, x25519_secret, key_option) =
         setup_kv_store(&validator_name, None).await?;
-    let cache: HashMap<String, Vec<u8>> = HashMap::new();
 
-    let app_state = AppState::new(
-        configuration.clone(),
-        kv_store.clone(),
-        sr25519_pair,
-        x25519_secret,
-        Arc::new(RwLock::new(cache)),
-    );
+    let app_state =
+        AppState::new(configuration.clone(), kv_store.clone(), sr25519_pair, x25519_secret);
 
     ensure!(
         setup_latest_block_number(&kv_store).await.is_ok(),

crates/threshold-signature-server/src/unsafe/api.rs

@@ -13,12 +13,13 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-use std::str;
-
+use crate::RequestLimitStorage;
 use axum::{extract::State, http::StatusCode, Json};
+use parity_scale_codec::Encode;
 use serde::{Deserialize, Serialize};
 #[cfg(test)]
 use serde_json::to_string;
+use std::str;
 
 use crate::AppState;
 
@@ -45,6 +46,12 @@ impl UnsafeQuery {
     }
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+pub struct UnsafeRequestLimitQuery {
+    pub key: String,
+    pub value: RequestLimitStorage,
+}
+
 /// Read a value from the encrypted KVDB.
 ///
 /// # Note
@@ -106,41 +113,41 @@ pub async fn put(State(app_state): State<AppState>, Json(key): Json<UnsafeQuery>
     }
 }
 
-/// Updates a value in the cache.
+/// Updates a value in the request_limit.
 ///
 /// # Note
 ///
 /// This should only be used for development purposes.
 #[tracing::instrument(
-    name = "Updating key from cache",
+    name = "Updating key from request_limit",
     skip_all,
     fields(key = key.key),
 )]
-pub async fn write_to_cache(
+pub async fn write_to_request_limit(
     State(app_state): State<AppState>,
-    Json(key): Json<UnsafeQuery>,
+    Json(key): Json<UnsafeRequestLimitQuery>,
 ) -> StatusCode {
-    tracing::trace!("Attempting to write value {:?} to cache", &key.value);
-    app_state.write_to_cache(key.key, key.value).unwrap();
+    tracing::trace!("Attempting to write value {:?} to request_limit", &key.value);
+    app_state.write_to_request_limit(key.key, key.value).unwrap();
     StatusCode::OK
 }
 
-/// Reads a value in the cache.
+/// Reads a value in the request_limit.
 ///
 /// # Note
 ///
 /// This should only be used for development purposes.
 #[tracing::instrument(
-    name = "Updating key from cache",
+    name = "Updating key from request_limit",
     skip_all,
     fields(key = key.key),
 )]
-pub async fn read_from_cache(
+pub async fn read_from_request_limit(
     State(app_state): State<AppState>,
     Json(key): Json<UnsafeQuery>,
 ) -> Vec<u8> {
     tracing::trace!("Attempting to read value {:?} to cache", &key.key);
-    app_state.read_from_cache(&key.key).unwrap().unwrap().to_vec()
+    app_state.read_from_request_limit(&key.key).unwrap().unwrap().encode()
 }
 
 /// Deletes any key from the KVDB.