feat: add request rate limiting

src/app.rs

@@ -1,11 +1,11 @@
 use crate::{
-    config::AppConfig, error::AppError, pokemon::PokeApiClient, service::PokedexService,
+    config::AppConfig, error::AppError, pokemon::PokeApiClient, rate_limit::RateLimiter,
-    telemetry::Telemetry, translation::TranslationClient,
+    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<PokedexService>,
     telemetry: Arc<Telemetry>,
+    rate_limiter: Arc<RateLimiter>,
 }
 
 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<AppState>) -> Result<impl IntoResponse, App
     Ok(([(header::CONTENT_TYPE, "text/plain; version=0.0.4")], body))
 }
 
+async fn enforce_rate_limit(
+    State(state): State<AppState>,
+    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<AppState>,
     req: Request,
@@ -120,8 +147,8 @@ struct HealthResponse {
 mod tests {
     use super::{AppState, create_app};
     use crate::{
-        pokemon::PokeApiClient, service::PokedexService, telemetry::Telemetry,
+        pokemon::PokeApiClient, rate_limit::RateLimiter, service::PokedexService,
-        translation::TranslationClient,
+        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"),
         )
     }
 

src/config.rs

@@ -13,6 +13,10 @@ pub struct AppConfig {
 
 impl AppConfig {
     pub fn from_env() -> Result<Self, ConfigError> {
+        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_per_second,
-            rate_limit_burst: env_or_parse("RATE_LIMIT_BURST", 40)?,
+            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<T>(name: &'static str, default: T) -> Result<T, ConfigError>
 where
     T: std::str::FromStr,

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;

src/rate_limit.rs

@@ -0,0 +1,90 @@
+use std::{sync::Mutex, time::Instant};
+
+#[derive(Debug)]
+pub struct RateLimiter {
+    bucket: Mutex<TokenBucket>,
+}
+
+impl RateLimiter {
+    pub fn new(refill_per_second: f64, burst: u32) -> Result<Self, RateLimitError> {
+        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());
+    }
+}