From cf55f50d1dd7abf6c80f0e373a4775803fec5f96 Mon Sep 17 00:00:00 2001 From: will <64425242+willser@users.noreply.github.com> Date: Mon, 30 Oct 2023 07:32:02 +0800 Subject: [PATCH] Add rate-limit example (#642) Co-authored-by: Rob Ede --- Cargo.lock | 107 ++++++++++++++++ Cargo.toml | 1 + middleware/middleware-rate-limit/Cargo.toml | 13 ++ middleware/middleware-rate-limit/README.md | 0 middleware/middleware-rate-limit/src/main.rs | 45 +++++++ .../middleware-rate-limit/src/rate_limit.rs | 121 ++++++++++++++++++ 6 files changed, 287 insertions(+) create mode 100644 middleware/middleware-rate-limit/Cargo.toml create mode 100644 middleware/middleware-rate-limit/README.md create mode 100644 middleware/middleware-rate-limit/src/main.rs create mode 100644 middleware/middleware-rate-limit/src/rate_limit.rs diff --git a/Cargo.lock b/Cargo.lock index d85f04d2..4f42ada9 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index b1cb7ed6..d091dbbc 100644 --- a/Cargo.toml +++ b/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", diff --git a/middleware/middleware-rate-limit/Cargo.toml b/middleware/middleware-rate-limit/Cargo.toml new file mode 100644 index 00000000..11bf5f5c --- /dev/null +++ b/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" diff --git a/middleware/middleware-rate-limit/README.md b/middleware/middleware-rate-limit/README.md new file mode 100644 index 00000000..e69de29b diff --git a/middleware/middleware-rate-limit/src/main.rs b/middleware/middleware-rate-limit/src/main.rs new file mode 100644 index 00000000..c2b1fc93 --- /dev/null +++ b/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 +} diff --git a/middleware/middleware-rate-limit/src/rate_limit.rs b/middleware/middleware-rate-limit/src/rate_limit.rs new file mode 100644 index 00000000..53514c6c --- /dev/null +++ b/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 { + service: S, + token_bucket: RefCell, +} + +impl Service for RateLimitService +where + S: Service, Error = Error>, + S::Future: 'static, + B: 'static, +{ + type Response = ServiceResponse>; + type Error = Error; + type Future = LocalBoxFuture<'static, Result>; + + 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 Transform for RateLimit +where + S: Service, Error = actix_web::Error>, + S::Future: 'static, + B: 'static, +{ + type Response = ServiceResponse>; + type Error = Error; + type Transform = RateLimitService; + type InitError = (); + type Future = Ready>; + + 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 + } + } +}