mirror of
https://github.com/actix/examples
synced 2024-11-23 22:41:07 +01:00
Add rate-limit example (#642)
Co-authored-by: Rob Ede <robjtede@icloud.com>
This commit is contained in:
will
GitHub
parent
c15a176b1c
commit
cf55f50d1d
107
Cargo.lock
generated
107
Cargo.lock
generated
@ -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"
|
||||
|
||||
@ -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",
|
||||
|
||||
13
middleware/middleware-rate-limit/Cargo.toml
Normal file
13
middleware/middleware-rate-limit/Cargo.toml
Normal file
@ -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"
|
||||
0
middleware/middleware-rate-limit/README.md
Normal file
0
middleware/middleware-rate-limit/README.md
Normal file
45
middleware/middleware-rate-limit/src/main.rs
Normal file
45
middleware/middleware-rate-limit/src/main.rs
Normal file
@ -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
|
||||
}
|
||||
121
middleware/middleware-rate-limit/src/rate_limit.rs
Normal file
121
middleware/middleware-rate-limit/src/rate_limit.rs
Normal file
@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user