diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,3 +20,7 @@ * 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. +## 0.2.0.0 -- 2025-08-24 + +* Second version. Changed Keter.RateLimiter.WAI module to remove multiple times computed data while working on the advice of @jappeace for keter integration and to prevent potential memory leak when keter reloads configuration. + diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -128,12 +128,13 @@ import Keter.RateLimiter.Cache (Algorithm(..)) import Keter.RateLimiter.IPZones (defaultIPZone) import Data.Text.Encoding (encodeUtf8) +import Network.HTTP.Types (hHost) main :: IO () main = do -- 1. Initialize environment with custom zone logic env <- initConfig $ \req -> - case requestHeaderHost req of + case lookup hHost (requestHeaders req) of Just "api.example.com" -> "api_zone" Just "admin.example.com" -> "admin_zone" _ -> defaultIPZone @@ -143,7 +144,7 @@ { throttleLimit = 1000 , throttlePeriod = 3600 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> fmap (encodeUtf8 . show) (remoteHost req) + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } @@ -151,7 +152,7 @@ { throttleLimit = 5 , throttlePeriod = 300 , throttleAlgorithm = TokenBucket - , throttleIdentifier = \req -> fmap (encodeUtf8 . show) (remoteHost req) + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 600 } @@ -170,18 +171,18 @@ ### 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 +- `IdIP` - Identify by client IP address +- `IdIPAndPath` - Identify by IP address and request path +- `IdIPAndUA` - Identify by IP address and User-Agent header +- `IdHeader headerName` - Identify by custom header value +- `IdCookie "session_id"` - Identify by cookie value +- `IdHeaderAndIP headerName` - 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 +- `ZoneDefault` - All requests use the same cache (no zone separation) +- `ZoneIP` - Separate zones by client IP address +- `ZoneHeader headerName` - Separate zones by custom header value ### Rate Limiting Algorithms @@ -233,12 +234,12 @@ ```haskell let config = RateLimiterConfig - { rlZoneBy = ZoneHeader "X-Tenant-ID" -- Separate by tenant + { rlZoneBy = ZoneHeader (hdr "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 "admin" 50 3600 SlidingWindow (IdHeader (hdr "X-Admin-Key")) Nothing , RLThrottle "lru_cache" 1000 60 TinyLRU IdIPAndPath Nothing ] } @@ -295,7 +296,7 @@ middleware <- buildRateLimiter config ``` -The old programmatic API is still fully supported for advanced use cases. +The old programmatic API is still fully supported for advanced use cases via `buildRateLimiterWithEnv` and related functions. ## License 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.2.0 +version: 0.2.0.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 @@ -2,13 +2,14 @@ {-# LANGUAGE TypeApplications #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE TupleSections #-} +{-# LANGUAGE DeriveGeneric #-} {-| Module : Keter.RateLimiter.WAI Description : WAI-compatible, plugin-friendly rate limiting middleware with IP-zone support License : MIT Maintainer : oleksandr.zhabenko@yahoo.com -Copyright : (c) 2025 Oleksandr Zhabenko +Copyright : (c) 2025 Oleksandr Zhabenko Stability : stable Portability : portable @@ -34,7 +35,7 @@ Key points ---------- -- Plugin-friendly construction: build an environment once (Env) from 'RateLimiterConfig' +- 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. @@ -44,7 +45,7 @@ - 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 +- 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. @@ -63,7 +64,7 @@ } @ -2) Build Env once and obtain a pure Middleware: +2) Build 'Env' once and obtain a pure 'Middleware': @ env <- buildEnvFromConfig cfg @@ -78,6 +79,74 @@ app = mw baseApplication @ +Usage patterns +-------------- + +__Declarative approach (recommended):__ + +@ +import Keter.RateLimiter.WAI +import Keter.RateLimiter.Cache (Algorithm(..)) + +main = do + let config = RateLimiterConfig + { rlZoneBy = ZoneIP + , rlThrottles = + [ RLThrottle "api" 100 3600 FixedWindow IdIP Nothing + ] + } + middleware <- buildRateLimiter config + let app = middleware baseApp + run 8080 app +@ + +__Programmatic approach (advanced):__ + +@ +import Keter.RateLimiter.WAI +import Keter.RateLimiter.Cache (Algorithm(..)) + +main = do + env <- initConfig (\\req -> "zone1") + let throttleConfig = ThrottleConfig + { throttleLimit = 100 + , throttlePeriod = 3600 + , throttleAlgorithm = FixedWindow + , throttleIdentifierBy = IdIP + , throttleTokenBucketTTL = Nothing + } + env' <- addThrottle env "api" throttleConfig + let middleware = buildRateLimiterWithEnv env' + app = middleware baseApp + run 8080 app +@ + +Configuration reference +----------------------- + +__Client identification strategies ('IdentifierBy'):__ + +- 'IdIP' - Identify by client IP address +- 'IdIPAndPath' - Identify by IP address and request path +- 'IdIPAndUA' - Identify by IP address and User-Agent header +- @'IdHeader' headerName@ - Identify by custom header value +- @'IdCookie' cookieName@ - Identify by cookie value +- @'IdHeaderAndIP' headerName@ - Identify by header value combined with IP + +__Zone derivation strategies ('ZoneBy'):__ + +- 'ZoneDefault' - All requests use the same cache (no zone separation) +- 'ZoneIP' - Separate zones by client IP address +- @'ZoneHeader' headerName@ - 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 + -} module Keter.RateLimiter.WAI @@ -92,10 +161,10 @@ , addThrottle -- * Middleware - , 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 + , 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 @@ -110,33 +179,42 @@ , fromHeaderName ) where +import Control.Concurrent.STM import Data.Aeson hiding (pairs) +import qualified Data.ByteString as S +import qualified Data.ByteString.Lazy as LBS +import Data.CaseInsensitive (mk, original) +import Data.Foldable (asum) +import Data.Hashable (Hashable(..)) +import qualified Data.HashMap.Strict as HM +import Data.Maybe (fromMaybe) import Data.Text (Text) import qualified Data.Text as Tx import qualified Data.Text.Encoding as TE import qualified Data.Text.Encoding.Error as TEE -import qualified Data.HashMap.Strict as HM -import Data.Foldable (asum) -import Network.Wai -import Network.HTTP.Types (status429, HeaderName, hCookie) -import Network.Socket (SockAddr(..)) -import Data.CaseInsensitive (mk, original) import GHC.Generics -import qualified Data.ByteString as S -import qualified Data.ByteString.Lazy as LBS -import Control.Concurrent.STM -import Keter.RateLimiter.Cache -import Keter.RateLimiter.IPZones (IPZoneIdentifier, defaultIPZone, ZoneSpecificCaches(..), createZoneCaches) -import qualified Keter.RateLimiter.SlidingWindow as SlidingWindow -import qualified Keter.RateLimiter.TokenBucket as TokenBucket -import qualified Keter.RateLimiter.LeakyBucket as LeakyBucket +import Network.HTTP.Types (HeaderName, hCookie, status429) +import Network.Socket (SockAddr (..)) +import Network.Wai +import qualified Web.Cookie as WC + +-- Import Cache with hiding Algorithm to avoid conflict, then import Algorithm explicitly +import Keter.RateLimiter.Cache hiding (Algorithm) +import Keter.RateLimiter.Cache (Algorithm(..)) import Keter.RateLimiter.CacheWithZone (allowFixedWindowRequest) +import Keter.RateLimiter.IPZones + ( IPZoneIdentifier + , ZoneSpecificCaches(..) + , createZoneCaches + , defaultIPZone + ) +import qualified Keter.RateLimiter.LeakyBucket as LeakyBucket import qualified Keter.RateLimiter.RequestUtils as RU +import qualified Keter.RateLimiter.SlidingWindow as SlidingWindow +import qualified Keter.RateLimiter.TokenBucket as TokenBucket import Data.TinyLRU (allowRequestTinyLRU) -import System.Clock (Clock(Monotonic), getTime) -import Data.Maybe (fromMaybe) +import System.Clock (Clock (Monotonic), getTime) import Data.Time.Clock.POSIX (getPOSIXTime) -import qualified Web.Cookie as WC -------------------------------------------------------------------------------- -- Configuration and Environment @@ -145,22 +223,25 @@ -- -- See 'RLThrottle' for the declarative counterpart. data ThrottleConfig = ThrottleConfig - { throttleLimit :: Int + { throttleLimit :: !Int -- ^ Maximum allowed requests per period. - , throttlePeriod :: Int + , throttlePeriod :: !Int -- ^ Period length in seconds. - , throttleAlgorithm :: Algorithm - -- ^ Throttling algorithm to use. - , throttleIdentifier :: Request -> IO (Maybe Text) - -- ^ Extracts an identifier (e.g., user, IP). Nothing => skip this throttle. - , throttleTokenBucketTTL :: Maybe Int + , throttleAlgorithm :: !Algorithm + -- ^ Which throttling algorithm to use. + , throttleIdentifierBy :: !IdentifierBy + -- ^ Declarative spec for extracting an identifier (e.g., IP, header, cookie). + -- At runtime we derive the extractor using 'mkIdentifier' and compute it + -- at most once per request per IdentifierBy group. If extraction yields + -- Nothing, this throttle does not apply to the request. + , throttleTokenBucketTTL :: !(Maybe Int) -- ^ Optional TTL (seconds) for TokenBucket entries. - } + } deriving (Show, Eq, Generic) -- | Thread-safe, shared state for rate limiting. -- --- Concurrency model --- ----------------- +-- = 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. @@ -212,87 +293,123 @@ attackMiddleware env app req respond = do blocked <- instrument env req if blocked - then respond $ responseLBS status429 [] (LBS.fromStrict $ TE.encodeUtf8 "Too Many Requests") + then respond $ responseLBS status429 [("Content-Type","text/plain; charset=utf-8")] + (LBS.fromStrict $ TE.encodeUtf8 "Too Many Requests") else app req respond -- | Inspect all active throttles in 'Env' for the given request. -- -- Returns True if the request should be blocked under any rule. -instrument - :: Env - -> Request - -> IO Bool +instrument :: Env -> Request -> IO Bool instrument env req = do throttles <- readTVarIO (envThrottles env) - let zone = envGetRequestIPZone env req - zoneCaches <- getOrCreateZoneCaches env zone - or <$> mapM (\(name, cfg) -> checkThrottle zoneCaches zone req name cfg) (HM.toList throttles) + if HM.null throttles + then pure False + else do + let zone = envGetRequestIPZone env req + caches <- getOrCreateZoneCaches env zone + let buckets = groupByIdentifier throttles + anyMHashMap + (\idBy group -> + case group of + [] -> pure False + ((_name0, _cfg0):_) -> do + -- Compute identifier once per IdentifierBy group + mIdent <- mkIdentifier idBy req + case mIdent of + Nothing -> pure False + Just ident -> + anyMList + (\(name, cfg) -> + checkThrottleWithIdent caches zone req name cfg (Just ident) + ) + group + ) + buckets --- | Check an individual throttle against a request. -checkThrottle :: ZoneSpecificCaches -> Text -> Request -> Text -> ThrottleConfig -> IO Bool -checkThrottle caches zone req throttleName cfg = do - mIdentifier <- throttleIdentifier cfg req +-- | Check an individual throttle with a precomputed identifier. +-- +-- True = block, False = allow. +checkThrottleWithIdent + :: ZoneSpecificCaches + -> Text -- ^ zone + -> Request + -> Text -- ^ throttle name + -> ThrottleConfig + -> Maybe Text -- ^ precomputed identifier + -> IO Bool +checkThrottleWithIdent caches zone _req throttleName cfg mIdentifier = case mIdentifier of - Nothing -> pure False - Just ident -> case throttleAlgorithm cfg of - FixedWindow -> - -- allowFixedWindowRequest cache throttleName zone ident limit period - not <$> allowFixedWindowRequest - (zscCounterCache caches) - throttleName - zone - ident - (throttleLimit cfg) - (throttlePeriod cfg) + Nothing -> pure False + Just ident -> + case throttleAlgorithm cfg of + -- Use unqualified Algorithm constructors since we imported them explicitly + FixedWindow -> + -- allowFixedWindowRequest cache throttleName zone ident limit period + not <$> allowFixedWindowRequest + (zscCounterCache caches) + throttleName + zone + ident + (throttleLimit cfg) + (throttlePeriod cfg) - SlidingWindow -> case zscTimestampCache caches of - Cache { cacheStore = TimestampStore tvar } -> - -- SlidingWindow.allowRequest timeNow tvar throttleName zone ident window limit - not <$> SlidingWindow.allowRequest - (realToFrac <$> getPOSIXTime) - tvar - throttleName - zone - ident - (throttlePeriod cfg) - (throttleLimit cfg) + SlidingWindow -> case zscTimestampCache caches of + Cache { cacheStore = TimestampStore tvar } -> + -- SlidingWindow.allowRequest timeNow tvar throttleName zone ident window limit + not <$> SlidingWindow.allowRequest + (realToFrac <$> getPOSIXTime) + tvar + throttleName + zone + ident + (throttlePeriod cfg) + (throttleLimit cfg) - TokenBucket -> do - let period = throttlePeriod cfg - limit = throttleLimit cfg - refillRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0 - ttl = fromMaybe 2 (throttleTokenBucketTTL cfg) - -- TokenBucket.allowRequest cache throttleName zone ident capacity refill expires - not <$> TokenBucket.allowRequest - (zscTokenBucketCache caches) - throttleName - zone - ident - (fromIntegral limit) - refillRate - (fromIntegral ttl) + TokenBucket -> do + let period = throttlePeriod cfg + limit = throttleLimit cfg + refillRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0 + ttl = fromMaybe 2 (throttleTokenBucketTTL cfg) + -- TokenBucket.allowRequest cache throttleName zone ident capacity refill expires + not <$> TokenBucket.allowRequest + (zscTokenBucketCache caches) + throttleName + zone + ident + (fromIntegral limit) + refillRate + (fromIntegral ttl) - LeakyBucket -> do - let period = throttlePeriod cfg - limit = throttleLimit cfg - leakRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0 - -- LeakyBucket.allowRequest cache throttleName zone ident capacity leakRate - not <$> LeakyBucket.allowRequest - (zscLeakyBucketCache caches) - throttleName - zone - ident - (fromIntegral limit) - leakRate + LeakyBucket -> do + let period = throttlePeriod cfg + limit = throttleLimit cfg + leakRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0 + -- LeakyBucket.allowRequest cache throttleName zone ident capacity leakRate + not <$> LeakyBucket.allowRequest + (zscLeakyBucketCache caches) + throttleName + zone + ident + (fromIntegral limit) + leakRate - TinyLRU -> do - now <- getTime Monotonic - case cacheStore (zscTinyLRUCache caches) of - TinyLRUStore tvar -> do - cache <- readTVarIO tvar - -- allowRequestTinyLRU now cache ident capacity periodSecs - not <$> atomically (allowRequestTinyLRU now cache ident (throttleLimit cfg) (throttlePeriod cfg)) + TinyLRU -> do + now <- getTime Monotonic + case cacheStore (zscTinyLRUCache caches) of + TinyLRUStore tvar -> do + cache <- readTVarIO tvar + -- allowRequestTinyLRU now cache ident capacity periodSecs + not <$> atomically (allowRequestTinyLRU now cache ident (throttleLimit cfg) (throttlePeriod cfg)) +-- | Backward-compatible entry that derives the identifier and delegates +-- to the precomputed path, ensuring no duplicate computation. +checkThrottle + :: ZoneSpecificCaches -> Text -> Request -> Text -> ThrottleConfig -> IO Bool +checkThrottle caches zone req throttleName cfg = do + mIdentifier <- mkIdentifier (throttleIdentifierBy cfg) req + checkThrottleWithIdent caches zone req throttleName cfg mIdentifier + -- | Reset all caches across all known zones. -- -- Useful in tests or administrative endpoints. @@ -345,6 +462,15 @@ | IdHeaderAndIP !HeaderName deriving (Show, Eq, Generic) +-- Manual Hashable instance since HeaderName doesn't have one +instance Hashable IdentifierBy where + hashWithSalt s IdIP = hashWithSalt s (0 :: Int) + hashWithSalt s (IdHeader h) = hashWithSalt s (1 :: Int, original h) + hashWithSalt s (IdCookie t) = hashWithSalt s (2 :: Int, t) + hashWithSalt s IdIPAndPath = hashWithSalt s (3 :: Int) + hashWithSalt s IdIPAndUA = hashWithSalt s (4 :: Int) + hashWithSalt s (IdHeaderAndIP h) = hashWithSalt s (5 :: Int, original h) + -- | How to derive IP zones from requests. data ZoneBy = ZoneDefault @@ -447,7 +573,7 @@ -- | Convenience: build an 'Env' from config and return the 'Middleware'. -- --- Suitable if you don’t need to retain the 'Env' for administrative operations. +-- Suitable if you don't need to retain the 'Env' for administrative operations. buildRateLimiter :: RateLimiterConfig -> IO Middleware buildRateLimiter cfg = buildRateLimiterWithEnv <$> buildEnvFromConfig cfg @@ -455,17 +581,15 @@ -- 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 - { throttleLimit = l - , throttlePeriod = p - , throttleAlgorithm = algo - , throttleIdentifier = mkIdentifier idBy - , throttleTokenBucketTTL = ttl - } - _ <- addThrottle env name cfg - pure () +registerThrottle :: Env -> RLThrottle -> IO Env +registerThrottle env (RLThrottle name l p algo idBy ttl) = + addThrottle env name ThrottleConfig + { throttleLimit = l + , throttlePeriod = p + , throttleAlgorithm = algo + , throttleIdentifierBy = idBy + , throttleTokenBucketTTL = ttl + } -- | Build a request-identifier function from a declarative spec. mkIdentifier :: IdentifierBy -> Request -> IO (Maybe Text) @@ -513,3 +637,26 @@ -- | Extract original bytes from a case-insensitive header name. fromHeaderName :: HeaderName -> S.ByteString fromHeaderName = original + +-------------------------------------------------------------------------------- +-- Internal helpers: grouping and traversal (to avoid duplicate work) + +type ThrottleName = Text +type Grouped = HM.HashMap IdentifierBy [(ThrottleName, ThrottleConfig)] + +-- | Group throttles by their IdentifierBy to compute the identifier once per group. +groupByIdentifier :: HM.HashMap ThrottleName ThrottleConfig -> Grouped +groupByIdentifier = + HM.foldlWithKey' step HM.empty + where + step acc name cfg = + HM.insertWith (++) (throttleIdentifierBy cfg) [(name, cfg)] acc + +anyMList :: (a -> IO Bool) -> [a] -> IO Bool +anyMList _ [] = pure False +anyMList f (x:xs) = do + b <- f x + if b then pure True else anyMList f xs + +anyMHashMap :: (k -> v -> IO Bool) -> HM.HashMap k v -> IO Bool +anyMHashMap f = anyMList (uncurry f) . HM.toList diff --git a/test/Keter/RateLimiter/LeakyBucketTests.hs b/test/Keter/RateLimiter/LeakyBucketTests.hs --- a/test/Keter/RateLimiter/LeakyBucketTests.hs +++ b/test/Keter/RateLimiter/LeakyBucketTests.hs @@ -47,7 +47,6 @@ import Network.HTTP.Types.Status (statusCode) import Network.Socket (SockAddr(..), tupleToHostAddress) import Keter.RateLimiter.Cache -import Keter.RateLimiter.RequestUtils import Keter.RateLimiter.WAI import Keter.RateLimiter.IPZones import Keter.RateLimiter.LeakyBucket (allowRequest) @@ -108,7 +107,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -126,7 +125,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -146,7 +145,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -164,7 +163,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -184,7 +183,7 @@ { throttleLimit = 2 -- capacity , throttlePeriod = 2 -- time to leak one request (leak rate is capacity / period) , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -209,7 +208,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -230,7 +229,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -249,7 +248,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle diff --git a/test/Keter/RateLimiter/SlidingWindowTests.hs b/test/Keter/RateLimiter/SlidingWindowTests.hs --- a/test/Keter/RateLimiter/SlidingWindowTests.hs +++ b/test/Keter/RateLimiter/SlidingWindowTests.hs @@ -46,7 +46,6 @@ import Keter.RateLimiter.Cache import Keter.RateLimiter.WAI import Keter.RateLimiter.SlidingWindow (allowRequest) -import Keter.RateLimiter.RequestUtils import Keter.RateLimiter.IPZones (defaultIPZone) import Data.CaseInsensitive (mk) import Data.IORef (newIORef, modifyIORef', readIORef) @@ -96,7 +95,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -114,7 +113,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -134,7 +133,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -152,7 +151,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -172,7 +171,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -193,7 +192,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -214,7 +213,7 @@ { throttleLimit = 2 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle @@ -240,7 +239,7 @@ { throttleLimit = 10 , throttlePeriod = 60 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } env' <- addThrottle env (T.pack "test_throttle") throttle diff --git a/test/Keter/RateLimiter/TokenBucketTests.hs b/test/Keter/RateLimiter/TokenBucketTests.hs --- a/test/Keter/RateLimiter/TokenBucketTests.hs +++ b/test/Keter/RateLimiter/TokenBucketTests.hs @@ -40,8 +40,8 @@ import Network.Socket (SockAddr(..), tupleToHostAddress) import Data.CaseInsensitive (mk) import qualified Data.Text.Encoding as TE -import Keter.RateLimiter.RequestUtils -import Keter.RateLimiter.WAI (ThrottleConfig(..), attackMiddleware, addThrottle, initConfig) +--import Keter.RateLimiter.RequestUtils +import Keter.RateLimiter.WAI (ThrottleConfig(..), IdentifierBy(..), attackMiddleware, addThrottle, initConfig) import Keter.RateLimiter.Cache (Algorithm(..)) import Keter.RateLimiter.IPZones (defaultIPZone) @@ -85,7 +85,7 @@ { throttleLimit = 3 , throttlePeriod = 10 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just ttlSeconds } env' <- addThrottle env (Text.pack "test_throttle") throttleConfig @@ -98,7 +98,7 @@ { throttleLimit = 2 , throttlePeriod = 20 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 2 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -122,7 +122,7 @@ { throttleLimit = 3 , throttlePeriod = 10 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -138,7 +138,7 @@ { throttleLimit = 2 , throttlePeriod = 5 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -158,7 +158,7 @@ { throttleLimit = 2 , throttlePeriod = 5 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -184,7 +184,7 @@ { throttleLimit = 10 , throttlePeriod = 10 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -208,7 +208,7 @@ { throttleLimit = 0 , throttlePeriod = 10 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -223,7 +223,7 @@ { throttleLimit = 1 , throttlePeriod = 10 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -240,7 +240,7 @@ { throttleLimit = 2 , throttlePeriod = 0 , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle @@ -262,7 +262,7 @@ { throttleLimit = 2 , throttlePeriod = 1 -- Changed from 0.1 to 1 (Int) , throttleAlgorithm = TokenBucket - , throttleIdentifier = byIP + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just 10 } env' <- addThrottle env (Text.pack "test_throttle") throttle 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 @@ -13,31 +13,34 @@ This module provides a comprehensive test suite for the WAI (Web Application Interface) middleware responsible for rate limiting. It uses the tasty, tasty-hunit, and tasty-quickcheck frameworks to define and run tests. The tests cover five distinct rate-limiting algorithms: - * Fixed Window - * Sliding Window - * Token Bucket - * Leaky Bucket - * TinyLRU +* Fixed Window +* Sliding Window +* Token Bucket +* Leaky Bucket +* TinyLRU + For each algorithm, the following scenarios are tested: - * Allowing requests under the defined limit. - * Blocking requests exceeding the defined limit. - * Correctly handling IPv4 and IPv6 addresses. - * Ensuring rate-limiting window resets correctly over time. - * Identifying clients using proxy headers like @x-forwarded-for@ and @x-real-ip@. - * Managing concurrent requests to prevent race conditions. - * Simulating high-volume concurrent requests to test DoS protection. +* Allowing requests under the defined limit. +* Blocking requests exceeding the defined limit. +* Correctly handling IPv4 and IPv6 addresses. +* Ensuring rate-limiting window resets correctly over time. +* Identifying clients using proxy headers like @x-forwarded-for@ and @x-real-ip@. +* Managing concurrent requests to prevent race conditions. +* Simulating high-volume concurrent requests to test DoS protection. + Additional tests cover: - * Configuration-driven middleware (buildRateLimiter). - * Multiple throttle rules simultaneously. - * Different identifier strategies (header, cookie, combined). - * Zone-based separation. - * JSON configuration parsing. - * Cache management functions. - * Error handling and edge cases. - * Property-based tests for robustness. +* Configuration-driven middleware (buildRateLimiter). +* Multiple throttle rules simultaneously. +* Different identifier strategies (header, cookie, combined). +* Zone-based separation. +* JSON configuration parsing. +* Cache management functions. +* Error handling and edge cases. +* Property-based tests for robustness. + The module defines helper functions to create mock requests and a mock application to isolate the middleware for testing. -} @@ -67,13 +70,12 @@ import Control.Concurrent.MVar import Control.Monad (replicateM) import Control.Monad.IO.Class (liftIO) -import Data.IORef import qualified Web.Cookie as WC import qualified Data.Text.Encoding.Error as TEE import Keter.RateLimiter.IPZones (defaultIPZone) import Keter.RateLimiter.WAI -import Keter.RateLimiter.RequestUtils import Keter.RateLimiter.Cache (Algorithm(..)) +import Text.Printf (printf) -- * Request Helpers @@ -93,13 +95,13 @@ mkRequestWithHeader :: Text -> Text -> Request mkRequestWithHeader name value = defaultRequest { requestHeaders = [(mk (TE.encodeUtf8 name), TE.encodeUtf8 value)] -} + } -- | Creates a mock 'Request' with a cookie header. mkRequestWithCookie :: Text -> Text -> Request mkRequestWithCookie name value = defaultRequest { requestHeaders = [(mk "Cookie", TE.encodeUtf8 $ name <> "=" <> value)] -} + } -- | Creates a mock 'Request' with an @x-forwarded-for@ header. mkRequestWithXFF :: Text -> Request @@ -114,8 +116,8 @@ extractCookieWC name raw = let pairs = WC.parseCookies raw in case lookup (TE.encodeUtf8 name) pairs of - Just v | not (BS.null v) -> Just (TE.decodeUtf8With TEE.lenientDecode v) - _ -> Nothing + Just v | not (BS.null v) -> Just (TE.decodeUtf8With TEE.lenientDecode v) + _ -> Nothing -- * Test Suite Definition @@ -144,15 +146,15 @@ ] where algorithmTestGroup algo = testGroup (show algo ++ " Algorithm") - [ testCase "Allows IPv4 requests below limit" $ testBelowLimit algo byIP mkIPv4Request - , testCase "Blocks IPv4 requests exceeding limit" $ testExceedLimit algo byIP mkIPv4Request - , testCase "Allows IPv6 requests below limit" $ testBelowLimit algo byIP mkIPv6Request - , testCase "Blocks IPv6 requests exceeding limit" $ testExceedLimit algo byIP mkIPv6Request - , testCase "Respects timing with IPv4" $ testTiming algo byIP - , testCase "Handles x-forwarded-for header for IPv4" $ testXFF algo byIP - , testCase "Handles x-real-ip header for IPv6" $ testRealIP algo byIP - , testCase "Handles concurrent requests" $ testConcurrent algo byIP - , testCase "Handles DoS-like concurrency" $ testDoS algo byIP + [ testCase "Allows IPv4 requests below limit" $ testBelowLimit algo IdIP mkIPv4Request + , testCase "Blocks IPv4 requests exceeding limit" $ testExceedLimit algo IdIP mkIPv4Request + , testCase "Allows IPv6 requests below limit" $ testBelowLimit algo IdIP mkIPv6Request + , testCase "Blocks IPv6 requests exceeding limit" $ testExceedLimit algo IdIP mkIPv6Request + , testCase "Respects timing with IPv4" $ testTiming algo IdIP + , testCase "Handles x-forwarded-for header for IPv4" $ testXFF algo IdIP + , testCase "Handles x-real-ip header for IPv6" $ testRealIP algo IdIP + , testCase "Handles concurrent requests" $ testConcurrent algo IdIP + , testCase "Handles DoS-like concurrency" $ testDoS algo IdIP ] -- | Test buildRateLimiter with various configurations. @@ -238,12 +240,12 @@ , testProperty "Header name round-trip" propHeaderNameRoundTrip , testProperty "IP extraction consistency" propIPExtraction , testProperty "Rate limiting monotonicity" propRateLimitingMonotonicity + , testProperty "Cookie extraction multiple" propCookieExtractionMultiple + , testProperty "Identifier independence" propIdentifierIndependence ] --- * Test Case Implementations - -- | Verifies that requests below the limit are allowed. -testBelowLimit :: Algorithm -> (Request -> IO (Maybe Text)) -> Request -> Assertion +testBelowLimit :: Algorithm -> IdentifierBy -> Request -> Assertion testBelowLimit algo identifier req = do env <- initConfig (const defaultIPZone) let throttle = ThrottleConfig 2 60 algo identifier (Just 3600) @@ -258,7 +260,7 @@ assertEqual "Second request status" status200 (simpleStatus response2) -- | Verifies that requests exceeding the limit are blocked. -testExceedLimit :: Algorithm -> (Request -> IO (Maybe Text)) -> Request -> Assertion +testExceedLimit :: Algorithm -> IdentifierBy -> Request -> Assertion testExceedLimit algo identifier req = do env <- initConfig (const defaultIPZone) let throttle = ThrottleConfig 2 60 algo identifier (Just 3600) @@ -273,7 +275,7 @@ assertEqual "Third request status" status429 (simpleStatus response3) -- | Verifies that the rate limit counter resets after the window period. -testTiming :: Algorithm -> (Request -> IO (Maybe Text)) -> Assertion +testTiming :: Algorithm -> IdentifierBy -> Assertion testTiming algo identifier = do env <- initConfig (const defaultIPZone) let throttle = ThrottleConfig 1 1 algo identifier (Just 3600) @@ -290,7 +292,7 @@ assertEqual "Second request status after reset" status200 (simpleStatus response2) -- | Verifies correct IP identification using x-forwarded-for header. -testXFF :: Algorithm -> (Request -> IO (Maybe Text)) -> Assertion +testXFF :: Algorithm -> IdentifierBy -> Assertion testXFF algo identifier = do env <- initConfig (const defaultIPZone) let throttle = ThrottleConfig 1 60 algo identifier (Just 3600) @@ -305,7 +307,7 @@ assertEqual "Second XFF request status" status429 (simpleStatus response2) -- | Verifies correct IP identification using x-real-ip header. -testRealIP :: Algorithm -> (Request -> IO (Maybe Text)) -> Assertion +testRealIP :: Algorithm -> IdentifierBy -> Assertion testRealIP algo identifier = do env <- initConfig (const defaultIPZone) let throttle = ThrottleConfig 1 60 algo identifier (Just 3600) @@ -320,7 +322,7 @@ assertEqual "Second Real-IP request status" status429 (simpleStatus response2) -- | Verifies behavior under moderate concurrent load. -testConcurrent :: Algorithm -> (Request -> IO (Maybe Text)) -> Assertion +testConcurrent :: Algorithm -> IdentifierBy -> Assertion testConcurrent algo identifier = do env <- initConfig (const defaultIPZone) let throttle = ThrottleConfig 5 60 algo identifier (Just 3600) @@ -331,12 +333,11 @@ result6 <- srequest $ SRequest mkIPv4Request "" return (responses, result6) (responses, response6) <- runSession session app - mapM_ (\(i, resp) -> assertEqual ("Request " ++ show i ++ " status") status200 (simpleStatus resp)) - (zip [1..5] responses) + mapM_ (\(i, resp) -> assertEqual ("Request " ++ show i ++ " status") status200 (simpleStatus resp)) (zip [1..5] responses) assertEqual "Sixth request status after limit" status429 (simpleStatus response6) -- | Simulates a DoS attack with high concurrency. -testDoS :: Algorithm -> (Request -> IO (Maybe Text)) -> Assertion +testDoS :: Algorithm -> IdentifierBy -> Assertion testDoS algo identifier = do env <- initConfig (const defaultIPZone) let throttle = ThrottleConfig 10 60 algo identifier (Just 3600) @@ -356,7 +357,7 @@ testBuildSingleThrottle :: Assertion testBuildSingleThrottle = do let config = RateLimiterConfig ZoneDefault - [ RLThrottle "api-limit" 5 60 FixedWindow IdIP Nothing ] + [ RLThrottle "api-limit" 5 60 FixedWindow IdIP Nothing ] middleware <- buildRateLimiter config let app = middleware mockApp let session = replicateM 6 (srequest $ SRequest mkIPv4Request "") @@ -370,9 +371,9 @@ testBuildMultipleThrottles :: Assertion testBuildMultipleThrottles = do let config = RateLimiterConfig ZoneDefault - [ RLThrottle "global-limit" 10 60 FixedWindow IdIP Nothing - , RLThrottle "api-limit" 3 60 SlidingWindow (IdHeader "X-API-Key") Nothing - ] + [ RLThrottle "global-limit" 10 60 FixedWindow IdIP Nothing + , RLThrottle "api-limit" 3 60 SlidingWindow (IdHeader "X-API-Key") Nothing + ] middleware <- buildRateLimiter config let app = middleware mockApp let requestWithApi = mkRequestWithHeader "X-API-Key" "test-key" @@ -390,7 +391,7 @@ testBuildWithZones :: Assertion testBuildWithZones = do let config = RateLimiterConfig (ZoneHeader "X-Zone") - [ RLThrottle "zone-limit" 2 60 FixedWindow IdIP Nothing ] + [ RLThrottle "zone-limit" 2 60 FixedWindow IdIP Nothing ] middleware <- buildRateLimiter config let app = middleware mockApp let zoneAReq = mkRequestWithHeader "X-Zone" "A" @@ -403,9 +404,7 @@ ra3 <- srequest $ SRequest zoneAReq "" return [ra1, ra2, rb1, rb2, ra3] responses <- runSession session app - assertEqual "Zone separation works" - [status200, status200, status200, status200, status429] - (map simpleStatus responses) + assertEqual "Zone separation works" [status200, status200, status200, status200, status429] (map simpleStatus responses) -- | Tests buildRateLimiter with an empty throttles list. testEmptyThrottles :: Assertion @@ -421,8 +420,8 @@ testMultipleSameAlgo :: Assertion testMultipleSameAlgo = do env <- initConfig (const defaultIPZone) - let throttle1 = ThrottleConfig 5 60 FixedWindow (mkIdentifier IdIP) Nothing - let throttle2 = ThrottleConfig 3 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle1 = ThrottleConfig 5 60 FixedWindow IdIP Nothing + let throttle2 = ThrottleConfig 3 60 FixedWindow IdIP Nothing env' <- addThrottle env "global" throttle1 >>= \e -> addThrottle e "strict" throttle2 let app = attackMiddleware env' mockApp let session = replicateM 4 (srequest $ SRequest mkIPv4Request "") @@ -434,8 +433,8 @@ testMultipleDiffAlgo :: Assertion testMultipleDiffAlgo = do env <- initConfig (const defaultIPZone) - let throttle1 = ThrottleConfig 10 60 FixedWindow (mkIdentifier IdIP) Nothing - let throttle2 = ThrottleConfig 5 60 TokenBucket (mkIdentifier IdIP) (Just 120) + let throttle1 = ThrottleConfig 10 60 FixedWindow IdIP Nothing + let throttle2 = ThrottleConfig 5 60 TokenBucket IdIP (Just 120) env' <- addThrottle env "fixed" throttle1 >>= \e -> addThrottle e "bucket" throttle2 let app = attackMiddleware env' mockApp let session = replicateM 6 (srequest $ SRequest mkIPv4Request "") @@ -447,8 +446,8 @@ testThrottlePriority :: Assertion testThrottlePriority = do env <- initConfig (const defaultIPZone) - let permissive = ThrottleConfig 1000 60 FixedWindow (mkIdentifier IdIP) Nothing - let restrictive = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIP) Nothing + let permissive = ThrottleConfig 1000 60 FixedWindow IdIP Nothing + let restrictive = ThrottleConfig 1 60 FixedWindow IdIP Nothing env' <- addThrottle env "permissive" permissive >>= \e -> addThrottle e "restrictive" restrictive let app = attackMiddleware env' mockApp let session = do @@ -462,8 +461,8 @@ testIndependentCounters :: Assertion testIndependentCounters = do env <- initConfig (const defaultIPZone) - let ipThrottle = ThrottleConfig 2 60 FixedWindow (mkIdentifier IdIP) Nothing - let headerThrottle = ThrottleConfig 2 60 FixedWindow (mkIdentifier (IdHeader "X-User-ID")) Nothing + let ipThrottle = ThrottleConfig 2 60 FixedWindow IdIP Nothing + let headerThrottle = ThrottleConfig 2 60 FixedWindow (IdHeader "X-User-ID") Nothing env' <- addThrottle env "ip" ipThrottle >>= \e -> addThrottle e "user" headerThrottle let app = attackMiddleware env' mockApp let userReq = mkRequestWithHeader "X-User-ID" "user123" @@ -480,7 +479,7 @@ testIdHeaderStrategy :: Assertion testIdHeaderStrategy = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 2 60 FixedWindow (mkIdentifier (IdHeader "X-Client-ID")) Nothing + let throttle = ThrottleConfig 2 60 FixedWindow (IdHeader "X-Client-ID") Nothing env' <- addThrottle env "header" throttle let app = attackMiddleware env' mockApp let client1Req = mkRequestWithHeader "X-Client-ID" "client1" @@ -501,7 +500,7 @@ testIdCookieStrategy :: Assertion testIdCookieStrategy = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier (IdCookie "session")) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow (IdCookie "session") Nothing env' <- addThrottle env "cookie" throttle let app = attackMiddleware env' mockApp let session1Req = mkRequestWithCookie "session" "sess123" @@ -520,7 +519,7 @@ testIdIPAndPathStrategy :: Assertion testIdIPAndPathStrategy = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIPAndPath) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow IdIPAndPath Nothing env' <- addThrottle env "ip-path" throttle let app = attackMiddleware env' mockApp let path1Req = mkIPv4Request { rawPathInfo = "/api/v1" } @@ -539,7 +538,7 @@ testIdIPAndUAStrategy :: Assertion testIdIPAndUAStrategy = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIPAndUA) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow IdIPAndUA Nothing env' <- addThrottle env "ip-ua" throttle let app = attackMiddleware env' mockApp let ua1Req = mkRequestWithHeader "User-Agent" "Browser/1.0" @@ -558,7 +557,7 @@ testIdHeaderAndIPStrategy :: Assertion testIdHeaderAndIPStrategy = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier (IdHeaderAndIP "X-Service")) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow (IdHeaderAndIP "X-Service") Nothing env' <- addThrottle env "header-ip" throttle let app = attackMiddleware env' mockApp let service1Req = mkRequestWithHeader "X-Service" "service1" @@ -577,7 +576,7 @@ testMissingIdentifiers :: Assertion testMissingIdentifiers = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier (IdHeader "Missing-Header")) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow (IdHeader "Missing-Header") Nothing env' <- addThrottle env "missing" throttle let app = attackMiddleware env' mockApp let session = replicateM 5 (srequest $ SRequest mkIPv4Request "") @@ -597,14 +596,15 @@ , ("malformed", Nothing) ] mapM_ (\(cookie, expected) -> do - let result = extractCookieWC "session" (TE.encodeUtf8 cookie) - assertEqual ("Cookie parsing: " <> T.unpack cookie) expected result) testCases + let result = extractCookieWC "session" (TE.encodeUtf8 cookie) + assertEqual ("Cookie parsing: " <> T.unpack cookie) expected result + ) testCases -- | Tests IP-based zone separation. testZoneIPSeparation :: Assertion testZoneIPSeparation = do let config = RateLimiterConfig ZoneIP - [ RLThrottle "ip-zone" 1 60 FixedWindow IdIP Nothing ] + [ RLThrottle "ip-zone" 1 60 FixedWindow IdIP Nothing ] middleware <- buildRateLimiter config let app = middleware mockApp let ip1Req = mkRequestWithXFF "192.168.1.1" @@ -624,7 +624,7 @@ testZoneHeaderSeparation :: Assertion testZoneHeaderSeparation = do let config = RateLimiterConfig (ZoneHeader "X-Tenant") - [ RLThrottle "tenant-limit" 1 60 FixedWindow IdIP Nothing ] + [ RLThrottle "tenant-limit" 1 60 FixedWindow IdIP Nothing ] middleware <- buildRateLimiter config let app = middleware mockApp let tenant1Req = mkRequestWithHeader "X-Tenant" "tenant1" @@ -643,7 +643,7 @@ testZoneCreation :: Assertion testZoneCreation = do env <- initConfig (\req -> maybe "default" TE.decodeUtf8 (lookup (mk "X-Zone") (requestHeaders req))) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow IdIP Nothing env' <- addThrottle env "test" throttle let zone1Req = mkRequestWithHeader "X-Zone" "zone1" let zone2Req = mkRequestWithHeader "X-Zone" "zone2" @@ -659,7 +659,7 @@ testDefaultZoneFallback :: Assertion testDefaultZoneFallback = do let config = RateLimiterConfig (ZoneHeader "Missing-Header") - [ RLThrottle "default-fallback" 1 60 FixedWindow IdIP Nothing ] + [ RLThrottle "default-fallback" 1 60 FixedWindow IdIP Nothing ] middleware <- buildRateLimiter config let app = middleware mockApp let session = do @@ -681,14 +681,16 @@ , ("{\"header\": \"X-API-Key\"}", Right (IdHeader (hdr "X-API-Key"))) , ("{\"cookie\": \"session\"}", Right (IdCookie "session")) , ("{\"header+ip\": \"X-User\"}", Right (IdHeaderAndIP (hdr "X-User"))) - , ("\"invalid\"", Left ("identifier_by must be one of:" :: String)) + , ("\"invalid\"", Left ("Error in $: identifier_by: 'ip' | 'ip+path' | 'ip+ua' | {header} | {cookie} | {header+ip}" :: String)) ] mapM_ (\(json, expected) -> do - let result = eitherDecode (LBS.fromStrict $ TE.encodeUtf8 json) :: Either String IdentifierBy - case (result, expected) of - (Right actual, Right expected') -> assertEqual ("Parse: " <> T.unpack json) expected' actual - (Left _, Left _) -> return () - _ -> assertFailure $ "Unexpected result for: " <> T.unpack json) testCases + let result = eitherDecode (LBS.fromStrict $ TE.encodeUtf8 json) :: Either String IdentifierBy + case (result, expected) of + (Right actual, Right expected') -> assertEqual ("Parse: " <> T.unpack json) expected' actual + (Left _, Left _) -> return () + (Left err, Right _) -> assertFailure $ "Parse failed unexpectedly for " <> T.unpack json <> ": " <> err + (Right res, Left _) -> assertFailure $ "Parse succeeded unexpectedly for " <> T.unpack json <> ": " <> show res + ) testCases -- | Tests parsing of ZoneBy JSON. testParseZoneBy :: Assertion @@ -697,14 +699,16 @@ [ ("\"default\"", Right ZoneDefault) , ("\"ip\"", Right ZoneIP) , ("{\"header\": \"X-Region\"}", Right (ZoneHeader (hdr "X-Region"))) - , ("\"invalid\"", Left ("zone_by must be" :: String)) + , ("\"invalid\"", Left ("Error in $: zone_by: 'default' | 'ip' | {header}" :: String)) ] mapM_ (\(json, expected) -> do - let result = eitherDecode (LBS.fromStrict $ TE.encodeUtf8 json) :: Either String ZoneBy - case (result, expected) of - (Right actual, Right expected') -> assertEqual ("Parse: " <> T.unpack json) expected' actual - (Left _, Left _) -> return () - _ -> assertFailure $ "Unexpected result for: " <> T.unpack json) testCases + let result = eitherDecode (LBS.fromStrict $ TE.encodeUtf8 json) :: Either String ZoneBy + case (result, expected) of + (Right actual, Right expected') -> assertEqual ("Parse: " <> T.unpack json) expected' actual + (Left _, Left _) -> return () + (Left err, Right _) -> assertFailure $ "Parse failed unexpectedly for " <> T.unpack json <> ": " <> err + (Right res, Left _) -> assertFailure $ "Parse succeeded unexpectedly for " <> T.unpack json <> ": " <> show res + ) testCases -- | Tests parsing of RLThrottle JSON. testParseRLThrottle :: Assertion @@ -743,7 +747,7 @@ testCacheResetAll :: Assertion testCacheResetAll = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow IdIP Nothing env' <- addThrottle env "test" throttle let app = attackMiddleware env' mockApp let session1 = srequest $ SRequest mkIPv4Request "" @@ -761,7 +765,7 @@ testZoneCacheIsolation :: Assertion testZoneCacheIsolation = do env <- initConfig (\req -> maybe "A" TE.decodeUtf8 (lookup (mk "X-Zone") (requestHeaders req))) - let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig 1 60 FixedWindow IdIP Nothing env' <- addThrottle env "test" throttle let app = attackMiddleware env' mockApp let zoneAReq = mkRequestWithHeader "X-Zone" "A" @@ -780,7 +784,7 @@ testMemoryCleanup :: Assertion testMemoryCleanup = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 100 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig 100 60 FixedWindow IdIP Nothing env' <- addThrottle env "test" throttle cacheResetAll env' cacheResetAll env' @@ -791,7 +795,7 @@ testZeroPeriod :: Assertion testZeroPeriod = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 10 0 TokenBucket (mkIdentifier IdIP) (Just 60) + let throttle = ThrottleConfig 10 0 TokenBucket IdIP (Just 60) env' <- addThrottle env "zero-period" throttle _ <- instrument env' mkIPv4Request return () @@ -800,7 +804,7 @@ testNegativeLimit :: Assertion testNegativeLimit = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig (-1) 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig (-1) 60 FixedWindow IdIP Nothing env' <- addThrottle env "negative" throttle _ <- instrument env' mkIPv4Request return () @@ -809,7 +813,7 @@ testLargeNumbers :: Assertion testLargeNumbers = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig (maxBound :: Int) (maxBound :: Int) FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig (maxBound :: Int) (maxBound :: Int) FixedWindow IdIP Nothing env' <- addThrottle env "large" throttle blocked <- instrument env' mkIPv4Request assertEqual "Large numbers handled" False blocked @@ -818,7 +822,7 @@ testMalformedRequests :: Assertion testMalformedRequests = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 5 60 FixedWindow (mkIdentifier (IdHeader "X-Malformed")) Nothing + let throttle = ThrottleConfig 5 60 FixedWindow (IdHeader "X-Malformed") Nothing env' <- addThrottle env "malformed" throttle let malformedReq = defaultRequest { requestHeaders = [(mk "X-Malformed", "\xFF\xFE\xFD")] } _ <- instrument env' malformedReq @@ -828,7 +832,7 @@ testConcurrentSafety :: Assertion testConcurrentSafety = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 100 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig 100 60 FixedWindow IdIP Nothing env' <- addThrottle env "concurrent" throttle results <- newMVar [] let worker :: Integer -> IO () @@ -847,7 +851,7 @@ testHighThroughputSingle :: Assertion testHighThroughputSingle = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 1000 60 FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig 1000 60 FixedWindow IdIP Nothing env' <- addThrottle env "throughput" throttle let app = attackMiddleware env' mockApp let session = replicateM 500 (srequest $ SRequest mkIPv4Request "") @@ -859,7 +863,7 @@ testManyClients :: Assertion testManyClients = do env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 2 60 FixedWindow (mkIdentifier (IdHeader "X-Client-ID")) Nothing + let throttle = ThrottleConfig 2 60 FixedWindow (IdHeader "X-Client-ID") Nothing env' <- addThrottle env "many-clients" throttle let app = attackMiddleware env' mockApp let makeClientRequest :: Integer -> Request @@ -874,21 +878,22 @@ testAlgorithmPerformance = do let algorithms = [FixedWindow, SlidingWindow, TokenBucket, LeakyBucket, TinyLRU] results <- mapM (\algo -> do - env <- initConfig (const defaultIPZone) - let throttle = ThrottleConfig 100 60 algo (mkIdentifier IdIP) (Just 120) - env' <- addThrottle env ("perf-" <> T.pack (show algo)) throttle - let start = (0 :: Integer) - mapM_ (\_ -> instrument env' mkIPv4Request) [1..100 :: Integer] - let end = (1 :: Integer) - return (algo, end - start)) algorithms + env <- initConfig (const defaultIPZone) + let throttle = ThrottleConfig 100 60 algo IdIP (Just 120) + env' <- addThrottle env ("perf-" <> T.pack (show algo)) throttle + let start = (0 :: Integer) + mapM_ (\_ -> instrument env' mkIPv4Request) [1..100 :: Integer] + let end = (1 :: Integer) + return (algo, end - start) + ) algorithms assertEqual "All algorithms tested" (length algorithms) (length results) -- | Tests memory usage with many zones. testManyZones :: Assertion testManyZones = do env <- initConfig (\req -> - maybe "default" TE.decodeUtf8 (lookup (mk "X-Zone-ID") (requestHeaders req))) - let throttle = ThrottleConfig 5 60 FixedWindow (mkIdentifier IdIP) Nothing + maybe "default" TE.decodeUtf8 (lookup (mk "X-Zone-ID") (requestHeaders req))) + let throttle = ThrottleConfig 5 60 FixedWindow IdIP Nothing env' <- addThrottle env "zones" throttle let makeZoneRequest :: Integer -> Request makeZoneRequest i = mkRequestWithHeader "X-Zone-ID" (T.pack $ "zone" <> show i) @@ -898,54 +903,151 @@ assertBool "Many zones created" (zoneCount > 10) assertBool "Reasonable zone count" (zoneCount <= 51) --- * Property-Based Tests +-- * Property-Based Tests (Fixed Version) --- | Generates valid token characters per RFC 6265 (simplified). -validTokenChar :: Gen Char -validTokenChar = elements $ ['!'..'~'] >>= \c -> - if c `elem` [';', ',', '=', ' '] then [] else [c] +-- | Generates valid cookie value characters (excluding problematic ones for Web.Cookie) +validCookieValueChar :: Gen Char +validCookieValueChar = elements $ concat + [ ['!'] + , ['#'..'&'] + , ['('..'/'] + , ['0'..'9'] + , [':'] -- Split the range to exclude ';' + , ['<'..'@'] + , ['A'..'Z'] + , ['['..'`'] + , ['a'..'z'] + , ['{'..'~'] + ] + -- Excludes: '"', ';', ',', '=', ' ', '\t', '\n', '\r' and control chars --- | Generates a valid token text. -tokenText :: Gen Text -tokenText = T.pack <$> listOf1 validTokenChar +-- | Generates valid cookie name characters (stricter than values) +validCookieNameChar :: Gen Char +validCookieNameChar = elements $ ['A'..'Z'] ++ ['a'..'z'] ++ ['0'..'9'] ++ ['!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', '^', '_', '`', '|', '~'] --- | Generates a valid cookie value text. -cookieValueText :: Gen Text -cookieValueText = T.pack <$> listOf1 validTokenChar +-- | Generates a valid cookie name. +cookieName :: Gen Text +cookieName = T.pack <$> listOf1 validCookieNameChar --- | Tests cookie extraction properties. +-- | Generates a valid cookie value (no quotes, semicolons, etc.) +cookieValue :: Gen Text +cookieValue = T.pack <$> listOf1 validCookieValueChar + +-- | Tests cookie extraction properties with proper cookie format. propCookieExtraction :: Property propCookieExtraction = - forAll tokenText $ \cookieName -> - forAll cookieValueText $ \cookieValue -> - let header = TE.encodeUtf8 (cookieName <> "=" <> cookieValue) - extracted = extractCookieWC cookieName header - in extracted === Just cookieValue + forAll cookieName $ \name -> + forAll cookieValue $ \value -> + let header = TE.encodeUtf8 (name <> "=" <> value) + extracted = extractCookieWC name header + in counterexample ("Header: " <> show header <> ", Expected: " <> show value <> ", Got: " <> show extracted) $ + extracted === Just value --- | Tests header name round-trip. +-- | Property test for cookie extraction with multiple cookies +propCookieExtractionMultiple :: Property +propCookieExtractionMultiple = + forAll cookieName $ \targetName -> + forAll cookieValue $ \targetValue -> + forAll (listOf ((,) <$> cookieName <*> cookieValue)) $ \otherCookies -> + let allCookies = (targetName, targetValue) : otherCookies + headerValue = T.intercalate "; " $ map (\(n, v) -> n <> "=" <> v) allCookies + header = TE.encodeUtf8 headerValue + extracted = extractCookieWC targetName header + in counterexample ("Header: " <> show headerValue <> ", Expected: " <> show targetValue <> ", Got: " <> show extracted) $ + extracted === Just targetValue + +-- | Tests header name round-trip with valid header characters only. propHeaderNameRoundTrip :: Property -propHeaderNameRoundTrip = property $ \headerText -> - let originalTxt = T.pack headerText +propHeaderNameRoundTrip = + forAll (listOf1 validHeaderChar) $ \headerChars -> + let originalTxt = T.pack headerChars headerName = hdr originalTxt roundTrip = TE.decodeUtf8 (fromHeaderName headerName) - in not (T.null originalTxt) ==> roundTrip === originalTxt + in counterexample ("Original: " <> show originalTxt <> ", RoundTrip: " <> show roundTrip) $ + roundTrip === originalTxt + where + -- Valid HTTP header name characters per RFC 7230 + validHeaderChar :: Gen Char + validHeaderChar = elements $ concat + [ ['!'] + , ['#'..'\''] + , ['*', '+', '-', '.'] + , ['0'..'9'] + , ['A'..'Z'] + , ['^'..'z'] + , ['|', '~'] + ] --- | Tests IP extraction consistency. +-- | Tests IP extraction consistency with valid IP formats. propIPExtraction :: Property -propIPExtraction = property $ \ipStr -> - let ip = T.pack ipStr - req = mkRequestWithXFF ip +propIPExtraction = + forAll genValidIP $ \ip -> + let req = mkRequestWithXFF ip extracted = getClientIPPure req expected = T.takeWhile (/= ',') ip - in not (T.null ip) ==> extracted === expected + in counterexample ("IP: " <> show ip <> ", Expected: " <> show expected <> ", Got: " <> show extracted) $ + extracted === expected + where + genValidIP = oneof [genIPv4, genIPv6, genIPv4List] --- | Tests rate limiting monotonicity. + genIPv4 = do + a <- choose (1, 255) :: Gen Int + b <- choose (0, 255) + c <- choose (0, 255) + d <- choose (0, 255) + -- Corrected: Convert each number to Text before intercalating. + return $ T.intercalate "." $ map (T.pack . show) [a, b, c, d] + + genIPv6 = do + segments <- replicateM 8 (choose (0, 65535 :: Int)) + -- Corrected: Convert each formatted hex string to Text before intercalating. + return $ T.intercalate ":" $ map (T.pack . printf "%x") segments + + genIPv4List = do + ips <- listOf1 genIPv4 + return $ T.intercalate ", " ips + +-- | Tests rate limiting monotonicity - more requests should never result in fewer blocks. propRateLimitingMonotonicity :: Property -propRateLimitingMonotonicity = property $ \limit period -> - limit > 0 && period > 0 ==> monadicIO $ do +propRateLimitingMonotonicity = + forAll (choose (1, 100)) $ \limit -> + forAll (choose (1, 3600)) $ \period -> + monadicIO $ do env <- run $ initConfig (const defaultIPZone) - let throttle = ThrottleConfig limit period FixedWindow (mkIdentifier IdIP) Nothing + let throttle = ThrottleConfig limit period FixedWindow IdIP Nothing env' <- run $ addThrottle env "prop" throttle - results <- run $ mapM (\_ -> instrument env' mkIPv4Request) [1..limit] - let blockedCount = length $ filter id results - Test.QuickCheck.Monadic.assert (blockedCount < limit) + -- Test with exactly limit requests + results1 <- run $ mapM (\_ -> instrument env' mkIPv4Request) [1..limit] + let blockedCount1 = length $ filter id results1 + -- Reset and test with limit + 1 requests + run $ cacheResetAll env' + results2 <- run $ mapM (\_ -> instrument env' mkIPv4Request) [1..limit + 1] + let blockedCount2 = length $ filter id results2 + -- Monotonicity: more requests should result in more (or equal) blocks + Test.QuickCheck.Monadic.assert (blockedCount2 >= blockedCount1) + -- First batch should allow all requests within limit + Test.QuickCheck.Monadic.assert (blockedCount1 == 0) + -- Second batch should block at least one request + Test.QuickCheck.Monadic.assert (blockedCount2 >= 1) + +-- | Tests that different identifiers are treated independently +propIdentifierIndependence :: Property +propIdentifierIndependence = + -- Removed unused cookieName generators + forAll cookieValue $ \cookieValue1 -> + forAll cookieValue $ \cookieValue2 -> + -- This precondition ensures the two identifiers are actually different. + cookieValue1 /= cookieValue2 ==> monadicIO $ do + env <- run $ initConfig (const defaultIPZone) + let throttle = ThrottleConfig 1 60 FixedWindow (IdCookie "session") Nothing + env' <- run $ addThrottle env "test" throttle + let req1 = mkRequestWithCookie "session" cookieValue1 + let req2 = mkRequestWithCookie "session" cookieValue2 + -- Both different session values should be allowed initially + blocked1 <- run $ instrument env' req1 + blocked2 <- run $ instrument env' req2 + Test.QuickCheck.Monadic.assert (not blocked1) + Test.QuickCheck.Monadic.assert (not blocked2) + -- Second request with same session should be blocked + blocked1_2 <- run $ instrument env' req1 + Test.QuickCheck.Monadic.assert blocked1_2 diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -74,7 +74,7 @@ import Keter.RateLimiter.Cache ( Cache(..), InMemoryStore(..), Algorithm(..), makeCacheKey, readCache, createInMemoryStore, newCache, incrementCache, writeCache, deleteCache ) import Keter.RateLimiter.CacheWithZone ( incStoreWithZone, writeCacheWithZone, readCacheWithZone ) import Keter.RateLimiter.IPZones (ZoneSpecificCaches(..), IPZoneIdentifier, defaultIPZone) -import Keter.RateLimiter.WAI ( Env, ThrottleConfig(..), initConfig, addThrottle, instrument, envZoneCachesMap, cacheResetAll ) +import Keter.RateLimiter.WAI ( Env, ThrottleConfig(..), initConfig, addThrottle, instrument, envZoneCachesMap, cacheResetAll, IdentifierBy(..) ) import qualified Keter.RateLimiter.RequestUtils as RequestUtils import System.IO.Unsafe (unsafePerformIO) @@ -108,14 +108,15 @@ -- Note: Uses 'RequestUtils.getClientIP' which respects 'x-real-ip' and 'x-forwarded-for' headers. mainTestGetRequestIPZone :: Request -> IPZoneIdentifier mainTestGetRequestIPZone req = - let ip = unsafePerformIO $ RequestUtils.getClientIP req + let ip = unsafePerformIO $ RequestUtils.byIP req in case ip of - _ | ip == ipForZoneA -> testIPZoneA - | ip == ipForZoneB -> testIPZoneB - | ip == genericTestIP -> defaultIPZone - | ip == genericTestIP2 -> defaultIPZone - | ip == ipForDefaultZone -> defaultIPZone - | otherwise -> defaultIPZone + Just val + | val == ipForZoneA -> testIPZoneA + | val == ipForZoneB -> testIPZoneB + | val == genericTestIP -> defaultIPZone + | val == genericTestIP2 -> defaultIPZone + | val == ipForDefaultZone -> defaultIPZone + _ -> defaultIPZone -- | A simple helper to extract the request path as 'Text'. getRequestPath :: Request -> Text @@ -222,7 +223,7 @@ { throttleLimit = 3 , throttlePeriod = 10 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "test-throttle" throttleConfig @@ -239,7 +240,7 @@ { throttleLimit = 3 , throttlePeriod = 10 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "test-throttle" throttleConfig @@ -263,7 +264,7 @@ { throttleLimit = 3 , throttlePeriod = 1 , throttleAlgorithm = TokenBucket - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Just ttlSeconds } envWithThrottle <- addThrottle env "test-throttle" throttleConfig @@ -286,7 +287,7 @@ { throttleLimit = 3 , throttlePeriod = 10 , throttleAlgorithm = LeakyBucket - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "test-throttle" throttleConfig @@ -315,18 +316,14 @@ { throttleLimit = 2 , throttlePeriod = 10 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> do - ip <- RequestUtils.getClientIP req - return $ if getRequestPath req == "\"/login\"" - then Just (ip <> Text.pack ":" <> getRequestPath req) - else Nothing + , throttleIdentifierBy = IdIPAndPath , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "login-throttle" throttleConfig let loginRequests = replicate 3 $ makeRequest genericTestIP "/login" homeRequests = replicate 3 $ makeRequest genericTestIP "/home" expectedLoginResponses = replicate 2 "Success" ++ ["Too Many Requests"] - expectedHomeResponses = replicate 3 "Success" + expectedHomeResponses = replicate 2 "Success" ++ ["Too Many Requests"] executeRequests envWithThrottle loginRequests expectedLoginResponses executeRequests envWithThrottle homeRequests expectedHomeResponses @@ -339,18 +336,14 @@ { throttleLimit = 5 , throttlePeriod = 10 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } loginThrottleConfig = ThrottleConfig { throttleLimit = 2 , throttlePeriod = 10 , throttleAlgorithm = SlidingWindow - , throttleIdentifier = \req -> do - ip <- RequestUtils.getClientIP req - return $ if getRequestPath req == "\"/login\"" - then Just (ip <> Text.pack ":" <> getRequestPath req) - else Nothing + , throttleIdentifierBy = IdIPAndPath , throttleTokenBucketTTL = Nothing } envWithIpThrottle <- addThrottle env "ip-throttle" ipThrottleConfig @@ -375,7 +368,7 @@ { throttleLimit = 2 , throttlePeriod = 1 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle1 <- addThrottle env1 "test-throttle" throttleConfig @@ -399,7 +392,7 @@ { throttleLimit = limit , throttlePeriod = periodSec , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "reset-throttle" throttleConfig @@ -419,7 +412,7 @@ Nothing -> assertFailure "Zone caches not found" >> return undefined let cache :: Cache (InMemoryStore 'FixedWindow) cache = zscCounterCache zoneCaches - key = makeCacheKey "reset-throttle" FixedWindow testIPZoneA (unsafePerformIO $ RequestUtils.getClientIP req) + key = makeCacheKey "reset-throttle" FixedWindow testIPZoneA (case unsafePerformIO $ RequestUtils.byIP req of Just ip -> ip; Nothing -> "") mVal <- readCache cache key :: IO (Maybe Int) when (isJust mVal) $ putStrLn $ "Cache value before third request: " ++ show mVal -- Try manually deleting if it exists @@ -451,7 +444,7 @@ { throttleLimit = limit , throttlePeriod = period , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "test-throttle" throttleConfig @@ -472,7 +465,7 @@ { throttleLimit = 1 , throttlePeriod = 10 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "test-throttle" throttleConfig @@ -491,7 +484,7 @@ { throttleLimit = 1 , throttlePeriod = 10 , throttleAlgorithm = FixedWindow - , throttleIdentifier = \req -> Just <$> RequestUtils.getClientIP req + , throttleIdentifierBy = IdIP , throttleTokenBucketTTL = Nothing } envWithThrottle <- addThrottle env "test-throttle" throttleConfig diff --git a/test/TinyLRUTests.hs b/test/TinyLRUTests.hs --- a/test/TinyLRUTests.hs +++ b/test/TinyLRUTests.hs @@ -316,7 +316,7 @@ , testCase "Concurrent Reset and Access" $ do cache <- createTinyLRU 3 now <- getTime Monotonic - results <- replicateConcurrently 50 $ do + _ <- replicateConcurrently 50 $ do -- Fixed: Add explicit type annotation for randomRIO action <- randomRIO (0, 1 :: Int) if action == 0