Giter Club home page Giter Club logo

tower-governor's Introduction

A Tower service and layer that provides a rate-limiting backend by governor. Based heavily on the work done for actix-governor. Works with Axum, Hyper, Tonic, and anything else based on Tower!

Features:

  • Rate limit requests based on peer IP address, IP address headers, globally, or via custom keys
  • Custom traffic limiting criteria per second, or to certain bursts
  • Simple to use
  • High customizability
  • High performance
  • Robust yet flexible API

How does it work?

Each governor middleware has a configuration that stores a quota. The quota specifies how many requests can be sent from an IP address before the middleware starts blocking further requests.

For example if the quota allowed ten requests a client could send a burst of ten requests in short time before the middleware starts blocking.

Once at least one element of the quota was used the elements of the quota will be replenished after a specified period.

For example if this period was 2 seconds and the quota was empty it would take 2 seconds to replenish one element of the quota. This means you could send one request every two seconds on average.

If there was a quota that allowed ten requests with the same period a client could again send a burst of ten requests and then had to wait two seconds before sending further requests or 20 seconds before the full quota would be replenished and he could send another burst.

Example

use axum::{error_handling::HandleErrorLayer, routing::get, BoxError, Router};
use std::net::SocketAddr;
use std::time::Duration;
use tokio::net::TcpListener;
use tower::ServiceBuilder;
use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};

async fn hello() -> &'static str {
   "Hello world"
}

#[tokio::main]
async fn main() {
   // Configure tracing if desired
   // construct a subscriber that prints formatted traces to stdout
   let subscriber = tracing_subscriber::FmtSubscriber::new();
   // use that subscriber to process traces emitted after this point
   tracing::subscriber::set_global_default(subscriber).unwrap();

   // Allow bursts with up to five requests per IP address
   // and replenishes one element every two seconds
   // We Box it because Axum 0.6 requires all Layers to be Clone
   // and thus we need a static reference to it
   let governor_conf = Box::new(
       GovernorConfigBuilder::default()
           .per_second(2)
           .burst_size(5)
           .finish()
           .unwrap(),
   );

   let governor_limiter = governor_conf.limiter().clone();
   let interval = Duration::from_secs(60);
   // a separate background task to clean up
   std::thread::spawn(move || {
       loop {
           std::thread::sleep(interval);
           tracing::info!("rate limiting storage size: {}", governor_limiter.len());
           governor_limiter.retain_recent();
       }
   });

   // build our application with a route
   let app = Router::new()
       // `GET /` goes to `root`
       .route("/", get(hello))
       .layer(GovernorLayer {
           // We can leak this because it is created once and then
           config: Box::leak(governor_conf),
       });

   // run our app with hyper
   // `axum::Server` is a re-export of `hyper::Server`
   let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
   tracing::debug!("listening on {}", addr);
   let listener = TcpListener::bind(addr).await.unwrap();
   axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
       .await
       .unwrap();
}

Configuration presets

Instead of using the configuration builder you can use predefined presets.

  • GovernorConfig::default(): The default configuration which is suitable for most services. Allows bursts with up to eight requests and replenishes one element after 500ms, based on peer IP.

  • GovernorConfig::secure(): A default configuration for security related services. Allows bursts with up to two requests and replenishes one element after four seconds, based on peer IP.

For example the secure configuration can be used as a short version of this code:

use tower_governor::governor::GovernorConfigBuilder;

let config = GovernorConfigBuilder::default()
    .per_second(4)
    .burst_size(2)
    .finish()
    .unwrap();

Customize rate limiting key

By default, rate limiting is done using the peer IP address (i.e. the IP address of the HTTP client that requested your app: either your user or a reverse proxy, depending on your deployment setup). You can configure a different behavior which:

  1. can be useful in itself
  2. allows you to setup multiple instances of this middleware based on different keys (for example, if you want to apply rate limiting with different rates on IP and API keys at the same time)

