Add rate-limit example (#642)

Co-authored-by: Rob Ede <robjtede@icloud.com>

Cargo.lock

@@ -139,6 +139,18 @@ dependencies = [
  "pin-project-lite 0.2.13",
 ]
 
+[[package]]
+name = "actix-governor"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "46ff2d40f2bc627b8054c5e20fa6b0b0cf9428699b54bd41634e9ae3098ad555"
+dependencies = [
+ "actix-http",
+ "actix-web",
+ "futures 0.3.28",
+ "governor",
+]
+
 [[package]]
 name = "actix-http"
 version = "3.4.0"
@@ -2603,6 +2615,19 @@ dependencies = [
  "syn 2.0.38",
 ]
 
+[[package]]
+name = "dashmap"
+version = "5.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856"
+dependencies = [
+ "cfg-if 1.0.0",
+ "hashbrown 0.14.0",
+ "lock_api 0.4.10",
+ "once_cell",
+ "parking_lot_core 0.9.8",
+]
+
 [[package]]
 name = "data-encoding"
 version = "2.4.0"
@@ -3546,6 +3571,24 @@ dependencies = [
  "walkdir",
 ]
 
+[[package]]
+name = "governor"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c390a940a5d157878dd057c78680a33ce3415bcd05b4799509ea44210914b4d5"
+dependencies = [
+ "cfg-if 1.0.0",
+ "dashmap",
+ "futures 0.3.28",
+ "futures-timer",
+ "no-std-compat",
+ "nonzero_ext",
+ "parking_lot 0.12.1",
+ "quanta",
+ "rand 0.8.5",
+ "smallvec 1.11.0",
+]
+
 [[package]]
 name = "graceful-shutdown"
 version = "0.2.0"
@@ -4566,6 +4609,15 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
 
+[[package]]
+name = "mach"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b823e83b2affd8f40a9ee8c29dbc56404c1e34cd2710921f2801e2cf29527afa"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "markup5ever"
 version = "0.11.0"
@@ -4693,6 +4745,18 @@ dependencies = [
  "rustls-pemfile",
 ]
 
+[[package]]
+name = "middleware-rate-limit"
+version = "1.0.0"
+dependencies = [
+ "actix-governor",
+ "actix-web",
+ "chrono",
+ "env_logger",
+ "futures-util",
+ "log",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.17"
@@ -5053,6 +5117,12 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e4a24736216ec316047a1fc4252e27dabb04218aa4a3f37c6e7ddbf1f9782b54"
 
+[[package]]
+name = "no-std-compat"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c"
+
 [[package]]
 name = "nom"
 version = "7.1.3"
@@ -5063,6 +5133,12 @@ dependencies = [
  "minimal-lexical",
 ]
 
+[[package]]
+name = "nonzero_ext"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21"
+
 [[package]]
 name = "notify"
 version = "5.2.0"
@@ -5780,6 +5856,22 @@ dependencies = [
  "syn 1.0.109",
 ]
 
+[[package]]
+name = "quanta"
+version = "0.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "20afe714292d5e879d8b12740aa223c6a88f118af41870e8b6196e39a02238a8"
+dependencies = [
+ "crossbeam-utils 0.8.16",
+ "libc",
+ "mach",
+ "once_cell",
+ "raw-cpuid",
+ "wasi 0.10.2+wasi-snapshot-preview1",
+ "web-sys",
+ "winapi 0.3.9",
+]
+
 [[package]]
 name = "quick-error"
 version = "1.2.3"
@@ -5904,6 +5996,15 @@ dependencies = [
  "rand_core 0.5.1",
 ]
 
+[[package]]
+name = "raw-cpuid"
+version = "10.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c297679cb867470fa8c9f67dbba74a78d78e3e98d7cf2b08d6d71540f797332"
+dependencies = [
+ "bitflags 1.3.2",
+]
+
 [[package]]
 name = "redis"
 version = "0.23.3"
@@ -8647,6 +8748,12 @@ version = "0.9.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519"
 
+[[package]]
+name = "wasi"
+version = "0.10.2+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6"
+
 [[package]]
 name = "wasi"
 version = "0.11.0+wasi-snapshot-preview1"

Cargo.toml

@@ -45,6 +45,7 @@ members = [
    "middleware/middleware-ext-mut",
    "middleware/middleware-http-to-https",
    "middleware/middleware",
+   "middleware/middleware-rate-limit",
    "protobuf",
    "run-in-thread",
    "server-sent-events",

middleware/middleware-rate-limit/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "middleware-rate-limit"
+version = "1.0.0"
+edition = "2021"
+
+[dependencies]
+actix-web.workspace = true
+log.workspace = true
+env_logger.workspace = true
+futures-util.workspace = true
+chrono.workspace = true
+
+actix-governor = "0.4"

middleware/middleware-rate-limit/src/main.rs

@@ -0,0 +1,45 @@
+use std::io;
+
+use actix_governor::{Governor, GovernorConfigBuilder};
+use actix_web::{
+    middleware,
+    web::{self},
+    App, HttpResponse, HttpServer,
+};
+
+mod rate_limit;
+
+async fn index() -> HttpResponse {
+    HttpResponse::Ok().body("succeed")
+}
+
+#[actix_web::main]
+async fn main() -> io::Result<()> {
+    env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));
+
+    let limit_cfg = GovernorConfigBuilder::default()
+        .per_second(10)
+        .burst_size(2)
+        .finish()
+        .unwrap();
+
+    log::info!("starting HTTP server at http://localhost:8080");
+
+    HttpServer::new(move || {
+        App::new()
+            .wrap(middleware::Logger::default())
+            .service(
+                web::resource("/test/governor")
+                    .wrap(Governor::new(&limit_cfg))
+                    .route(web::get().to(index)),
+            )
+            .service(
+                web::resource("/test/simple")
+                    .wrap(rate_limit::RateLimit::new(2))
+                    .route(web::get().to(index)),
+            )
+    })
+    .bind(("127.0.0.1", 8080))?
+    .run()
+    .await
+}

middleware/middleware-rate-limit/src/rate_limit.rs

@@ -0,0 +1,121 @@
+use std::cell::RefCell;
+use std::cmp::min;
+use std::future::{ready, Ready};
+
+use actix_web::body::EitherBody;
+use actix_web::{
+    dev,
+    dev::{Service, ServiceRequest, ServiceResponse, Transform},
+    Error, HttpResponse,
+};
+use chrono::{Local, NaiveDateTime};
+use futures_util::future::LocalBoxFuture;
+
+#[doc(hidden)]
+pub struct RateLimitService<S> {
+    service: S,
+    token_bucket: RefCell<TokenBucket>,
+}
+
+impl<S, B> Service<ServiceRequest> for RateLimitService<S>
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    dev::forward_ready!(service);
+
+    fn call(&self, req: ServiceRequest) -> Self::Future {
+        log::info!("request is passing through the AddMsg middleware");
+
+        req.uri().path();
+        // if be limited
+        if !self.token_bucket.borrow_mut().allow_query() {
+            return Box::pin(async {
+                Ok(req.into_response(
+                    HttpResponse::TooManyRequests()
+                        .body("")
+                        .map_into_right_body(),
+                ))
+            });
+        }
+
+        let fut = self.service.call(req);
+        Box::pin(async move { fut.await.map(ServiceResponse::map_into_left_body) })
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct RateLimit {
+    // limit in 10s
+    limit: u64,
+}
+
+impl RateLimit {
+    pub fn new(limit: u64) -> Self {
+        Self { limit }
+    }
+}
+
+impl<S, B> Transform<S, ServiceRequest> for RateLimit
+where
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = actix_web::Error>,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<EitherBody<B>>;
+    type Error = Error;
+    type Transform = RateLimitService<S>;
+    type InitError = ();
+    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        ready(Ok(RateLimitService {
+            service,
+            token_bucket: RefCell::new(TokenBucket::new(self.limit)),
+        }))
+    }
+}
+
+struct TokenBucket {
+    // limit in ten sec
+    limit: u64,
+    last_query_time: NaiveDateTime,
+    // max query number in ten sec,in this case equal limit
+    capacity: u64,
+    // numbers of token,default equal capacity
+    tokens: u64,
+}
+
+impl TokenBucket {
+    fn new(limit: u64) -> Self {
+        TokenBucket {
+            limit,
+            last_query_time: Default::default(),
+            capacity: limit,
+            tokens: 0,
+        }
+    }
+
+    fn allow_query(&mut self) -> bool {
+        let current_time = Local::now().naive_local();
+
+        let time_elapsed = (current_time.timestamp() - self.last_query_time.timestamp()) as u64;
+
+        let tokens_to_add = time_elapsed * self.limit / 10;
+
+        self.tokens = min(self.tokens + tokens_to_add, self.capacity);
+
+        if self.tokens > 0 {
+            self.last_query_time = current_time;
+            self.tokens -= 1;
+            true
+        } else {
+            false
+        }
+    }
+}