diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,3 +16,7 @@ * First version revised C. Added new functions to support more middleware functionality for better keter integration. Changed api for some functions because of the necessity to take into account also throttle names for every request (especially important for multiple throttles per zone/user etc). Improved documentation. Added more tests and some dependencies (most of them are already dependencies — at least indirect — of keter) and removed not used any more dependency on containers. +## 0.1.2.0 -- 2025-08-18 + +* First version revised D. Changed Keter.RateLimiter.WAI module to more declaraive approach recommended by @jappeace for keter integration. Reflected the changes in the README.md file. + diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -2,33 +2,35 @@ **keter-rate-limiting-plugin** is a modern, high-performance, and highly customizable rate-limiting plugin for [Keter](https://github.com/snoyberg/keter). It addresses [issue \#301](https://github.com/snoyberg/keter/issues/301) and brings robust, production-grade request throttling to Haskell web applications, featuring efficient in-memory caching with HashMap-based lookups and IP zone isolation. -This library is inspired by [rack-attack](https://github.com/rack/rack-attack) and and [Ruby on Rails](https://github.com/rails/rails) (for Keter.RateLimiter.Notifications) and provides a powerful middleware for Keter-managed applications, though it can be integrated with any WAI-compatible Haskell web stack. +This library is inspired by [rack-attack](https://github.com/rack/rack-attack) and [Ruby on Rails](https://github.com/rails/rails) (for Keter.RateLimiter.Notifications) and provides a powerful middleware for Keter-managed applications, though it can be integrated with any WAI-compatible Haskell web stack. ## Features - - **Five window algorithms**: - - Fixed Window - - Sliding Window - - Token Bucket - - Leaky Bucket - - TinyLRU (Least Recently Used) - - **IP Zone Support**: Isolate caches and throttling policies per IP zone, customer segment, or any other logical grouping with efficient HashMap-based zone lookups. - - **Flexible Throttle Configuration**: Set limits, periods, algorithms, and unique identifiers on a per-throttle basis, stored in optimized HashMap structures. - - **WAI Middleware**: Integrates seamlessly as a middleware into any WAI application. - - **Convenient and Customizable API**: - - Use simple wrappers for common scenarios with automatic key composition. - - Or, for advanced use, fully control cache key structure and throttling logic. - - **Memory-efficient**: Designed for large-scale, high-traffic deployments with automatic cleanup of expired entries and HashMap-based O(1) average-case lookups. - - **Easy Integration**: Minimal code changes are required to get started. +- **Five window algorithms**: + - Fixed Window + - Sliding Window + - Token Bucket + - Leaky Bucket + - TinyLRU (Least Recently Used) +- **IP Zone Support**: Isolate caches and throttling policies per IP zone, customer segment, or any other logical grouping with efficient HashMap-based zone lookups. +- **Declarative Configuration**: Define throttling rules using JSON/YAML configuration with automatic serialization support. +- **Flexible Client Identification**: Multiple strategies for identifying clients (IP, headers, cookies, combinations). +- **Configurable Zone Derivation**: Flexible strategies for deriving IP zones from requests. +- **WAI Middleware**: Integrates seamlessly as a middleware into any WAI application. +- **Convenient and Customizable API**: + - Use declarative configuration for common scenarios with automatic setup. + - Or, for advanced use, fully control cache key structure and throttling logic. +- **Memory-efficient**: Designed for large-scale, high-traffic deployments with automatic cleanup of expired entries and HashMap-based O(1) average-case lookups. +- **Easy Integration**: Minimal code changes are required to get started. ## Why Use This Plugin? - - **Scalability**: Per-zone caches with HashMap-based storage and flexible throttling allow you to scale from single-user apps to multi-tenant platforms. - - **Performance**: The in-memory backend is built on efficient STM-based containers with HashMap optimizations for high-concurrency workloads. - - **Security**: Protects your application from abusive clients and denial-of-service attacks. - - **Flexibility**: Choose between the convenience of wrappers and the full customizability of manual key management. - - **Production-Ready**: Inspired by industry-standard tools, thoroughly documented, and designed for reliability with efficient data structures. - - **Open Source**: MIT licensed and community-friendly. +- **Scalability**: Per-zone caches with HashMap-based storage and flexible throttling allow you to scale from single-user apps to multi-tenant platforms. +- **Performance**: The in-memory backend is built on efficient STM-based containers with HashMap optimizations for high-concurrency workloads. +- **Security**: Protects your application from abusive clients and denial-of-service attacks. +- **Flexibility**: Choose between declarative configuration and full programmatic customization. +- **Production-Ready**: Inspired by industry-standard tools, thoroughly documented, and designed for reliability with efficient data structures. +- **Open Source**: MIT licensed and community-friendly. ## Installation @@ -52,18 +54,18 @@ ## Quick Start -The following example sets up a simple WAI application with a single rate-limiting rule: 10 requests per 60 seconds from a given IP address. It also demonstrates assigning requests to different IP zones using efficient HashMap-based zone resolution. +### Declarative Configuration (Recommended) +The recommended approach uses declarative configuration that can be loaded from JSON or YAML files: + ```haskell {-# LANGUAGE OverloadedStrings #-} import Keter.RateLimiter.WAI import Keter.RateLimiter.Cache (Algorithm(..)) -import Keter.RateLimiter.IPZones (defaultIPZone) -import Network.Wai (Request, responseLBS, Application) +import Network.Wai (responseLBS, Application) import Network.HTTP.Types (status200) import Network.Wai.Handler.Warp (run) -import Data.Text.Encoding (encodeUtf8) -- A simple application that runs behind the middleware. myApp :: Application @@ -71,32 +73,124 @@ main :: IO () main = do - -- 1. Initialize the rate limiter configuration. - -- This function determines the IP Zone for each request. - -- Here, we route traffic from "127.0.0.1" to "local_zone" and all other traffic to the default zone. - -- The environment uses HashMap-based storage for efficient zone cache lookups. - env <- initConfig (\req -> if requestHeaderHost req == Just "127.0.0.1" then "local_zone" else defaultIPZone) + -- 1. Define declarative configuration + let config = RateLimiterConfig + { rlZoneBy = ZoneIP -- Separate zones by client IP + , rlThrottles = + [ RLThrottle "api" 100 3600 FixedWindow IdIP Nothing -- 100 requests/hour by IP + , RLThrottle "login" 5 300 TokenBucket IdIP (Just 600) -- 5 login attempts/5min by IP with 10min idle timeout + ] + } - -- 2. Define a throttle rule. - let ipThrottle = ThrottleConfig - { throttleLimit = 10 - , throttlePeriod = 60 + -- 2. Build middleware from configuration + middleware <- buildRateLimiter config + + -- 3. Apply middleware to your application + let appWithMiddleware = middleware myApp + + putStrLn "Server starting on port 8080..." + run 8080 appWithMiddleware +``` + +### JSON Configuration + +You can also load configuration from JSON files: + +```json +{ + "zone_by": "ip", + "throttles": [ + { + "name": "api", + "limit": 100, + "period": 3600, + "algorithm": "fixed_window", + "identifier_by": "ip" + }, + { + "name": "login", + "limit": 5, + "period": 300, + "algorithm": "token_bucket", + "identifier_by": "ip", + "token_bucket_ttl": 600 + } + ] +} +``` + +### Advanced Programmatic Configuration + +For more control, you can build the environment programmatically: + +```haskell +import Keter.RateLimiter.WAI +import Keter.RateLimiter.Cache (Algorithm(..)) +import Keter.RateLimiter.IPZones (defaultIPZone) +import Data.Text.Encoding (encodeUtf8) + +main :: IO () +main = do + -- 1. Initialize environment with custom zone logic + env <- initConfig $ \req -> + case requestHeaderHost req of + Just "api.example.com" -> "api_zone" + Just "admin.example.com" -> "admin_zone" + _ -> defaultIPZone + + -- 2. Add throttle configurations + let apiThrottle = ThrottleConfig + { throttleLimit = 1000 + , throttlePeriod = 3600 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> fmap (encodeUtf8 . show) (remoteHost req) -- Identify requests by IP address - , throttleTokenBucketTTL = Nothing -- Not used for FixedWindow + , throttleIdentifier = \req -> fmap (encodeUtf8 . show) (remoteHost req) + , throttleTokenBucketTTL = Nothing } - -- 3. Add the throttle rule to the environment. - -- The throttle is stored in a HashMap for O(1) average-case retrieval. - env' <- addThrottle env "req/ip" ipThrottle + let loginThrottle = ThrottleConfig + { throttleLimit = 5 + , throttlePeriod = 300 + , throttleAlgorithm = TokenBucket + , throttleIdentifier = \req -> fmap (encodeUtf8 . show) (remoteHost req) + , throttleTokenBucketTTL = Just 600 + } - -- 4. Wrap your application with the middleware. - let appWithMiddleware = attackMiddleware env' myApp + env' <- addThrottle env "api" apiThrottle + env'' <- addThrottle env' "login" loginThrottle + -- 3. Create middleware + let middleware = buildRateLimiterWithEnv env'' + appWithMiddleware = middleware myApp + putStrLn "Server starting on port 8080..." run 8080 appWithMiddleware ``` +## Configuration Reference + +### Client Identification Strategies (`IdentifierBy`) + +- `"ip"` - Identify by client IP address +- `"ip+path"` - Identify by IP address and request path +- `"ip+ua"` - Identify by IP address and User-Agent header +- `{"header": "X-API-Key"}` - Identify by custom header value +- `{"cookie": "session_id"}` - Identify by cookie value +- `{"header+ip": "X-User-ID"}` - Identify by header value combined with IP + +### Zone Derivation Strategies (`ZoneBy`) + +- `"default"` - All requests use the same cache (no zone separation) +- `"ip"` - Separate zones by client IP address +- `{"header": "X-Tenant-ID"}` - Separate zones by custom header value + +### Rate Limiting Algorithms + +- **`FixedWindow`** - Traditional fixed-window counting +- **`SlidingWindow`** - Precise sliding-window with timestamp tracking +- **`TokenBucket`** - Allow bursts up to capacity, refill over time +- **`LeakyBucket`** - Smooth rate limiting with configurable leak rate +- **`TinyLRU`** - Least-recently-used eviction for memory efficiency + ## Example Usage ### Using the Convenient API @@ -135,24 +229,19 @@ mVal <- readCache cache customKey :: IO (Maybe Int) ``` -### Token Bucket Example (with TTL) - -The Token Bucket algorithm allows for bursts of traffic. You can also specify a TTL for how long an idle bucket remains in memory. The throttle configuration is efficiently stored and retrieved using HashMap-based lookups. +### Multi-Algorithm Configuration Example ```haskell -import Keter.RateLimiter.WAI -import Keter.RateLimiter.Cache (Algorithm(..)) - -let tokenBucketThrottle = ThrottleConfig - { throttleLimit = 100 -- Bucket capacity - , throttlePeriod = 60 -- Refills 100 tokens over 60 seconds - , throttleAlgorithm = TokenBucket - , throttleIdentifier = \req -> getAuthToken req -- A function to get a user's API token - , throttleTokenBucketTTL = Just 7200 -- Purge idle buckets after 2 hours +let config = RateLimiterConfig + { rlZoneBy = ZoneHeader "X-Tenant-ID" -- Separate by tenant + , rlThrottles = + [ RLThrottle "api_burst" 100 60 TokenBucket IdIP (Just 300) + , RLThrottle "api_sustained" 1000 3600 FixedWindow IdIP Nothing + , RLThrottle "login" 5 300 LeakyBucket IdIP Nothing + , RLThrottle "admin" 50 3600 SlidingWindow (IdHeader "X-Admin-Key") Nothing + , RLThrottle "lru_cache" 1000 60 TinyLRU IdIPAndPath Nothing + ] } - --- env' <- addThrottle env "api/token" tokenBucketThrottle --- The throttle will be stored in the environment's HashMap for efficient retrieval ``` ## Performance Characteristics @@ -183,17 +272,36 @@ ## When to Use This Library - - You need robust and efficient request throttling for your Haskell web application. - - You want to protect your service from abuse and DoS attacks. - - You require per-zone or per-user isolation of throttling policies with efficient lookups. - - You value both convenience and the ability to customize behavior as needed. - - You need high-performance rate limiting that can scale to handle large numbers of concurrent requests and zones. +- You need robust and efficient request throttling for your Haskell web application. +- You want to protect your service from abuse and DoS attacks. +- You require per-zone or per-user isolation of throttling policies with efficient lookups. +- You value both declarative configuration and the ability to customize behavior as needed. +- You need high-performance rate limiting that can scale to handle large numbers of concurrent requests and zones. +## Migration from Earlier Versions + +If you're upgrading from an earlier version that used the programmatic API, the declarative configuration approach is now recommended: + +**Old approach:** +```haskell +env <- initConfig getZoneFunction +env' <- addThrottle env "api" throttleConfig +let middleware = attackMiddleware env' +``` + +**New recommended approach:** +```haskell +let config = RateLimiterConfig { ... } +middleware <- buildRateLimiter config +``` + +The old programmatic API is still fully supported for advanced use cases. + ## License MIT License © 2025 Oleksandr Zhabenko ## References - - [rack-attack (Ruby)](https://github.com/rack/rack-attack) - - [keter (Haskell)](https://github.com/snoyberg/keter) +- [rack-attack (Ruby)](https://github.com/rack/rack-attack) +- [keter (Haskell)](https://github.com/snoyberg/keter) diff --git a/keter-rate-limiting-plugin.cabal b/keter-rate-limiting-plugin.cabal --- a/keter-rate-limiting-plugin.cabal +++ b/keter-rate-limiting-plugin.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: keter-rate-limiting-plugin -version: 0.1.1.0 +version: 0.1.2.0 synopsis: Simple Keter rate limiting plugin. description: A modern, high-performance, and highly customisable rate limiting plugin for keter. diff --git a/src/Keter/RateLimiter/WAI.hs b/src/Keter/RateLimiter/WAI.hs --- a/src/Keter/RateLimiter/WAI.hs +++ b/src/Keter/RateLimiter/WAI.hs @@ -5,10 +5,10 @@ {-| Module : Keter.RateLimiter.WAI -Description : WAI-compatible rate limiting middleware with IP zone support -Copyright : (c) 2025 Oleksandr Zhabenko +Description : WAI-compatible, plugin-friendly rate limiting middleware with IP-zone support License : MIT Maintainer : oleksandr.zhabenko@yahoo.com +Copyright : (c) 2025 Oleksandr Zhabenko Stability : stable Portability : portable @@ -19,25 +19,65 @@ Oleksandr Zhabenko added several implementations of the window algorithm: tinyLRU, sliding window, token bucket window, leaky bucket window alongside with the initial count algorithm using AI chatbots. IP Zone functionality added to allow separate caches per IP zone. -This implementation is released under the MIT License. +Overview +======== -This module provides WAI middleware that implements multiple IP-zone-specific -rate limiting strategies including: +This module provides WAI middleware for declarative, IP-zone-aware rate limiting with +multiple algorithms: -* Fixed Window -* Sliding Window -* Token Bucket -* Leaky Bucket -* TinyLRU +- Fixed Window +- Sliding Window +- Token Bucket +- Leaky Bucket +- TinyLRU -It supports multiple throttles applied simultaneously (e.g., per endpoint or -user group), each with its own algorithm, rate, and identifier logic. +Key points +---------- -Rate limiting state is tracked separately per IP zone using the -'ZoneSpecificCaches' abstraction with efficient HashMap-based lookups. +- Plugin-friendly construction: build an environment once (Env) from 'RateLimiterConfig' + and produce a pure WAI 'Middleware'. This matches common WAI patterns and avoids + per-request setup or global mutable state. -Use 'attackMiddleware' to enforce throttling in your WAI application. +- Concurrency model: all shared structures inside 'Env' use STM 'TVar', not 'IORef'. + This ensures thread-safe updates under GHC's lightweight (green) threads. +- Zone-specific caches: per-IP-zone caches are stored in a HashMap keyed by zone + identifiers. Zones are derived from a configurable strategy ('ZoneBy'), with a default. + +- No global caches in Keter: you can build one Env per compiled middleware chain + and cache that chain externally (e.g., per-vhost + middleware-list), preserving + counters/windows across requests. + +Quick start +----------- + +1) Declarative configuration (e.g., parsed from JSON/YAML): + +@ +let cfg = RateLimiterConfig + { rlZoneBy = ZoneDefault + , rlThrottles = + [ RLThrottle "api" 1000 3600 FixedWindow IdIP Nothing + , RLThrottle "login" 5 300 TokenBucket IdIP (Just 600) + ] + } +@ + +2) Build Env once and obtain a pure Middleware: + +@ +env <- buildEnvFromConfig cfg +let mw = buildRateLimiterWithEnv env +app = mw baseApplication +@ + +Alternatively: + +@ +mw <- buildRateLimiter cfg -- convenience: Env creation + Middleware +app = mw baseApplication +@ + -} module Keter.RateLimiter.WAI @@ -50,13 +90,18 @@ , RateLimiterConfig(..) , initConfig , addThrottle + -- * Middleware - , attackMiddleware + , attackMiddleware -- low-level: apply throttling with an existing Env + , buildRateLimiter -- convenience: build Env from config, return Middleware + , buildRateLimiterWithEnv -- preferred: pure Middleware from a pre-built Env + , buildEnvFromConfig -- build Env once from RateLimiterConfig + -- * Manual Control & Inspection , instrument , cacheResetAll - -- * Functions for Middleware - , buildRateLimiter + + -- * Helpers for configuration , registerThrottle , mkIdentifier , mkZoneFn @@ -79,7 +124,6 @@ import GHC.Generics import qualified Data.ByteString as S import qualified Data.ByteString.Lazy as LBS -import Data.IORef import Control.Concurrent.STM import Keter.RateLimiter.Cache import Keter.RateLimiter.IPZones (IPZoneIdentifier, defaultIPZone, ZoneSpecificCaches(..), createZoneCaches) @@ -95,213 +139,149 @@ import qualified Web.Cookie as WC -------------------------------------------------------------------------------- +-- Configuration and Environment --- | Configuration for a single throttle rule. +-- | Runtime throttle parameters assembled from declarative configuration. -- --- Defines the parameters for a specific rate limiting rule including --- the algorithm to use, rate limits, and identifier extraction logic. +-- See 'RLThrottle' for the declarative counterpart. data ThrottleConfig = ThrottleConfig { throttleLimit :: Int - -- ^ Maximum number of requests allowed per period. + -- ^ Maximum allowed requests per period. , throttlePeriod :: Int - -- ^ Period over which requests are counted (in seconds). + -- ^ Period length in seconds. , throttleAlgorithm :: Algorithm - -- ^ Algorithm used for throttling (e.g., 'FixedWindow', 'LeakyBucket'). + -- ^ Throttling algorithm to use. , throttleIdentifier :: Request -> IO (Maybe Text) - -- ^ Function to extract an identifier (e.g., user or IP) from the request. - -- Returns 'Nothing' if throttling should be skipped. + -- ^ Extracts an identifier (e.g., user, IP). Nothing => skip this throttle. , throttleTokenBucketTTL :: Maybe Int - -- ^ Optional TTL (seconds) for TokenBucket entries. If not provided, defaults to 2. + -- ^ Optional TTL (seconds) for TokenBucket entries. } --- | Global environment for throttling, including zone-specific cache HashMap and throttle configurations. --- --- The environment uses efficient HashMap-based lookups for both IP zone caches and --- throttle configurations, providing O(1) average-case performance for cache and --- throttle rule retrieval. --- --- === HashMap Usage --- --- * /envZoneCachesMap/: HashMap from IP zones to their corresponding rate limiter caches --- * /envThrottles/: HashMap of throttle rules indexed by their unique names --- --- === Concurrency Model +-- | Thread-safe, shared state for rate limiting. -- --- Both HashMaps are wrapped in 'IORef' for thread-safe atomic updates while --- maintaining efficient read access patterns typical in web middleware. +-- Concurrency model +-- ----------------- +-- - Uses 'TVar' from STM for in-memory HashMaps. +-- - Safe for green-threaded request handlers. +-- - No global variables: construct 'Env' in your wiring/bootstrap and reuse it. data Env = Env - { envZoneCachesMap :: IORef (HM.HashMap IPZoneIdentifier ZoneSpecificCaches) - -- ^ HashMap from IP zones to rate limiter caches for O(1) zone cache lookup. - , envThrottles :: IORef (HM.HashMap Text ThrottleConfig) - -- ^ HashMap of registered throttle rules by name for O(1) throttle rule retrieval. + { envZoneCachesMap :: TVar (HM.HashMap IPZoneIdentifier ZoneSpecificCaches) + -- ^ Per-zone caches for all algorithms. + , envThrottles :: TVar (HM.HashMap Text ThrottleConfig) + -- ^ Named throttle configurations. , envGetRequestIPZone :: Request -> IPZoneIdentifier - -- ^ Function to derive an IP zone from the incoming WAI request. + -- ^ Function deriving the IP zone for a given request. } --------------------------------------------------------------------------------- - --- | Initialize the rate-limiter environment with a default zone and no throttles. --- --- Creates a fresh environment with: --- --- * Empty HashMap for zone-specific caches (initialized with default zone only) --- * Empty HashMap for throttle configurations --- * Custom zone derivation function --- --- === Performance Characteristics +-- | Initialize an empty environment with a zone-derivation function. -- --- * Initial HashMap creation: O(1) --- * Default zone cache creation: O(1) --- * Thread-safe IORef initialization: O(1) +-- Populates the default zone lazily as needed; a default cache is allocated +-- immediately for the default zone to keep fast-path lookups cheap. initConfig - :: (Request -> IPZoneIdentifier) - -- ^ Function to extract an IP zone label from the request. + :: (Request -> IPZoneIdentifier) -- ^ Request -> zone label -> IO Env initConfig getIPZone = do defaultCaches <- createZoneCaches - zoneCachesMap <- newIORef $ HM.singleton defaultIPZone defaultCaches - throttles <- newIORef HM.empty - return $ Env zoneCachesMap throttles getIPZone + zoneCachesMap <- newTVarIO $ HM.singleton defaultIPZone defaultCaches + throttles <- newTVarIO HM.empty + pure $ Env zoneCachesMap throttles getIPZone --- | Register a new named throttle rule in the environment. --- --- Adds a throttle configuration to the environment's HashMap-based registry. --- If a throttle with the same name already exists, it will be replaced. --- --- === HashMap Update +-- | Add or replace a named throttle configuration. -- --- * Average case: O(1) insertion into throttles HashMap --- * Atomic update via 'modifyIORef'' for thread safety --- * Returns updated environment for method chaining +-- STM-backed insertion for concurrency safety. addThrottle :: Env - -- ^ Environment to update -> Text - -- ^ Throttle rule name (must be unique) -> ThrottleConfig - -- ^ Throttle configuration -> IO Env - -- ^ Updated environment addThrottle env name config = do - modifyIORef' (envThrottles env) $ HM.insert name config - return env + atomically $ modifyTVar' (envThrottles env) $ HM.insert name config + pure env --- | WAI middleware that applies registered throttles from the environment. --- --- If any throttle blocks the request, a 429 response is returned. --- Otherwise, the request is forwarded to the application. --- --- ==== __Examples__ +-------------------------------------------------------------------------------- +-- Middleware (application of throttles) + +-- | Low-level middleware: apply throttling using an existing 'Env'. -- --- @ --- -- Basic setup with IP-based rate limiting --- env <- initConfig (const defaultIPZone) --- let throttleConfig = ThrottleConfig --- { throttleLimit = 100 --- , throttlePeriod = 3600 -- 1 hour --- , throttleAlgorithm = FixedWindow --- , throttleIdentifier = RU.byIP --- , throttleTokenBucketTTL = Nothing --- } --- env' <- addThrottle env \"api-limit\" throttleConfig --- --- -- Use as WAI middleware --- app = attackMiddleware env' baseApplication --- @ +-- If any throttle denies the request, a 429 response is returned. +-- Otherwise, 'app' is invoked. attackMiddleware :: Env - -- ^ Rate limiting environment -> Application - -- ^ Base WAI application -> Application - -- ^ Rate-limited WAI application attackMiddleware env app req respond = do blocked <- instrument env req if blocked then respond $ responseLBS status429 [] (LBS.fromStrict $ TE.encodeUtf8 "Too Many Requests") else app req respond --- | Inspect all throttle rules and apply them to the incoming request. --- --- Returns 'True' if any rule blocks the request. --- --- === HashMap Processing +-- | Inspect all active throttles in 'Env' for the given request. -- --- * O(1) throttles HashMap lookup via 'readIORef' --- * O(n) processing of all throttle rules where n = number of throttles --- * O(1) average case zone cache lookup in IP zone HashMap +-- Returns True if the request should be blocked under any rule. instrument :: Env - -- ^ Rate limiting environment -> Request - -- ^ Incoming WAI request -> IO Bool - -- ^ 'True' if request should be blocked instrument env req = do - throttles <- readIORef (envThrottles env) + throttles <- readTVarIO (envThrottles env) let zone = envGetRequestIPZone env req zoneCaches <- getOrCreateZoneCaches env zone - -- Check all throttles and collect their block/allow decisions, passing throttle names - anyBlocked <- or <$> mapM (\(name, config) -> checkThrottle zoneCaches zone req name config) (HM.toList throttles) - return anyBlocked + or <$> mapM (\(name, cfg) -> checkThrottle zoneCaches zone req name cfg) (HM.toList throttles) --- | Internal function to check a single throttle rule against a request. --- --- This function is called from 'instrument' and has access to throttle names --- for proper cache key construction. +-- | Check an individual throttle against a request. checkThrottle :: ZoneSpecificCaches -> Text -> Request -> Text -> ThrottleConfig -> IO Bool -checkThrottle caches zone req throttleName config = do - mIdentifier <- throttleIdentifier config req +checkThrottle caches zone req throttleName cfg = do + mIdentifier <- throttleIdentifier cfg req case mIdentifier of - Nothing -> return False - Just identifier -> case throttleAlgorithm config of + Nothing -> pure False + Just ident -> case throttleAlgorithm cfg of FixedWindow -> - -- allowFixedWindowRequest cache throttleName ipZone userKey limit period + -- allowFixedWindowRequest cache throttleName zone ident limit period not <$> allowFixedWindowRequest - (zscCounterCache caches) - throttleName - zone - identifier - (throttleLimit config) - (throttlePeriod config) + (zscCounterCache caches) + throttleName + zone + ident + (throttleLimit cfg) + (throttlePeriod cfg) SlidingWindow -> case zscTimestampCache caches of Cache { cacheStore = TimestampStore tvar } -> - -- SlidingWindow.allowRequest getTimeNow stmMapTVar throttleName ipZone userKey windowSize limit + -- SlidingWindow.allowRequest timeNow tvar throttleName zone ident window limit not <$> SlidingWindow.allowRequest (realToFrac <$> getPOSIXTime) tvar throttleName zone - identifier - (throttlePeriod config) - (throttleLimit config) + ident + (throttlePeriod cfg) + (throttleLimit cfg) TokenBucket -> do - let period = throttlePeriod config - limit = throttleLimit config + let period = throttlePeriod cfg + limit = throttleLimit cfg refillRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0 - ttl = fromMaybe 2 (throttleTokenBucketTTL config) - -- TokenBucket.allowRequest cache throttleName ipZone userKey capacity refillRate expiresIn + ttl = fromMaybe 2 (throttleTokenBucketTTL cfg) + -- TokenBucket.allowRequest cache throttleName zone ident capacity refill expires not <$> TokenBucket.allowRequest (zscTokenBucketCache caches) throttleName zone - identifier + ident (fromIntegral limit) refillRate (fromIntegral ttl) LeakyBucket -> do - let period = throttlePeriod config - limit = throttleLimit config + let period = throttlePeriod cfg + limit = throttleLimit cfg leakRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0 - -- LeakyBucket.allowRequest cache throttleName ipZone userKey capacity leakRate + -- LeakyBucket.allowRequest cache throttleName zone ident capacity leakRate not <$> LeakyBucket.allowRequest (zscLeakyBucketCache caches) throttleName zone - identifier + ident (fromIntegral limit) leakRate @@ -310,117 +290,82 @@ case cacheStore (zscTinyLRUCache caches) of TinyLRUStore tvar -> do cache <- readTVarIO tvar - not <$> atomically (allowRequestTinyLRU now cache identifier (throttleLimit config) (throttlePeriod config)) + -- allowRequestTinyLRU now cache ident capacity periodSecs + not <$> atomically (allowRequestTinyLRU now cache ident (throttleLimit cfg) (throttlePeriod cfg)) --- | Reset all caches for all IP zones tracked in the environment. --- --- Iterates through the HashMap of zone caches and resets each zone's --- rate limiting state across all algorithms. --- --- === HashMap Processing +-- | Reset all caches across all known zones. -- --- * O(1) HashMap read via 'readIORef' --- * O(n) iteration over all zones where n = number of active zones --- * Each zone reset involves clearing multiple algorithm-specific caches +-- Useful in tests or administrative endpoints. cacheResetAll :: Env -> IO () cacheResetAll env = do - zoneCachesMap <- readIORef (envZoneCachesMap env) - mapM_ (resetZoneCaches . snd) (HM.toList zoneCachesMap) + zoneCachesMap <- readTVarIO (envZoneCachesMap env) + mapM_ (resetZone . snd) (HM.toList zoneCachesMap) where - resetZoneCaches :: ZoneSpecificCaches -> IO () - resetZoneCaches caches = do + resetZone :: ZoneSpecificCaches -> IO () + resetZone caches = do cacheReset (zscCounterCache caches) cacheReset (zscTimestampCache caches) cacheReset (zscTokenBucketCache caches) cacheReset (zscLeakyBucketCache caches) cacheReset (zscTinyLRUCache caches) --- | Retrieve or create the set of caches for a specific IP zone. --- --- Uses HashMap-based zone cache lookup with atomic fallback creation --- for new zones. Ensures thread-safe zone cache initialization. --- --- === HashMap Operations +-- | Retrieve or create caches for a given IP zone. -- --- * O(1) average case lookup in zone caches HashMap --- * O(1) average case insertion for new zones via 'insertWith' --- * Atomic 'modifyIORef'' prevents race conditions during zone creation --- * Uses 'insertWith' with preference to existing caches (first writer wins) +-- Ensures a single writer initializes a new zone; readers see either the +-- existing or newly-inserted caches. getOrCreateZoneCaches :: Env - -- ^ Environment containing zone cache HashMap -> IPZoneIdentifier - -- ^ IP zone identifier -> IO ZoneSpecificCaches - -- ^ Zone-specific rate limiting caches getOrCreateZoneCaches env zone = do - readIORef (envZoneCachesMap env) >>= \m -> - case HM.lookup zone m of - Just caches -> return caches - Nothing -> do - newCaches <- createZoneCaches - atomicModifyIORef' (envZoneCachesMap env) $ \currentMap -> - let updatedMap = HM.insertWith (\_ old -> old) zone newCaches currentMap - in (updatedMap, updatedMap HM.! zone) - + m <- readTVarIO (envZoneCachesMap env) + case HM.lookup zone m of + Just caches -> pure caches + Nothing -> do + newCaches <- createZoneCaches + atomically $ do + m0 <- readTVar (envZoneCachesMap env) + case HM.lookup zone m0 of + Just existing -> pure existing + Nothing -> do + let m1 = HM.insert zone newCaches m0 + writeTVar (envZoneCachesMap env) m1 + pure newCaches -------------------------------------------------------------------------------- --- * Configuration Data Types +-- Declarative configuration types --- | Specification for extracting client identifiers from requests. --- --- Determines how clients are identified for rate limiting purposes. --- Different strategies allow grouping by IP, headers, cookies, or --- combinations thereof. --- --- @since 0.1.1.0 +-- | How to identify clients for throttling. data IdentifierBy - = IdIP -- ^ Group by client IP address - | IdHeader !HeaderName -- ^ Group by specific HTTP header value - | IdCookie !Text -- ^ Group by cookie value from Cookie header - | IdIPAndPath -- ^ Group by IP address + request path - | IdIPAndUA -- ^ Group by IP address + User-Agent header - | IdHeaderAndIP !HeaderName -- ^ Group by header value + IP address + = IdIP + | IdHeader !HeaderName + | IdCookie !Text + | IdIPAndPath + | IdIPAndUA + | IdHeaderAndIP !HeaderName deriving (Show, Eq, Generic) --- | Specification for how to derive IP zones from requests. --- --- IP zones allow rate limiting caches to be separated by logical or --- geographic boundaries. Each zone maintains independent rate limiting --- state using separate HashMap entries, enabling different policies for --- different request sources. --- --- @since 0.1.1.0 +-- | How to derive IP zones from requests. data ZoneBy - = ZoneDefault -- ^ All requests use default zone - | ZoneIP -- ^ Derive zone from client IP - | ZoneHeader !HeaderName -- ^ Derive zone from header value + = ZoneDefault + | ZoneIP + | ZoneHeader !HeaderName deriving (Show, Eq, Generic) --- | Complete specification for a single rate limiting rule. --- --- Combines algorithm choice, rate parameters, and client identification --- strategy into a single declarative configuration. --- --- @since 0.1.1.0 +-- | Declarative throttle rule (parsed from JSON/YAML). data RLThrottle = RLThrottle - { rlName :: !Text -- ^ Unique throttle identifier - , rlLimit :: !Int -- ^ Maximum requests per period - , rlPeriod :: !Int -- ^ Time period in seconds - , rlAlgo :: !Algorithm -- ^ Rate limiting algorithm - , rlIdBy :: !IdentifierBy -- ^ Request identifier strategy - , rlTokenBucketTTL :: !(Maybe Int) -- ^ TTL for TokenBucket (seconds) + { rlName :: !Text + , rlLimit :: !Int + , rlPeriod :: !Int + , rlAlgo :: !Algorithm + , rlIdBy :: !IdentifierBy + , rlTokenBucketTTL :: !(Maybe Int) } deriving (Show, Eq, Generic) --- | Complete rate limiter configuration combining zone strategy and throttles. --- --- Top-level configuration structure that defines both how to partition --- requests into zones and what rate limiting rules to apply. --- --- @since 0.1.1.0 +-- | Top-level configuration: zone strategy and throttle rules. data RateLimiterConfig = RateLimiterConfig - { rlZoneBy :: !ZoneBy -- ^ Zone derivation strategy - , rlThrottles :: ![RLThrottle] -- ^ List of throttle rules + { rlZoneBy :: !ZoneBy + , rlThrottles :: ![RLThrottle] } deriving (Show, Eq, Generic) instance FromJSON IdentifierBy where @@ -432,7 +377,7 @@ , IdCookie <$> o .: "cookie" , IdHeaderAndIP . hdr <$> o .: "header+ip" ] - parseJSON _ = fail "identifier_by must be one of: 'ip' | 'ip+path' | 'ip+ua' | {header: ...} | {cookie: ...} | {header+ip: ...}" + parseJSON _ = fail "identifier_by: 'ip' | 'ip+path' | 'ip+ua' | {header} | {cookie} | {header+ip}" instance ToJSON IdentifierBy where toJSON IdIP = String "ip" @@ -446,7 +391,7 @@ parseJSON (String "default") = pure ZoneDefault parseJSON (String "ip") = pure ZoneIP parseJSON (Object o) = ZoneHeader . hdr <$> o .: "header" - parseJSON _ = fail "zone_by must be 'default' | 'ip' | {header: ...}" + parseJSON _ = fail "zone_by: 'default' | 'ip' | {header}" instance ToJSON ZoneBy where toJSON ZoneDefault = String "default" @@ -481,53 +426,35 @@ object [ "zone_by" .= zb, "throttles" .= ths ] -------------------------------------------------------------------------------- --- * Functions for Middleware +-- Public builders (preferred wiring API) --- | Convert Text header name to WAI HeaderName. +-- | Build 'Env' once from a declarative 'RateLimiterConfig'. -- --- @since 0.1.1.0 -hdr :: Text -> HeaderName -hdr = mk . TE.encodeUtf8 +-- Use this at wiring time; the returned 'Env' is stable and reused across requests. +buildEnvFromConfig :: RateLimiterConfig -> IO Env +buildEnvFromConfig (RateLimiterConfig zb ths) = do + let zoneFn = mkZoneFn zb + env <- initConfig zoneFn + mapM_ (registerThrottle env) ths + pure env --- | Extract the original ByteString from a case-insensitive HeaderName. +-- | Produce a pure 'Middleware' from an existing 'Env'. -- --- @since 0.1.1.0 -fromHeaderName :: HeaderName -> S.ByteString -fromHeaderName = original +-- This is the recommended way to integrate with WAI/Keter: the middleware is +-- a pure function, while the state is already encapsulated in 'Env'. +buildRateLimiterWithEnv :: Env -> Middleware +buildRateLimiterWithEnv = attackMiddleware --- | Build a complete WAI middleware from a rate limiter configuration. --- --- This is the primary entry point for declarative rate limiter setup. --- It creates an environment, registers all throttles, and returns --- ready-to-use WAI middleware. --- --- ==== __Examples__ --- --- @ --- -- JSON-driven configuration --- let config = RateLimiterConfig --- { rlZoneBy = ZoneDefault --- , rlThrottles = --- [ RLThrottle \"api\" 1000 3600 FixedWindow IdIP Nothing --- , RLThrottle \"login\" 5 300 TokenBucket IdIP (Just 600) --- ] --- } --- --- middleware <- buildRateLimiter config --- app = middleware baseApplication --- @ +-- | Convenience: build an 'Env' from config and return the 'Middleware'. -- --- @since 0.1.1.0 +-- Suitable if you don’t need to retain the 'Env' for administrative operations. buildRateLimiter :: RateLimiterConfig -> IO Middleware -buildRateLimiter (RateLimiterConfig zb ths) = do - let zoneFn = mkZoneFn zb - env <- initConfig zoneFn - mapM_ (registerThrottle env) ths - pure (attackMiddleware env) +buildRateLimiter cfg = buildRateLimiterWithEnv <$> buildEnvFromConfig cfg --- | Register a single throttle rule in an existing environment. --- --- @since 0.1.1.0 +-------------------------------------------------------------------------------- +-- Helper functions for configuration + +-- | Register a single throttle rule into an 'Env'. registerThrottle :: Env -> RLThrottle -> IO () registerThrottle env (RLThrottle name l p algo idBy ttl) = do let cfg = ThrottleConfig @@ -540,9 +467,7 @@ _ <- addThrottle env name cfg pure () --- | Create an identifier extraction function from declarative specification. --- --- @since 0.1.1.0 +-- | Build a request-identifier function from a declarative spec. mkIdentifier :: IdentifierBy -> Request -> IO (Maybe Text) mkIdentifier IdIP = RU.byIP mkIdentifier IdIPAndPath = RU.byIPAndPath @@ -551,31 +476,22 @@ mkIdentifier (IdCookie name) = \req -> pure $ cookieLookupText name req mkIdentifier (IdHeaderAndIP h) = RU.byHeaderAndIP h --- | Internal helper: cookie lookup via Web.Cookie with empty-value rejection (matches previous semantics). +-- | Cookie lookup via Web.Cookie; ignores empty values. cookieLookupText :: Text -> Request -> Maybe Text cookieLookupText n req = do raw <- lookup hCookie (requestHeaders req) let pairs = WC.parseCookies raw v <- lookup (TE.encodeUtf8 n) pairs - if S.null v - then Nothing - else Just (TE.decodeUtf8With TEE.lenientDecode v) + if S.null v then Nothing else Just (TE.decodeUtf8With TEE.lenientDecode v) --- | Create a zone derivation function from declarative specification. --- --- @since 0.1.1.0 +-- | Derive IP zone function from a declarative spec. mkZoneFn :: ZoneBy -> (Request -> IPZoneIdentifier) mkZoneFn ZoneDefault = const defaultIPZone mkZoneFn ZoneIP = getClientIPPure mkZoneFn (ZoneHeader h) = \req -> maybe defaultIPZone (TE.decodeUtf8With TEE.lenientDecode) (lookup h (requestHeaders req)) --- | Extract client IP address from WAI request with header precedence. --- --- Checks headers in order of precedence: X-Forwarded-For, X-Real-IP, --- then falls back to socket address. Handles IPv4, IPv6, and Unix sockets. --- --- @since 0.1.1.0 +-- | Extract client IP with header precedence: X-Forwarded-For, X-Real-IP, then socket. getClientIPPure :: Request -> IPZoneIdentifier getClientIPPure req = let safeDecode = TE.decodeUtf8With TEE.lenientDecode @@ -589,3 +505,11 @@ SockAddrInet _ addr -> RU.ipv4ToString addr SockAddrInet6 _ _ addr _ -> RU.ipv6ToString addr SockAddrUnix path -> Tx.pack path + +-- | Construct a case-insensitive header name from Text. +hdr :: Text -> HeaderName +hdr = mk . TE.encodeUtf8 + +-- | Extract original bytes from a case-insensitive header name. +fromHeaderName :: HeaderName -> S.ByteString +fromHeaderName = original diff --git a/test/Keter/RateLimiter/WAITests.hs b/test/Keter/RateLimiter/WAITests.hs --- a/test/Keter/RateLimiter/WAITests.hs +++ b/test/Keter/RateLimiter/WAITests.hs @@ -55,6 +55,7 @@ import Network.HTTP.Types import Network.Socket (SockAddr(..), tupleToHostAddress) import Data.Text (Text) +import Control.Concurrent.STM (readTVarIO) import qualified Data.Text as T import qualified Data.Text.Encoding as TE import qualified Data.ByteString as BS @@ -646,11 +647,11 @@ env' <- addThrottle env "test" throttle let zone1Req = mkRequestWithHeader "X-Zone" "zone1" let zone2Req = mkRequestWithHeader "X-Zone" "zone2" - zoneCaches <- readIORef (envZoneCachesMap env) + zoneCaches <- readTVarIO (envZoneCachesMap env) initialSize <- return $ HM.size zoneCaches _ <- instrument env' zone1Req _ <- instrument env' zone2Req - zoneCaches' <- readIORef (envZoneCachesMap env') + zoneCaches' <- readTVarIO (envZoneCachesMap env') finalSize <- return $ HM.size zoneCaches' assertBool "New zones created" (finalSize > initialSize) @@ -771,7 +772,7 @@ return [ra1, rb1] responses <- runSession session app assertEqual "Both zones populated" [status200, status200] (map simpleStatus responses) - zoneCaches <- readIORef (envZoneCachesMap env') + zoneCaches <- readTVarIO (envZoneCachesMap env') let zoneCount = HM.size zoneCaches assertBool "Multiple zones created" (zoneCount >= 2) @@ -892,7 +893,7 @@ let makeZoneRequest :: Integer -> Request makeZoneRequest i = mkRequestWithHeader "X-Zone-ID" (T.pack $ "zone" <> show i) mapM_ (\i -> instrument env' (makeZoneRequest i)) [1..50 :: Integer] - zoneCaches <- readIORef (envZoneCachesMap env') + zoneCaches <- readTVarIO (envZoneCachesMap env') let zoneCount = HM.size zoneCaches assertBool "Many zones created" (zoneCount > 10) assertBool "Reasonable zone count" (zoneCount <= 51) diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -59,7 +59,8 @@ import Control.Concurrent (threadDelay) import Control.Monad (when) import Data.Maybe (isJust) -import Data.IORef (readIORef) +--import Data.IORef (readIORef) +import Control.Concurrent.STM (readTVarIO) import Data.Cache (purgeExpired) import Test.Tasty (TestTree, defaultMain, testGroup, after, DependencyType(..)) import Test.Tasty.HUnit (testCase, assertEqual, assertFailure) @@ -412,7 +413,7 @@ -- Wait for period + buffer threadDelay 1_500_000 -- 1.5 seconds to ensure expiration -- Debug cache state - cachesMap <- readIORef (envZoneCachesMap envWithThrottle) + cachesMap <- readTVarIO (envZoneCachesMap envWithThrottle) zoneCaches <- case HashMap.lookup testIPZoneA cachesMap of Just caches -> return caches Nothing -> assertFailure "Zone caches not found" >> return undefined