Rate limiting com tower-governor ou custom middleware

1. Introdução ao Rate Limiting em Rust

Rate limiting é uma técnica fundamental para proteger APIs contra abusos, garantir qualidade de serviço e prevenir ataques de negação de serviço. Em sistemas assíncronos Rust, onde concorrência é a norma, controlar a taxa de requisições torna-se ainda mais crítico.

Os algoritmos mais comuns incluem:

  • Janela fixa (Fixed Window): Divide o tempo em janelas e limita requisições por janela. Simples, mas pode permitir picos nos limites das janelas.
  • Sliding Window: Suaviza o comportamento da janela fixa usando timestamps contínuos.
  • Token Bucket: Um balde com tokens que se regeneram a uma taxa constante. Permite rajadas controladas.
  • Leaky Bucket: Similar ao token bucket, mas processa requisições a uma taxa constante, descartando excessos.

O ecossistema Tower fornece uma abstração poderosa para middleware em Rust, permitindo compor comportamentos como rate limiting de forma modular e reutilizável.

2. Configurando o Ambiente com Tower

Para começar, adicione as dependências necessárias ao seu Cargo.toml:

[dependencies]
axum = "0.7"
tower = "0.4"
tower-http = "0.5"
tokio = { version = "1", features = ["full"] }
serde = { version = "1", features = ["derive"] }

Um servidor básico com camadas de middleware:

use axum::{Router, routing::get, response::Json};
use tower_http::trace::TraceLayer;
use std::net::SocketAddr;

async fn health_check() -> Json<&'static str> {
    Json("OK")
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/health", get(health_check))
        .layer(TraceLayer::new_for_http());

    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    println!("Servidor rodando em {}", addr);

    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();
}

3. Implementação com tower-governor

O crate tower-governor implementa o algoritmo Generic Cell Rate Algorithm (GCRA), similar ao leaky bucket. Adicione ao Cargo.toml:

tower-governor = "0.4"

Configuração básica com limite por IP:

use tower_governor::{GovernorLayer, GovernorConfigBuilder, key_extractor::PeerIpKeyExtractor};
use std::time::Duration;

fn rate_limiter_config() -> GovernorLayer<PeerIpKeyExtractor> {
    let config = GovernorConfigBuilder::default()
        .per_second(10)        // 10 requisições por segundo
        .burst_size(20)        // permite rajadas de até 20 requisições
        .key_extractor(PeerIpKeyExtractor)
        .finish()
        .unwrap();

    GovernorLayer { config }
}

Aplicando ao router:

let app = Router::new()
    .route("/api/data", get(get_data))
    .layer(rate_limiter_config());

Para rate limiting por chave de API:

use tower_governor::key_extractor::{KeyExtractor, RequestKey};
use axum::http::Request;
use std::sync::Arc;

struct ApiKeyExtractor;

impl KeyExtractor for ApiKeyExtractor {
    type Key = String;

    fn extract<B>(&self, req: &Request<B>) -> Option<Self::Key> {
        req.headers()
            .get("X-API-Key")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string())
    }
}

let config = GovernorConfigBuilder::default()
    .per_minute(60)
    .key_extractor(ApiKeyExtractor)
    .finish()
    .unwrap();

4. Custom Middleware de Rate Limiting

Implementar um middleware customizado oferece controle total sobre o comportamento. Aqui está uma implementação com armazenamento em memória:

use tower::{Layer, Service};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use axum::http::{Request, StatusCode, Response};
use std::time::{Duration, Instant};

#[derive(Clone)]
struct RateLimiter {
    state: Arc<Mutex<HashMap<String, (u32, Instant)>>>,
    max_requests: u32,
    window: Duration,
}

impl RateLimiter {
    fn new(max_requests: u32, window: Duration) -> Self {
        RateLimiter {
            state: Arc::new(Mutex::new(HashMap::new())),
            max_requests,
            window,
        }
    }

    fn check(&self, key: &str) -> bool {
        let mut state = self.state.lock().unwrap();
        let now = Instant::now();

        // Limpeza de entradas expiradas
        state.retain(|_, (_, time)| now.duration_since(*time) < self.window);

        let entry = state.entry(key.to_string()).or_insert((0, now));
        entry.1 = now;

        if entry.0 >= self.max_requests {
            false
        } else {
            entry.0 += 1;
            true
        }
    }
}

impl<S, B> Layer<S> for RateLimiter {
    type Service = RateLimitService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        RateLimitService {
            inner,
            limiter: self.clone(),
        }
    }
}

struct RateLimitService<S> {
    inner: S,
    limiter: RateLimiter,
}

impl<S, B> Service<Request<B>> for RateLimitService<S>
where
    S: Service<Request<B>, Response = Response<axum::body::Body>>,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request<B>) -> Self::Future {
        let ip = req.headers()
            .get("X-Forwarded-For")
            .and_then(|v| v.to_str().ok())
            .unwrap_or("unknown")
            .to_string();

        if !self.limiter.check(&ip) {
            let response = Response::builder()
                .status(StatusCode::TOO_MANY_REQUESTS)
                .header("Retry-After", "1")
                .body(axum::body::Body::from("Rate limit exceeded"))
                .unwrap();
            return Box::pin(async { Ok(response) });
        }

        let future = self.inner.call(req);
        Box::pin(async move { future.await })
    }
}

