Rate Limiter

Middleware that provides flow control functionality.

Main Features

  • RateIssuer provides an abstraction of the assigned key value to identify the visitor's identity. RemoteIpIssuer is an implementation of it that can determine the visitor based on the requested IP address. Eq + Send + Sync + 'static` constraint types can be used as keys.

  • RateGuard provides an abstraction for the flow control algorithm. By default, two implementations of fixed window (FixedGuard) and sliding window (SlidingGuard) are implemented.

  • RateStore provides access to data. MokaStore is a built-in moka-based memory cache implementation. You can also define your own implementation.

  • RateLimiter is a structure that implements Handler, and there is also a skipper field inside, which can be specified to skip certain requests that do not require caching. By default, none_skipper will be used to not skip any requests.

  • QuotaGetter provides the abstraction of quota acquisition, which can obtain a quota object according to the visitor's Key, which means that we can configure the user quota and other information into the database, change it dynamically, and acquire it dynamically.

  • Appilty to add ratelimit headers to response with RateLimiter::add_headersopen in new window instance function, will add the following headers to the response:

    HeaderDescription
    X-RateLimit-LimitThe maximum number of requests that the consumer is permitted to make per quota period.
    X-RateLimit-RemainingThe number of requests remaining in the current rate limit window.
    X-RateLimit-ResetThe time at which the current rate limit window resets in UTC epoch seconds.

Example

Use static quota

use salvo::prelude::*;
use salvo::rate_limiter::{BasicQuota, FixedGuard, MokaStore, RateLimiter, RemoteIpIssuer};

#[handler]
async fn hello() -> &'static str {
    "Hello World"
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt().init();

    let limiter = RateLimiter::new(
        FixedGuard::new(),
        MokaStore::new(),
        RemoteIpIssuer,
        BasicQuota::per_second(1),
    );
    let router = Router::with_hoop(limiter).get(hello);
    let acceptor = TcpListener::new("127.0.0.1:5800").bind().await;
    Server::new(acceptor).serve(router).await;
}
[package]
name = "example-rate-limiter-static"
version = "0.1.0"
edition = "2021"
publish = false

[dependencies]
salvo = { workspace = true, features = ["rate-limiter"]}
tokio = { version = "1", features = ["macros"] }
tracing = "0.1"
tracing-subscriber = "0.3"

Use dynamic quota

use std::borrow::Borrow;
use std::collections::HashMap;
use std::hash::Hash;

use once_cell::sync::Lazy;
use salvo::prelude::*;
use salvo::rate_limiter::{
    CelledQuota, MokaStore, QuotaGetter, RateIssuer, RateLimiter, SlidingGuard,
};
use salvo::Error;

static USER_QUOTAS: Lazy<HashMap<String, CelledQuota>> = Lazy::new(|| {
    let mut map = HashMap::new();
    map.insert("user1".into(), CelledQuota::per_second(1, 1));
    map.insert("user2".into(), CelledQuota::set_seconds(1, 1, 5));
    map.insert("user3".into(), CelledQuota::set_seconds(1, 1, 10));
    map
});

struct UserIssuer;
impl RateIssuer for UserIssuer {
    type Key = String;
    async fn issue(&self, req: &mut Request, _depot: &Depot) -> Option<Self::Key> {
        req.query::<Self::Key>("user")
    }
}

struct CustomQuotaGetter;
impl QuotaGetter<String> for CustomQuotaGetter {
    type Quota = CelledQuota;
    type Error = Error;

    async fn get<Q>(&self, key: &Q) -> Result<Self::Quota, Self::Error>
    where
        String: Borrow<Q>,
        Q: Hash + Eq + Sync,
    {
        USER_QUOTAS
            .get(key)
            .cloned()
            .ok_or_else(|| Error::other("user not found"))
    }
}

#[handler]
async fn limited() -> &'static str {
    "Limited page"
}
#[handler]
async fn home() -> Text<&'static str> {
    Text::Html(HOME_HTML)
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt().init();

    let limiter = RateLimiter::new(
        SlidingGuard::new(),
        MokaStore::new(),
        UserIssuer,
        CustomQuotaGetter,
    );
    let router = Router::new()
        .get(home)
        .push(Router::with_path("limited").hoop(limiter).get(limited));
    let acceptor = TcpListener::new("127.0.0.1:5800").bind().await;
    Server::new(acceptor).serve(router).await;
}

static HOME_HTML: &str = r#"
<!DOCTYPE html>
<html>
    <head>
        <title>Rate Limiter Dynmaic</title>
    </head>
    <body>
        <h2>Rate Limiter Dynamic</h2>
        <p>
            This example shows how to set limit for different users. 
        </p>
        <p>
            <a href="/limited?user=user1" target="_blank">Limited page for user1: 1/second</a>
        </p>
        <p>
            <a href="/limited?user=user2" target="_blank">Limited page for user2: 1/5seconds</a>
        </p>
        <p>
            <a href="/limited?user=user3" target="_blank">Limited page for user3: 1/10seconds</a>
        </p>
    </body>
</html>
"#;
[package]
name = "example-rate-limiter-dynamic"
version = "0.1.0"
edition = "2021"
publish = false

[dependencies]
salvo = { workspace = true, features = ["rate-limiter"] }
tokio = { version = "1", features = ["macros"] }
tracing = "0.1"
tracing-subscriber = "0.3"
once_cell = "1"