benwis / tower-governor

Rate Limiting middleware for Tower/Axum/Tonic/Hyper utilizing the governor crate
MIT License
174 stars 29 forks source link
axum governor rust tonic tower

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:

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 = Arc::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 {
        config: 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()`](https://docs.rs/tower_governor/latest/tower_governor/governor/struct.GovernorConfig.html#method.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()`](https://docs.rs/tower_governor/latest/tower_governor/governor/struct.GovernorConfig.html#method.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:

 ```rust
 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 std::sync::Arc;
    
    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(axum::body::Body::from("mock response"))) 
    });
    
    let config = Arc::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 })
    .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 and retry-after headers are being sent. If you want to add 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

  3. 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.

  4. 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!