When you first start thinking about rate limiting, it's tempting to think of it as a simple counting problem. You'll probably tell yourself, "I'll just keep a HashMap of IP addresses and a counter in my app state. Every time a request comes in, I'll increment the number. If it's over 100, I'll return a 429."

The "Simple HashMap" Trap

The problem is that Axum is asynchronous and multi-threaded. If you try to wrap a HashMap in a Mutex and update it inside a handler or middleware, you aren't just solving a counting problem; you're creating a massive synchronization bottleneck.

// DON'T DO THIS in a high-traffic app
async fn rate_limit_middleware(
    State(state): State<Arc<Mutex<HashMap<IpAddr, u32>>>>,
    req: Request,
    next: Next,
) -> Response {
    let mut counts = state.lock().unwrap(); // Every single request blocks here
    let ip = req.extensions().get<ConnectInfo<SocketAddr>>().map(|i| i.0.ip());
    // ... increment and check ...
    next.run(req).await
}

In the example above, every single request to your server—regardless of whether it's from the same user or a completely different one—must wait in line for that one Mutex lock. You've effectively turned your high-performance asynchronous server into a single-threaded one. Under a real DDoS attack (the exact time you actually need a rate limiter), your server will spend more time fighting over the lock than actually processing requests.

Using Atomic State and Fixed Windows

To do this right, we need to minimize the time we hold locks or avoid them entirely for the "hot path." For a simple in-memory limiter, I recommend starting with a "Fixed Window" approach using a thread-safe store. Instead of one giant lock, we can use a DashMap (a concurrent hash map) or, for even better performance, a dedicated crate like governor. But since we're learning how this works under the hood, let's implement a basic version using std::sync::atomic for the counters.

We'll create a state object that tracks the current "window" (e.g., the current minute) and the request count for each IP. If the window changes, we reset the count.

use axum::{
    extract::{ConnectInfo, State},
    http::{Request, StatusCode},
    middleware::Next,
    response::Response,
};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::net::IpAddr;
use std::time::{SystemTime, UNIX_EPOCH};

struct RateLimitState {
    // We use a Mutex here for the Map, but the values inside could be Atomics
    // For this simple example, we'll keep it straightforward.
    requests: Mutex<HashMap<(IpAddr, u64), u32>>,
}

async fn rate_limiter(
    State(state): State<Arc<RateLimitState>>,
    ConnectInfo(addr): ConnectInfo<std::net::SocketAddr>,
    req: Request,
    next: Next,
) -> Result<Response, StatusCode> {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs() / 60; // Window = 1 minute

    let ip = addr.ip();
    
    {
        let mut map = state.requests.lock().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
        let count = map.entry((ip, now)).or_insert(0);
        *count += 1;

        if *count > 100 {
            return Err(StatusCode::TOO_MANY_REQUESTS);
        }
    } // Lock is dropped here immediately

    Ok(next.run(req).await)
}

Integrating the Middleware into the Router

The beauty of middleware::from_fn is that it allows us to treat our rate limiter as a wrapper. I usually apply this globally, but you can also apply it to specific "expensive" routes (like /api/search or /api/login) to protect them without affecting static assets.

Keep in mind: the ConnectInfo extractor requires you to use a specific transport layer in your server setup (like Router::into_make_service_with_connect_info), otherwise the middleware will panic because it can't find the IP address of the requester. I've seen plenty of developers spend an hour wondering why their middleware is crashing, only to realize they forgot to enable the connection info at the listener level.

let state = Arc::new(RateLimitState {
    requests: Mutex::new(HashMap::new()),
});

let app = Router::new()
    .route("/", get(|| async { "Hello!" }))
    .layer(axum::middleware::from_fn_with_state(state.clone(), rate_limiter));

// IMPORTANT: You must use into_make_service_with_connect_info
let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(
    listener, 
    app.into_make_service_with_connect_info<std::net::SocketAddr>()
).await.unwrap();

📋 Practical Task

Exercise: Implementing a Tiered Rate Limiter for API Keys

In a real-world scenario, you rarely rate limit by IP alone. Usually, you have different tiers (e.g., "Free" users get 10 requests/min, "Premium" users get 1000 requests/min).

Your Task: Modify the rate limiter middleware implemented in the lesson to support API Key-based limiting instead of IP-based limiting.

• Create a UserTier enum with Free and Premium variants.
• Create a mock UserDatabase (a simple HashMap<String, UserTier>) that maps API keys to their tier.
• Update the middleware to extract an x-api-key header from the request.
• If the key is missing, return 401 Unauthorized.
• If the key exists, look up their tier and apply the corresponding limit (10 for Free, 1000 for Premium).
• Ensure the lock on the request map is held for the shortest time possible.