From c64b738117cdfede32c8f5629d00f5b34a869008 Mon Sep 17 00:00:00 2001 From: Alberto Moretti Date: Sat, 13 Jun 2026 06:14:47 +0200 Subject: [PATCH] feat: add request rate limiting --- src/app.rs | 74 ++++++++++++++++++++++++++++++++++---- src/config.rs | 18 ++++++++-- src/lib.rs | 1 + src/rate_limit.rs | 90 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 174 insertions(+), 9 deletions(-) create mode 100644 src/rate_limit.rs diff --git a/src/app.rs b/src/app.rs index e1fe0aa..f715410 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1,11 +1,11 @@ use crate::{ - config::AppConfig, error::AppError, pokemon::PokeApiClient, service::PokedexService, - telemetry::Telemetry, translation::TranslationClient, + config::AppConfig, error::AppError, pokemon::PokeApiClient, rate_limit::RateLimiter, + service::PokedexService, telemetry::Telemetry, translation::TranslationClient, }; use axum::{ Json, Router, extract::{MatchedPath, Path, Request, State}, - http::header, + http::{StatusCode, header}, middleware, response::{IntoResponse, Response}, routing::get, @@ -23,14 +23,16 @@ use tower_http::{ pub struct AppState { service: Arc, telemetry: Arc, + rate_limiter: Arc, } impl AppState { #[must_use] - pub fn new(service: PokedexService, telemetry: Telemetry) -> Self { + pub fn new(service: PokedexService, telemetry: Telemetry, rate_limiter: RateLimiter) -> Self { Self { service: Arc::new(service), telemetry: Arc::new(telemetry), + rate_limiter: Arc::new(rate_limiter), } } @@ -40,7 +42,12 @@ impl AppState { let translations = TranslationClient::new(&config.translations_base_url, config.request_timeout)?; let service = PokedexService::new(pokemon, translations, telemetry.metrics()); - Ok(Self::new(service, telemetry)) + Ok(Self::new( + service, + telemetry, + RateLimiter::new(config.rate_limit_per_second, config.rate_limit_burst) + .map_err(|_| AppError::Internal)?, + )) } } @@ -60,6 +67,10 @@ pub fn create_app(state: AppState) -> Router { .route("/pokemon/{name}", get(get_pokemon)) .route("/pokemon/translated/{name}", get(get_translated_pokemon)) .with_state(state.clone()) + .layer(middleware::from_fn_with_state( + state.clone(), + enforce_rate_limit, + )) .layer(middleware::from_fn_with_state(state, record_http_metrics)) .layer(tracing_layers) } @@ -73,6 +84,22 @@ async fn metrics(State(state): State) -> Result, + req: Request, + next: middleware::Next, +) -> Response { + if state.rate_limiter.try_acquire() { + next.run(req).await + } else { + ( + StatusCode::TOO_MANY_REQUESTS, + Json(serde_json::json!({ "error": "rate limit exceeded" })), + ) + .into_response() + } +} + async fn record_http_metrics( State(state): State, req: Request, @@ -120,8 +147,8 @@ struct HealthResponse { mod tests { use super::{AppState, create_app}; use crate::{ - pokemon::PokeApiClient, service::PokedexService, telemetry::Telemetry, - translation::TranslationClient, + pokemon::PokeApiClient, rate_limit::RateLimiter, service::PokedexService, + telemetry::Telemetry, translation::TranslationClient, }; use axum::{ body::{Body, to_bytes}, @@ -256,6 +283,38 @@ mod tests { assert!(body.contains("upstream_client_requests")); } + #[tokio::test] + async fn rate_limit_rejects_requests_over_the_configured_burst() { + let pokeapi = MockServer::start().await; + let translations = MockServer::start().await; + let telemetry = Telemetry::new("test-service"); + let app = create_app(AppState::new( + PokedexService::new( + PokeApiClient::new(pokeapi.uri(), Duration::from_secs(1)) + .expect("client should build"), + TranslationClient::new(translations.uri(), Duration::from_secs(1)) + .expect("client should build"), + telemetry.metrics(), + ), + telemetry, + RateLimiter::new(0.01, 1).expect("rate limiter should build"), + )); + + let first = app + .clone() + .oneshot(request("/health")) + .await + .expect("request should be handled"); + let second = app + .oneshot(request("/health")) + .await + .expect("request should be handled"); + + assert_eq!(first.status(), StatusCode::OK); + assert_eq!(second.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(json_body(second).await["error"], "rate limit exceeded"); + } + #[tokio::test] async fn pokemon_endpoint_maps_missing_pokemon_to_404() { let pokeapi = MockServer::start().await; @@ -308,6 +367,7 @@ mod tests { telemetry.metrics(), ), telemetry, + RateLimiter::new(100.0, 100).expect("rate limiter should build"), ) } diff --git a/src/config.rs b/src/config.rs index 72847fc..37afd67 100644 --- a/src/config.rs +++ b/src/config.rs @@ -13,6 +13,10 @@ pub struct AppConfig { impl AppConfig { pub fn from_env() -> Result { + let rate_limit_per_second = env_or_parse("RATE_LIMIT_PER_SECOND", 20.0)?; + let rate_limit_burst = env_or_parse("RATE_LIMIT_BURST", 40)?; + validate_rate_limit(rate_limit_per_second, rate_limit_burst)?; + Ok(Self { bind_addr: env_or_parse("BIND_ADDR", default_bind_addr())?, pokeapi_base_url: env_or_string("POKEAPI_BASE_URL", "https://pokeapi.co/api/v2"), @@ -21,8 +25,8 @@ impl AppConfig { "https://funtranslations.mercxry.me", ), request_timeout: Duration::from_secs(env_or_parse("REQUEST_TIMEOUT_SECONDS", 5)?), - rate_limit_per_second: env_or_parse("RATE_LIMIT_PER_SECOND", 20.0)?, - rate_limit_burst: env_or_parse("RATE_LIMIT_BURST", 40)?, + rate_limit_per_second, + rate_limit_burst, service_name: env_or_string("OTEL_SERVICE_NAME", env!("CARGO_PKG_NAME")), }) } @@ -30,6 +34,8 @@ impl AppConfig { #[derive(Debug, thiserror::Error)] pub enum ConfigError { + #[error("rate limit values must be positive")] + InvalidRateLimit, #[error("invalid value for {name}: {source}")] InvalidValue { name: &'static str, @@ -45,6 +51,14 @@ fn default_bind_addr() -> SocketAddr { SocketAddr::from(([0, 0, 0, 0], 5000)) } +const fn validate_rate_limit(per_second: f64, burst: u32) -> Result<(), ConfigError> { + if per_second.is_sign_positive() && burst > 0 { + Ok(()) + } else { + Err(ConfigError::InvalidRateLimit) + } +} + fn env_or_parse(name: &'static str, default: T) -> Result where T: std::str::FromStr, diff --git a/src/lib.rs b/src/lib.rs index b397c01..5b9b9c2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ pub mod config; pub mod domain; pub mod error; pub mod pokemon; +pub mod rate_limit; pub mod service; pub mod telemetry; pub mod translation; diff --git a/src/rate_limit.rs b/src/rate_limit.rs new file mode 100644 index 0000000..242d1c9 --- /dev/null +++ b/src/rate_limit.rs @@ -0,0 +1,90 @@ +use std::{sync::Mutex, time::Instant}; + +#[derive(Debug)] +pub struct RateLimiter { + bucket: Mutex, +} + +impl RateLimiter { + pub fn new(refill_per_second: f64, burst: u32) -> Result { + if !refill_per_second.is_sign_positive() || burst == 0 { + return Err(RateLimitError::InvalidConfiguration); + } + + Ok(Self { + bucket: Mutex::new(TokenBucket { + capacity: f64::from(burst), + tokens: f64::from(burst), + refill_per_second, + last_refill: Instant::now(), + }), + }) + } + + #[must_use] + pub fn try_acquire(&self) -> bool { + self.bucket + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .try_acquire() + } +} + +#[derive(Debug, thiserror::Error)] +pub enum RateLimitError { + #[error("rate limit configuration must be positive")] + InvalidConfiguration, +} + +#[derive(Debug)] +struct TokenBucket { + capacity: f64, + tokens: f64, + refill_per_second: f64, + last_refill: Instant, +} + +impl TokenBucket { + fn try_acquire(&mut self) -> bool { + self.refill(); + if self.tokens >= 1.0 { + self.tokens -= 1.0; + true + } else { + false + } + } + + fn refill(&mut self) { + let elapsed = self.last_refill.elapsed().as_secs_f64(); + self.tokens = self + .capacity + .min(self.tokens + elapsed.mul_add(self.refill_per_second, 0.0)); + self.last_refill = Instant::now(); + } +} + +#[cfg(test)] +mod tests { + use super::RateLimiter; + use std::time::Duration; + + #[test] + fn allows_requests_up_to_the_burst_capacity() { + let limiter = RateLimiter::new(100.0, 2).expect("rate limiter should build"); + + assert!(limiter.try_acquire()); + assert!(limiter.try_acquire()); + assert!(!limiter.try_acquire()); + } + + #[test] + fn refills_tokens_over_time() { + let limiter = RateLimiter::new(100.0, 1).expect("rate limiter should build"); + + assert!(limiter.try_acquire()); + assert!(!limiter.try_acquire()); + std::thread::sleep(Duration::from_millis(20)); + assert!(limiter.try_acquire()); + } +}