5. Rate Limiting Distribuído com Redis

Para ambientes com múltiplas instâncias, Redis oferece armazenamento compartilhado:

redis = { version = "0.24", features = ["tokio-comp"] }

Implementação de sliding window counter com sorted sets:

use redis::AsyncCommands;
use std::time::{SystemTime, UNIX_EPOCH};

async fn check_rate_limit(conn: &mut redis::aio::Connection, key: &str, limit: u32, window_secs: u64) -> bool {
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_millis() as u64;

    let window_start = now - (window_secs * 1000);

    // Remove entradas antigas
    let _: () = conn.zrembyscore(key, 0, window_start).await.unwrap();

    // Conta requisições atuais
    let count: u32 = conn.zcard(key).await.unwrap();

    if count >= limit {
        false
    } else {
        let _: () = conn.zadd(key, now, now).await.unwrap();
        let _: () = conn.expire(key, window_secs as usize).await.unwrap();
        true
    }
}

6. Headers de Rate Limiting e Respostas HTTP

Headers padronizados informam clientes sobre limites:

use axum::http::HeaderValue;

fn add_rate_limit_headers(response: &mut Response<Body>, limit: u32, remaining: u32, reset: u64) {
    response.headers_mut().insert(
        "X-RateLimit-Limit",
        HeaderValue::from(limit),
    );
    response.headers_mut().insert(
        "X-RateLimit-Remaining",
        HeaderValue::from(remaining),
    );
    response.headers_mut().insert(
        "X-RateLimit-Reset",
        HeaderValue::from_str(&reset.to_string()).unwrap(),
    );
}

Tratamento de excesso com status 429:

use axum::response::IntoResponse;

struct RateLimitExceeded {
    retry_after: u64,
}

impl IntoResponse for RateLimitExceeded {
    fn into_response(self) -> Response<Body> {
        Response::builder()
            .status(StatusCode::TOO_MANY_REQUESTS)
            .header("Retry-After", self.retry_after.to_string())
            .header("Content-Type", "application/json")
            .body(Body::from(
                serde_json::json!({
                    "error": "rate_limit_exceeded",
                    "message": "Muitas requisições. Tente novamente mais tarde.",
                    "retry_after": self.retry_after
                }).to_string()
            ))
            .unwrap()
    }
}

7. Testes de Rate Limiting

Testes com tokio-test para verificar comportamento:

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::Request;
    use tower::ServiceExt;

    #[tokio::test]
    async fn test_rate_limiting() {
        let limiter = RateLimiter::new(3, Duration::from_secs(10));
        let mut service = limiter.layer(
            tower::service_fn(|req: Request<Body>| async {
                Ok::<_, std::convert::Infallible>(Response::new(Body::from("OK")))
            })
        );

        // 3 requisições devem passar
        for _ in 0..3 {
            let req = Request::builder()
                .header("X-Forwarded-For", "192.168.1.1")
                .body(Body::empty())
                .unwrap();
            let resp = service.ready().await.unwrap().call(req).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK);
        }

        // A quarta deve ser bloqueada
        let req = Request::builder()
            .header("X-Forwarded-For", "192.168.1.1")
            .body(Body::empty())
            .unwrap();
        let resp = service.ready().await.unwrap().call(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
    }
}

8. Boas Práticas e Considerações de Performance

Escolha entre tower-governor vs custom middleware:

  • Use tower-governor quando precisar de uma solução robusta e testada com algoritmo GCRA
  • Prefira custom middleware quando necessitar de lógica específica de negócio ou integração com sistemas legados

Impacto na performance:

  • Armazenamento em memória é mais rápido que Redis, mas não escala horizontalmente
  • Para alta concorrência, considere usar dashmap em vez de Mutex<HashMap>
  • Implemente limpeza periódica de cache para evitar vazamento de memória

Monitoramento:

use tracing::{info, warn};

fn log_rate_limit_event(ip: &str, endpoint: &str, blocked: bool) {
    if blocked {
        warn!("Rate limit bloqueado para IP: {} no endpoint: {}", ip, endpoint);
    } else {
        info!("Requisição permitida para IP: {} no endpoint: {}", ip, endpoint);
    }
}

Considerações finais:

  • Sempre use rate limiting em endpoints críticos (login, API keys, uploads)
  • Combine rate limiting com autenticação para maior granularidade
  • Documente os limites para consumidores da sua API
  • Implemente backoff exponencial no cliente para lidar com respostas 429

Rate limiting é uma camada essencial de defesa e qualidade de serviço. Com Tower e as ferramentas do ecossistema Rust, você pode implementar soluções elegantes e performáticas.

Referências