Merge pull request #41 from NishantJoshi00/feat/rate-limit-delete

Cargo.toml

@@ -8,8 +8,9 @@ default-run = "locker"
 
 [features]
 default = []
-release = ["kms", "middleware", "key_custodian"]
+release = ["kms", "middleware", "key_custodian", "limit"]
 kms = ["dep:aws-config", "dep:aws-sdk-kms"]
+limit = []
 middleware = []
 key_custodian = []
 
@@ -30,7 +31,7 @@ once_cell = "1.18.0"
 tokio = { version = "1.33.0", features = ["macros", "rt-multi-thread"] }
 axum = "0.6.20"
 hyper = "0.14.27"
-tower = "0.4.13"
+tower = { version = "0.4.13", features = ["limit", "buffer", "load-shed"] }
 tower-http = { version = "0.4.4", features = ["trace"] }
 
 

config.example.toml

@@ -0,0 +1,30 @@
+[log.console]
+enabled = true # To enable logging in console
+level = "DEBUG" # level to be set for the logging framework
+log_format = "default" # format to be used for logging default | json
+
+[server]
+host = "127.0.0.1" # The host that the server should be exposed to 
+port = 8080 # The port where the server should be hosted on
+
+[limit]
+request_count = 1 # The requests per duration
+duration = 60 # duration to rate limit the delete api (in sec)
+
+[database]
+username = "sam" # username for the database
+password = "damn" # password of the database
+host = "localhost" # the host where the database is hosted on
+port = 5432 # the port of the database
+dbname = "locker" # the name of the database where the cards are stored
+
+[secrets]
+tenant = "hyperswitch" # the tenant that we are currently configured for
+master_key = "feffe9928665731c6d6a8f9467308308feffe9928665731c6d6a8f9467308308" # master key that is used for database encryption this could be aes encrypted by key custodian
+
+tenant_public_key = ""  # The tenant public key to which the communication will be established with
+locker_private_key = "" # the locker private key to used used and the private key present with the tenant
+
+[kms]
+region = "us-west-2" # the kms details needed to perform kms decryption of the above mentioned environment variables
+key_id = "abc"

config/development.toml

@@ -7,6 +7,10 @@ log_format = "default"
 host = "127.0.0.1"
 port = 8080
 
+[limit]
+request_count = 1
+duration = 60
+
 [database]
 username = "sam"
 password = "damn"

docs/guides/setup.md

@@ -132,6 +132,9 @@ LOCKER__DATABASE__HOST=
 LOCKER__DATABASE__PORT=
 LOCKER__DATABASE__DBNAME=
 
+LOCKER__LIMIT__REQUEST_COUNT=100
+LOCKER__LIMIT__DURATION=60
+
 LOCKER__SECRETS__TENANT=
 LOCKER__SECRETS__MASTER_KEY=
 LOCKER__SECRETS__LOCKER_PRIVATE_KEY=

src/app.rs

@@ -99,14 +99,14 @@ where
         .nest(
             "/data",
             routes::data::serve(
-                #[cfg(feature = "middleware")]
+                #[cfg(any(feature = "middleware", feature = "limit"))]
                 state.clone(),
             ),
         )
         .nest(
             "/cards",
             routes::data::serve(
-                #[cfg(feature = "middleware")]
+                #[cfg(any(feature = "middleware", feature = "limit"))]
                 state.clone(),
             ),
         )

src/config.rs

@@ -12,7 +12,18 @@ pub struct Config {
     #[cfg(feature = "kms")]
     pub kms: kms::KmsConfig,
     pub log: Log,
+    #[cfg(feature = "limit")]
+    pub limit: Limit,
 }
+
+#[cfg(feature = "limit")]
+#[derive(Clone, serde::Deserialize, Debug)]
+pub struct Limit {
+    pub request_count: u64,
+    pub duration: u64, // in sec
+    pub buffer_size: Option<usize>,
+}
+
 #[derive(Clone, serde::Deserialize, Debug)]
 pub struct Server {
     pub host: String,

src/routes/data.rs

@@ -1,4 +1,6 @@
-use axum::{extract, routing::post, Json};
+use axum::{
+    error_handling::HandleErrorLayer, extract, response::IntoResponse, routing::post, Json,
+};
 
 #[cfg(feature = "middleware")]
 use axum::middleware;
@@ -21,14 +23,39 @@ use self::types::Validation;
 mod transformers;
 pub mod types;
 
+const BUFFER_LIMIT: usize = 1024;
+
+async fn ratelimit_err_handler(_: axum::BoxError) -> impl IntoResponse {
+    (hyper::StatusCode::TOO_MANY_REQUESTS, "Rate Limit Applied")
+}
+
 ///
 /// Function for creating the server that is specifically handling the cards api
 ///
 #[allow(clippy::let_and_return)]
-pub fn serve(#[cfg(feature = "middleware")] state: AppState) -> axum::Router<AppState> {
+pub fn serve(
+    #[cfg(any(feature = "middleware", feature = "limit"))] state: AppState,
+) -> axum::Router<AppState> {
+    #[cfg(feature = "limit")]
+    let ratelimit_middleware = tower::ServiceBuilder::new()
+        .layer(HandleErrorLayer::new(ratelimit_err_handler))
+        .buffer(state.config.limit.buffer_size.unwrap_or(BUFFER_LIMIT))
+        .load_shed()
+        .rate_limit(
+            state.config.limit.request_count,
+            std::time::Duration::from_secs(state.config.limit.duration),
+        )
+        .into_inner();
+
+    #[cfg(feature = "limit")]
+    let delete_route = post(delete_card).layer(ratelimit_middleware);
+
+    #[cfg(not(feature = "limit"))]
+    let delete_route = post(delete_card);
+
     let router = axum::Router::new()
+        .route("/delete", delete_route)
         .route("/add", post(add_card))
-        .route("/delete", post(delete_card))
         .route("/retrieve", post(retrieve_card));
 
     #[cfg(feature = "middleware")]