This is achieved by defining a [KeyExtractor] and giving it to a [Governor] instance. Three ready-to-use key extractors are provided:

  • [PeerIpKeyExtractor]: this is the default, it uses the peer IP address of the request.
  • [SmartIpKeyExtractor]: Looks for common IP identification headers usually provided by reverse proxies in order(x-forwarded-for,x-real-ip, forwarded) and falls back to the peer IP address.
  • [GlobalKeyExtractor]: uses the same key for all incoming requests

Check out the custom_key_bearer example for more information.

Crate feature flags

tower-governor uses feature flags to reduce the amount of compiled code and it is possible to enable certain features over others. Below is a list of the available feature flags:

  • axum: Enables support for axum web framework
  • tracing: Enables tracing output for this middleware

Example for no-default-features

  • Disabling default feature will change behavior of [PeerIpKeyExtractor] and [SmartIpKeyExtractor]: These two key extractors will expect SocketAddr type from Request's Extensions.
  • Fail to provide valid SocketAddr could result in [GovernorError::UnableToExtractKey].

Cargo.toml

[dependencies]
tower-governor = { version = "0.3", default-features = false }

main.rs

use std::{convert::Infallible, net::SocketAddr};

use http::{Request, Response};
use tower::{service_fn, ServiceBuilder, ServiceExt};
use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};
# async fn service() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {

// service function expecting rate limiting by governor.
let service = service_fn(|_: Request<()>| async { 
   Ok::<_, Infallible>(Response::new(String::from("mock response"))) 
});

let config = Box::new(GovernorConfigBuilder::default().finish().unwrap());

// build service with governor layer
let service = ServiceBuilder::new()
   // the caller of service must provide SocketAddr to governor layer middleware
   .map_request(|(mut req, addr): (Request<()>, SocketAddr)| {
       // insert SocketAddr to request's extensions and governor is expecting it.
       req.extensions_mut().insert(addr);
       req
   })
   .layer(GovernorLayer { config: Box::leak(config) })
   .service(service);

// mock client socket addr and http request.
let addr = "127.0.0.1:12345".parse().unwrap();
let req = Request::default();

// execute service
service.oneshot((req, addr)).await?;
# Ok(())
# }

Add x-ratelimit headers

By default, x-ratelimit-after is enabled but if you want to enable x-ratelimit-limit, x-ratelimit-whitelisted and x-ratelimit-remaining use the .use_headers() method on your GovernorConfig.

Error Handling

This crate surfaces a GovernorError with suggested headers, and includes GovernorConfigBuilder::error_handler method that will turn those errors into a Response. Feel free to provide your own error handler that takes in [GovernorError] and returns a Response.

Common pitfalls

  1. Do not construct the same configuration multiple times, unless explicitly wanted! This will create an independent rate limiter for each configuration! Instead pass the same configuration reference into Governor::new(), like it is described in the example.

  2. Be careful to create your server with .into_make_service_with_connect_info::<SocketAddr> instead of .into_make_service() if you are using the default PeerIpKeyExtractor. Otherwise there will be no peer ip address for Tower to find!

tower-governor's People

Contributors

atrox avatar benwis avatar ernestas-poskus avatar fakeshadow avatar genericnerd avatar maksimryndin avatar untitaker avatar

Stargazers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

Watchers

 avatar  avatar  avatar

tower-governor's Issues

Relax From<String> constraint on the Response?

Hi Ben,

Thanks for the amazing library,

Is it possible to relax the From constraint on the Response for the error_handler ?

I am having a custom json body that I am constructing from a Error response type with into_response implementation that make it hard to use error_handler with custom response types.

rethink error handling.

The current way of handling error is both aggressive and bloated at the same time.

Imagine two use cases of this crate:

  • in some generic tower service tree that is not axum
use http::{Request, Response};
use tower::{ServiceBuilder, ServiceExt};
use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};

#[derive(Debug, thiserror::Error)]
#[error("MyError")]
struct MyError;

