Implementing API Rate Limiting in Rust

shuttle_dev

Shuttle

Posted on February 22, 2024

Implementing API Rate Limiting in Rust

Hello world! We’re going to talk about implementing rate limiting for your API in Rust. When it comes to services in production, you want to ensure that bad actors aren’t abusing your APIs - this is where API rate limiting comes in.

For this tutorial we will be implementing the “sliding window” algorithm by having a dynamic period to check request histories over as well, using a basic in-memory hashmap to store users and their request times. We will also look at using tower-governor to configure rate limiting for you.

Implementing a naive sliding window rate limiter

To figure out how we can do this from the ground up, let’s write a naive sliding window IP based rate limiter from scratch.

To get started, we’re going to initialise a regular project using cargo init and follow the prompt, picking Axum as our framework of choice.

We’re going to need some extra dependencies, so let’s install them with this shell snippet:

cargo add serde@1.0.196 -F derive
cargo add chrono@0.4.34 -F serde,clock
Enter fullscreen mode Exit fullscreen mode

We’ll declare a new struct that holds a HashMap of IpAddr keys with the values being Vec<DateTime<Utc>> (a Vector of UTC-timezone timestamps).

use std::sync::{Arc, Mutex};
use std::collections::HashMap;
use std::net::IpAddr;
use chrono::{DateTime, Utc};

// This will be the request limit (per minute) for a user to access an endpoint
// If the user attempts to go beyond this limit, we should return an error
const REQUEST_LIMIT: usize = 120;

#[derive(Clone, Default)]
pub struct RateLimiter {
    requests: Arc<Mutex<HashMap<IpAddr, Vec<DateTime<Utc>>>>>,
}
Enter fullscreen mode Exit fullscreen mode

To get started, we’ll want to lock our hashmap by using .lock() which gives us write access. Then we’ll want to check whether or not the hashmap contains a key containing the IP address we want to check for with the .entry() function, then modify it by retaining valid timestamps and pushing a new entry depending on whether or not the length is under the request limit. We then check if the entry length is a higher length than the request limit - if so, return an error; if not, return Ok(()).

impl RateLimiter {
    fn check_if_rate_limited(&self, ip_addr: IpAddr) -> Result<(), String> {
        // we only want to keep timestamps from up to 60 seconds ago
        let throttle_time_limit = Utc::now() - std::time::Duration::from_secs(60);

        let mut requests_hashmap = self.requests.lock().unwrap();

        let mut requests_for_ip = requests_hashmap
            // grab the entry here and allow us to modify it in place
            .entry(ip_addr)
            // if the entry is empty, insert a vec with the current timestamp
            .or_insert(Vec::new());

        requests_for_ip.retain(|x| x.to_utc() > throttle_time_limit);
        requests_for_ip.push(Utc::now());

        if requests_for_ip.len() > REQUEST_LIMIT {
            return Err("IP is rate limited :(".to_string());
        }

        Ok(())
    }
}
Enter fullscreen mode Exit fullscreen mode

Here is a basic example of how you might use this:

use std::net::Ipv4Addr;

fn main() {
let rate_limiter = RateLimiter::default();

    let localhost_v4 = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));

    // here we request 120 times - our request limit
    for _ in 1..80 {
    assert!(rate_limiter.check_if_rate_limited(localhost_v4).is_ok())
    }

    // wait 30 seconds
    std::thread::sleep(std::time::Duration::from_secs(30));

    // make another 40 requests here to satisfy request quota
    for _ in 1..40 {
    assert!(rate_limiter.check_if_rate_limited(localhost_v4).is_ok())
    }

    // wait another 30 seconds
    std::thread::sleep(std::time::Duration::from_secs(30));

    // now we can make another 80 requests
    for _ in 1..80 {
    assert!(rate_limiter.check_if_rate_limited(localhost_v4).is_ok())
    }
}
Enter fullscreen mode Exit fullscreen mode

From here if we wanted to extend this to work with Axum, we could. However, production-ready rate-limiting systems are typically much more advanced than this. We’ll be discussing how you can utilize crates for rate limiting below, including usage of user-based rate limiting.

Implementing user-based rate limiting

For external-facing websites without a login, IP addresses are the only thing you can use (besides browser information) to track users. However, it can be much more useful to rate limit based on authenticated users rather than IP addresses. While working with IP addresses, you may run into the following issues:

  • Multiple users may have the same IP address
  • Users can simply change the IP address they use if you block them (via proxy or other methods)

Working with user-based rate limiting allows us to solve these issues. While users can have more than one IP address, we can assign it all to the same user.

Getting started

To initialise our web service we'll use cargo shuttle init (requires cargo-shuttle installed) to create our project, making sure to pick Axum as the framework.

Before adding the rate limiter itself, we’re going to create a custom header key! This will be used in routes where we require user authentication. We can also use the header when implementing our custom key extractor for the rate limiter later on. We’ll want to start by adding axum-extra with the typed-header feature:

cargo add axum-extra -F typed-header
Enter fullscreen mode Exit fullscreen mode

Next we’ll want to create a struct that will hold a String and implement the axum_extra re-export of headers::Header. You can see the Header implementation below, where it decodes the value by iterating over HeaderValue and creates the CustomHeader struct.

We can start by defining a HeaderName:

static X: HeaderName = HeaderName::from_static("x-custom-key");
static CUSTOM_HEADER: &HeaderName = &X;

