feat: apply rate limiting to the REST server

Signed-off-by: ljedrz <[email protected]>

node/rest/Cargo.toml

@@ -79,6 +79,10 @@ version = "0.3"
 [dependencies.tokio]
 version = "1"
 
+[dependencies.tower]
+version = "0.4"
+features = ["buffer", "limit"]
+
 [dependencies.tower-http]
 version = "0.4"
 features = ["cors", "trace"]

node/rest/src/lib.rs

@@ -33,18 +33,21 @@ use snarkvm::{
 
 use anyhow::Result;
 use axum::{
+    error_handling::HandleErrorLayer,
     extract::{ConnectInfo, DefaultBodyLimit, Path, Query, State},
     http::{header::CONTENT_TYPE, Method, Request, StatusCode},
     middleware,
     middleware::Next,
     response::Response,
     routing::{get, post},
+    BoxError,
     Json,
 };
 use axum_extra::response::ErasedJson;
 use parking_lot::Mutex;
-use std::{net::SocketAddr, str::FromStr, sync::Arc};
+use std::{net::SocketAddr, str::FromStr, sync::Arc, time::Duration};
 use tokio::task::JoinHandle;
+use tower::{buffer::BufferLayer, limit::RateLimitLayer, ServiceBuilder};
 use tower_http::{
     cors::{Any, CorsLayer},
     trace::TraceLayer,
@@ -99,6 +102,18 @@ impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
             .allow_methods([Method::GET, Method::POST, Method::OPTIONS])
             .allow_headers([CONTENT_TYPE]);
 
+        let rate_limiter = ServiceBuilder::new()
+            .layer(HandleErrorLayer::new(|err: BoxError| async move {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Unhandled error: {}", err),
+                )
+            }))
+            // Use a buffer layer to allow rate limiting.
+            .layer(BufferLayer::new(1024))
+            // Apply a rate limit of 5 requests/s.
+            .layer(RateLimitLayer::new(5, Duration::from_secs(1)));
+
         let router = {
             axum::Router::new()
 
@@ -149,6 +164,8 @@ impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
             .layer(cors)
             // Cap body size at 10MB.
             .layer(DefaultBodyLimit::max(10 * 1024 * 1024))
+            // Apply a rate limiting layer.
+            .layer(rate_limiter)
             // JWT auth.
             // .layer(middleware::from_fn(auth_middleware))
         };