#[tokio::main]
async fn main() {
    let governor_conf = Box::new(
        GovernorConfigBuilder::default()
            .per_second(60)
            .burst_size(10)
            .finish()
            .unwrap(),
    );

    // MyError type is lost because of governor layer and dynamic type casting is forced.
    let _e = ServiceBuilder::new()
        .layer(GovernorLayer {
            config: Box::leak(governor_conf),
        })
        .service_fn(|_: Request<()>| async { Err::<Response<String>, _>(MyError) })
        .oneshot(Request::default())
        .await
        .unwrap_err();
}
  • in axum the caller is forced to write HandleErrorLayer.
Router::new() .route(..).layer(GovernorLayer { .. }) // ideally this is what I want to write for a middleware

In other word the middleware force user to do runtime type casting no matter where it's used.

To make the API simpler I would suggest make use of the Response<Body> type for everything the governor would produce and just pass through the inner Layer's Service::Error type. For example:

impl<K, S, ReqBody, ResBody> Service<Request<ReqBody>>
    for Governor<K, StateInformationMiddleware, S>
where
    K: KeyExtractor,
    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
    // convert error string to ResBody. this works because axum has impl From<String> for axum::body::Body
    ResBody: From<String>, 
{
    type Response = S::Response;
    // forward error type for S 
    type Error = S::Error; 
    type Future = ResponseFuture<S::Future>;

This change would resolve the case 2 issue completely and no more HandleErrorLayer is needed for axum. For case 1 it keeps it's type error and need an extra impl From<String> for it's body type. Which is much easier because convert a String to http body is trivial in most case.

I can make an PR for this change if it's desired.

Limit an I/O throughput?

Hi, does tower-governor support limiting IO throughput?
For example, instead of limiting the number of requests per IP, I would like to limit the AsyncWrite/Read instance to 1000 kbps.

Error reported on axum version 0.6.1

Hi, first of all thank you for such a useful crate, when I upgraded axum to 0.6.1 it apparently had a compatibility issue, see the log below for details:

error[E0277]: the trait bound `GovernorLayer<'_, PeerIpKeyExtractor, governor::middleware::StateInformationMiddleware>: Clone` is not satisfied
   --> src/main.rs:50:9
    |
49  |       let app = routes::build().layer(Extension(database)).layer(
    |                                                            ----- required by a bound introduced by this call
50  | /         ServiceBuilder::new()
51  | |             .layer(HandleErrorLayer::new(|err: BoxError| async move {
52  | |                 (
53  | |                     StatusCode::INTERNAL_SERVER_ERROR,
...   |
58  | |                 config: &governor_conf,
59  | |             }),
    | |______________^ the trait `Clone` is not implemented for `GovernorLayer<'_, PeerIpKeyExtractor, governor::middleware::StateInformationMiddleware>`
    |
    = note: required for `GovernorLayer<'_, PeerIpKeyExtractor, governor::middleware::StateInformationMiddleware>` to implement `Clone`
    = note: 2 redundant requirements hidden
    = note: required for `ServiceBuilder<Stack<GovernorLayer<'_, PeerIpKeyExtractor, governor::middleware::StateInformationMiddleware>, Stack<..., ...>>>` to implement `Clone`
note: required by a bound in `Router::<S, B>::layer`
   --> /Users/wangeguo/.cargo/registry/src/mirrors.ustc.edu.cn-12df342d903acd47/axum-0.6.1/src/routing/mod.rs:316:30
    |
316 |         L: Layer<Route<B>> + Clone + Send + 'static,
    |                              ^^^^^ required by this bound in `Router::<S, B>::layer`
help: consider borrowing here
    |
50  |         &ServiceBuilder::new()
    |         +

For more information about this error, try `rustc --explain E0277`.

Referencing a Governor in a struct

Hi! I'm having a bit of trouble referencing a Governor in my own struct due to the lack of visibility to the RateLimitingMiddleware. Are you able to help with how I'm meant to provide a Governor to a struct?

Any example using tonic?

Trying to use this with tonic:

pub fn ip_rate_limiter() -> GovernorLayer<'static, PeerIpKeyExtractor, NoOpMiddleware> {
    // Allow bursts with up to five requests per IP address
    // and replenishes one element every two seconds
    // We Box it because Axum 0.6 requires all Layers to be Clone
    // and thus we need a static reference to it
    let governor_conf = Box::new(
        GovernorConfigBuilder::default()
            .per_second(2)
            .burst_size(5)
            .finish()
            .unwrap(),
    );

    let governor_limiter = governor_conf.limiter().clone();
    let interval = Duration::from_secs(60);
    // a separate background task to clean up
    std::thread::spawn(move || loop {
        std::thread::sleep(interval);
        tracing::info!("rate limiting storage size: {}", governor_limiter.len());
        governor_limiter.retain_recent();
    });
    GovernorLayer {
        // We can leak this because it is created once and then
        config: Box::leak(governor_conf),
    }
}
    let middlewares = tower::ServiceBuilder::new()
        // .timeout(Duration::from_secs(10))
        // .concurrency_limit(1_000)
        .layer(ip_rate_limiter())
        .into_inner();

    info!("starting grpc server at {socket_addr}");
    Server::builder()
        .layer(middlewares)
        .await
        .serve(socket_addr)
        .expect("server to start");

I get the following error:

error[E0277]: the trait bound `tower_governor::governor::Governor<tower_governor::key_extractor::PeerIpKeyExtractor, governor::middleware::NoOpMiddleware<governor::clock::quanta::QuantaInstant>, Routes>: Service<request::Request<tonic::transport::Body>>` is not satisfied
   --> src/auth_service/src/main.rs:245:10
    |
245 |         .serve(socket_addr)
    |          ^^^^^ the trait `Service<request::Request<tonic::transport::Body>>` is not implemented for `tower_governor::governor::Governor<tower_governor::key_extractor::PeerIpKeyExtractor, governor::middleware::NoOpMiddleware<governor::clock::quanta::QuantaInstant>, Routes>`
    |
    = help: the following other types implement trait `Service<Request>`:
              tower_governor::governor::Governor<K, governor::middleware::NoOpMiddleware, S>
              tower_governor::governor::Governor<K, governor::middleware::StateInformationMiddleware, S>
    = note: required for `InterceptedService<tower_governor::governor::Governor<tower_governor::key_extractor::PeerIpKeyExtractor, governor::middleware::NoOpMiddleware<governor::clock::quanta::QuantaInstant>, Routes>, HealthChecker>` to implement `Service<request::Request<tonic::transport::Body>>`


consider making axum an optional dependency.

right now axum::extract::ConnectInfo type is needed for providing default connection info(SocketAddr) to governor. This means user must carry axum as dependency and manual provide it's connect info when they are outside the context of it.

I would suggest making axum support an optional feature and be opt-out. This way the crate would keep compat with current feature flags and disable axum support when needed.

There are multiple approaches to this change and personally I find this could be the most simple one:

#[cfg(feature = "axum")]
fn maybe_connect_info<T>(req: &Request<T>) -> Option<IpAddr> {
    req.extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|ConnectInfo(addr)| addr.ip())
}

// for non axum case just extract SocketAddr type from extension.
#[cfg(not(feature = "axum"))]
fn maybe_connect_info<T>(req: &Request<T>) -> Option<IpAddr> {
    req.extensions()
        .get::<SocketAddr>()
        .map(|addr| addr.ip())
}

impl KeyExtractor for PeerIpKeyExtractor {
    fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, GovernorError> {
        maybe_connect_info(req).ok_or(GovernorError::UnableToExtractKey)
    }
}

impl KeyExtractor for SmartIpKeyExtractor {
    fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, GovernorError> {
        let headers = req.headers();

        maybe_x_forwarded_for(headers)
            .or_else(|| maybe_x_real_ip(headers))
            .or_else(|| maybe_forwarded(headers))
            .or_else(|| maybe_connect_info(req))
            .ok_or(GovernorError::UnableToExtractKey)
    }
}

// teach user to do this pattern when in non axum case:
ServiceBuilder::new()
    .map_request(<AddSocketAddrToExtension>)
    .layer(<Governor>)

status code 429 will be fiddled to 200 if no bufferlayer is used

this is my main.rs:

#[tokio::main]
async fn main() {
    // ... my own implementation of initialization ...
    let governor = Box::new(
        GovernorConfigBuilder::default()
            .per_second(10)
            .burst_size(20)
            .use_headers()
            .finish()
            .unwrap(),
    );
    let app = Router::new()
        .route("/test", get(test))
        // ... my own middlewares and shared states ...
        .layer(
            ServiceBuilder::new()
                .layer(HandleErrorLayer::new(|e: BoxError| async move {
                    display_error(e)
                }))
                .layer(BufferLayer::new(1024)) // remove this will cause 429 too many requests fiddled to 200
                .layer(GovernorLayer {
                    config: Box::leak(governor),
                }),
        );

    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    tracing::info!("Listening on {}", addr);
    axum::Server::bind(&addr)
        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
        .await
        .unwrap();

if I put bufferlayer there everything works perfectly fine, if remove it then issue occurs, the 429 responses will be changed to 200, with an empty response body and no headers even if I enabled use_headers() when I set up governor config

I doubt this could be related to axum itself, but I'm not sure yet what happened

my Cargo.toml dependencies:

[dependencies]
dotenv = "0.15"
axum = { version = "0.6.18", features = ["ws", "json", "headers", "http2"] }
tower = { version = "0.4.13", features = ["buffer"] }
tower-http = { version = "0.4.3", features = ["fs", "trace", "util", "request-id", "sensitive-headers"], default-features = false }
tokio = { version = "1.28", features = ["full"] }
futures-util = { default-features = false, version = "0.3.28" }
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
tracing = { version = "0.1.37", features = ["attributes"] }
serde = { version = "1.0.164", features = ["derive"] }
serde_json = "1"
once_cell = "1"
postcard = { version = "1", features = ["use-std"] }
chrono = { version = "0.4", features = ["clock", "serde"] }
derivative = "2"
jsonwebtoken = "8"
tower_governor = { version = "*", features = ["tracing"] }
governor = "0.5.1"
http = "0.2.9"

Error: undeclared crate or module `axum`

Hi!
Thank you for your great crate!

I tried it with tonic and got an error:

error[E0433]: failed to resolve: use of undeclared crate or module `axum`
 --> /home/ben/.cargo/registry/src/index.crates.io-6f17d22bba15001f/tower_governor-0.3.2/src/governor.rs:5:5 

 | use axum::body::Body;
 |     ^^^^ use of undeclared crate or module `axum`

Remove `axum` feature in the documentation

I have tried to install the package with the axum feature but I get this error from rust analyzer

Cargo.toml

tower_governor = { version = "0.3.0", features = ["tracing", "axum"] }
the package `server` depends on `tower_governor`, with features: `axum` but `tower_governor` does not have these features.
 It has a required dependency with that name, but only optional dependencies can be used as features.

The documentation is misleading because the packages doesn't have feature axum but it says it does.

Cannot create a function with `RateLimitingMiddleware`

Hi, I'm trying to create a function that returns a Box<GovernorConfig<PeerIpKeyExtractor, StateInformationMiddleware>> but it seems that GovernorConfig in the root crate is private, and StateInformationMiddleware is also private.

In my axum project I have a main function that joins all the routes from different modules. I want to configure only a small set of routes with rate limiting, so I want to be able to create the rate limiting on demand.
Could those structs be made public?
Thanks

Unable To Extract Key!

The problem was I missed that I had to start the axum server with

 .serve(router.into_make_service_with_connect_info::<SocketAddr>())

all good now

feat req: Allow other "identifiers" than IP

Imagine a corp environment or environments where many requests come from the same IP but sourcing from various entities. In that case I'd like to rate limit against the basic auth user e.g.

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    ๐Ÿ–– Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. ๐Ÿ“Š๐Ÿ“ˆ๐ŸŽ‰

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google โค๏ธ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.