pub struct CustomHeader(String);

impl CustomHeader {
    pub fn key(self) -> String {
        self.0
    }
}
Enter fullscreen mode Exit fullscreen mode

Now that we’ve defined our custom header name (which will be used as the header key), we can implement axum_extra::headers::Header for CustomHeader:


impl Header for CustomHeader {
    fn name() -> &'static HeaderName {
        CUSTOM_HEADER
    }

    fn decode<'i, I>(values: &mut I) -> Result<Self, axum_extra::headers::Error>
    where
        I: Iterator<Item = &'i HeaderValue>,
    {
        let value = values
            .next()
            .ok_or_else(axum_extra::headers::Error::invalid)?;

        Ok(CustomHeader(value.to_str().unwrap().to_owned()))
    }

    fn encode<E>(&self, values: &mut E)
    where
        E: Extend<HeaderValue>,
    {
        let s = &self.0;

        let value = HeaderValue::from_str(s).unwrap();

        values.extend(std::iter::once(value));
    }
}
Enter fullscreen mode Exit fullscreen mode

To use CustomHeader as an Axum extractor, we need to wrap it in TypedHeader like so:

async fn register(
    TypedHeader(header): TypedHeader<CustomHeader>,
) -> impl IntoResponse {
    // .. your code goes here
}
Enter fullscreen mode Exit fullscreen mode

This is all well and good, but how does this relate to rate limiting?

While we can use this in middleware, a better alternative solution would be to use tower_governor. This crate is a Tower service backed by the governor crate (a crate for regulating data with rate limiting) and makes it much easier to implement rate limiting, The crate uses the Generic Cell Rate Algorithm (GCRA) which is a much more sophisticated version of a leaky bucket. You can read much more about GCRA here.

To get started, we’ll add the crate to our Rust program:

cargo add tower-governor
Enter fullscreen mode Exit fullscreen mode

When we want to add it to our main function, we can do it by using GovernorConfigBuilder and then adding it into GovernorLayer. Note that while GovernorConfigBuilder doesn’t implement Clone, adding a Tower service layer requires it to implement Clone. This means that we need to box the config builder and then later on, we can use Box::leak to leak the box to get a &'static lifetime GovernorConfig for usage with our axum::Router:

use auxm::{Router, routing::get};
use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};

#[shuttle_runtime::main]
async fn main() -> shuttle_axum::ShuttleAxum {
    let governor_conf = Box::new(
        GovernorConfigBuilder::default()
            .per_second(2)
            .burst_size(5)
            .finish()
            .unwrap(),
    );

    let router = Router::new()
        .route("/", get(hello_world))
        .layer(GovernorLayer {
            // We can leak this because it is created once and then never needs to be destructed
            config: Box::leak(governor_conf),
        });

    Ok(router.into())
}
Enter fullscreen mode Exit fullscreen mode

By default, GovernorConfigBuilder uses a type called PeerIpKeyExtractor which attempts to grab the IP key of a connecting client. However, to use our header as the extracted key we can implement tower_governor::key_extractor::KeyExtractor. To do this, we’ll use a unit struct as when we add it to GovernorConfigBuilder later on, there aren’t currently any extra variables we need:

use tower_governor::GovernorError;
use axum::http::Request;

#[derive(Clone)]
pub struct CustomHeaderExtractor;

impl KeyExtractor for CustomHeaderExtractor {
    type Key = String;

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

        match headers.get(CUSTOM_HEADER) {
            Some(res) => {
                let res = res.to_str()
                    .map_err(|_| GovernorError::UnableToExtractKey)?;

                Ok(res.to_owned())
            },
            None => Err(GovernorError::UnableToExtractKey)
        }
    }
}
Enter fullscreen mode Exit fullscreen mode

This allows us to add CustomHeaderExtractor to our GovernorConfigBuilder in our main function.

let governor_conf = Box::new(
    GovernorConfigBuilder::default()
        .per_second(2)
        .burst_size(5)
        .key_extractor(CustomHeaderExtractor)
        .finish()
        .unwrap(),
);
Enter fullscreen mode Exit fullscreen mode

When a user attempts to access any route that is layered with the GovernorLayer, now it’ll attempt to get a header with the header name x-custom-key - if it’s not present, the route will return an error. Here we have set the limit to allow users to send 5 requests every 2 seconds.

Note that in the builder, the per_second() function tells us exactly how many seconds the interval will be between replenishing the quota and burst_size tells us what the quota is before tower-governor will start blocking requests from a given IP address (or API key, in our case). We can also additionally set per_millisecond() and per_nanosecond() parameters so that if you want to replenish the quota every half a second for example, you can use per_millisecond(500) in the builder.

Deploying

Now that we’re done, you can deploy using cargo shuttle deploy (add --ad if on a dirty Git branch) and watch the magic happen. Once finished, Shuttle will output the details of your deployment in the terminal.

Finishing Up

Thanks for reading! With this guide, implementing rate limiting in a Rust web service should be much easier to do. Productionizing Rust web services has never been easier!

Read more:

  • Try out implementing JWT authentication here.
  • Learn more about the Tracing ecosystem for logging here.
💖 💪 🙅 🚩
shuttle_dev
Shuttle

Posted on February 22, 2024

Join Our Newsletter. No Spam, Only the good stuff.

Sign up to receive the latest update from our blog.

Related