keter-rate-limiting-plugin 0.1.0.2 → 0.1.1.0
raw patch · 22 files changed
+2135/−888 lines, 22 filesdep +QuickCheckdep +cookiedep +tasty-quickcheckdep ~containersPVP: major bump suggested
API removals or changes: PVP suggests a major version bump
Dependencies added: QuickCheck, cookie, tasty-quickcheck
Dependency ranges changed: containers
API changes (from Hackage documentation)
- Keter.RateLimiter.WAI: envZoneCachesMap :: Env -> IORef (Map IPZoneIdentifier ZoneSpecificCaches)
+ Keter.RateLimiter.Cache: algoToText :: Algorithm -> Text
+ Keter.RateLimiter.Cache: parseAlgoText :: Text -> Parser Algorithm
+ Keter.RateLimiter.WAI: Env :: IORef (HashMap IPZoneIdentifier ZoneSpecificCaches) -> IORef (HashMap Text ThrottleConfig) -> (Request -> IPZoneIdentifier) -> Env
+ Keter.RateLimiter.WAI: IdCookie :: !Text -> IdentifierBy
+ Keter.RateLimiter.WAI: IdHeader :: !HeaderName -> IdentifierBy
+ Keter.RateLimiter.WAI: IdHeaderAndIP :: !HeaderName -> IdentifierBy
+ Keter.RateLimiter.WAI: IdIP :: IdentifierBy
+ Keter.RateLimiter.WAI: IdIPAndPath :: IdentifierBy
+ Keter.RateLimiter.WAI: IdIPAndUA :: IdentifierBy
+ Keter.RateLimiter.WAI: RLThrottle :: !Text -> !Int -> !Int -> !Algorithm -> !IdentifierBy -> !Maybe Int -> RLThrottle
+ Keter.RateLimiter.WAI: RateLimiterConfig :: !ZoneBy -> ![RLThrottle] -> RateLimiterConfig
+ Keter.RateLimiter.WAI: ZoneDefault :: ZoneBy
+ Keter.RateLimiter.WAI: ZoneHeader :: !HeaderName -> ZoneBy
+ Keter.RateLimiter.WAI: ZoneIP :: ZoneBy
+ Keter.RateLimiter.WAI: [envGetRequestIPZone] :: Env -> Request -> IPZoneIdentifier
+ Keter.RateLimiter.WAI: [envThrottles] :: Env -> IORef (HashMap Text ThrottleConfig)
+ Keter.RateLimiter.WAI: [envZoneCachesMap] :: Env -> IORef (HashMap IPZoneIdentifier ZoneSpecificCaches)
+ Keter.RateLimiter.WAI: [rlAlgo] :: RLThrottle -> !Algorithm
+ Keter.RateLimiter.WAI: [rlIdBy] :: RLThrottle -> !IdentifierBy
+ Keter.RateLimiter.WAI: [rlLimit] :: RLThrottle -> !Int
+ Keter.RateLimiter.WAI: [rlName] :: RLThrottle -> !Text
+ Keter.RateLimiter.WAI: [rlPeriod] :: RLThrottle -> !Int
+ Keter.RateLimiter.WAI: [rlThrottles] :: RateLimiterConfig -> ![RLThrottle]
+ Keter.RateLimiter.WAI: [rlTokenBucketTTL] :: RLThrottle -> !Maybe Int
+ Keter.RateLimiter.WAI: [rlZoneBy] :: RateLimiterConfig -> !ZoneBy
+ Keter.RateLimiter.WAI: buildRateLimiter :: RateLimiterConfig -> IO Middleware
+ Keter.RateLimiter.WAI: data IdentifierBy
+ Keter.RateLimiter.WAI: data RLThrottle
+ Keter.RateLimiter.WAI: data RateLimiterConfig
+ Keter.RateLimiter.WAI: data ZoneBy
+ Keter.RateLimiter.WAI: fromHeaderName :: HeaderName -> ByteString
+ Keter.RateLimiter.WAI: getClientIPPure :: Request -> IPZoneIdentifier
+ Keter.RateLimiter.WAI: hdr :: Text -> HeaderName
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.FromJSON.FromJSON Keter.RateLimiter.WAI.IdentifierBy
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.FromJSON.FromJSON Keter.RateLimiter.WAI.RLThrottle
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.FromJSON.FromJSON Keter.RateLimiter.WAI.RateLimiterConfig
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.FromJSON.FromJSON Keter.RateLimiter.WAI.ZoneBy
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.ToJSON.ToJSON Keter.RateLimiter.WAI.IdentifierBy
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.ToJSON.ToJSON Keter.RateLimiter.WAI.RLThrottle
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.ToJSON.ToJSON Keter.RateLimiter.WAI.RateLimiterConfig
+ Keter.RateLimiter.WAI: instance Data.Aeson.Types.ToJSON.ToJSON Keter.RateLimiter.WAI.ZoneBy
+ Keter.RateLimiter.WAI: instance GHC.Classes.Eq Keter.RateLimiter.WAI.IdentifierBy
+ Keter.RateLimiter.WAI: instance GHC.Classes.Eq Keter.RateLimiter.WAI.RLThrottle
+ Keter.RateLimiter.WAI: instance GHC.Classes.Eq Keter.RateLimiter.WAI.RateLimiterConfig
+ Keter.RateLimiter.WAI: instance GHC.Classes.Eq Keter.RateLimiter.WAI.ZoneBy
+ Keter.RateLimiter.WAI: instance GHC.Generics.Generic Keter.RateLimiter.WAI.IdentifierBy
+ Keter.RateLimiter.WAI: instance GHC.Generics.Generic Keter.RateLimiter.WAI.RLThrottle
+ Keter.RateLimiter.WAI: instance GHC.Generics.Generic Keter.RateLimiter.WAI.RateLimiterConfig
+ Keter.RateLimiter.WAI: instance GHC.Generics.Generic Keter.RateLimiter.WAI.ZoneBy
+ Keter.RateLimiter.WAI: instance GHC.Show.Show Keter.RateLimiter.WAI.IdentifierBy
+ Keter.RateLimiter.WAI: instance GHC.Show.Show Keter.RateLimiter.WAI.RLThrottle
+ Keter.RateLimiter.WAI: instance GHC.Show.Show Keter.RateLimiter.WAI.RateLimiterConfig
+ Keter.RateLimiter.WAI: instance GHC.Show.Show Keter.RateLimiter.WAI.ZoneBy
+ Keter.RateLimiter.WAI: mkIdentifier :: IdentifierBy -> Request -> IO (Maybe Text)
+ Keter.RateLimiter.WAI: mkZoneFn :: ZoneBy -> Request -> IPZoneIdentifier
+ Keter.RateLimiter.WAI: registerThrottle :: Env -> RLThrottle -> IO ()
- Keter.RateLimiter.Cache: makeCacheKey :: Algorithm -> Text -> Text -> Text
+ Keter.RateLimiter.Cache: makeCacheKey :: Text -> Algorithm -> Text -> Text -> Text
- Keter.RateLimiter.CacheWithZone: allowFixedWindowRequest :: Cache (InMemoryStore 'FixedWindow) -> Text -> Text -> Int -> Int -> IO Bool
+ Keter.RateLimiter.CacheWithZone: allowFixedWindowRequest :: Cache (InMemoryStore 'FixedWindow) -> Text -> Text -> Text -> Int -> Int -> IO Bool
- Keter.RateLimiter.CacheWithZone: deleteCacheWithZone :: CacheStore store v IO => Cache store -> Text -> Text -> IO ()
+ Keter.RateLimiter.CacheWithZone: deleteCacheWithZone :: CacheStore store v IO => Cache store -> Text -> Text -> Text -> IO ()
- Keter.RateLimiter.CacheWithZone: incStoreWithZone :: (CacheStore store v IO, FromJSON v, ToJSON v, Num v, Ord v) => Cache store -> Text -> Text -> Int -> IO v
+ Keter.RateLimiter.CacheWithZone: incStoreWithZone :: (CacheStore store v IO, FromJSON v, ToJSON v, Num v, Ord v) => Cache store -> Text -> Text -> Text -> Int -> IO v
- Keter.RateLimiter.CacheWithZone: readCacheWithZone :: CacheStore store v IO => Cache store -> Text -> Text -> IO (Maybe v)
+ Keter.RateLimiter.CacheWithZone: readCacheWithZone :: CacheStore store v IO => Cache store -> Text -> Text -> Text -> IO (Maybe v)
- Keter.RateLimiter.CacheWithZone: writeCacheWithZone :: CacheStore store v IO => Cache store -> Text -> Text -> v -> Int -> IO ()
+ Keter.RateLimiter.CacheWithZone: writeCacheWithZone :: CacheStore store v IO => Cache store -> Text -> Text -> Text -> v -> Int -> IO ()
- Keter.RateLimiter.LeakyBucket: allowRequest :: MonadIO m => Cache (InMemoryStore 'LeakyBucket) -> Text -> Text -> Int -> Double -> m Bool
+ Keter.RateLimiter.LeakyBucket: allowRequest :: MonadIO m => Cache (InMemoryStore 'LeakyBucket) -> Text -> Text -> Text -> Int -> Double -> m Bool
- Keter.RateLimiter.SlidingWindow: allowRequest :: IO Double -> TVar (Map Text [Double]) -> Text -> Text -> Int -> Int -> IO Bool
+ Keter.RateLimiter.SlidingWindow: allowRequest :: IO Double -> TVar (Map Text [Double]) -> Text -> Text -> Text -> Int -> Int -> IO Bool
- Keter.RateLimiter.TokenBucket: allowRequest :: MonadIO m => Cache (InMemoryStore 'TokenBucket) -> Text -> Text -> Int -> Double -> Int -> m Bool
+ Keter.RateLimiter.TokenBucket: allowRequest :: MonadIO m => Cache (InMemoryStore 'TokenBucket) -> Text -> Text -> Text -> Int -> Double -> Int -> m Bool
Files
- CHANGELOG.md +4/−0
- README.md +28/−13
- keter-rate-limiting-plugin.cabal +7/−4
- src/Data/TinyLRU.hs +70/−62
- src/Keter/RateLimiter/AutoPurge.hs +63/−43
- src/Keter/RateLimiter/Cache.hs +147/−36
- src/Keter/RateLimiter/CacheWithZone.hs +72/−29
- src/Keter/RateLimiter/IPZones.hs +26/−16
- src/Keter/RateLimiter/LeakyBucket.hs +60/−49
- src/Keter/RateLimiter/Notifications.hs +20/−4
- src/Keter/RateLimiter/RequestUtils.hs +24/−18
- src/Keter/RateLimiter/SlidingWindow.hs +89/−49
- src/Keter/RateLimiter/TokenBucket.hs +55/−54
- src/Keter/RateLimiter/TokenBucketWorker.hs +13/−12
- src/Keter/RateLimiter/Types.hs +25/−24
- src/Keter/RateLimiter/WAI.hs +438/−71
- test/Keter/RateLimiter/IPZonesTests.hs +11/−11
- test/Keter/RateLimiter/LeakyBucketTests.hs +9/−7
- test/Keter/RateLimiter/SlidingWindowTests.hs +4/−4
- test/Keter/RateLimiter/WAITests.hs +938/−349
- test/Main.hs +13/−16
- test/TinyLRUTests.hs +19/−17
CHANGELOG.md view
@@ -12,3 +12,7 @@ * First version revised B. Fixed some documentation issues. +## 0.1.1.0 -- 2025-08-11++* 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. +
README.md view
@@ -1,6 +1,6 @@ # keter-rate-limiting-plugin -**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 and IP zone isolation.+**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. @@ -12,22 +12,22 @@ - 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.- - **Flexible Throttle Configuration**: Set limits, periods, algorithms, and unique identifiers on a per-throttle basis.+ - **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.+ - **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 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 for high-concurrency workloads.+ - **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.+ - **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,7 +52,7 @@ ## 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.+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. ```haskell {-# LANGUAGE OverloadedStrings #-}@@ -74,6 +74,7 @@ -- 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) -- 2. Define a throttle rule.@@ -86,6 +87,7 @@ } -- 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 -- 4. Wrap your application with the middleware.@@ -99,7 +101,7 @@ ### Using the Convenient API -The `CacheWithZone` module provides helpers that automatically compose cache keys from the algorithm, zone, and user key, simplifying common use cases.+The `CacheWithZone` module provides helpers that automatically compose cache keys from the algorithm, zone, and user key, simplifying common use cases while leveraging efficient HashMap-based zone lookups. ```haskell import Keter.RateLimiter.Cache@@ -112,12 +114,13 @@ -- Increment a counter for a user in a specific zone. -- The key "rate_limiter:zoneX:userX" is created automatically. -- The request is allowed if the count is within the limit.+-- Zone lookup uses HashMap for O(1) average performance. isAllowed <- allowFixedWindowRequest cache "zoneX" "userX" 100 3600 -- 100 requests per hour ``` ### Using the Customizable API -For more complex scenarios, you can manually construct cache keys and interact directly with the `Cache` module. This gives you full control over the key structure.+For more complex scenarios, you can manually construct cache keys and interact directly with the `Cache` module. This gives you full control over the key structure while still benefiting from HashMap-optimized storage. ```haskell import Keter.RateLimiter.Cache@@ -134,7 +137,7 @@ ### 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 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. ```haskell import Keter.RateLimiter.WAI@@ -149,11 +152,22 @@ } -- env' <- addThrottle env "api/token" tokenBucketThrottle+-- The throttle will be stored in the environment's HashMap for efficient retrieval ``` +## Performance Characteristics++This library is optimized for high-performance scenarios:++- **HashMap-based zone caches**: O(1) average-case lookup for IP zone cache resolution+- **HashMap-based throttle storage**: O(1) average-case retrieval of throttle configurations+- **STM-based concurrent access**: Thread-safe operations with minimal contention+- **Memory-efficient algorithms**: Automatic cleanup of expired entries across all rate limiting algorithms+- **Scalable architecture**: Designed to handle thousands of concurrent requests with minimal overhead+ ## Testing -This package includes an extensive test suite covering all supported rate-limiting algorithms, IP zone isolation, and cache management.+This package includes an extensive test suite covering all supported rate-limiting algorithms, IP zone isolation, cache management, and HashMap-based performance optimizations. To run the tests: @@ -171,8 +185,9 @@ - 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.+ - 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. ## License
keter-rate-limiting-plugin.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: keter-rate-limiting-plugin-version: 0.1.0.2+version: 0.1.1.0 synopsis: Simple Keter rate limiting plugin. description: A modern, high-performance, and highly customisable rate limiting plugin for keter.@@ -18,7 +18,7 @@ build-type: Simple extra-doc-files: CHANGELOG.md, README.md extra-source-files: README.md-tested-with: GHC ==9.2.8, GHC ==9.4.8, GHC ==9.6.5+tested-with: GHC ==9.2.8, GHC ==9.4.8, GHC ==9.6.5, GHC ==9.8.4 common warnings ghc-options: -Wall@@ -53,12 +53,12 @@ case-insensitive >=1.2.1.0, aeson >= 1.4, cache >= 0.1,- containers >= 0.6, time >= 1.9, focus >=1.0.3, stm >=2.5.0 && <2.6, stm-containers >=1.2 && <2, deepseq,+ cookie, list-t >=1.0.5 && <2, network >=3.1.2 && <3.2 || ^>=3.2.0, unordered-containers >=0.2.17 && <0.3,@@ -101,13 +101,16 @@ tasty >= 1.4 && < 1.6, filepath >=1.4.2 && <1.6, tasty-hunit >= 0.10 && < 0.11,- containers ^>= 0.6.4 || ^>= 0.7,+ QuickCheck,+ tasty-quickcheck, text >= 1.2.5 && < 3.0,+ cookie, random >=1.2.1 && <1.4, stm >=2.5.0 && <2.6, stm-containers >=1.2 && <2, cache >= 0.1, containers >= 0.6,+ unordered-containers >=0.2.17 && <0.3, time >= 1.9, hashable >= 1.4.2.0 && < 2, clock >= 0.8.3 && < 1,
src/Data/TinyLRU.hs view
@@ -145,8 +145,11 @@ -- lruCache <- atomically $ initTinyLRU 1000 -- @ ----- @param cap The maximum number of items the cache can hold. Must be > 0.--- @return An 'STM' action that yields a new, empty 'TinyLRUCache'.+-- The capacity must be greater than 0.+--+-- * 'cap' - The maximum number of items the cache can hold. Must be > 0.+-- +-- Returns an 'STM' action that yields a new, empty 'TinyLRUCache'. initTinyLRU :: Int -> STM (TinyLRUCache s) initTinyLRU cap = do cache <- Map.new@@ -160,9 +163,10 @@ -- | Checks if a cache node is expired relative to the current time. ----- @param now The current 'TimeSpec'.--- @param node The 'LRUNode' to check.--- @return 'True' if the node is expired, 'False' otherwise.+-- * 'now' - The current 'TimeSpec'.+-- * 'node' - The 'LRUNode' to check.+--+-- Returns 'True' if the node is expired, 'False' otherwise. isExpired :: TimeSpec -> LRUNode s -> Bool isExpired now node = case nodeExpiry node of@@ -173,8 +177,8 @@ -- It correctly updates the 'nodePrev' and 'nodeNext' pointers of the -- neighboring nodes and the list's 'lruHead' and 'lruTail' pointers if necessary. ----- @param listTVar The 'TVar' of the 'LRUList' from which to remove the node.--- @param nodeTVar The 'TVar' of the 'LRUNode' to remove.+-- * 'listTVar' - The 'TVar' of the 'LRUList' from which to remove the node.+-- * 'nodeTVar' - The 'TVar' of the 'LRUNode' to remove. removeNode :: TVar (LRUList s) -> TVar (LRUNode s) -> STM () removeNode listTVar nodeTVar = do node <- readTVar nodeTVar@@ -194,12 +198,13 @@ -- | Creates a new node and adds it to the front (most recently used position) -- of the cache's linked list. ----- @param now The current 'TimeSpec', used for calculating expiry.--- @param ttl The time-to-live in seconds. A value `<= 0` means it never expires.--- @param cache The 'TinyLRUCache' instance.--- @param key The key for the new entry.--- @param value The 'ByteString' value for the new entry.--- @return The 'TVar' of the newly created 'LRUNode'.+-- * 'now' - The current 'TimeSpec', used for calculating expiry.+-- * 'ttl' - The time-to-live in seconds. A value <= 0 means it never expires.+-- * 'cache' - The 'TinyLRUCache' instance.+-- * 'key' - The key for the new entry.+-- * 'value' - The 'ByteString' value for the new entry.+--+-- Returns the 'TVar' of the newly created 'LRUNode'. addToFront :: TimeSpec -> Int -> TinyLRUCache s -> Text -> ByteString -> STM (TVar (LRUNode s)) addToFront now ttl cache key value = do let expiry = mkExpiry now ttl@@ -218,8 +223,8 @@ -- | Moves an existing node to the front of the linked list, marking it as the -- most recently used. This is a core operation for the LRU logic. ----- @param listTVar The 'TVar' of the 'LRUList'.--- @param nodeTVar The 'TVar' of the 'LRUNode' to move.+-- * 'listTVar' - The 'TVar' of the 'LRUList'.+-- * 'nodeTVar' - The 'TVar' of the 'LRUNode' to move. moveToFront :: TVar (LRUList s) -> TVar (LRUNode s) -> STM () moveToFront listTVar nodeTVar = do list <- readTVar listTVar@@ -240,7 +245,7 @@ -- | Evicts the least recently used item from the cache. This involves removing -- the tail of the linked list and deleting the corresponding entry from the hash map. ----- @param cache The 'TinyLRUCache' to perform eviction on.+-- * 'cache' - The 'TinyLRUCache' to perform eviction on. evictLRU :: TinyLRUCache s -> STM () evictLRU cache = do list <- readTVar (lruList cache)@@ -256,8 +261,8 @@ -- | Deletes an entry from the cache by its key. -- This removes the item from both the internal map and the linked list. ----- @param key The key of the item to delete.--- @param cache The 'TinyLRUCache' instance.+-- * 'key' - The key of the item to delete.+-- * 'cache' - The 'TinyLRUCache' instance. deleteKey :: Text -> TinyLRUCache s -> STM () deleteKey key cache = do mNodeTVar <- Map.lookup key (lruCache cache)@@ -280,21 +285,22 @@ -- -- The logic is as follows: ----- 1. It first cleans up any expired items in the cache.--- 2. It looks for the key.--- 3. If the key exists and is not expired, it moves the item to the front (as it's now--- the most recently used) and returns its value wrapped in 'Just'.--- 4. If the key does not exist, or if it exists but has expired, it inserts the--- provided new value. If the cache is full, it evicts the least recently used--- item before insertion. It then returns the new value wrapped in 'Just'.--- 5. If the key is invalid (empty or too long), it returns 'Nothing'.+-- 1. It first cleans up any expired items in the cache.+-- 2. It looks for the key.+-- 3. If the key exists and is not expired, it moves the item to the front (as it's now+-- the most recently used) and returns its value wrapped in 'Just'.+-- 4. If the key does not exist, or if it exists but has expired, it inserts the+-- provided new value. If the cache is full, it evicts the least recently used+-- item before insertion. It then returns the new value wrapped in 'Just'.+-- 5. If the key is invalid (empty or longer than 256 characters), it returns 'Nothing'. ----- @param now The current 'TimeSpec', for expiry checks.--- @param key The key to look up. Length must be between 1 and 256.--- @param val The value to insert if the key is not found. It must have 'ToJSON' and 'FromJSON' instances.--- @param ttl The time-to-live in seconds for the new entry if it's inserted.--- @param cache The 'TinyLRUCache' instance.--- @return An 'STM' action that yields 'Just' the value (either existing or newly inserted), or 'Nothing' if the key is invalid.+-- * 'now' - The current 'TimeSpec', for expiry checks.+-- * 'key' - The key to look up. Length must be between 1 and 256 characters.+-- * 'val' - The value to insert if the key is not found. It must have 'ToJSON' and 'FromJSON' instances.+-- * 'ttl' - The time-to-live in seconds for the new entry if it's inserted.+-- * 'cache' - The 'TinyLRUCache' instance.+--+-- Returns an 'STM' action that yields 'Just' the value (either existing or newly inserted), or 'Nothing' if the key is invalid. access :: forall a s. (FromJSON a, ToJSON a) => TimeSpec -> Text -> a -> Int -> TinyLRUCache s -> STM (Maybe a) access now key val ttl cache | T.null key || T.length key > 256 = return Nothing@@ -331,21 +337,22 @@ -- -- The logic is as follows: ----- 1. It first cleans up any expired items in the cache.--- 2. It looks for the key.--- 3. If the key exists, it updates the value and expiry time, and moves it to the front.--- 4. If the key does not exist, it inserts the new value. If the cache is full, it--- evicts the least recently used item first.--- 5. If the key is invalid (empty or too long), it returns 'Nothing'.+-- 1. It first cleans up any expired items in the cache.+-- 2. It looks for the key.+-- 3. If the key exists, it updates the value and expiry time, and moves it to the front.+-- 4. If the key does not exist, it inserts the new value. If the cache is full, it+-- evicts the least recently used item first.+-- 5. If the key is invalid (empty or longer than 256 characters), it returns 'Nothing'. -- -- Unlike 'access', this function /always/ writes the provided value. ----- @param now The current 'TimeSpec', for expiry calculations.--- @param key The key to update or insert. Length must be between 1 and 256.--- @param val The new value to write. It must have 'ToJSON' and 'FromJSON' instances.--- @param ttl The new time-to-live in seconds for the entry.--- @param cache The 'TinyLRUCache' instance.--- @return An 'STM' action that yields 'Just' the value that was written, or 'Nothing' if the key is invalid.+-- * 'now' - The current 'TimeSpec', for expiry calculations.+-- * 'key' - The key to update or insert. Length must be between 1 and 256 characters.+-- * 'val' - The new value to write. It must have 'ToJSON' and 'FromJSON' instances.+-- * 'ttl' - The new time-to-live in seconds for the entry.+-- * 'cache' - The 'TinyLRUCache' instance.+--+-- Returns an 'STM' action that yields 'Just' the value that was written, or 'Nothing' if the key is invalid. updateValue :: forall a s. (FromJSON a, ToJSON a) => TimeSpec -> Text -> a -> Int -> TinyLRUCache s -> STM (Maybe a) updateValue now key val ttl cache | T.null key || T.length key > 256 = return Nothing@@ -373,7 +380,7 @@ -- | Resets the cache, removing all entries. ----- @param cache The 'TinyLRUCache' to reset.+-- * 'cache' - The 'TinyLRUCache' to reset. resetTinyLRU :: TinyLRUCache s -> STM () resetTinyLRU cache = do Map.reset (lruCache cache)@@ -384,20 +391,21 @@ -- -- The logic is as follows: ----- 1. It looks for the key. The value associated with the key is treated as a counter ('Int').--- 2. If the entry for the key exists and is not expired:--- * If the counter is less than the 'limit', it increments the counter,--- refreshes the expiry time, and returns 'True' (request allowed).--- * If the counter has reached the 'limit', it does nothing and returns 'False' (request denied).--- 3. If the entry does not exist or has expired, it creates a new entry with the--- counter set to 1 and returns 'True'.+-- 1. It looks for the key. The value associated with the key is treated as a counter ('Int').+-- 2. If the entry for the key exists and is not expired:+-- * If the counter is less than the 'limit', it increments the counter,+-- refreshes the expiry time, and returns 'True' (request allowed).+-- * If the counter has reached the 'limit', it does nothing and returns 'False' (request denied).+-- 3. If the entry does not exist or has expired, it creates a new entry with the+-- counter set to 1 and returns 'True'. ----- @param now The current 'TimeSpec'.--- @param cache The 'TinyLRUCache' instance.--- @param key The key to identify the entity being rate-limited (e.g., a user ID or IP address).--- @param limit The maximum number of requests allowed.--- @param period The time period in seconds for the rate limit window.--- @return An 'STM' action yielding 'True' if the request is allowed, 'False' otherwise.+-- * 'now' - The current 'TimeSpec'.+-- * 'cache' - The 'TinyLRUCache' instance.+-- * 'key' - The key to identify the entity being rate-limited (e.g., a user ID or IP address).+-- * 'limit' - The maximum number of requests allowed.+-- * 'period' - The time period in seconds for the rate limit window.+--+-- Returns an 'STM' action yielding 'True' if the request is allowed, 'False' otherwise. allowRequestTinyLRU :: TimeSpec -> TinyLRUCache s -> Text -> Int -> Int -> STM Bool allowRequestTinyLRU now cache key limit period | T.null key = return False@@ -443,8 +451,8 @@ -- | Low-level function to remove a node from the cache's list. -- Alias for 'removeNode'. Most users should use 'deleteKey' instead. ----- @param cache The 'TinyLRUCache' instance.--- @param nodeTVar The 'TVar' of the 'LRUNode' to remove.+-- * 'cache' - The 'TinyLRUCache' instance.+-- * 'nodeTVar' - The 'TVar' of the 'LRUNode' to remove. removeNodeFromCache :: TinyLRUCache s -> TVar (LRUNode s) -> STM () removeNodeFromCache cache = removeNode (lruList cache) @@ -452,7 +460,7 @@ -- Alias for 'moveToFront'. Most users should use 'access' or 'updateValue' which -- call this internally. ----- @param cache The 'TinyLRUCache' instance.--- @param nodeTVar The 'TVar' of the 'LRUNode' to move.+-- * 'cache' - The 'TinyLRUCache' instance.+-- * 'nodeTVar' - The 'TVar' of the 'LRUNode' to move. moveToFrontInCache :: TinyLRUCache s -> TVar (LRUNode s) -> STM () moveToFrontInCache cache = moveToFront (lruList cache)
src/Keter/RateLimiter/AutoPurge.hs view
@@ -350,19 +350,25 @@ -- __Memory Impact:__ Prevents unbounded memory growth in applications with dynamic rate limiting keys. -- -- __Performance:__ O(n) scan of all buckets, but typically runs infrequently during low-traffic periods.+--+-- * 'stmMap' - STM map containing token bucket entries keyed by identifier.+-- Typically uses API keys, user IDs, or IP addresses as keys.+-- Map is scanned atomically during each purge cycle.+-- * 'intervalSeconds' - Purge interval in seconds. How often to scan for expired buckets.+-- Balance between cleanup responsiveness and CPU usage.+-- Recommended: 300-3600 seconds depending on traffic patterns.+-- * 'ttlSeconds' - TTL (time-to-live) in seconds. Buckets unused for this duration+-- are considered expired and eligible for removal.+-- Should be much larger than typical request intervals.+-- Recommended: 3600-86400 seconds.+--+-- Returns thread ID of the background purge thread.+-- Can be used with 'killThread' to stop purging. startCustomPurgeTokenBucket- :: StmMap.Map Text TokenBucketEntry -- ^ STM map containing token bucket entries keyed by identifier.- -- Typically uses API keys, user IDs, or IP addresses as keys.- -- Map is scanned atomically during each purge cycle.- -> Integer -- ^ Purge interval in seconds. How often to scan for expired buckets.- -- Balance between cleanup responsiveness and CPU usage.- -- Recommended: 300-3600 seconds depending on traffic patterns.- -> Integer -- ^ TTL (time-to-live) in seconds. Buckets unused for this duration- -- are considered expired and eligible for removal.- -- Should be much larger than typical request intervals.- -- Recommended: 3600-86400 seconds.- -> IO ThreadId -- ^ Returns thread ID of the background purge thread.- -- Can be used with 'killThread' to stop purging.+ :: StmMap.Map Text TokenBucketEntry+ -> Integer+ -> Integer+ -> IO ThreadId startCustomPurgeTokenBucket stmMap intervalSeconds ttlSeconds = startCustomPurge (\entry -> do TokenBucketState _ lastT <- readTVar (tbeState entry)@@ -418,18 +424,24 @@ -- __Precision:__ Sub-second timestamp accuracy for fine-grained rate control. -- -- __Cleanup Behavior:__ More aggressive cleanup suitable for high-frequency, short-lived connections.+--+-- * 'stmMap' - STM map containing leaky bucket entries keyed by identifier.+-- Keys typically represent connections, streams, or data flows.+-- All entries are scanned atomically during purge operations.+-- * 'intervalSeconds' - Purge interval in seconds. Frequency of expired entry cleanup.+-- For high-frequency applications, shorter intervals (60-600s)+-- provide more responsive cleanup.+-- * 'ttlSeconds' - TTL in seconds. Buckets inactive for this duration are purged.+-- Should account for typical connection/stream lifetime patterns.+-- Shorter TTL (900-3600s) suitable for transient connections.+--+-- Returns thread ID for background purge thread management.+-- Thread runs continuously until explicitly terminated. startCustomPurgeLeakyBucket- :: StmMap.Map Text LeakyBucketEntry -- ^ STM map containing leaky bucket entries keyed by identifier.- -- Keys typically represent connections, streams, or data flows.- -- All entries are scanned atomically during purge operations.- -> Integer -- ^ Purge interval in seconds. Frequency of expired entry cleanup.- -- For high-frequency applications, shorter intervals (60-600s)- -- provide more responsive cleanup.- -> Integer -- ^ TTL in seconds. Buckets inactive for this duration are purged.- -- Should account for typical connection/stream lifetime patterns.- -- Shorter TTL (900-3600s) suitable for transient connections.- -> IO ThreadId -- ^ Thread ID for background purge thread management.- -- Thread runs continuously until explicitly terminated.+ :: StmMap.Map Text LeakyBucketEntry+ -> Integer+ -> Integer+ -> IO ThreadId startCustomPurgeLeakyBucket stmMap intervalSeconds ttlSeconds = startCustomPurge (\entry -> do LeakyBucketState _ lastT <- readTVar (lbeState entry)@@ -520,28 +532,36 @@ -- __Extensibility:__ Custom delete actions enable complex cleanup scenarios. -- -- __Reliability:__ Robust error handling ensures continuous operation.+--+-- * 'getTimestamp' - Timestamp extraction function. Called for each entry+-- to determine last activity time. Should be fast and+-- side-effect free. Returns Unix timestamp as 'Double'+-- for sub-second precision.+-- * 'deleteAction' - Custom delete action executed for expired entries.+-- Receives both key and entry for context. Should+-- handle resource cleanup and map removal atomically.+-- Can perform logging, resource release, etc.+-- * 'stmMap' - STM map to purge. Entries are identified by 'Text' keys+-- and can be any type 'entry'. Map is scanned completely+-- during each purge cycle for expired entries.+-- * 'intervalSeconds' - Purge interval in seconds. Controls how frequently+-- the purge operation runs. Shorter intervals provide+-- more responsive cleanup but increase CPU usage.+-- * 'ttlSeconds' - TTL (time-to-live) in seconds. Entries with timestamps+-- older than (currentTime - ttlSeconds) are considered+-- expired and eligible for removal.+--+-- Returns thread ID of the purge thread. Thread runs+-- indefinitely until killed. Can be used for thread+-- management and monitoring. startCustomPurge :: forall entry.- (entry -> STM Double) -- ^ Timestamp extraction function. Called for each entry- -- to determine last activity time. Should be fast and- -- side-effect free. Returns Unix timestamp as 'Double'- -- for sub-second precision.- -> (Text -> entry -> STM ()) -- ^ Custom delete action executed for expired entries.- -- Receives both key and entry for context. Should- -- handle resource cleanup and map removal atomically.- -- Can perform logging, resource release, etc.- -> StmMap.Map Text entry -- ^ STM map to purge. Entries are identified by 'Text' keys- -- and can be any type 'entry'. Map is scanned completely- -- during each purge cycle for expired entries.- -> Integer -- ^ Purge interval in seconds. Controls how frequently- -- the purge operation runs. Shorter intervals provide- -- more responsive cleanup but increase CPU usage.- -> Integer -- ^ TTL (time-to-live) in seconds. Entries with timestamps- -- older than (currentTime - ttlSeconds) are considered- -- expired and eligible for removal.- -> IO ThreadId -- ^ Returns thread ID of the purge thread. Thread runs- -- indefinitely until killed. Can be used for thread- -- management and monitoring.+ (entry -> STM Double)+ -> (Text -> entry -> STM ())+ -> StmMap.Map Text entry+ -> Integer+ -> Integer+ -> IO ThreadId startCustomPurge getTimestamp deleteAction stmMap intervalSeconds ttlSeconds = do purgeSignal <- newMVar () forkIO $ forever $ do
src/Keter/RateLimiter/Cache.hs view
@@ -107,11 +107,11 @@ -- Type-safe cache creation tokenCache <- do- store <- createInMemoryStore @'TokenBucket+ store <- createInMemoryStore \@'TokenBucket return $ newCache TokenBucket store fixedWindowCache <- do- store <- createInMemoryStore @'FixedWindow + store <- createInMemoryStore \@'FixedWindow return $ newCache FixedWindow store @ @@ -130,7 +130,7 @@ @ -- Composite key generation-let userKey = makeCacheKey TokenBucket "zone1" "user456"+let userKey = makeCacheKey "throttle1" TokenBucket "zone1" "user456" writeCache tokenCache userKey state 7200 -- Cleanup and reset@@ -211,6 +211,9 @@ -- * Entry Creation , createTokenBucketEntry , createLeakyBucketEntry+ -- * Algorithm utilities+ , algoToText+ , parseAlgoText ) where import Control.Concurrent.STM@@ -219,8 +222,10 @@ import Control.Monad (void, forever) import Control.Concurrent (forkIO, threadDelay) import Data.Aeson (ToJSON, FromJSON, decodeStrict, encode)+import Data.Aeson.Types (Parser) import qualified Data.ByteString.Lazy as LBS import Data.Text (Text)+import qualified Data.Text as Tx (unpack, toLower) import qualified Data.TinyLRU as TinyLRU import Data.Text.Encoding (encodeUtf8, decodeUtf8') import qualified Data.Cache as C@@ -305,18 +310,18 @@ -- -- ==== Type Parameters ----- * @store@: The storage backend type (e.g., 'InMemoryStore', Redis, etc.)--- * The algorithm is captured in the store type for type safety+-- The 'store' parameter represents the storage backend type (e.g., 'InMemoryStore', Redis, etc.)+-- The algorithm is captured in the store type for type safety -- -- ==== Example Usage -- -- @ -- -- Create a token bucket cache--- tokenStore <- createInMemoryStore @'TokenBucket+-- tokenStore <- createInMemoryStore \@'TokenBucket -- let tokenCache = Cache TokenBucket tokenStore -- -- -- Create a fixed window cache--- counterStore <- createInMemoryStore @'FixedWindow+-- counterStore <- createInMemoryStore \@'FixedWindow -- let counterCache = Cache FixedWindow counterStore -- @ data Cache store = Cache@@ -521,11 +526,11 @@ -- -- @ -- -- Compiler ensures correct store type--- tokenStore <- createStore @'TokenBucket -- Creates TokenBucketStore--- counterStore <- createStore @'FixedWindow -- Creates CounterStore+-- tokenStore <- createStore \@'TokenBucket -- Creates TokenBucketStore+-- counterStore <- createStore \@'FixedWindow -- Creates CounterStore -- -- -- This would be a compile error:--- -- tokenStore <- createStore @'SlidingWindow -- Type mismatch!+-- -- tokenStore <- createStore \@'SlidingWindow -- Type mismatch! -- @ -- -- ==== Automatic Services@@ -637,12 +642,12 @@ -- -- @ -- -- These are all valid and type-safe:--- tokenStore <- createInMemoryStore @'TokenBucket--- counterStore <- createInMemoryStore @'FixedWindow--- lruStore <- createInMemoryStore @'TinyLRU+-- tokenStore <- createInMemoryStore \@'TokenBucket+-- counterStore <- createInMemoryStore \@'FixedWindow+-- lruStore <- createInMemoryStore \@'TinyLRU -- -- -- This would be a compile error (typo):--- -- badStore <- createInMemoryStore @'TokenBuckett -- Not a valid algorithm+-- -- badStore <- createInMemoryStore \@'TokenBuckett -- Not a valid algorithm -- @ -- -- ==== Background Services@@ -660,8 +665,8 @@ -- -- main = do -- -- Create stores for different algorithms--- tokenStore <- createInMemoryStore @'TokenBucket--- slidingStore <- createInMemoryStore @'SlidingWindow+-- tokenStore <- createInMemoryStore \@'TokenBucket+-- slidingStore <- createInMemoryStore \@'SlidingWindow -- -- -- Use in cache creation -- let tokenCache = newCache TokenBucket tokenStore@@ -687,7 +692,7 @@ -- -- @ -- -- Create a token bucket cache with in-memory storage--- store <- createInMemoryStore @'TokenBucket+-- store <- createInMemoryStore \@'TokenBucket -- let cache = newCache TokenBucket store -- -- -- Later operations use the unified cache interface@@ -1109,8 +1114,8 @@ -- -- @ -- -- Start worker for streaming rate limiter--- startLeakyBucketWorker stateVar queue 100 2.0 "stream123"--- -- Capacity: 100 requests, Leak rate: 2 requests/second+-- startLeakyBucketWorker stateVar queue 100 2.0+-- -- Capacity: 100 requests, Leak rate: 2 requests\/second -- @ startLeakyBucketWorker :: TVar LeakyBucketState -- ^ Shared bucket state@@ -1198,7 +1203,7 @@ newCache <- C.newCache Nothing atomically $ writeTVar tvar newCache --- | Compose a unique cache key from algorithm, IP zone, and user identifier.+-- | Compose a unique cache key from throttle name, algorithm, IP zone, and user identifier. -- -- Creates hierarchical cache keys that prevent collisions and enable -- efficient organization of rate limiting data. The key format follows@@ -1207,17 +1212,7 @@ -- ==== Key Format -- -- @--- "algorithm_prefix:ip_zone:user_key"--- @------ ==== Examples------ @--- makeCacheKey TokenBucket "api" "user123"--- -- Result: "token_bucket:api:user123"------ makeCacheKey FixedWindow "web" "192.168.1.1" --- -- Result: "rate_limiter:web:192.168.1.1"+-- \<algorithm>:\<throttleName>:\<ipZone>:\<userKey> -- @ -- -- ==== Use Cases@@ -1233,8 +1228,124 @@ -- * __Query Efficiency__: Pattern-based queries and cleanup -- * __Debugging__: Clear key structure aids troubleshooting -- * __Monitoring__: Easy to aggregate metrics by zone or user type-makeCacheKey :: Algorithm -- ^ Rate limiting algorithm for prefix- -> Text -- ^ IP zone or service identifier (e.g., "api", "web", "zone1")- -> Text -- ^ User key (e.g., API key, user ID, IP address)- -> Text -- ^ Complete hierarchical cache key-makeCacheKey algo ipZone userKey = algorithmPrefix algo <> ":" <> ipZone <> ":" <> userKey+--+-- ==== Example+--+-- @+-- -- Create key for API rate limiting+-- let key = makeCacheKey "api_limit" TokenBucket "us-east-1" "user123"+-- -- Result: "TokenBucket:api_limit:us-east-1:user123"+--+-- -- Create key for login attempts+-- let key = makeCacheKey "login_attempts" FixedWindow "global" "192.168.1.1"+-- -- Result: "FixedWindow:login_attempts:global:192.168.1.1"+-- @+makeCacheKey :: Text -- ^ Throttle name (unique per rule)+ -> Algorithm -- ^ Rate limiting algorithm+ -> Text -- ^ IP zone or region identifier+ -> Text -- ^ User key (user ID, IP address, API key, etc.)+ -> Text -- ^ Complete hierarchical cache key+makeCacheKey throttleName algo ipZone userKey =+ algoToText algo <> ":" <> throttleName <> ":" <> ipZone <> ":" <> userKey++-- | Convert an Algorithm value to its canonical Text representation.+--+-- This function provides a standardized textual representation of rate limiting+-- algorithms for serialization, logging, and configuration files. Each algorithm+-- has exactly one canonical PascalCase representation matching the constructor names.+--+-- ==== Output Format+--+-- @+-- FixedWindow → "FixedWindow"+-- SlidingWindow → "SlidingWindow" +-- TokenBucket → "TokenBucket"+-- LeakyBucket → "LeakyBucket"+-- TinyLRU → "TinyLRU"+-- @+--+-- ==== Examples+--+-- @+-- -- Basic usage+-- algoToText TokenBucket -- "TokenBucket"+--+-- -- Configuration serialization+-- data RateLimitConfig = RateLimitConfig+-- { algorithm :: Text, limit :: Int, window :: Int }+-- +-- config = RateLimitConfig (algoToText SlidingWindow) 1000 3600+--+-- -- Logging+-- logInfo $ "Using " <> algoToText currentAlgo <> " rate limiting"+-- @+--+-- This function is pure, thread-safe, and forms a reversible pair with 'parseAlgoText'.+-- Time complexity: O(1), Space complexity: O(1).+algoToText :: Algorithm -- ^ The algorithm to convert+ -> Text -- ^ Canonical text representation+algoToText FixedWindow = "FixedWindow"+algoToText SlidingWindow = "SlidingWindow"+algoToText TokenBucket = "TokenBucket"+algoToText LeakyBucket = "LeakyBucket"+algoToText TinyLRU = "TinyLRU"++-- | Parse a Text representation back into an Algorithm value.+--+-- Provides flexible parsing with case-insensitive matching and hyphenated format support.+-- Returns an Aeson Parser for seamless JSON integration with rich error reporting.+--+-- ==== Supported Formats+--+-- @+-- Algorithm │ Accepted Inputs (case-insensitive)+-- ──────────────┼───────────────────────────────────+-- FixedWindow │ "fixedwindow", "fixed-window"+-- SlidingWindow │ "slidingwindow", "sliding-window"+-- TokenBucket │ "tokenbucket", "token-bucket"+-- LeakyBucket │ "leakybucket", "leaky-bucket"+-- TinyLRU │ "tinylru", "tiny-lru"+-- @+--+-- ==== Examples+--+-- @+-- -- JSON parsing+-- instance FromJSON RateLimitConfig where+-- parseJSON = withObject "RateLimitConfig" $ \\o ->+-- RateLimitConfig <$> (o .: "algorithm" >>= parseAlgoText)+-- <*> o .: "limit"+--+-- -- Configuration file (accepts flexible formats)+-- parseAlgoText "token-bucket" -- Success TokenBucket+-- parseAlgoText "SLIDING-WINDOW" -- Success SlidingWindow+-- parseAlgoText "invalid" -- Error "Unknown algorithm: invalid"+-- @+--+-- ==== Error Handling+--+-- Calls 'fail' with descriptive error messages for invalid inputs. In aeson's+-- Parser context, this provides rich error reporting with parsing context.+--+-- ==== Round-trip Property+--+-- @+-- forall algo. parseAlgoText (algoToText algo) == Success algo+-- @+--+-- This function is pure, thread-safe, with O(1) time complexity after O(n) normalization.+parseAlgoText :: Text -- ^ Text representation to parse+ -> Parser Algorithm -- ^ Parsed algorithm or error+parseAlgoText t =+ case Tx.toLower t of+ "fixedwindow" -> pure FixedWindow+ "fixed-window" -> pure FixedWindow+ "slidingwindow" -> pure SlidingWindow+ "sliding-window" -> pure SlidingWindow+ "tokenbucket" -> pure TokenBucket+ "token-bucket" -> pure TokenBucket+ "leakybucket" -> pure LeakyBucket+ "leaky-bucket" -> pure LeakyBucket+ "tinylru" -> pure TinyLRU+ "tiny-lru" -> pure TinyLRU+ _ -> fail ("Unknown algorithm: " <> Tx.unpack t)
src/Keter/RateLimiter/CacheWithZone.hs view
@@ -11,7 +11,7 @@ Stability : stable Portability : portable -This module provides utility functions for interacting with `Cache` instances+This module provides utility functions for interacting with 'Cache' instances in a zone-aware manner. Each key is built from an algorithm-specific prefix, an IP zone label, and a user-specific identifier. This allows different logical groups (IP zones) to maintain independent rate limiting state even if@@ -48,25 +48,24 @@ -- zone and user key. The request is allowed if the resulting count -- does not exceed the limit for the given period. ----- == Parameters------ [@cache@] The FixedWindow cache to use.--- [@ipZone@] IP zone or logical tenant label.--- [@userKey@] Unique per-client identifier (e.g., IP or token).--- [@limit@] Maximum number of allowed requests in the given period.--- [@period@] Time window in seconds.------ == Returns------ A boolean indicating whether the request is allowed.+-- The function takes a FixedWindow cache, throttle name, IP zone, user key,+-- request limit, and time period. It returns 'True' if the request should+-- be allowed, 'False' otherwise. ----- == Example+-- ==== __Examples__ ----- > allowed <- allowFixedWindowRequest myCache "zone1" "userA" 10 60+-- @+-- allowed <- allowFixedWindowRequest myCache \"throttle1\" \"zone1\" \"userA\" 10 60+-- if allowed+-- then putStrLn \"Request allowed\"+-- else putStrLn \"Request denied - rate limit exceeded\"+-- @ allowFixedWindowRequest :: Cache (InMemoryStore 'FixedWindow) -- ^ FixedWindow-based cache handle -> Text+ -- ^ Throttle name+ -> Text -- ^ IP zone or tenant key -> Text -- ^ User key@@ -76,8 +75,8 @@ -- ^ Period in seconds -> IO Bool -- ^ Whether the request is allowed-allowFixedWindowRequest cache ipZone userKey limit period = do- count <- incStoreWithZone cache ipZone userKey period+allowFixedWindowRequest cache throttleName ipZone userKey limit period = do+ count <- incStoreWithZone cache throttleName ipZone userKey period return $ count <= limit --------------------------------------------------------------------------------@@ -85,16 +84,25 @@ -- | Increment a cache entry using a composed key derived from the zone and user key. -- -- This is typically used in fixed-window rate limiters to track the number--- of requests within a time period.+-- of requests within a time period. The function creates a composite cache key+-- from the throttle name, algorithm type, IP zone, and user key, then increments+-- the stored counter value. ----- == Requirements+-- The stored type must be JSON-serializable and support numeric operations. ----- The stored type @v@ must be JSON-serializable and support numeric operations.+-- ==== __Examples__+--+-- @+-- newCount <- incStoreWithZone cache \"api-throttle\" \"zone-a\" \"user123\" 3600+-- print newCount -- Prints the incremented counter value+-- @ incStoreWithZone :: (CacheStore store v IO, FromJSON v, ToJSON v, Num v, Ord v) => Cache store -- ^ Cache with algorithm-specific store -> Text+ -- ^ Throttle name+ -> Text -- ^ IP zone -> Text -- ^ User key@@ -102,37 +110,61 @@ -- ^ TTL (seconds) -> IO v -- ^ New counter value-incStoreWithZone cache ipZone userKey expiresIn = do- let key = makeCacheKey (cacheAlgorithm cache) ipZone userKey+incStoreWithZone cache throttleName ipZone userKey expiresIn = do+ let key = makeCacheKey throttleName (cacheAlgorithm cache) ipZone userKey prefix = algorithmPrefix (cacheAlgorithm cache) incStore (cacheStore cache) prefix key expiresIn -- | Read a cache entry using a composed key. ----- Useful for debugging or non-mutating cache inspection.+-- Useful for debugging or non-mutating cache inspection. The function+-- constructs the appropriate cache key from the provided components and+-- retrieves the stored value if it exists.+--+-- ==== __Examples__+--+-- @+-- maybeValue <- readCacheWithZone cache \"throttle1\" \"zone1\" \"user456\"+-- case maybeValue of+-- Just val -> print val+-- Nothing -> putStrLn \"No entry found\"+-- @ readCacheWithZone :: (CacheStore store v IO) => Cache store -- ^ Cache handle -> Text+ -- ^ Throttle name+ -> Text -- ^ IP zone -> Text -- ^ User key -> IO (Maybe v) -- ^ Optional stored value-readCacheWithZone cache ipZone userKey = do- let key = makeCacheKey (cacheAlgorithm cache) ipZone userKey+readCacheWithZone cache throttleName ipZone userKey = do+ let key = makeCacheKey throttleName (cacheAlgorithm cache) ipZone userKey prefix = algorithmPrefix (cacheAlgorithm cache) readStore (cacheStore cache) prefix key -- | Write a value into the cache using a zone-aware key. -- -- Typically used for manual updates, testing, or initializing state.+-- The function creates a composite key and stores the provided value+-- with the specified expiration time.+--+-- ==== __Examples__+--+-- @+-- writeCacheWithZone cache \"throttle1\" \"zone1\" \"user789\" (42 :: Int) 1800+-- putStrLn \"Value written to cache\"+-- @ writeCacheWithZone :: (CacheStore store v IO) => Cache store -- ^ Cache handle -> Text+ -- ^ Throttle name+ -> Text -- ^ IP zone -> Text -- ^ User key@@ -141,24 +173,35 @@ -> Int -- ^ Expiration in seconds -> IO ()-writeCacheWithZone cache ipZone userKey val expiresIn = do- let key = makeCacheKey (cacheAlgorithm cache) ipZone userKey+writeCacheWithZone cache throttleName ipZone userKey val expiresIn = do+ let key = makeCacheKey throttleName (cacheAlgorithm cache) ipZone userKey prefix = algorithmPrefix (cacheAlgorithm cache) writeStore (cacheStore cache) prefix key val expiresIn -- | Delete an entry from the cache using a zone-aware key. ----- Can be used to forcibly reset rate limiter state or clean up.+-- Can be used to forcibly reset rate limiter state or clean up expired+-- entries manually. The function constructs the appropriate cache key+-- and removes the corresponding entry.+--+-- ==== __Examples__+--+-- @+-- deleteCacheWithZone cache \"throttle1\" \"zone1\" \"user123\"+-- putStrLn \"Cache entry deleted\"+-- @ deleteCacheWithZone :: (CacheStore store v IO) => Cache store -- ^ Cache handle -> Text+ -- ^ Throttle name+ -> Text -- ^ IP zone -> Text -- ^ User key -> IO ()-deleteCacheWithZone cache ipZone userKey = do- let key = makeCacheKey (cacheAlgorithm cache) ipZone userKey+deleteCacheWithZone cache throttleName ipZone userKey = do+ let key = makeCacheKey throttleName (cacheAlgorithm cache) ipZone userKey prefix = algorithmPrefix (cacheAlgorithm cache) deleteStore (cacheStore cache) prefix key
src/Keter/RateLimiter/IPZones.hs view
@@ -17,7 +17,7 @@ ensures multi-tenant systems can rate-limit different clients or groups in isolation. -The primary structure here is `ZoneSpecificCaches`, which contains multiple+The primary structure here is 'ZoneSpecificCaches', which contains multiple caches per zone. Utility functions allow dynamic creation, reset, and lookup of caches for specific zones. @@ -61,7 +61,7 @@ -- | Type alias representing an identifier for an IP zone. -- -- This is used as a logical namespace or grouping key for assigning and isolating rate limiters.--- Examples: `"default"`, `"zone-a"`, or `"192.168.1.0/24"`.+-- Examples: @\"default\"@, @\"zone-a\"@, or @\"192.168.1.0\/24\"@. type IPZoneIdentifier = Text -- | The default IP zone identifier used when no specific zone is assigned.@@ -89,13 +89,17 @@ -- | Create a new set of caches for a single IP zone. ----- Each algorithm receives its own store. For `LeakyBucket`, a background+-- Each algorithm receives its own store. For 'LeakyBucket', a background -- cleanup thread is also started to remove inactive entries periodically. ----- == Example+-- The cleanup thread runs every 60 seconds and removes entries older than 2 hours. ----- > zoneCaches <- createZoneCaches--- > cacheReset (zscTokenBucketCache zoneCaches)+-- ==== __Examples__+--+-- @+-- zoneCaches <- createZoneCaches+-- cacheReset (zscTokenBucketCache zoneCaches)+-- @ createZoneCaches :: IO ZoneSpecificCaches createZoneCaches = do counterStore <- createInMemoryStore @'FixedWindow@@ -127,9 +131,11 @@ -- -- Clears all internal state, including token counts, timestamps, and queues. ----- == Example+-- ==== __Examples__ ----- > resetSingleZoneCaches zoneCaches+-- @+-- resetSingleZoneCaches zoneCaches+-- @ resetSingleZoneCaches :: ZoneSpecificCaches -> IO () resetSingleZoneCaches zsc = do cacheReset (zscCounterCache zsc)@@ -142,9 +148,11 @@ -- -- This is useful when only one type of rate limiter needs a reset. ----- == Example+-- ==== __Examples__ ----- > resetZoneCache zoneCaches TokenBucket+-- @+-- resetZoneCache zoneCaches TokenBucket+-- @ resetZoneCache :: ZoneSpecificCaches -> Algorithm -> IO () resetZoneCache zsc algorithm = case algorithm of FixedWindow -> cacheReset (zscCounterCache zsc)@@ -155,14 +163,16 @@ -- | Convert a socket address into an IP zone identifier. ----- IPv4 addresses are rendered using `fromHostAddress`. IPv6 addresses are+-- IPv4 addresses are rendered using 'fromHostAddress'. IPv6 addresses are -- expanded and zero-padded for consistency. Any unknown or unsupported--- address formats fall back to the 'defaultIPZone'.+-- address formats fall back to 'defaultIPZone'. ----- == Example+-- ==== __Examples__ ----- > zone <- sockAddrToIPZone (SockAddrInet 80 0x7f000001)--- > print zone -- "127.0.0.1"+-- @+-- zone <- sockAddrToIPZone (SockAddrInet 80 0x7f000001)+-- print zone -- \"127.0.0.1\"+-- @ sockAddrToIPZone :: SockAddr -> IO Text sockAddrToIPZone (SockAddrInet _ hostAddr) = do let ip = fromHostAddress hostAddr@@ -173,4 +183,4 @@ w3 `shiftR` 16, w3 .&. 0xFFFF, w4 `shiftR` 16, w4 .&. 0xFFFF] where showHexWord n = let s = showHex n "" in if length s < 4 then replicate (4 - length s) '0' ++ s else s-sockAddrToIPZone _ = return "default"+sockAddrToIPZone _ = return defaultIPZone
src/Keter/RateLimiter/LeakyBucket.hs view
@@ -1,5 +1,3 @@--- file: Keter.RateLimiter.LeakyBucket.hs- {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE LambdaCase #-}@@ -17,7 +15,7 @@ This module implements the core logic for the Leaky Bucket rate-limiting algorithm. The primary goal of this algorithm is to smooth out bursts of requests into a steady, predictable flow. It is conceptually similar to a bucket with a hole in the bottom. -Incoming requests are like water being added to the bucket. The bucket has a finite 'capacity'. If a request arrives when the bucket is full, it is rejected (it "spills over"). The hole in the bucket allows requests to be processed (or "leak out") at a constant 'leakRate'.+Incoming requests are like water being added to the bucket. The bucket has a finite capacity. If a request arrives when the bucket is full, it is rejected (it \"spills over\"). The hole in the bucket allows requests to be processed (or \"leak out\") at a constant leak rate. This implementation uses a dedicated worker thread for each bucket (e.g., for each unique user or IP address) to process requests from a queue. This ensures that processing is serialized and state is managed safely under concurrent access. The first request for a given key will spawn the worker, which then serves all subsequent requests for that same key. -}@@ -45,72 +43,88 @@ -- -- A dedicated worker thread is responsible for processing the queue for each bucket. This function ensures -- that a worker is started for a new bucket but avoids starting duplicate workers. The final result--- (`True` or `False`) is communicated back to the caller once the worker has processed the request.+-- ('True' or 'False') is communicated back to the caller once the worker has processed the request. ----- ==== __Example__+-- === Algorithm Flow ----- > {-# LANGUAGE OverloadedStrings #-}--- > import Keter.RateLimiter.LeakyBucket (allowRequest)--- > import Keter.RateLimiter.Cache--- >--- > main :: IO ()--- > main = do--- > -- Initialize the cache store for the LeakyBucket algorithm.--- > leakyStore <- createStore @'LeakyBucket--- > let cache = newCache LeakyBucket leakyStore--- >--- > -- Simulate a request for a user from a specific IP zone.--- > -- Capacity: 10 requests (burst limit)--- > -- Leak Rate: 1 request per second--- > isAllowed <- allowRequest cache "us-east-1" "user-123" 10 1.0--- >--- > if isAllowed--- > then putStrLn "Request is allowed."--- > else putStrLn "Request is blocked (429 Too Many Requests)."+-- 1. __Validation__: Check if capacity is valid (> 0)+-- 2. __Key Construction__: Build composite cache key from throttle name, IP zone, and user key+-- 3. __Bucket Management__: Find existing bucket or create new one atomically+-- 4. __Request Queuing__: Add request to bucket's processing queue+-- 5. __Worker Management__: Start worker thread for new buckets (one-time operation)+-- 6. __Response__: Wait for worker to process request and return result --+-- === Concurrency Model+--+-- * Each bucket has its own worker thread for serialized processing+-- * Multiple clients can safely call this function concurrently+-- * STM ensures atomic bucket creation and state management+-- * Workers are started lazily on first request per bucket+--+-- ==== __Examples__+--+-- @+-- -- Basic usage: 10 request capacity, 1 request per second leak rate+-- isAllowed <- allowRequest cache \"api-throttle\" \"us-east-1\" \"user-123\" 10 1.0+-- if isAllowed+-- then putStrLn \"Request is allowed.\"+-- else putStrLn \"Request is blocked (429 Too Many Requests).\"+-- @+--+-- @+-- -- High-throughput API with burst tolerance+-- let capacity = 100 -- Allow burst of up to 100 requests+-- leakRate = 10.0 -- Process 10 requests per second steadily+-- +-- result <- allowRequest cache \"high-volume\" \"zone-premium\" \"client-456\" capacity leakRate+-- when result $ processApiRequest+-- @+--+-- @+-- -- Rate limiting for different service tiers+-- let (cap, rate) = case userTier of+-- Premium -> (50, 5.0) -- 50 burst, 5\/sec sustained+-- Standard -> (20, 2.0) -- 20 burst, 2\/sec sustained +-- Basic -> (5, 0.5) -- 5 burst, 1 per 2 seconds+-- +-- allowed <- allowRequest cache \"tiered-api\" zone userId cap rate+-- @+--+-- /Thread Safety:/ This function is fully thread-safe and can be called+-- concurrently from multiple threads.+--+-- /Performance:/ New buckets incur a one-time worker thread creation cost.+-- Subsequent requests are queued with minimal overhead. allowRequest :: MonadIO m => Cache (InMemoryStore 'LeakyBucket)- -- ^ The cache containing the state of all leaky buckets.+ -- ^ Leaky bucket cache instance -> Text- -- ^ The IP-based zone for segmenting the cache (e.g., "eu-central-1").+ -- ^ Throttle name (logical grouping identifier) -> Text- -- ^ A unique key identifying the user or client (e.g., an IP address or API key).+ -- ^ IP zone identifier for multi-tenant isolation+ -> Text+ -- ^ User key (unique client identifier) -> Int- -- ^ The bucket's capacity (the maximum burst size). If a request arrives when the bucket is full, it will be rejected.+ -- ^ Bucket capacity (maximum queued requests, must be > 0) -> Double- -- ^ The rate at which the bucket "leaks" (in requests per second). This determines the sustained throughput.+ -- ^ Leak rate in requests per second (must be positive) -> m Bool- -- ^ Returns 'True' if the request is allowed, 'False' otherwise.-allowRequest cache ipZone userKey capacity leakRate = liftIO $ do- -- Immediately reject if capacity is non-positive, as it's an invalid configuration.+ -- ^ 'True' if request is allowed, 'False' if bucket is full+allowRequest cache throttleName ipZone userKey capacity leakRate = liftIO $ do if capacity <= 0 then pure False else do now <- realToFrac <$> getPOSIXTime- let fullKey = makeCacheKey LeakyBucket ipZone userKey+ let fullKey = makeCacheKey throttleName LeakyBucket ipZone userKey LeakyBucketStore tvBuckets = cacheStore cache replyVar <- newEmptyTMVarIO-- -- Create a template for a new bucket entry. Its state is neutral (level 0). newEntry <- createLeakyBucketEntry (LeakyBucketState 0 now)-- -- Atomically get the existing bucket for the key or create a new one.- -- This 'focus' operation provides an atomic get-or-create pattern. entry <- atomically $ do buckets <- readTVar tvBuckets StmMap.focus- -- If the key is missing, create and insert the newEntry. (F.cases (newEntry, F.Set newEntry)- -- If the key exists, just return the existing entry. (\existing -> (existing, F.Leave))) fullKey buckets-- -- All requests, new or existing, queue up to be processed by the worker.- -- This serializes access to the bucket's state. atomically $ writeTQueue (lbeQueue entry) replyVar-- -- The first request to arrive will acquire the lock and start the worker.- -- Subsequent requests for the same bucket won't be able to acquire the lock- -- and will simply wait for the running worker to process them. started <- atomically $ tryPutTMVar (lbeWorkerLock entry) () when started $ startLeakyBucketWorker@@ -118,7 +132,4 @@ (lbeQueue entry) capacity leakRate-- -- All callers block here, waiting for the worker to process their specific- -- request and provide a boolean result in the reply variable. atomically $ takeTMVar replyVar
src/Keter/RateLimiter/Notifications.hs view
@@ -52,7 +52,7 @@ import Network.Wai import Network.Socket import Network.Wai.Handler.Warp (run)-import Network.HTTP.Types (status429)+import Network.HTTP.Types (status200, status429) import Control.Monad.IO.Class (liftIO) -- A hypothetical function called when a client hits a rate limit.@@ -87,6 +87,8 @@ @ import qualified Data.Text.IO as TIO import Control.Exception (try, SomeException)+import Data.Time.Clock (getCurrentTime)+import Data.Time.Format (defaultTimeLocale, formatTime) -- A notifier that attempts to write to a file, falling back to console on error fileNotifier :: FilePath -> Notifier@@ -144,10 +146,17 @@ YYYY-MM-DD HH:MM:SS - throttleName action item (limit: N) @ +When the throttle name is empty, it is omitted from the output:++@+YYYY-MM-DD HH:MM:SS - action item (limit: N)+@+ Examples: @ 2025-01-30 13:45:12 - loginAttempts blocked "user123" (limit: 5) 2025-01-30 13:45:13 - api-global blocked "192.168.1.100" (limit: 1000)+2025-01-30 13:45:14 - blocked "item" (limit: 50) @ Note that textual items are quoted due to the use of 'show' for conversion.@@ -175,7 +184,7 @@ * **IPv6 addresses**: @2001:0db8:0000:0000:0000:0000:0000:0001:443@ * **Unix sockets**: Port information is omitted * **Empty query strings**: Only the path is included-* **Empty throttle names**: The throttle name is simply omitted from the output+* **Non-empty query strings**: Included with the leading '?' character == Thread-safety @@ -195,6 +204,8 @@ you may want to wrap these notifiers with appropriate error handling: @+import Control.Exception (try, SomeException)+ safeNotifier :: Notifier -> Notifier safeNotifier baseNotifier = baseNotifier { notifierAction = \\throttle act item limit -> do@@ -232,13 +243,18 @@ Example test notifier: @+import Data.IORef+ createTestNotifier :: IORef [Text] -> Notifier createTestNotifier outputRef = Notifier { notifierName = "test" , notifierAction = \\throttle act item limit -> do now <- getCurrentTime- let message = -- format message as needed- modifyIORef' outputRef (message :)+ let timestamp = formatTime defaultTimeLocale "%Y-%m-%d %H:%M:%S" now+ parts = [throttle, act, item, T.concat ["(limit: ", T.pack (show limit), ")"]]+ message = T.intercalate " " $ filter (not . T.null) parts+ fullMessage = T.pack timestamp <> " - " <> message+ modifyIORef' outputRef (fullMessage :) } @
src/Keter/RateLimiter/RequestUtils.hs view
@@ -19,7 +19,7 @@ 1. /Zero/ allocation whenever the value is already available in the request record (e.g. @rawPathInfo@ or @requestMethod@ are reused verbatim).-2. No reverse DNS or other network round-trips – the functions are pure and+2. No reverse DNS or other network round-trips -- the functions are pure and fast. 3. Header names are handled case-insensitively via the 'Data.CaseInsensitive.CI' type.@@ -31,6 +31,7 @@ import Control.Monad.IO.Class (liftIO) import Network.Wai import Network.Wai.Handler.Warp (run)+import Network.HTTP.Types (status200) import Keter.RateLimiter.RequestUtils (byIPAndPath) import Data.Text.IO as TIO @@ -79,6 +80,7 @@ import Data.Text (Text) import qualified Data.Text as T import qualified Data.Text.Encoding as TE+import qualified Data.Text.Encoding.Error as TEE import Network.Wai (Request) import qualified Network.Wai as WAI import Network.Socket@@ -87,7 +89,7 @@ , HostAddress6 , hostAddressToTuple )-import Network.HTTP.Types.Header (HeaderName)+import Network.HTTP.Types.Header (HeaderName, hHost, hUserAgent) import Data.Bits ((.&.), shiftR) import Data.CaseInsensitive (mk) import Numeric (showHex)@@ -108,7 +110,7 @@ in T.intercalate "." (map (T.pack . show) [o1, o2, o3, o4]) -- | Render an IPv6 'HostAddress6' as eight 16-bit hex blocks separated--- by ‘:’. Each block is zero-padded to four characters. This rendering+-- by ':'. Each block is zero-padded to four characters. This rendering -- is canonical but not compressed (e.g., it does not use @::@). -- -- The function is micro-optimised to avoid lists and string formatting functions.@@ -150,33 +152,34 @@ -- Unix sockets are represented by their file path. getClientIP :: Request -> IO Text getClientIP req = do- let ipTxt = case lookup (mk "x-forwarded-for") (WAI.requestHeaders req) of- Just xff -> T.takeWhile (/= ',') $ TE.decodeUtf8 xff+ let safeDecode = TE.decodeUtf8With TEE.lenientDecode+ ipTxt = case lookup (mk "x-forwarded-for") (WAI.requestHeaders req) of+ Just xff -> T.takeWhile (/= ',') . safeDecode $ xff Nothing -> case lookup (mk "x-real-ip") (WAI.requestHeaders req) of- Just rip -> TE.decodeUtf8 rip+ Just rip -> safeDecode rip Nothing -> case WAI.remoteHost req of SockAddrInet _ addr -> ipv4ToString addr SockAddrInet6 _ _ addr _ -> ipv6ToString addr SockAddrUnix path -> T.pack path pure ipTxt --- | Extracts the raw path info from the request and decodes it as UTF-8 'Text'.--- This is equivalent to @'TE.decodeUtf8' . 'WAI.rawPathInfo'@.+-- | Extracts the raw path info from the request and decodes it using lenient UTF-8 'Text'.+-- This is equivalent to @'TE.decodeUtf8With' 'TEE.lenientDecode' . 'WAI.rawPathInfo'@. getRequestPath :: Request -> Text-getRequestPath = TE.decodeUtf8 . WAI.rawPathInfo+getRequestPath = TE.decodeUtf8With TEE.lenientDecode . WAI.rawPathInfo -- | Extracts the HTTP request method (e.g., @"GET"@, @"POST"@) and returns it as a 'Text' value.--- This is equivalent to @'TE.decodeUtf8' . 'WAI.requestMethod'@.+-- This is equivalent to @'TE.decodeUtf8' . 'WAI.requestMethod'@ (methods are ASCII). getRequestMethod :: Request -> Text getRequestMethod = TE.decodeUtf8 . WAI.requestMethod --- | Extracts the value of the @Host@ header, if present.+-- | Extracts the value of the @Host@ header, if present, using lenient UTF-8 decoding. getRequestHost :: Request -> Maybe Text-getRequestHost = fmap TE.decodeUtf8 . WAI.requestHeaderHost+getRequestHost = fmap (TE.decodeUtf8With TEE.lenientDecode) . lookup hHost . WAI.requestHeaders --- | Extracts the value of the @User-Agent@ header, if present.+-- | Extracts the value of the @User-Agent@ header, if present, using lenient UTF-8 decoding. getRequestUserAgent :: Request -> Maybe Text-getRequestUserAgent = fmap TE.decodeUtf8 . WAI.requestHeaderUserAgent+getRequestUserAgent = fmap (TE.decodeUtf8With TEE.lenientDecode) . lookup hUserAgent . WAI.requestHeaders ---------------------------------------------------------------------- -- Composite key builders@@ -184,7 +187,8 @@ -- | Creates a request key based solely on the client's IP address. ----- The result is always 'Just' a value, as an IP or equivalent is always available.+-- This function always succeeds and returns a 'Just' value, as every WAI+-- request has an associated socket address (IPv4, IPv6, or Unix socket). -- -- ==== __Example__ --@@ -198,7 +202,8 @@ -- separated by a colon. -- -- This is useful for rate-limiting access to specific endpoints rather than--- penalizing a client for all of its requests.+-- penalizing a client for all of its requests. This function always succeeds+-- since both IP and path are always available. -- -- ==== __Example__ --@@ -231,7 +236,8 @@ -- | Builds a key from an arbitrary header and the client IP, joined by a colon. ----- Header lookup is case-insensitive. Returns 'Nothing' if the header is absent.+-- Header lookup is case-insensitive. Returns 'Nothing' if the specified header+-- is absent from the request. -- -- ==== __Example__ --@@ -245,4 +251,4 @@ byHeaderAndIP headerName req = do ip <- getClientIP req let mVal = lookup headerName (WAI.requestHeaders req)- pure $ fmap (\hv -> ip <> ":" <> TE.decodeUtf8 hv) mVal+ pure $ fmap (\hv -> ip <> ":" <> TE.decodeUtf8With TEE.lenientDecode hv) mVal
src/Keter/RateLimiter/SlidingWindow.hs view
@@ -13,7 +13,7 @@ -- -- This module provides an implementation of the /Sliding Window Counter/ -- algorithm for rate limiting. It is implemented using STM primitives and--- integrates with the `Keter.RateLimiter.Cache` key structure.+-- integrates with the "Keter.RateLimiter.Cache" key structure. -- -- A sliding window allows a fixed number of requests within a moving time -- interval (window). It tracks individual request timestamps and filters them@@ -22,11 +22,13 @@ -- -- == Example usage ----- > let getTimeNow = realToFrac <$> getPOSIXTime--- > result <- allowRequest getTimeNow myMap "zone1" "user42" 60 100--- > when result (putStrLn "Request allowed")+-- @+-- let getTimeNow = realToFrac \<$\> getPOSIXTime+-- result <- allowRequest getTimeNow myMap \"zone1\" \"user42\" 60 100+-- when result (putStrLn \"Request allowed\")+-- @ ----- This example checks whether "user42" from "zone1" can make a request,+-- This example checks whether @\"user42\"@ from @\"zone1\"@ can make a request, -- allowing up to 100 requests in a 60-second window. module Keter.RateLimiter.SlidingWindow@@ -50,42 +52,80 @@ -- the request is denied. -- -- This version is /thread-safe/ and uses STM to avoid race conditions under--- concurrency. The timestamp list is updated atomically using `StmMap.focus`.+-- concurrency. The timestamp list is updated atomically using 'StmMap.focus'. ----- == Parameters+-- === Algorithm Steps ----- [@getTimeNow@] Function that returns the current time as a `Double` (in seconds).--- [@stmMapTVar@] STM container storing per-client timestamp lists.--- [@ipZone@] A textual label identifying the group or IP zone.--- [@userKey@] Unique identifier for the client (IP, token, etc).--- [@windowSize@] Sliding window size in seconds.--- [@limit@] Maximum allowed requests within the time window.+-- 1. __Time Acquisition__: Get current timestamp using provided time function+-- 2. __Key Construction__: Build composite cache key from throttle name, IP zone, and user key+-- 3. __Atomic Update__: Use STM Focus to atomically read, filter, and update timestamp list+-- 4. __Window Filtering__: Remove timestamps older than the sliding window+-- 5. __Limit Check__: Allow request if filtered list length is below limit+-- 6. __State Update__: Add current timestamp to list if request is allowed ----- == Returns+-- === Memory Management ----- A `Bool` indicating whether the request is allowed (`True`) or rate-limited (`False`).+-- The algorithm automatically cleans up old timestamps and removes empty entries+-- from the map, ensuring efficient memory usage for long-running applications. ----- == Example+-- ==== __Examples__ ----- > getTime <- realToFrac <$> getPOSIXTime--- > allowed <- allowRequest getTime myStore "zone" "user" 10 5--- > if allowed then serve else reject+-- @+-- -- Basic API rate limiting: 100 requests per minute+-- getTime <- realToFrac \<$\> getPOSIXTime+-- allowed <- allowRequest (pure getTime) stmMapTVar \"api\" \"zone1\" \"user123\" 60 100+-- if allowed +-- then processApiRequest+-- else sendRateLimitError+-- @+--+-- @+-- -- High-frequency trading API: 1000 requests per second+-- let getTimeIO = realToFrac \<$\> getPOSIXTime+-- result <- allowRequest getTimeIO storage \"hft\" \"premium\" \"trader456\" 1 1000+-- when result $ executeTrade+-- @+--+-- @+-- -- Multi-tier rate limiting with different windows+-- let (windowSecs, maxReqs) = case userTier of+-- Premium -> (60, 1000) -- 1000 requests per minute+-- Standard -> (60, 100) -- 100 requests per minute +-- Free -> (3600, 50) -- 50 requests per hour+--+-- allowed <- allowRequest getTime storage \"tiered\" zone userId windowSecs maxReqs+-- @+--+-- /Thread Safety:/ All operations are atomic via STM. Multiple threads can+-- safely call this function concurrently for the same or different keys.+--+-- /Performance:/ Time complexity is O(n) where n is the number of timestamps+-- within the sliding window. Space complexity is O(k*n) where k is the number+-- of active clients. allowRequest- :: IO Double -- ^ Time retrieval function (returns time in seconds)- -> TVar (StmMap.Map Text [Double]) -- ^ STM Map container (wrapped in a TVar)- -> Text -- ^ Zone identifier (or grouping key)- -> Text -- ^ User (or client) key- -> Int -- ^ Sliding window size (in seconds)- -> Int -- ^ Request limit within the window- -> IO Bool -- ^ Returns True if the request is allowed, False if rate-limited-allowRequest getTimeNow stmMapTVar ipZone userKey windowSize limit = do+ :: IO Double+ -- ^ Action to get current time as fractional seconds since epoch+ -> TVar (StmMap.Map Text [Double])+ -- ^ STM map storing per-client timestamp lists+ -> Text+ -- ^ Throttle name (logical grouping identifier)+ -> Text+ -- ^ IP zone identifier for multi-tenant isolation+ -> Text+ -- ^ User key (unique client identifier)+ -> Int+ -- ^ Sliding window size in seconds (must be positive)+ -> Int+ -- ^ Maximum allowed requests within the time window (must be positive)+ -> IO Bool+ -- ^ 'True' if request is allowed, 'False' if rate-limited+allowRequest getTimeNow stmMapTVar throttleName ipZone userKey windowSize limit = do now <- getTimeNow- let key = makeCacheKey SlidingWindow ipZone userKey+ let key = makeCacheKey throttleName SlidingWindow ipZone userKey windowSizeD = fromIntegral windowSize atomically $ do stmMap <- readTVar stmMapTVar- -- Use StmMap.focus for an atomic read-modify-write operation on the map entry StmMap.focus (updateTimestamps now windowSizeD limit) key stmMap --------------------------------------------------------------------------------@@ -97,42 +137,42 @@ -- is appended to the list. If the list becomes empty, the entry is removed -- from the map. ----- This ensures memory-efficient cleanup of old clients.+-- This ensures memory-efficient cleanup of old clients and automatic garbage+-- collection of inactive entries. ----- == Parameters+-- === Behavior Details ----- [@now@] The current time in fractional seconds.--- [@windowSize@] Duration of the sliding window.--- [@limit@] Maximum allowed request count in the window.+-- * If the client has no prior entries, the request is allowed and a new list is created.+-- * If the client exists, the list is trimmed to remove stale timestamps.+-- * The request is allowed only if the filtered list has fewer than the specified limit.+-- * Empty timestamp lists are automatically removed from the map. ----- == Behavior+-- === Algorithm Complexity ----- - If the client has no prior entries, the request is allowed and a new list is created.--- - If the client exists, the list is trimmed, and the request is allowed only if--- the resulting list has fewer than @limit@ entries.+-- * /Time/: O(n) where n is the number of timestamps in the current window+-- * /Space/: O(1) additional space per operation (filtering is done in-place conceptually) updateTimestamps- :: Double -- ^ Current time- -> Double -- ^ Sliding window size (seconds)- -> Int -- ^ Request limit+ :: Double+ -- ^ Current time in fractional seconds+ -> Double + -- ^ Sliding window duration in seconds+ -> Int+ -- ^ Request limit within the window -> Focus.Focus [Double] STM Bool+ -- ^ STM Focus that returns whether the request is allowed updateTimestamps now windowSize limit = Focus.Focus- -- Case 1: The key does not exist.- -- This is the first request. It's allowed. Insert a new list with the current timestamp. (do+ -- New client: allow first request and create timestamp list let newList = [now] pure (True, Focus.Set newList) )- -- Case 2: The key exists.- -- Update the list and check the limit. (\currentTimestamps -> do- -- Filter out timestamps older than the window+ -- Existing client: filter old timestamps and check limit let freshTimestamps = filter (\t -> now - t <= windowSize) currentTimestamps- -- Check if the request is allowed let allowed = length freshTimestamps < limit- -- If allowed, add the new timestamp. Otherwise, keep the old list. let updatedTimestamps = if allowed then now : freshTimestamps else freshTimestamps - -- If the list becomes empty after filtering, remove it from the map.+ -- Clean up empty entries to prevent memory leaks if null updatedTimestamps then pure (allowed, Focus.Remove) else pure (allowed, Focus.Set updatedTimestamps)
src/Keter/RateLimiter/TokenBucket.hs view
@@ -11,23 +11,25 @@ -- Portability : portable -- -- This module provides a rate limiter based on the /Token Bucket/ algorithm.--- It integrates with the `Keter.RateLimiter.Cache` infrastructure and uses STM+-- It integrates with the "Keter.RateLimiter.Cache" infrastructure and uses STM -- and worker threads to manage refill and request allowance. ----- The token bucket algorithm allows for a configurable burst size (`capacity`)+-- The token bucket algorithm allows for a configurable burst size (capacity) -- and replenishes tokens over time at a fixed rate. If a request is made and -- a token is available, the request is allowed and a token is consumed. -- Otherwise, the request is denied. -- -- == Example usage ----- > import Keter.RateLimiter.TokenBucket (allowRequest)--- > --- > allowed <- allowRequest cache "zone1" "user123" 10 2.5 60--- > when allowed $ doSomething+-- @+-- import Keter.RateLimiter.TokenBucket (allowRequest)+-- +-- allowed <- allowRequest cache \"zone1\" \"user123\" 10 2.5 60+-- when allowed $ doSomething+-- @ ----- This call checks if a request by "user123" in "zone1" is allowed, given a--- bucket with a capacity of 10, a refill rate of 2.5 tokens/second, and a TTL of 60 seconds.+-- This call checks if a request by @\"user123\"@ in @\"zone1\"@ is allowed, given a+-- bucket with a capacity of 10, a refill rate of 2.5 tokens\/second, and a TTL of 60 seconds. module Keter.RateLimiter.TokenBucket ( -- * Request Evaluation@@ -66,103 +68,102 @@ -- -- The token bucket is defined by: ----- - @capacity@: maximum number of tokens in the bucket (i.e., max burst size).--- - @refillRate@: tokens added per second (can be fractional).--- - @expiresIn@: TTL in seconds; determines how long idle buckets live.+-- * /capacity/: maximum number of tokens in the bucket (i.e., max burst size).+-- * /refillRate/: tokens added per second (can be fractional).+-- * /expiresIn/: TTL in seconds; determines how long idle buckets live. ----- == Parameters+-- The function performs the following steps: ----- [@cache@] The configured cache for storing per-user token bucket state.--- [@ipZone@] A label identifying the IP zone (e.g., region or tenant).--- [@userKey@] A unique identifier for the user (e.g., IP address or token).--- [@capacity@] Maximum tokens the bucket can hold.--- [@refillRate@] Token refill rate (tokens per second).--- [@expiresIn@] Time-to-live for idle entries in seconds.+-- 1. Validates that TTL meets the minimum threshold+-- 2. Creates or retrieves the token bucket for the given key+-- 3. For new buckets: starts a worker thread and allows the first request+-- 4. For existing buckets: queues the request and waits for the worker's response ----- == Returns+-- ==== __Examples__ ----- A boolean in an arbitrary `MonadIO` context indicating whether the request is allowed.+-- @+-- -- Allow 100 requests per minute with burst capacity of 10+-- let capacity = 10+-- refillRate = 100.0 \/ 60.0 -- ~1.67 tokens per second+-- ttl = 300 -- 5 minutes TTL ----- == Example+-- result <- allowRequest cache \"api-throttle\" \"192.168.1.1\" \"user456\" capacity refillRate ttl+-- if result+-- then putStrLn \"Request allowed\"+-- else putStrLn \"Request denied - rate limit exceeded\"+-- @ ----- > allowed <- allowRequest cache "zoneA" "192.168.0.1" 5 1.0 60--- > when allowed $ putStrLn "Proceeding with request..."+-- @+-- -- High-frequency API with small bursts+-- allowed <- allowRequest cache \"fast-api\" \"zone-premium\" \"client789\" 5 10.0 120+-- @+--+-- /Thread Safety:/ This function is thread-safe and can be called concurrently+-- from multiple threads for the same or different keys.+--+-- /Performance:/ For new buckets, there's a one-time setup cost of starting+-- a worker thread. Subsequent requests are processed asynchronously with+-- minimal blocking. allowRequest :: MonadIO m => Cache (InMemoryStore 'TokenBucket)- -- ^ Token bucket cache backend+ -- ^ Token bucket cache instance -> Text- -- ^ IP zone (e.g. region or customer identifier)+ -- ^ Throttle name (logical grouping identifier) -> Text- -- ^ User key (e.g. IP address or API token)+ -- ^ IP zone identifier+ -> Text+ -- ^ User key (unique client identifier) -> Int- -- ^ Bucket capacity (max number of tokens)+ -- ^ Bucket capacity (maximum tokens, must be positive) -> Double- -- ^ Refill rate (tokens per second)+ -- ^ Refill rate in tokens per second (must be positive, can be fractional) -> Int- -- ^ TTL (seconds) for the bucket state+ -- ^ TTL in seconds (must be >= 'minTTL') -> m Bool-allowRequest cache ipZone userKey capacity refillRate expiresIn = liftIO $+ -- ^ 'True' if request is allowed, 'False' if denied+allowRequest cache throttleName ipZone userKey capacity refillRate expiresIn = liftIO $ if expiresIn < minTTL then do- putStrLn $- "TokenBucket: Request denied due to invalid TTL: "- ++ show expiresIn pure False else do now <- floor <$> getPOSIXTime- let key = makeCacheKey (cacheAlgorithm cache) ipZone userKey+ let key = makeCacheKey throttleName (cacheAlgorithm cache) ipZone userKey TokenBucketStore tvBuckets = cacheStore cache replyVar <- newEmptyMVar-- ----------------------------------------------------------------------- -- 1. Obtain (or create) bucket entry and enqueue the request.- -- Create the entry in IO, then pass it into the STM transaction. newEntryInitialState <- createTokenBucketEntry (TokenBucketState (capacity - 1) now) (wasNew, entry) <- atomically $ do buckets <- readTVar tvBuckets- -- Use F.Focus directly to allow STM actions in the handler, bypassing F.cases. (wasNewEntry, ent) <- StmMap.focus (F.Focus- -- Handler for when the key is NOT found (the "Nothing" case) (pure ((True, newEntryInitialState), F.Set newEntryInitialState))- -- Handler for when the key IS found (the "Just" case) (\existingEnt -> do- -- This handler can now perform STM actions. workerLockEmpty <- isEmptyTMVar (tbeWorkerLock existingEnt) if workerLockEmpty- then pure ((True, newEntryInitialState), F.Set newEntryInitialState) -- Replace dead entry- else pure ((False, existingEnt), F.Leave) -- Keep existing entry+ then pure ((True, newEntryInitialState), F.Set newEntryInitialState)+ else pure ((False, existingEnt), F.Leave) ) ) key buckets pure (wasNewEntry, ent)-- ----------------------------------------------------------------------- -- 2. Spawn a worker once for a fresh bucket. if wasNew then- -- For a new bucket, the first request is allowed only if there is capacity. if capacity > 0 then do workerReadyVar <- atomically newEmptyTMVar- atomically $ putTMVar (tbeWorkerLock entry) () -- Mark worker lock as taken- -- Start the worker with ready synchronization+ atomically $ putTMVar (tbeWorkerLock entry) () startTokenBucketWorker (tbeState entry) (tbeQueue entry) capacity refillRate workerReadyVar- -- Wait for the worker to signal it's ready before proceeding atomically $ takeTMVar workerReadyVar pure True- else do- -- If capacity is 0, no request can ever be allowed.+ else pure False else do- -- For existing buckets, enqueue the request and wait for response atomically $ writeTQueue (tbeQueue entry) replyVar result <- takeMVar replyVar pure result
src/Keter/RateLimiter/TokenBucketWorker.hs view
@@ -47,7 +47,7 @@ requestQueue <- newTBroadcastTQueueIO readySignal <- newEmptyTMVarIO --- Start worker: 100 token capacity, 10 tokens/second refill rate+-- Start worker: 100 token capacity, 10 tokens\/second refill rate startTokenBucketWorker initialState requestQueue 100 10.0 readySignal -- Wait for worker to be ready@@ -90,13 +90,13 @@ -- Each request is handled atomically: the bucket state is read, tokens are refilled -- based on elapsed time, a token is consumed if available, and the new state is written back. ----- ==== Worker Lifecycle+-- === Worker Lifecycle -- -- 1. __Startup__: Worker thread is forked and signals readiness via 'TMVar' -- 2. __Processing Loop__: Worker waits for requests, processes them atomically -- 3. __Response__: Results are sent back to clients via 'MVar' ----- ==== Token Refill Algorithm+-- === Token Refill Algorithm -- -- Tokens are refilled using the formula: --@@ -105,11 +105,12 @@ -- @ -- -- This ensures:+-- -- * Tokens are added proportionally to elapsed time -- * Bucket capacity is never exceeded -- * Sub-second precision for refill calculations ----- ==== Atomic Request Processing+-- === Atomic Request Processing -- -- Each request is processed in a single STM transaction that: --@@ -118,9 +119,9 @@ -- 3. Computes available tokens after refill -- 4. Attempts to consume one token if available -- 5. Updates bucket state with new token count and timestamp--- 6. Returns allow/deny decision+-- 6. Returns allow\/deny decision ----- ==== Error Handling+-- === Error Handling -- -- The worker is designed to be resilient: --@@ -128,12 +129,12 @@ -- * Negative elapsed time (clock adjustments) results in no refill -- * Worker continues running even if individual requests fail ----- ==== Example+-- ==== __Examples__ -- -- @--- -- Create a bucket for API rate limiting: 1000 requests/hour = ~0.278 req/sec+-- -- Create a bucket for API rate limiting: 1000 requests\/hour = ~0.278 req\/sec -- let capacity = 100 -- Allow bursts up to 100 requests--- refillRate = 1000.0 / 3600.0 -- 1000 requests per hour+-- refillRate = 1000.0 \/ 3600.0 -- 1000 requests per hour -- -- initialState <- newTVarIO $ TokenBucketState capacity now -- requestQueue <- newTBroadcastTQueueIO @@ -142,15 +143,15 @@ -- startTokenBucketWorker initialState requestQueue capacity refillRate readySignal -- @ ----- __Thread Safety:__ All state updates are atomic via STM transactions.+-- /Thread Safety:/ All state updates are atomic via STM transactions. ----- __Resource Usage:__ Creates one background thread that runs indefinitely.+-- /Resource Usage:/ Creates one background thread that runs indefinitely. startTokenBucketWorker :: TVar TokenBucketState -- ^ Shared bucket state (tokens + last update time). -- This 'TVar' is read and updated atomically by the worker. -> TQueue (MVar Bool) -- ^ Request queue containing 'MVar's for client responses. -- Clients place their response 'MVar' in this queue and wait- -- for the worker to write the allow/deny decision.+ -- for the worker to write the allow\/deny decision. -> Int -- ^ Maximum bucket capacity (maximum tokens that can be stored). -- This sets the upper limit for burst traffic handling. -- Must be positive.
src/Keter/RateLimiter/Types.hs view
@@ -22,7 +22,7 @@ Both types support JSON serialization for persistence and configuration purposes, with validation to ensure state consistency. -== Example Usage+= Example Usage @ -- Token bucket example@@ -51,19 +51,19 @@ -- Each incoming request consumes one or more tokens from the bucket. Tokens are replenished -- at a fixed rate. If insufficient tokens are available, the request is either delayed or rejected. ----- ==== Token Bucket Properties+-- == Token Bucket Properties -- -- * __Bursty Traffic__: Allows bursts of traffic up to the bucket capacity -- * __Rate Control__: Long-term average rate is controlled by token replenishment rate -- * __Memory Efficient__: Only requires tracking token count and last update time ----- ==== Use Cases+-- == Use Cases -- -- * API rate limiting with burst allowance -- * Network traffic shaping -- * Resource allocation with temporary overages ----- ==== Example+-- = Example -- -- @ -- -- Initial state with 50 tokens, last updated at Unix timestamp 1640995200@@ -88,7 +88,7 @@ -- -- Serializes the token bucket state to JSON format for persistence or network transmission. ----- ==== JSON Format+-- = JSON Format -- -- @ -- {@@ -102,10 +102,10 @@ -- -- Deserializes JSON to 'TokenBucketState' with the following validation rules: ----- * @tokens@ field must be non-negative (>= 0)--- * @lastUpdate@ field must be present and parseable as 'Int'+-- * 'tokens' field must be non-negative (>= 0)+-- * 'lastUpdate' field must be present and parseable as 'Int' ----- ==== Validation Examples+-- = Validation Examples -- -- @ -- -- Valid JSON@@ -117,7 +117,7 @@ -- -- Nothing -- @ ----- __Throws:__ Parse error if @tokens@ is negative or required fields are missing.+-- Throws a parse error if 'tokens' is negative or required fields are missing. instance FromJSON TokenBucketState where parseJSON = withObject "TokenBucketState" $ \o -> do tokens <- o .: "tokens"@@ -131,20 +131,20 @@ -- at a constant rate. Incoming requests add water to the bucket, and if the bucket -- overflows, requests are rejected. This provides smooth rate limiting without bursts. ----- ==== Leaky Bucket Properties+-- == Leaky Bucket Properties -- -- * __Smooth Rate__: Enforces a consistent output rate regardless of input bursts -- * __No Bursts__: Unlike token bucket, doesn't allow temporary rate exceedance -- * __Queue Modeling__: Can model request queuing with bucket level representing queue depth ----- ==== Use Cases+-- == Use Cases -- -- * Smooth traffic shaping for network connections -- * Audio/video streaming rate control -- * Database connection throttling -- * Prevention of thundering herd problems ----- ==== Mathematical Model+-- = Mathematical Model -- -- The bucket level changes according to: --@@ -153,12 +153,13 @@ -- @ -- -- Where:--- * @requestSize@ is the size of the incoming request--- * @drainRate@ is the constant drain rate (requests per second)--- * @timeDelta@ is the elapsed time since last update ----- ==== Example+-- * 'requestSize' is the size of the incoming request+-- * 'drainRate' is the constant drain rate (requests per second)+-- * 'timeDelta' is the elapsed time since last update --+-- = Example+-- -- @ -- -- Initial state: half-full bucket at timestamp 1640995200.5 -- let initialState = LeakyBucketState @@ -174,7 +175,7 @@ -- @ data LeakyBucketState = LeakyBucketState { level :: Double -- ^ Current fill level of the bucket (0.0 to capacity).- -- Represents the amount of "water" (pending requests)+ -- Represents the amount of \"water\" (pending requests) -- currently in the bucket. Higher values indicate -- more backpressure or pending work. , lastTime :: Double -- ^ Timestamp of last bucket update as Unix time with@@ -186,7 +187,7 @@ -- -- Serializes the leaky bucket state to JSON format with full double precision. ----- ==== JSON Format+-- = JSON Format -- -- @ -- {@@ -200,14 +201,14 @@ -- -- Deserializes JSON to 'LeakyBucketState' with the following validation rules: ----- * @level@ must be non-negative (>= 0.0)--- * @level@ must not exceed 1,000,000 (practical upper bound)--- * @lastTime@ must be present and parseable as 'Double'+-- * 'level' must be non-negative (>= 0.0)+-- * 'level' must not exceed 1,000,000 (practical upper bound)+-- * 'lastTime' must be present and parseable as 'Double' ----- The upper bound on @level@ prevents potential overflow issues and ensures+-- The upper bound on 'level' prevents potential overflow issues and ensures -- reasonable memory usage for bucket state tracking. ----- ==== Validation Examples+-- = Validation Examples -- -- @ -- -- Valid JSON@@ -223,7 +224,7 @@ -- -- Nothing -- @ ----- __Throws:__ Parse error if @level@ is negative, exceeds 1,000,000, or required fields are missing.+-- Throws a parse error if 'level' is negative, exceeds 1,000,000, or required fields are missing. instance FromJSON LeakyBucketState where parseJSON = withObject "LeakyBucketState" $ \o -> do level <- o .: "level"
src/Keter/RateLimiter/WAI.hs view
@@ -1,6 +1,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE DataKinds #-}+{-# LANGUAGE TupleSections #-} {-| Module : Keter.RateLimiter.WAI@@ -11,9 +12,9 @@ Stability : stable Portability : portable -This file is a ported to Haskell language code with some simplifications of rack-attack -https://github.com/rack/rack-attack/blob/main/lib/rack/attack.rb-and is based on the structure of the original code of +This file is a ported to Haskell language code with some simplifications of rack-attack+<https://github.com/rack/rack-attack/blob/main/lib/rack/attack.rb>+and is based on the structure of the original code of rack-attack, Copyright (c) 2016 by Kickstarter, PBC, under the MIT License. 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.@@ -23,26 +24,30 @@ This module provides WAI middleware that implements multiple IP-zone-specific rate limiting strategies including: -- 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. Rate limiting state is tracked separately per IP zone using the-'ZoneSpecificCaches' abstraction.+'ZoneSpecificCaches' abstraction with efficient HashMap-based lookups. -Use `attackMiddleware` to enforce throttling in your WAI application.+Use 'attackMiddleware' to enforce throttling in your WAI application. -} module Keter.RateLimiter.WAI ( -- * Environment & Configuration- Env+ Env(..) , ThrottleConfig(..)+ , IdentifierBy(..)+ , ZoneBy(..)+ , RLThrottle(..)+ , RateLimiterConfig(..) , initConfig , addThrottle -- * Middleware@@ -50,15 +55,29 @@ -- * Manual Control & Inspection , instrument , cacheResetAll- -- * Accessors- , envZoneCachesMap+ -- * Functions for Middleware+ , buildRateLimiter+ , registerThrottle+ , mkIdentifier+ , mkZoneFn+ , getClientIPPure+ , hdr+ , fromHeaderName ) where +import Data.Aeson hiding (pairs) import Data.Text (Text)+import qualified Data.Text as Tx import qualified Data.Text.Encoding as TE-import qualified Data.Map.Strict as Map+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)+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 Data.IORef import Control.Concurrent.STM@@ -68,13 +87,19 @@ import qualified Keter.RateLimiter.TokenBucket as TokenBucket import qualified Keter.RateLimiter.LeakyBucket as LeakyBucket import Keter.RateLimiter.CacheWithZone (allowFixedWindowRequest)+import qualified Keter.RateLimiter.RequestUtils as RU import Data.TinyLRU (allowRequestTinyLRU) import System.Clock (Clock(Monotonic), getTime) import Data.Maybe (fromMaybe) import Data.Time.Clock.POSIX (getPOSIXTime)+import qualified Web.Cookie as WC --------------------------------------------------------------------------------+ -- | Configuration for a single throttle rule.+--+-- Defines the parameters for a specific rate limiting rule including+-- the algorithm to use, rate limits, and identifier extraction logic. data ThrottleConfig = ThrottleConfig { throttleLimit :: Int -- ^ Maximum number of requests allowed per period.@@ -84,17 +109,31 @@ -- ^ Algorithm used for throttling (e.g., 'FixedWindow', 'LeakyBucket'). , 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.+ -- Returns 'Nothing' if throttling should be skipped. , throttleTokenBucketTTL :: Maybe Int -- ^ Optional TTL (seconds) for TokenBucket entries. If not provided, defaults to 2. } --- | Global environment for throttling, including zone-specific cache maps and throttle configurations.+-- | 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+--+-- Both HashMaps are wrapped in 'IORef' for thread-safe atomic updates while+-- maintaining efficient read access patterns typical in web middleware. data Env = Env- { envZoneCachesMap :: IORef (Map.Map IPZoneIdentifier ZoneSpecificCaches)- -- ^ Map from IP zones to rate limiter caches.- , envThrottles :: IORef (Map.Map Text ThrottleConfig)- -- ^ Registered throttle rules by name.+ { 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. , envGetRequestIPZone :: Request -> IPZoneIdentifier -- ^ Function to derive an IP zone from the incoming WAI request. }@@ -102,34 +141,80 @@ -------------------------------------------------------------------------------- -- | 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+--+-- * Initial HashMap creation: O(1)+-- * Default zone cache creation: O(1)+-- * Thread-safe IORef initialization: O(1) initConfig :: (Request -> IPZoneIdentifier) -- ^ Function to extract an IP zone label from the request. -> IO Env initConfig getIPZone = do defaultCaches <- createZoneCaches- zoneCachesMap <- newIORef $ Map.singleton defaultIPZone defaultCaches- throttles <- newIORef Map.empty+ zoneCachesMap <- newIORef $ HM.singleton defaultIPZone defaultCaches+ throttles <- newIORef HM.empty return $ 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+--+-- * Average case: O(1) insertion into throttles HashMap+-- * Atomic update via 'modifyIORef'' for thread safety+-- * Returns updated environment for method chaining addThrottle :: Env- -> Text -- ^ Throttle rule name (must be unique).- -> ThrottleConfig -- ^ Throttle configuration.+ -- ^ 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) $ Map.insert name config+ modifyIORef' (envThrottles env) $ HM.insert name config return 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__+--+-- @+-- -- 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+-- @ 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@@ -139,66 +224,108 @@ -- | Inspect all throttle rules and apply them to the incoming request. -- -- Returns 'True' if any rule blocks the request.+--+-- === HashMap Processing+--+-- * 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 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) let zone = envGetRequestIPZone env req zoneCaches <- getOrCreateZoneCaches env zone- anyBlocked <- or <$> mapM (checkThrottle zoneCaches zone req) (Map.elems throttles)+ -- 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- where- checkThrottle :: ZoneSpecificCaches -> Text -> Request -> ThrottleConfig -> IO Bool- checkThrottle caches zone req config = do- mIdentifier <- throttleIdentifier config req- case mIdentifier of- Nothing -> return False- Just identifier -> case throttleAlgorithm config of- FixedWindow ->- not <$> allowFixedWindowRequest- (zscCounterCache caches) zone identifier- (throttleLimit config) (throttlePeriod config) - SlidingWindow -> case zscTimestampCache caches of- Cache { cacheStore = TimestampStore tvar } ->- not <$> SlidingWindow.allowRequest- (realToFrac <$> getPOSIXTime)- tvar zone identifier- (throttlePeriod config) (throttleLimit config)+-- | 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.+checkThrottle :: ZoneSpecificCaches -> Text -> Request -> Text -> ThrottleConfig -> IO Bool+checkThrottle caches zone req throttleName config = do+ mIdentifier <- throttleIdentifier config req+ case mIdentifier of+ Nothing -> return False+ Just identifier -> case throttleAlgorithm config of+ FixedWindow ->+ -- allowFixedWindowRequest cache throttleName ipZone userKey limit period+ not <$> allowFixedWindowRequest+ (zscCounterCache caches)+ throttleName+ zone+ identifier+ (throttleLimit config)+ (throttlePeriod config) - TokenBucket -> do- let period = throttlePeriod config- limit = throttleLimit config- refillRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0- ttl = fromMaybe 2 (throttleTokenBucketTTL config)- not <$> TokenBucket.allowRequest- (zscTokenBucketCache caches)- zone identifier- limit refillRate ttl+ SlidingWindow -> case zscTimestampCache caches of+ Cache { cacheStore = TimestampStore tvar } ->+ -- SlidingWindow.allowRequest getTimeNow stmMapTVar throttleName ipZone userKey windowSize limit+ not <$> SlidingWindow.allowRequest+ (realToFrac <$> getPOSIXTime)+ tvar+ throttleName+ zone+ identifier+ (throttlePeriod config)+ (throttleLimit config) - LeakyBucket -> do- let period = throttlePeriod config- limit = throttleLimit config- leakRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0- not <$> LeakyBucket.allowRequest- (zscLeakyBucketCache caches)- zone identifier- limit leakRate+ TokenBucket -> do+ let period = throttlePeriod config+ limit = throttleLimit config+ 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+ not <$> TokenBucket.allowRequest+ (zscTokenBucketCache caches)+ throttleName+ zone+ identifier+ (fromIntegral limit)+ refillRate+ (fromIntegral ttl) - TinyLRU -> do- now <- getTime Monotonic- case cacheStore (zscTinyLRUCache caches) of- TinyLRUStore tvar -> do- cache <- readTVarIO tvar- not <$> atomically (allowRequestTinyLRU now cache identifier (throttleLimit config) (throttlePeriod config))+ LeakyBucket -> do+ let period = throttlePeriod config+ limit = throttleLimit config+ leakRate = if period > 0 then fromIntegral limit / fromIntegral period else 0.0+ -- LeakyBucket.allowRequest cache throttleName ipZone userKey capacity leakRate+ not <$> LeakyBucket.allowRequest+ (zscLeakyBucketCache caches)+ throttleName+ zone+ identifier+ (fromIntegral limit)+ leakRate + TinyLRU -> do+ now <- getTime Monotonic+ case cacheStore (zscTinyLRUCache caches) of+ TinyLRUStore tvar -> do+ cache <- readTVarIO tvar+ not <$> atomically (allowRequestTinyLRU now cache identifier (throttleLimit config) (throttlePeriod config))+ -- | 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+--+-- * 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 cacheResetAll :: Env -> IO () cacheResetAll env = do zoneCachesMap <- readIORef (envZoneCachesMap env)- mapM_ (resetZoneCaches . snd) (Map.toList zoneCachesMap)+ mapM_ (resetZoneCaches . snd) (HM.toList zoneCachesMap) where resetZoneCaches :: ZoneSpecificCaches -> IO () resetZoneCaches caches = do@@ -209,16 +336,256 @@ 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+--+-- * 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) 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 Map.lookup zone m of+ case HM.lookup zone m of Just caches -> return caches Nothing -> do newCaches <- createZoneCaches atomicModifyIORef' (envZoneCachesMap env) $ \currentMap ->- let updatedMap = Map.insertWith (\_ old -> old) zone newCaches currentMap- in (updatedMap, updatedMap Map.! zone)+ let updatedMap = HM.insertWith (\_ old -> old) zone newCaches currentMap+ in (updatedMap, updatedMap HM.! zone)+++--------------------------------------------------------------------------------+-- * Configuration Data 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+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+ 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+data ZoneBy+ = ZoneDefault -- ^ All requests use default zone+ | ZoneIP -- ^ Derive zone from client IP+ | ZoneHeader !HeaderName -- ^ Derive zone from header value+ 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+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)+ } 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+data RateLimiterConfig = RateLimiterConfig+ { rlZoneBy :: !ZoneBy -- ^ Zone derivation strategy+ , rlThrottles :: ![RLThrottle] -- ^ List of throttle rules+ } deriving (Show, Eq, Generic)++instance FromJSON IdentifierBy where+ parseJSON (String "ip") = pure IdIP+ parseJSON (String "ip+path") = pure IdIPAndPath+ parseJSON (String "ip+ua") = pure IdIPAndUA+ parseJSON (Object o) =+ asum [ IdHeader . hdr <$> o .: "header"+ , 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: ...}"++instance ToJSON IdentifierBy where+ toJSON IdIP = String "ip"+ toJSON IdIPAndPath = String "ip+path"+ toJSON IdIPAndUA = String "ip+ua"+ toJSON (IdHeader h) = object ["header" .= TE.decodeUtf8 (fromHeaderName h)]+ toJSON (IdCookie c) = object ["cookie" .= c]+ toJSON (IdHeaderAndIP h) = object ["header+ip" .= TE.decodeUtf8 (fromHeaderName h)]++instance FromJSON ZoneBy where+ 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: ...}"++instance ToJSON ZoneBy where+ toJSON ZoneDefault = String "default"+ toJSON ZoneIP = String "ip"+ toJSON (ZoneHeader h) = object ["header" .= TE.decodeUtf8 (fromHeaderName h)]++instance FromJSON RLThrottle where+ parseJSON = withObject "throttle" $ \o -> do+ n <- o .: "name"+ l <- o .: "limit"+ p <- o .: "period"+ at <- o .: "algorithm" >>= parseAlgoText+ idb <- o .: "identifier_by"+ ttl <- o .:? "token_bucket_ttl"+ pure (RLThrottle n l p at idb ttl)++instance ToJSON RLThrottle where+ toJSON (RLThrottle n l p a idb ttl) =+ object [ "name" .= n, "limit" .= l, "period" .= p+ , "algorithm" .= algoToText a, "identifier_by" .= idb+ , "token_bucket_ttl" .= ttl+ ]++instance FromJSON RateLimiterConfig where+ parseJSON = withObject "rate-limiter" $ \o ->+ RateLimiterConfig+ <$> o .:? "zone_by" .!= ZoneDefault+ <*> o .: "throttles"++instance ToJSON RateLimiterConfig where+ toJSON (RateLimiterConfig zb ths) =+ object [ "zone_by" .= zb, "throttles" .= ths ]++--------------------------------------------------------------------------------+-- * Functions for Middleware++-- | Convert Text header name to WAI HeaderName.+--+-- @since 0.1.1.0+hdr :: Text -> HeaderName+hdr = mk . TE.encodeUtf8++-- | Extract the original ByteString from a case-insensitive HeaderName.+--+-- @since 0.1.1.0+fromHeaderName :: HeaderName -> S.ByteString+fromHeaderName = original++-- | 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+-- @+--+-- @since 0.1.1.0+buildRateLimiter :: RateLimiterConfig -> IO Middleware+buildRateLimiter (RateLimiterConfig zb ths) = do+ let zoneFn = mkZoneFn zb+ env <- initConfig zoneFn+ mapM_ (registerThrottle env) ths+ pure (attackMiddleware env)++-- | Register a single throttle rule in an existing environment.+--+-- @since 0.1.1.0+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 ()++-- | Create an identifier extraction function from declarative specification.+--+-- @since 0.1.1.0+mkIdentifier :: IdentifierBy -> Request -> IO (Maybe Text)+mkIdentifier IdIP = RU.byIP+mkIdentifier IdIPAndPath = RU.byIPAndPath+mkIdentifier IdIPAndUA = RU.byIPAndUserAgent+mkIdentifier (IdHeader h) = \req -> pure $ fmap (TE.decodeUtf8With TEE.lenientDecode) . lookup h . requestHeaders $ req+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).+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)++-- | Create a zone derivation function from declarative specification.+--+-- @since 0.1.1.0+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+getClientIPPure :: Request -> IPZoneIdentifier+getClientIPPure req =+ let safeDecode = TE.decodeUtf8With TEE.lenientDecode+ in case lookup (mk "x-forwarded-for") (requestHeaders req) of+ Just xff -> Tx.takeWhile (/= ',') $ safeDecode xff+ Nothing ->+ case lookup (mk "x-real-ip") (requestHeaders req) of+ Just rip -> safeDecode rip+ Nothing ->+ case remoteHost req of+ SockAddrInet _ addr -> RU.ipv4ToString addr+ SockAddrInet6 _ _ addr _ -> RU.ipv6ToString addr+ SockAddrUnix path -> Tx.pack path
test/Keter/RateLimiter/IPZonesTests.hs view
@@ -56,40 +56,40 @@ , testCase "incStoreWithZone with IPv4 address" $ do store <- createInMemoryStore @'FixedWindow let cache = newCache FixedWindow store- _ <- incStoreWithZone cache "192.168.1.100" "user123" 60- result <- readCacheWithZone cache "192.168.1.100" "user123" :: IO (Maybe Int)+ _ <- incStoreWithZone cache "test-throttle" "192.168.1.100" "user123" 60+ result <- readCacheWithZone cache "test-throttle" "192.168.1.100" "user123" :: IO (Maybe Int) assertEqual "readCacheWithZone should return Just 1" (Just 1) result , testCase "incStoreWithZone with IPv6 address" $ do store <- createInMemoryStore @'FixedWindow let cache = newCache FixedWindow store- _ <- incStoreWithZone cache "2001:0db8:0000:0000:0000:0000:0000:0001" "user123" 60- result <- readCacheWithZone cache "2001:0db8:0000:0000:0000:0000:0000:0001" "user123" :: IO (Maybe Int)+ _ <- incStoreWithZone cache "test-throttle" "2001:0db8:0000:0000:0000:0000:0000:0001" "user123" 60+ result <- readCacheWithZone cache "test-throttle" "2001:0db8:0000:0000:0000:0000:0000:0001" "user123" :: IO (Maybe Int) assertEqual "readCacheWithZone should return Just 1" (Just 1) result , testCase "incStoreWithZone increments existing counter" $ do store <- createInMemoryStore @'FixedWindow let cache = newCache FixedWindow store- _ <- incStoreWithZone cache "192.168.1.100" "user123" 60- result <- incStoreWithZone cache "192.168.1.100" "user123" 60+ _ <- incStoreWithZone cache "test-throttle" "192.168.1.100" "user123" 60+ result <- incStoreWithZone cache "test-throttle" "192.168.1.100" "user123" 60 assertEqual "incStoreWithZone should increment to 2" 2 result , testCase "resetSingleZoneCaches clears all caches" $ do zoneCaches <- createZoneCaches- _ <- incStoreWithZone (zscCounterCache zoneCaches) "192.168.1.100" "user123" 60+ _ <- incStoreWithZone (zscCounterCache zoneCaches) "test-throttle" "192.168.1.100" "user123" 60 _ <- resetSingleZoneCaches zoneCaches- result <- readCacheWithZone (zscCounterCache zoneCaches) "192.168.1.100" "user123" :: IO (Maybe Int)+ result <- readCacheWithZone (zscCounterCache zoneCaches) "test-throttle" "192.168.1.100" "user123" :: IO (Maybe Int) assertEqual "Should not retrieve counter after reset" Nothing result , testCase "incStoreWithZone with expired entry" $ do store <- createInMemoryStore @'FixedWindow let cache = newCache FixedWindow store -- Set a short TTL of 1 second.- _ <- incStoreWithZone cache "192.168.1.100" "user123" 1+ _ <- incStoreWithZone cache "test-throttle" "192.168.1.100" "user123" 1 -- Wait for the entry to expire. liftIO $ threadDelay (2 * 1000000) -- Wait for 2 seconds -- Increment again, which should create a new entry.- _ <- incStoreWithZone cache "192.168.1.100" "user123" 60- result <- readCacheWithZone cache "192.168.1.100" "user123" :: IO (Maybe Int)+ _ <- incStoreWithZone cache "test-throttle" "192.168.1.100" "user123" 60+ result <- readCacheWithZone cache "test-throttle" "192.168.1.100" "user123" :: IO (Maybe Int) assertEqual "Should retrieve new counter after expiration" (Just 1) result ]
test/Keter/RateLimiter/LeakyBucketTests.hs view
@@ -287,13 +287,14 @@ -- Start a purge thread for the LeakyBucketStore -- Removed the extra predicate argument, as startCustomPurgeLeakyBucket doesn't take it. _ <- startCustomPurgeLeakyBucket leakyBucketMap 60 7200- let ipZone = defaultIPZone+ let throttleName = "test_throttle" -- Add throttleName parameter+ ipZone = defaultIPZone userKey = "test_user" capacity = 2 leakRate = 1- allowed1 <- allowRequest cache ipZone userKey capacity leakRate+ allowed1 <- allowRequest cache throttleName ipZone userKey capacity leakRate assertBool "First request should be allowed" allowed1- allowed2 <- allowRequest cache ipZone userKey capacity leakRate+ allowed2 <- allowRequest cache throttleName ipZone userKey capacity leakRate assertBool "Second request should be allowed" allowed2 , testCase "Direct allowRequest denies request at capacity" $ do leakyBucketMap <- atomically StmMap.new@@ -303,14 +304,15 @@ -- Start a purge thread for the LeakyBucketStore -- Removed the extra predicate argument, as startCustomPurgeLeakyBucket doesn't take it. _ <- startCustomPurgeLeakyBucket leakyBucketMap 60 7200- let ipZone = defaultIPZone+ let throttleName = "test_throttle" -- Add throttleName parameter+ ipZone = defaultIPZone userKey = "test_user" capacity = 2 leakRate = 1- allowed1 <- allowRequest cache ipZone userKey capacity leakRate+ allowed1 <- allowRequest cache throttleName ipZone userKey capacity leakRate assertBool "First request should be allowed" allowed1- allowed2 <- allowRequest cache ipZone userKey capacity leakRate+ allowed2 <- allowRequest cache throttleName ipZone userKey capacity leakRate assertBool "Second request should be allowed" allowed2- allowed3 <- allowRequest cache ipZone userKey capacity leakRate+ allowed3 <- allowRequest cache throttleName ipZone userKey capacity leakRate assertBool "Third request should be denied" (not allowed3) ]
test/Keter/RateLimiter/SlidingWindowTests.hs view
@@ -277,10 +277,10 @@ let TimestampStore tvar = cacheStore cacheInstance -- First two calls should be allowed.- allowed1 <- allowRequest getFakeTime tvar (T.pack ipZone) (T.pack userKey) window limit- allowed2 <- allowRequest getFakeTime tvar (T.pack ipZone) (T.pack userKey) window limit+ allowed1 <- allowRequest getFakeTime tvar (T.pack "test_throttle") (T.pack ipZone) (T.pack userKey) window limit+ allowed2 <- allowRequest getFakeTime tvar (T.pack "test_throttle") (T.pack ipZone) (T.pack userKey) window limit -- Third call should be blocked as the window is full.- allowed3 <- allowRequest getFakeTime tvar (T.pack ipZone) (T.pack userKey) window limit+ allowed3 <- allowRequest getFakeTime tvar (T.pack "test_throttle") (T.pack ipZone) (T.pack userKey) window limit assertBool "First request should be allowed" allowed1 assertBool "Second request should be allowed" allowed2@@ -290,6 +290,6 @@ advanceTime 3.0 -- This new request should now be allowed.- allowed4 <- allowRequest getFakeTime tvar (T.pack ipZone) (T.pack userKey) window limit+ allowed4 <- allowRequest getFakeTime tvar (T.pack "test_throttle") (T.pack ipZone) (T.pack userKey) window limit assertBool "After advancing time, a new request should be allowed" allowed4 ]
test/Keter/RateLimiter/WAITests.hs view
@@ -10,352 +10,941 @@ Stability : experimental Portability : POSIX -This module provides a comprehensive test suite for the WAI (Web Application Interface) middleware responsible for rate limiting. It uses the tasty and tasty-hunit frameworks to define and run tests.--The tests cover five distinct rate-limiting algorithms:- * Fixed Window- * Sliding Window- * Token Bucket- * Leaky Bucket- * TinyLRU--For each algorithm, the following scenarios are tested:- * Allowing requests that are under the defined limit.- * Blocking requests that exceed the defined limit.- * Correctly handling both IPv4 and IPv6 addresses.- * Ensuring the 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 a high-volume of concurrent requests to test DoS (Denial of Service) protection.--The module defines several helper functions to create mock requests and a mock application to isolate the middleware for testing.--}-module Keter.RateLimiter.WAITests (- -- * Test Suite- tests-) where--import Test.Tasty-import Test.Tasty.HUnit-import Network.Wai-import Network.Wai.Test-import Network.HTTP.Types-import Network.Socket (SockAddr(..), tupleToHostAddress)-import Data.Text (Text)-import qualified Data.Text as T-import qualified Data.Text.Encoding as TE-import qualified Data.ByteString.Lazy as LBS-import Control.Concurrent (threadDelay)-import Control.Concurrent.MVar (newMVar, modifyMVar_)-import Control.Monad (replicateM)-import Control.Monad.IO.Class (liftIO)-import Data.CaseInsensitive (mk)-import Keter.RateLimiter.IPZones (defaultIPZone)-import Keter.RateLimiter.WAI-import Keter.RateLimiter.RequestUtils-import Keter.RateLimiter.Cache (Algorithm(..))---- * Request Helpers---- | Creates a mock 'Request' with a default IPv4 address (127.0.0.1).--- This is used to simulate a standard IPv4 client connection.-mkIPv4Request :: Request-mkIPv4Request = defaultRequest { remoteHost = SockAddrInet 0 (tupleToHostAddress (127, 0, 0, 1)) } -- 127.0.0.1---- | Creates a mock 'Request' with a default IPv6 address (::1).--- This is used to simulate a standard IPv6 client connection.-mkIPv6Request :: Request-mkIPv6Request = defaultRequest { remoteHost = SockAddrInet6 0 0 (0, 0, 0, 1) 0 } -- ::1---- | Creates a mock 'Request' containing an @x-forwarded-for@ header.--- This is used to test scenarios where the application is behind a proxy--- and the client's IP address is forwarded in this header.------ ==== __Example__--- > mkRequestWithXFF "192.168.1.1"-mkRequestWithXFF- :: Text -- ^ The IP address to set in the header.- -> Request-mkRequestWithXFF ip = defaultRequest { requestHeaders = [(mk "x-forwarded-for", TE.encodeUtf8 ip)] }---- | Creates a mock 'Request' containing an @x-real-ip@ header.--- This is another common header used by proxies to forward the original client IP.------ ==== __Example__--- > mkRequestWithRealIP "2001:db8::1"-mkRequestWithRealIP- :: Text -- ^ The IP address to set in the header.- -> Request-mkRequestWithRealIP ip = defaultRequest { requestHeaders = [(mk "x-real-ip", TE.encodeUtf8 ip)] }----- * Mock Application---- | A simple WAI 'Application' that always returns a 200 OK response.--- This serves as the downstream application in tests, allowing focus to be--- on the middleware's behavior.-mockApp :: Application-mockApp _ respond = respond $ responseLBS status200 [] (LBS.fromStrict $ TE.encodeUtf8 "OK")----- * Test Suite Definition---- | The main entry point for all tests in this module.--- It groups tests by the rate-limiting algorithm being tested.-tests :: TestTree-tests = testGroup "Rate Limiting Tests"- [ testGroup "Fixed Window Algorithm"- [ testCase "Allows IPv4 requests below limit" $ testBelowLimit FixedWindow byIP- , testCase "Blocks IPv4 requests exceeding limit" $ testExceedLimit FixedWindow byIP- , testCase "Allows IPv6 requests below limit" $ testBelowLimit FixedWindow byIP- , testCase "Blocks IPv6 requests exceeding limit" $ testExceedLimit FixedWindow byIP- , testCase "Respects fixed window timing with IPv4" $ testTiming FixedWindow byIP- , testCase "Handles x-forwarded-for header for IPv4" $ testXFF FixedWindow byIP- , testCase "Handles x-real-ip header for IPv6" $ testRealIP FixedWindow byIP- , testCase "Handles concurrent requests" $ testConcurrent FixedWindow byIP- , testCase "Handles DoS-like concurrency" $ testDoS FixedWindow byIP- ]- , testGroup "Sliding Window Algorithm"- [ testCase "Allows IPv4 requests below limit" $ testBelowLimit SlidingWindow byIP- , testCase "Blocks IPv4 requests exceeding limit" $ testExceedLimit SlidingWindow byIP- , testCase "Allows IPv6 requests below limit" $ testBelowLimit SlidingWindow byIP- , testCase "Blocks IPv6 requests exceeding limit" $ testExceedLimit SlidingWindow byIP- , testCase "Respects sliding window timing with IPv4" $ testTiming SlidingWindow byIP- , testCase "Handles x-forwarded-for header for IPv4" $ testXFF SlidingWindow byIP- , testCase "Handles x-real-ip header for IPv6" $ testRealIP SlidingWindow byIP- , testCase "Handles concurrent requests" $ testConcurrent SlidingWindow byIP- , testCase "Handles DoS-like concurrency" $ testDoS SlidingWindow byIP- ]- , testGroup "Token Bucket Algorithm"- [ testCase "Allows IPv4 requests below limit" $ testBelowLimit TokenBucket byIP- , testCase "Blocks IPv4 requests exceeding limit" $ testExceedLimit TokenBucket byIP- , testCase "Allows IPv6 requests below limit" $ testBelowLimit TokenBucket byIP- , testCase "Blocks IPv6 requests exceeding limit" $ testExceedLimit TokenBucket byIP- , testCase "Respects token bucket timing with IPv4" $ testTiming TokenBucket byIP- , testCase "Handles x-forwarded-for header for IPv4" $ testXFF TokenBucket byIP- , testCase "Handles x-real-ip header for IPv6" $ testRealIP TokenBucket byIP- , testCase "Handles concurrent requests" $ testConcurrent TokenBucket byIP- , testCase "Handles DoS-like concurrency" $ testDoS TokenBucket byIP- ]- , testGroup "Leaky Bucket Algorithm"- [ testCase "Allows IPv4 requests below limit" $ testBelowLimit LeakyBucket byIP- , testCase "Blocks IPv4 requests exceeding limit" $ testExceedLimit LeakyBucket byIP- , testCase "Allows IPv6 requests below limit" $ testBelowLimit LeakyBucket byIP- , testCase "Blocks IPv6 requests exceeding limit" $ testExceedLimit LeakyBucket byIP- , testCase "Respects leaky bucket timing with IPv4" $ testTiming LeakyBucket byIP- , testCase "Handles x-forwarded-for header for IPv4" $ testXFF LeakyBucket byIP- , testCase "Handles x-real-ip header for IPv6" $ testRealIP LeakyBucket byIP- , testCase "Handles concurrent requests" $ testConcurrent LeakyBucket byIP- , testCase "Handles DoS-like concurrency" $ testDoS LeakyBucket byIP- ]- , testGroup "TinyLRU Algorithm"- [ testCase "Allows IPv4 requests below limit" $ testBelowLimit TinyLRU byIP- , testCase "Blocks IPv4 requests exceeding limit" $ testExceedLimit TinyLRU byIP- , testCase "Allows IPv6 requests below limit" $ testBelowLimit TinyLRU byIP- , testCase "Blocks IPv6 requests exceeding limit" $ testExceedLimit TinyLRU byIP- , testCase "Respects TinyLRU timing with IPv4" $ testTiming TinyLRU byIP- , testCase "Handles x-forwarded-for header for IPv4" $ testXFF TinyLRU byIP- , testCase "Handles x-real-ip header for IPv6" $ testRealIP TinyLRU byIP- , testCase "Handles concurrent requests" $ testConcurrent TinyLRU byIP- , testCase "Handles DoS-like concurrency" $ testDoS TinyLRU byIP- ]- ]---- * Test Case Implementations---- | **Test Scenario**: Verifies that the middleware allows requests when the count is below the configured limit.--- It sends two requests to an endpoint with a limit of 2 and asserts that both receive a 200 OK status.-testBelowLimit- :: Algorithm -- ^ The rate-limiting algorithm to test.- -> (Request -> IO (Maybe Text)) -- ^ The function to identify the client.- -> Assertion-testBelowLimit algo identifier = do- env <- initConfig (const defaultIPZone)- let throttle = ThrottleConfig- { throttleLimit = 2- , throttlePeriod = 60- , throttleAlgorithm = algo- , throttleIdentifier = identifier- , throttleTokenBucketTTL = Just 3600- }- env' <- addThrottle env (T.pack "test_throttle") throttle- let app = attackMiddleware env' mockApp- let session = do- result1 <- srequest $ SRequest mkIPv4Request ""- result2 <- srequest $ SRequest mkIPv4Request ""- return (result1, result2)- (response1, response2) <- runSession session app- assertEqual "First request status" status200 (simpleStatus response1)- assertEqual "Second request status" status200 (simpleStatus response2)---- | **Test Scenario**: Verifies that the middleware blocks requests once the limit is exceeded.--- It sends three requests to an endpoint with a limit of 2. It asserts that the first two succeed (200 OK)--- and the third is blocked (429 Too Many Requests).-testExceedLimit- :: Algorithm -- ^ The rate-limiting algorithm to test.- -> (Request -> IO (Maybe Text)) -- ^ The function to identify the client.- -> Assertion-testExceedLimit algo identifier = do- env <- initConfig (const defaultIPZone)- let throttle = ThrottleConfig- { throttleLimit = 2- , throttlePeriod = 60- , throttleAlgorithm = algo- , throttleIdentifier = identifier- , throttleTokenBucketTTL = Just 3600- }- env' <- addThrottle env (T.pack "test_throttle") throttle- let app = attackMiddleware env' mockApp- let session = do- _ <- srequest $ SRequest mkIPv4Request ""- _ <- srequest $ SRequest mkIPv4Request ""- result3 <- srequest $ SRequest mkIPv4Request ""- return result3- response3 <- runSession session app- assertEqual "Third request status" status429 (simpleStatus response3)---- | **Test Scenario**: Verifies that the rate limit counter resets after the configured window period.--- It sets a limit of 1 request per 1-second window. It sends one request, waits for 2 seconds,--- then sends a second request. It asserts that both requests are successful (200 OK).-testTiming- :: Algorithm -- ^ The rate-limiting algorithm to test.- -> (Request -> IO (Maybe Text)) -- ^ The function to identify the client.- -> Assertion-testTiming algo identifier = do- env <- initConfig (const defaultIPZone)- let throttle = ThrottleConfig- { throttleLimit = 1- , throttlePeriod = 1- , throttleAlgorithm = algo- , throttleIdentifier = identifier- , throttleTokenBucketTTL = Just 3600- }- env' <- addThrottle env (T.pack "test_throttle") throttle- let app = attackMiddleware env' mockApp- lock <- newMVar ()- let session = do- result1 <- srequest $ SRequest mkIPv4Request ""- liftIO $ modifyMVar_ lock $ \_ -> threadDelay 2000000 >> return () -- Wait 2s- result2 <- srequest $ SRequest mkIPv4Request ""- return (result1, result2)- (response1, response2) <- runSession session app- assertEqual "First request status" status200 (simpleStatus response1)- assertEqual "Second request status after reset" status200 (simpleStatus response2)---- | **Test Scenario**: Verifies correct IP identification using the @x-forwarded-for@ header.--- It sends two requests from the same forwarded IP address with a limit of 1.--- It asserts that the first request succeeds (200 OK) and the second is blocked (429 Too Many Requests).-testXFF- :: Algorithm -- ^ The rate-limiting algorithm to test.- -> (Request -> IO (Maybe Text)) -- ^ The function to identify the client.- -> Assertion-testXFF algo identifier = do- env <- initConfig (const defaultIPZone)- let throttle = ThrottleConfig- { throttleLimit = 1- , throttlePeriod = 60- , throttleAlgorithm = algo- , throttleIdentifier = identifier- , throttleTokenBucketTTL = Just 3600- }- env' <- addThrottle env (T.pack "test_throttle") throttle- let app = attackMiddleware env' mockApp- let session = do- result1 <- srequest $ SRequest (mkRequestWithXFF "192.168.1.1") ""- result2 <- srequest $ SRequest (mkRequestWithXFF "192.168.1.1") ""- return (result1, result2)- (response1, response2) <- runSession session app- assertEqual "First XFF request status" status200 (simpleStatus response1)- assertEqual "Second XFF request status" status429 (simpleStatus response2)---- | **Test Scenario**: Verifies correct IP identification using the @x-real-ip@ header.--- It sends two requests from the same real IP address with a limit of 1.--- It asserts that the first request succeeds (200 OK) and the second is blocked (429 Too Many Requests).-testRealIP- :: Algorithm -- ^ The rate-limiting algorithm to test.- -> (Request -> IO (Maybe Text)) -- ^ The function to identify the client.- -> Assertion-testRealIP algo identifier = do- env <- initConfig (const defaultIPZone)- let throttle = ThrottleConfig- { throttleLimit = 1- , throttlePeriod = 60- , throttleAlgorithm = algo- , throttleIdentifier = identifier- , throttleTokenBucketTTL = Just 3600- }- env' <- addThrottle env (T.pack "test_throttle") throttle- let app = attackMiddleware env' mockApp- let session = do- result1 <- srequest $ SRequest (mkRequestWithRealIP "::1") ""- result2 <- srequest $ SRequest (mkRequestWithRealIP "::1") ""- return (result1, result2)- (response1, response2) <- runSession session app- assertEqual "First Real-IP request status" status200 (simpleStatus response1)- assertEqual "Second Real-IP request status" status429 (simpleStatus response2)---- | **Test Scenario**: Verifies the middleware's behavior under moderate concurrent load.--- It sends 5 concurrent requests to an endpoint with a limit of 5. It then sends a sixth request.--- It asserts that the first 5 requests succeed (200 OK) and the sixth is blocked (429 Too Many Requests).-testConcurrent- :: Algorithm -- ^ The rate-limiting algorithm to test.- -> (Request -> IO (Maybe Text)) -- ^ The function to identify the client.- -> Assertion-testConcurrent algo identifier = do- env <- initConfig (const defaultIPZone)- let throttle = ThrottleConfig- { throttleLimit = 5- , throttlePeriod = 60- , throttleAlgorithm = algo- , throttleIdentifier = identifier- , throttleTokenBucketTTL = Just 3600- }- env' <- addThrottle env (T.pack "test_throttle") throttle- let app = attackMiddleware env' mockApp- - -- Execute concurrent requests sequentially within the session- let session = do- responses <- replicateM 5 (srequest $ SRequest mkIPv4Request "")- result6 <- srequest $ SRequest mkIPv4Request ""- return (responses, result6)- (responses, response6) <- runSession session app- - -- Check that first 5 requests succeeded- mapM_ (\(i, resp) -> assertEqual ("Request " ++ show i ++ " status") status200 (simpleStatus resp))- (zip [1..5] responses)- -- Check that 6th request was throttled- assertEqual "Sixth request status after limit" status429 (simpleStatus response6)---- | **Test Scenario**: Simulates a Denial-of-Service (DoS) attack with high concurrency.--- It sends 15 concurrent requests to an endpoint with a limit of 10.--- It asserts that some requests were successful, some were blocked, and the total processed--- matches the number sent, ensuring the server remains stable.-testDoS- :: Algorithm -- ^ The rate-limiting algorithm to test.- -> (Request -> IO (Maybe Text)) -- ^ The function to identify the client.- -> Assertion-testDoS algo identifier = do- env <- initConfig (const defaultIPZone)- let throttle = ThrottleConfig- { throttleLimit = 10- , throttlePeriod = 60- , throttleAlgorithm = algo- , throttleIdentifier = identifier- , throttleTokenBucketTTL = Just 3600- }- env' <- addThrottle env (T.pack "test_throttle") throttle- let app = attackMiddleware env' mockApp- - -- Execute many requests sequentially to simulate DoS scenario- let session = do- responses <- replicateM 15 (srequest $ SRequest mkIPv4Request "")- return responses- responses <- runSession session app- - -- Count how many succeeded (should be <= 10)- let successCount = length $ filter (\resp -> simpleStatus resp == status200) responses- let throttledCount = length $ filter (\resp -> simpleStatus resp == status429) responses- - assertBool "Some requests should succeed" (successCount > 0)- assertBool "Some requests should be throttled" (throttledCount > 0)- assertEqual "Total requests processed" 15 (successCount + throttledCount)+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++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.++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.++The module defines helper functions to create mock requests and a mock application to isolate the middleware for testing.+-}++module Keter.RateLimiter.WAITests (+ -- * Test Suite+ tests+) where++import Test.Tasty+import Test.Tasty.HUnit+import Test.Tasty.QuickCheck+import Test.QuickCheck.Monadic+import Network.Wai+import Network.Wai.Test+import Network.HTTP.Types+import Network.Socket (SockAddr(..), tupleToHostAddress)+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Text.Encoding as TE+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import qualified Data.HashMap.Strict as HM+import Data.Aeson hiding (pairs)+import Data.CaseInsensitive (mk)+import Control.Concurrent (threadDelay, forkIO)+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(..))++-- * Request Helpers++-- | A simple WAI 'Application' that always returns a 200 OK response.+mockApp :: Application+mockApp _ respond = respond $ responseLBS status200 [] (LBS.fromStrict $ TE.encodeUtf8 "OK")++-- | Creates a mock 'Request' with a default IPv4 address (127.0.0.1).+mkIPv4Request :: Request+mkIPv4Request = defaultRequest { remoteHost = SockAddrInet 0 (tupleToHostAddress (127, 0, 0, 1)) }++-- | Creates a mock 'Request' with a default IPv6 address (::1).+mkIPv6Request :: Request+mkIPv6Request = defaultRequest { remoteHost = SockAddrInet6 0 0 (0, 0, 0, 1) 0 }++-- | Creates a mock 'Request' with a specific header.+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+mkRequestWithXFF ip = defaultRequest { requestHeaders = [(mk "x-forwarded-for", TE.encodeUtf8 ip)] }++-- | Creates a mock 'Request' with an @x-real-ip@ header.+mkRequestWithRealIP :: Text -> Request+mkRequestWithRealIP ip = defaultRequest { requestHeaders = [(mk "x-real-ip", TE.encodeUtf8 ip)] }++-- | Extracts a cookie value using Web.Cookie, ignoring empty values.+extractCookieWC :: Text -> BS.ByteString -> Maybe Text+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++-- * Test Suite Definition++tests :: TestTree+tests = testGroup "Rate Limiting Tests"+ [ algorithmTests+ , configurationTests+ , multipleThrottleTests+ , identifierStrategyTests+ , zoneBasedTests+ , jsonConfigTests+ , cacheManagementTests+ , errorHandlingTests+ , performanceTests+ , propertyBasedTests+ ]++-- | Tests for each rate-limiting algorithm across various scenarios.+algorithmTests :: TestTree+algorithmTests = testGroup "Algorithm-Specific Tests"+ [ algorithmTestGroup FixedWindow+ , algorithmTestGroup SlidingWindow+ , algorithmTestGroup TokenBucket+ , algorithmTestGroup LeakyBucket+ , algorithmTestGroup TinyLRU+ ]+ 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+ ]++-- | Test buildRateLimiter with various configurations.+configurationTests :: TestTree+configurationTests = testGroup "Configuration-Driven Middleware"+ [ testCase "buildRateLimiter with single throttle" testBuildSingleThrottle+ , testCase "buildRateLimiter with multiple throttles" testBuildMultipleThrottles+ , testCase "buildRateLimiter with different zones" testBuildWithZones+ , testCase "Empty throttles list" testEmptyThrottles+ ]++-- | Test multiple throttles running simultaneously.+multipleThrottleTests :: TestTree+multipleThrottleTests = testGroup "Multiple Throttle Rules"+ [ testCase "Multiple throttles with same algorithm" testMultipleSameAlgo+ , testCase "Multiple throttles with different algorithms" testMultipleDiffAlgo+ , testCase "Throttle priority and interaction" testThrottlePriority+ , testCase "Independent throttle counters" testIndependentCounters+ ]++-- | Test different identifier strategies.+identifierStrategyTests :: TestTree+identifierStrategyTests = testGroup "Identifier Strategies"+ [ testCase "IdHeader strategy" testIdHeaderStrategy+ , testCase "IdCookie strategy" testIdCookieStrategy+ , testCase "IdIPAndPath strategy" testIdIPAndPathStrategy+ , testCase "IdIPAndUA strategy" testIdIPAndUAStrategy+ , testCase "IdHeaderAndIP strategy" testIdHeaderAndIPStrategy+ , testCase "Missing header/cookie handling" testMissingIdentifiers+ , testCase "Cookie parsing edge cases" testCookieParsing+ ]++-- | Test zone-based separation.+zoneBasedTests :: TestTree+zoneBasedTests = testGroup "Zone-Based Separation"+ [ testCase "ZoneIP separation" testZoneIPSeparation+ , testCase "ZoneHeader separation" testZoneHeaderSeparation+ , testCase "Zone creation and cleanup" testZoneCreation+ , testCase "Default zone fallback" testDefaultZoneFallback+ ]++-- | Test JSON configuration parsing.+jsonConfigTests :: TestTree+jsonConfigTests = testGroup "JSON Configuration"+ [ testCase "Parse IdentifierBy JSON" testParseIdentifierBy+ , testCase "Parse ZoneBy JSON" testParseZoneBy+ , testCase "Parse RLThrottle JSON" testParseRLThrottle+ , testCase "Parse RateLimiterConfig JSON" testParseRateLimiterConfig+ , testCase "Invalid JSON handling" testInvalidJSON+ ]++-- | Test cache management functions.+cacheManagementTests :: TestTree+cacheManagementTests = testGroup "Cache Management"+ [ testCase "cacheResetAll functionality" testCacheResetAll+ , testCase "Zone cache isolation" testZoneCacheIsolation+ , testCase "Memory cleanup after reset" testMemoryCleanup+ ]++-- | Test error handling and edge cases.+errorHandlingTests :: TestTree+errorHandlingTests = testGroup "Error Handling & Edge Cases"+ [ testCase "Zero period handling" testZeroPeriod+ , testCase "Negative limit handling" testNegativeLimit+ , testCase "Very large numbers" testLargeNumbers+ , testCase "Malformed requests" testMalformedRequests+ , testCase "Concurrent access safety" testConcurrentSafety+ ]++-- | Performance and stress tests.+performanceTests :: TestTree+performanceTests = testGroup "Performance Tests"+ [ testCase "High throughput single client" testHighThroughputSingle+ , testCase "Many unique clients" testManyClients+ , testCase "Algorithm performance comparison" testAlgorithmPerformance+ , testCase "Memory usage with many zones" testManyZones+ ]++-- | Property-based tests using QuickCheck.+propertyBasedTests :: TestTree+propertyBasedTests = testGroup "Property-Based Tests"+ [ testProperty "Cookie extraction properties" propCookieExtraction+ , testProperty "Header name round-trip" propHeaderNameRoundTrip+ , testProperty "IP extraction consistency" propIPExtraction+ , testProperty "Rate limiting monotonicity" propRateLimitingMonotonicity+ ]++-- * Test Case Implementations++-- | Verifies that requests below the limit are allowed.+testBelowLimit :: Algorithm -> (Request -> IO (Maybe Text)) -> Request -> Assertion+testBelowLimit algo identifier req = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 2 60 algo identifier (Just 3600)+ env' <- addThrottle env "test_throttle" throttle+ let app = attackMiddleware env' mockApp+ let session = do+ result1 <- srequest $ SRequest req ""+ result2 <- srequest $ SRequest req ""+ return (result1, result2)+ (response1, response2) <- runSession session app+ assertEqual "First request status" status200 (simpleStatus response1)+ assertEqual "Second request status" status200 (simpleStatus response2)++-- | Verifies that requests exceeding the limit are blocked.+testExceedLimit :: Algorithm -> (Request -> IO (Maybe Text)) -> Request -> Assertion+testExceedLimit algo identifier req = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 2 60 algo identifier (Just 3600)+ env' <- addThrottle env "test_throttle" throttle+ let app = attackMiddleware env' mockApp+ let session = do+ _ <- srequest $ SRequest req ""+ _ <- srequest $ SRequest req ""+ result3 <- srequest $ SRequest req ""+ return result3+ response3 <- runSession session app+ 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 algo identifier = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 1 algo identifier (Just 3600)+ env' <- addThrottle env "test_throttle" throttle+ let app = attackMiddleware env' mockApp+ lock <- newMVar ()+ let session = do+ result1 <- srequest $ SRequest mkIPv4Request ""+ liftIO $ modifyMVar_ lock $ \_ -> threadDelay 2000000 >> return () -- Wait 2s+ result2 <- srequest $ SRequest mkIPv4Request ""+ return (result1, result2)+ (response1, response2) <- runSession session app+ assertEqual "First request status" status200 (simpleStatus response1)+ 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 algo identifier = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 algo identifier (Just 3600)+ env' <- addThrottle env "test_throttle" throttle+ let app = attackMiddleware env' mockApp+ let session = do+ result1 <- srequest $ SRequest (mkRequestWithXFF "192.168.1.1") ""+ result2 <- srequest $ SRequest (mkRequestWithXFF "192.168.1.1") ""+ return (result1, result2)+ (response1, response2) <- runSession session app+ assertEqual "First XFF request status" status200 (simpleStatus response1)+ 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 algo identifier = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 algo identifier (Just 3600)+ env' <- addThrottle env "test_throttle" throttle+ let app = attackMiddleware env' mockApp+ let session = do+ result1 <- srequest $ SRequest (mkRequestWithRealIP "::1") ""+ result2 <- srequest $ SRequest (mkRequestWithRealIP "::1") ""+ return (result1, result2)+ (response1, response2) <- runSession session app+ assertEqual "First Real-IP request status" status200 (simpleStatus response1)+ assertEqual "Second Real-IP request status" status429 (simpleStatus response2)++-- | Verifies behavior under moderate concurrent load.+testConcurrent :: Algorithm -> (Request -> IO (Maybe Text)) -> Assertion+testConcurrent algo identifier = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 5 60 algo identifier (Just 3600)+ env' <- addThrottle env "test_throttle" throttle+ let app = attackMiddleware env' mockApp+ let session = do+ responses <- replicateM 5 (srequest $ SRequest mkIPv4Request "")+ 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)+ assertEqual "Sixth request status after limit" status429 (simpleStatus response6)++-- | Simulates a DoS attack with high concurrency.+testDoS :: Algorithm -> (Request -> IO (Maybe Text)) -> Assertion+testDoS algo identifier = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 10 60 algo identifier (Just 3600)+ env' <- addThrottle env "test_throttle" throttle+ let app = attackMiddleware env' mockApp+ let session = do+ responses <- replicateM 15 (srequest $ SRequest mkIPv4Request "")+ return responses+ responses <- runSession session app+ let successCount = length $ filter (\resp -> simpleStatus resp == status200) responses+ let throttledCount = length $ filter (\resp -> simpleStatus resp == status429) responses+ assertBool "Some requests should succeed" (successCount > 0)+ assertBool "Some requests should be throttled" (throttledCount > 0)+ assertEqual "Total requests processed" 15 (successCount + throttledCount)++-- | Tests buildRateLimiter with a single throttle.+testBuildSingleThrottle :: Assertion+testBuildSingleThrottle = do+ let config = RateLimiterConfig ZoneDefault+ [ RLThrottle "api-limit" 5 60 FixedWindow IdIP Nothing ]+ middleware <- buildRateLimiter config+ let app = middleware mockApp+ let session = replicateM 6 (srequest $ SRequest mkIPv4Request "")+ responses <- runSession session app+ let successCount = length $ filter (\r -> simpleStatus r == status200) responses+ let throttledCount = length $ filter (\r -> simpleStatus r == status429) responses+ assertEqual "Success count" 5 successCount+ assertEqual "Throttled count" 1 throttledCount++-- | Tests buildRateLimiter with multiple throttles.+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+ ]+ middleware <- buildRateLimiter config+ let app = middleware mockApp+ let requestWithApi = mkRequestWithHeader "X-API-Key" "test-key"+ let session = do+ r1 <- srequest $ SRequest requestWithApi ""+ r2 <- srequest $ SRequest requestWithApi ""+ r3 <- srequest $ SRequest requestWithApi ""+ r4 <- srequest $ SRequest requestWithApi ""+ return [r1, r2, r3, r4]+ responses <- runSession session app+ assertEqual "First 3 API requests" [status200, status200, status200] (map simpleStatus $ take 3 responses)+ assertEqual "4th API request blocked" status429 (simpleStatus $ responses !! 3)++-- | Tests buildRateLimiter with different zones.+testBuildWithZones :: Assertion+testBuildWithZones = do+ let config = RateLimiterConfig (ZoneHeader "X-Zone")+ [ RLThrottle "zone-limit" 2 60 FixedWindow IdIP Nothing ]+ middleware <- buildRateLimiter config+ let app = middleware mockApp+ let zoneAReq = mkRequestWithHeader "X-Zone" "A"+ let zoneBReq = mkRequestWithHeader "X-Zone" "B"+ let session = do+ ra1 <- srequest $ SRequest zoneAReq ""+ ra2 <- srequest $ SRequest zoneAReq ""+ rb1 <- srequest $ SRequest zoneBReq ""+ rb2 <- srequest $ SRequest zoneBReq ""+ 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)++-- | Tests buildRateLimiter with an empty throttles list.+testEmptyThrottles :: Assertion+testEmptyThrottles = do+ let config = RateLimiterConfig ZoneDefault []+ middleware <- buildRateLimiter config+ let app = middleware mockApp+ let session = replicateM 10 (srequest $ SRequest mkIPv4Request "")+ responses <- runSession session app+ assertEqual "All requests succeed" 10 (length $ filter (\r -> simpleStatus r == status200) responses)++-- | Tests multiple throttles with the same algorithm.+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+ env' <- addThrottle env "global" throttle1 >>= \e -> addThrottle e "strict" throttle2+ let app = attackMiddleware env' mockApp+ let session = replicateM 4 (srequest $ SRequest mkIPv4Request "")+ responses <- runSession session app+ let successCount = length $ filter (\r -> simpleStatus r == status200) responses+ assertEqual "Limited by stricter rule" 3 successCount++-- | Tests multiple throttles with different algorithms.+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)+ env' <- addThrottle env "fixed" throttle1 >>= \e -> addThrottle e "bucket" throttle2+ let app = attackMiddleware env' mockApp+ let session = replicateM 6 (srequest $ SRequest mkIPv4Request "")+ responses <- runSession session app+ let successCount = length $ filter (\r -> simpleStatus r == status200) responses+ assertEqual "Multiple algorithms interact" 5 successCount++-- | Tests throttle priority and interaction.+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+ env' <- addThrottle env "permissive" permissive >>= \e -> addThrottle e "restrictive" restrictive+ let app = attackMiddleware env' mockApp+ let session = do+ r1 <- srequest $ SRequest mkIPv4Request ""+ r2 <- srequest $ SRequest mkIPv4Request ""+ return [r1, r2]+ responses <- runSession session app+ assertEqual "Most restrictive rule wins" [status200, status429] (map simpleStatus responses)++-- | Tests independent throttle counters.+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+ env' <- addThrottle env "ip" ipThrottle >>= \e -> addThrottle e "user" headerThrottle+ let app = attackMiddleware env' mockApp+ let userReq = mkRequestWithHeader "X-User-ID" "user123"+ let session = do+ ri1 <- srequest $ SRequest mkIPv4Request ""+ ri2 <- srequest $ SRequest mkIPv4Request ""+ ru1 <- srequest $ SRequest userReq ""+ ru2 <- srequest $ SRequest userReq ""+ return [ri1, ri2, ru1, ru2]+ responses <- runSession session app+ assertEqual "Independent counters" [status200, status200, status200, status200] (map simpleStatus responses)++-- | Tests header-based identifier strategy.+testIdHeaderStrategy :: Assertion+testIdHeaderStrategy = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 2 60 FixedWindow (mkIdentifier (IdHeader "X-Client-ID")) Nothing+ env' <- addThrottle env "header" throttle+ let app = attackMiddleware env' mockApp+ let client1Req = mkRequestWithHeader "X-Client-ID" "client1"+ let client2Req = mkRequestWithHeader "X-Client-ID" "client2"+ let session = do+ r1 <- srequest $ SRequest client1Req ""+ r2 <- srequest $ SRequest client1Req ""+ r3 <- srequest $ SRequest client2Req ""+ r4 <- srequest $ SRequest client2Req ""+ r5 <- srequest $ SRequest client1Req ""+ return [r1, r2, r3, r4, r5]+ responses <- runSession session app+ assertEqual "Header-based identification"+ [status200, status200, status200, status200, status429]+ (map simpleStatus responses)++-- | Tests cookie-based identifier strategy.+testIdCookieStrategy :: Assertion+testIdCookieStrategy = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier (IdCookie "session")) Nothing+ env' <- addThrottle env "cookie" throttle+ let app = attackMiddleware env' mockApp+ let session1Req = mkRequestWithCookie "session" "sess123"+ let session2Req = mkRequestWithCookie "session" "sess456"+ let session = do+ r1 <- srequest $ SRequest session1Req ""+ r2 <- srequest $ SRequest session2Req ""+ r3 <- srequest $ SRequest session1Req ""+ return [r1, r2, r3]+ responses <- runSession session app+ assertEqual "Cookie-based identification"+ [status200, status200, status429]+ (map simpleStatus responses)++-- | Tests IP+Path identifier strategy.+testIdIPAndPathStrategy :: Assertion+testIdIPAndPathStrategy = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIPAndPath) Nothing+ env' <- addThrottle env "ip-path" throttle+ let app = attackMiddleware env' mockApp+ let path1Req = mkIPv4Request { rawPathInfo = "/api/v1" }+ let path2Req = mkIPv4Request { rawPathInfo = "/api/v2" }+ let session = do+ r1 <- srequest $ SRequest path1Req ""+ r2 <- srequest $ SRequest path2Req ""+ r3 <- srequest $ SRequest path1Req ""+ return [r1, r2, r3]+ responses <- runSession session app+ assertEqual "IP+Path identification"+ [status200, status200, status429]+ (map simpleStatus responses)++-- | Tests IP+UserAgent identifier strategy.+testIdIPAndUAStrategy :: Assertion+testIdIPAndUAStrategy = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIPAndUA) Nothing+ env' <- addThrottle env "ip-ua" throttle+ let app = attackMiddleware env' mockApp+ let ua1Req = mkRequestWithHeader "User-Agent" "Browser/1.0"+ let ua2Req = mkRequestWithHeader "User-Agent" "Browser/2.0"+ let session = do+ r1 <- srequest $ SRequest ua1Req ""+ r2 <- srequest $ SRequest ua2Req ""+ r3 <- srequest $ SRequest ua1Req ""+ return [r1, r2, r3]+ responses <- runSession session app+ assertEqual "IP+UserAgent identification"+ [status200, status200, status429]+ (map simpleStatus responses)++-- | Tests Header+IP identifier strategy.+testIdHeaderAndIPStrategy :: Assertion+testIdHeaderAndIPStrategy = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier (IdHeaderAndIP "X-Service")) Nothing+ env' <- addThrottle env "header-ip" throttle+ let app = attackMiddleware env' mockApp+ let service1Req = mkRequestWithHeader "X-Service" "service1"+ let service2Req = mkRequestWithHeader "X-Service" "service2"+ let session = do+ r1 <- srequest $ SRequest service1Req ""+ r2 <- srequest $ SRequest service2Req ""+ r3 <- srequest $ SRequest service1Req ""+ return [r1, r2, r3]+ responses <- runSession session app+ assertEqual "Header+IP identification"+ [status200, status200, status429]+ (map simpleStatus responses)++-- | Tests handling of missing identifiers.+testMissingIdentifiers :: Assertion+testMissingIdentifiers = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier (IdHeader "Missing-Header")) Nothing+ env' <- addThrottle env "missing" throttle+ let app = attackMiddleware env' mockApp+ let session = replicateM 5 (srequest $ SRequest mkIPv4Request "")+ responses <- runSession session app+ let successCount = length $ filter (\r -> simpleStatus r == status200) responses+ assertEqual "Missing identifiers bypass throttling" 5 successCount++-- | Tests cookie parsing edge cases.+testCookieParsing :: Assertion+testCookieParsing = do+ let testCases =+ [ ("session=abc123", Just "abc123")+ , ("session=abc123; other=value", Just "abc123")+ , ("other=value; session=def456", Just "def456")+ , ("session=; other=value", Nothing)+ , ("other=value", Nothing)+ , ("malformed", Nothing)+ ]+ mapM_ (\(cookie, expected) -> do+ 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 ]+ middleware <- buildRateLimiter config+ let app = middleware mockApp+ let ip1Req = mkRequestWithXFF "192.168.1.1"+ let ip2Req = mkRequestWithXFF "192.168.1.2"+ let session = do+ r1 <- srequest $ SRequest ip1Req ""+ r2 <- srequest $ SRequest ip2Req ""+ r3 <- srequest $ SRequest ip1Req ""+ r4 <- srequest $ SRequest ip2Req ""+ return [r1, r2, r3, r4]+ responses <- runSession session app+ assertEqual "IP zone separation"+ [status200, status200, status429, status429]+ (map simpleStatus responses)++-- | Tests header-based zone separation.+testZoneHeaderSeparation :: Assertion+testZoneHeaderSeparation = do+ let config = RateLimiterConfig (ZoneHeader "X-Tenant")+ [ RLThrottle "tenant-limit" 1 60 FixedWindow IdIP Nothing ]+ middleware <- buildRateLimiter config+ let app = middleware mockApp+ let tenant1Req = mkRequestWithHeader "X-Tenant" "tenant1"+ let tenant2Req = mkRequestWithHeader "X-Tenant" "tenant2"+ let session = do+ r1 <- srequest $ SRequest tenant1Req ""+ r2 <- srequest $ SRequest tenant2Req ""+ r3 <- srequest $ SRequest tenant1Req ""+ return [r1, r2, r3]+ responses <- runSession session app+ assertEqual "Header zone separation"+ [status200, status200, status429]+ (map simpleStatus responses)++-- | Tests zone creation and cleanup.+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+ env' <- addThrottle env "test" throttle+ let zone1Req = mkRequestWithHeader "X-Zone" "zone1"+ let zone2Req = mkRequestWithHeader "X-Zone" "zone2"+ zoneCaches <- readIORef (envZoneCachesMap env)+ initialSize <- return $ HM.size zoneCaches+ _ <- instrument env' zone1Req+ _ <- instrument env' zone2Req+ zoneCaches' <- readIORef (envZoneCachesMap env')+ finalSize <- return $ HM.size zoneCaches'+ assertBool "New zones created" (finalSize > initialSize)++-- | Tests default zone fallback.+testDefaultZoneFallback :: Assertion+testDefaultZoneFallback = do+ let config = RateLimiterConfig (ZoneHeader "Missing-Header")+ [ RLThrottle "default-fallback" 1 60 FixedWindow IdIP Nothing ]+ middleware <- buildRateLimiter config+ let app = middleware mockApp+ let session = do+ r1 <- srequest $ SRequest mkIPv4Request ""+ r2 <- srequest $ SRequest mkIPv4Request ""+ return [r1, r2]+ responses <- runSession session app+ assertEqual "Default zone fallback"+ [status200, status429]+ (map simpleStatus responses)++-- | Tests parsing of IdentifierBy JSON.+testParseIdentifierBy :: Assertion+testParseIdentifierBy = do+ let testCases =+ [ ("\"ip\"", Right IdIP)+ , ("\"ip+path\"", Right IdIPAndPath)+ , ("\"ip+ua\"", Right IdIPAndUA)+ , ("{\"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))+ ]+ 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++-- | Tests parsing of ZoneBy JSON.+testParseZoneBy :: Assertion+testParseZoneBy = do+ let testCases =+ [ ("\"default\"", Right ZoneDefault)+ , ("\"ip\"", Right ZoneIP)+ , ("{\"header\": \"X-Region\"}", Right (ZoneHeader (hdr "X-Region")))+ , ("\"invalid\"", Left ("zone_by must be" :: 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++-- | Tests parsing of RLThrottle JSON.+testParseRLThrottle :: Assertion+testParseRLThrottle = do+ let json = "{\"name\":\"test\",\"limit\":100,\"period\":3600,\"algorithm\":\"FixedWindow\",\"identifier_by\":\"ip\"}"+ let result = eitherDecode (LBS.fromStrict $ TE.encodeUtf8 json) :: Either String RLThrottle+ case result of+ Right throttle -> do+ assertEqual "Name" "test" (rlName throttle)+ assertEqual "Limit" 100 (rlLimit throttle)+ assertEqual "Period" 3600 (rlPeriod throttle)+ assertEqual "Algorithm" FixedWindow (rlAlgo throttle)+ Left err -> assertFailure $ "Parse failed: " <> err++-- | Tests parsing of RateLimiterConfig JSON.+testParseRateLimiterConfig :: Assertion+testParseRateLimiterConfig = do+ let json = "{\"zone_by\":\"default\",\"throttles\":[{\"name\":\"test\",\"limit\":100,\"period\":3600,\"algorithm\":\"FixedWindow\",\"identifier_by\":\"ip\"}]}"+ let result = eitherDecode (LBS.fromStrict $ TE.encodeUtf8 json) :: Either String RateLimiterConfig+ case result of+ Right config -> do+ assertEqual "Zone by" ZoneDefault (rlZoneBy config)+ assertEqual "Throttles length" 1 (length $ rlThrottles config)+ Left err -> assertFailure $ "Parse failed: " <> err++-- | Tests handling of invalid JSON.+testInvalidJSON :: Assertion+testInvalidJSON = do+ let invalidJson = "{\"invalid\": true}"+ let result = eitherDecode (LBS.fromStrict $ TE.encodeUtf8 invalidJson) :: Either String RateLimiterConfig+ case result of+ Left _ -> return ()+ Right _ -> assertFailure "Should have failed to parse invalid JSON"++-- | Tests cacheResetAll functionality.+testCacheResetAll :: Assertion+testCacheResetAll = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1 60 FixedWindow (mkIdentifier IdIP) Nothing+ env' <- addThrottle env "test" throttle+ let app = attackMiddleware env' mockApp+ let session1 = srequest $ SRequest mkIPv4Request ""+ resp1 <- runSession session1 app+ assertEqual "First request succeeds" status200 (simpleStatus resp1)+ let session2 = srequest $ SRequest mkIPv4Request ""+ resp2 <- runSession session2 app+ assertEqual "Second request blocked" status429 (simpleStatus resp2)+ cacheResetAll env'+ let session3 = srequest $ SRequest mkIPv4Request ""+ resp3 <- runSession session3 app+ assertEqual "Request succeeds after reset" status200 (simpleStatus resp3)++-- | Tests zone cache isolation.+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+ env' <- addThrottle env "test" throttle+ let app = attackMiddleware env' mockApp+ let zoneAReq = mkRequestWithHeader "X-Zone" "A"+ let zoneBReq = mkRequestWithHeader "X-Zone" "B"+ let session = do+ ra1 <- srequest $ SRequest zoneAReq ""+ rb1 <- srequest $ SRequest zoneBReq ""+ return [ra1, rb1]+ responses <- runSession session app+ assertEqual "Both zones populated" [status200, status200] (map simpleStatus responses)+ zoneCaches <- readIORef (envZoneCachesMap env')+ let zoneCount = HM.size zoneCaches+ assertBool "Multiple zones created" (zoneCount >= 2)++-- | Tests memory cleanup after reset.+testMemoryCleanup :: Assertion+testMemoryCleanup = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 100 60 FixedWindow (mkIdentifier IdIP) Nothing+ env' <- addThrottle env "test" throttle+ cacheResetAll env'+ cacheResetAll env'+ _ <- instrument env' mkIPv4Request+ return ()++-- | Tests handling of zero period.+testZeroPeriod :: Assertion+testZeroPeriod = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 10 0 TokenBucket (mkIdentifier IdIP) (Just 60)+ env' <- addThrottle env "zero-period" throttle+ _ <- instrument env' mkIPv4Request+ return ()++-- | Tests handling of negative limit.+testNegativeLimit :: Assertion+testNegativeLimit = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig (-1) 60 FixedWindow (mkIdentifier IdIP) Nothing+ env' <- addThrottle env "negative" throttle+ _ <- instrument env' mkIPv4Request+ return ()++-- | Tests handling of very large numbers.+testLargeNumbers :: Assertion+testLargeNumbers = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig (maxBound :: Int) (maxBound :: Int) FixedWindow (mkIdentifier IdIP) Nothing+ env' <- addThrottle env "large" throttle+ blocked <- instrument env' mkIPv4Request+ assertEqual "Large numbers handled" False blocked++-- | Tests handling of malformed requests.+testMalformedRequests :: Assertion+testMalformedRequests = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 5 60 FixedWindow (mkIdentifier (IdHeader "X-Malformed")) Nothing+ env' <- addThrottle env "malformed" throttle+ let malformedReq = defaultRequest { requestHeaders = [(mk "X-Malformed", "\xFF\xFE\xFD")] }+ _ <- instrument env' malformedReq+ return ()++-- | Tests concurrent access safety.+testConcurrentSafety :: Assertion+testConcurrentSafety = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 100 60 FixedWindow (mkIdentifier IdIP) Nothing+ env' <- addThrottle env "concurrent" throttle+ results <- newMVar []+ let worker :: Integer -> IO ()+ worker i = do+ if i `mod` 10 == 0+ then cacheResetAll env'+ else do+ blocked <- instrument env' mkIPv4Request+ modifyMVar_ results (return . (blocked:))+ mapM_ (\i -> forkIO (worker i)) [1..50 :: Integer]+ threadDelay 100000+ finalResults <- readMVar results+ assertBool "Concurrent operations completed" (length finalResults > 0)++-- | Tests high throughput for a single client.+testHighThroughputSingle :: Assertion+testHighThroughputSingle = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 1000 60 FixedWindow (mkIdentifier IdIP) Nothing+ env' <- addThrottle env "throughput" throttle+ let app = attackMiddleware env' mockApp+ let session = replicateM 500 (srequest $ SRequest mkIPv4Request "")+ responses <- runSession session app+ let successCount = length $ filter (\r -> simpleStatus r == status200) responses+ assertBool "High throughput handled" (successCount > 0)++-- | Tests handling of many unique clients.+testManyClients :: Assertion+testManyClients = do+ env <- initConfig (const defaultIPZone)+ let throttle = ThrottleConfig 2 60 FixedWindow (mkIdentifier (IdHeader "X-Client-ID")) Nothing+ env' <- addThrottle env "many-clients" throttle+ let app = attackMiddleware env' mockApp+ let makeClientRequest :: Integer -> Request+ makeClientRequest i = mkRequestWithHeader "X-Client-ID" (T.pack $ "client" <> show i)+ let session = mapM (\i -> srequest $ SRequest (makeClientRequest i) "") [1..100 :: Integer]+ responses <- runSession session app+ let successCount = length $ filter (\r -> simpleStatus r == status200) responses+ assertEqual "Many clients handled" 100 successCount++-- | Tests performance across different algorithms.+testAlgorithmPerformance :: Assertion+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+ 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+ env' <- addThrottle env "zones" throttle+ 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')+ let zoneCount = HM.size zoneCaches+ assertBool "Many zones created" (zoneCount > 10)+ assertBool "Reasonable zone count" (zoneCount <= 51)++-- * Property-Based Tests++-- | Generates valid token characters per RFC 6265 (simplified).+validTokenChar :: Gen Char+validTokenChar = elements $ ['!'..'~'] >>= \c ->+ if c `elem` [';', ',', '=', ' '] then [] else [c]++-- | Generates a valid token text.+tokenText :: Gen Text+tokenText = T.pack <$> listOf1 validTokenChar++-- | Generates a valid cookie value text.+cookieValueText :: Gen Text+cookieValueText = T.pack <$> listOf1 validTokenChar++-- | Tests cookie extraction properties.+propCookieExtraction :: Property+propCookieExtraction =+ forAll tokenText $ \cookieName ->+ forAll cookieValueText $ \cookieValue ->+ let header = TE.encodeUtf8 (cookieName <> "=" <> cookieValue)+ extracted = extractCookieWC cookieName header+ in extracted === Just cookieValue++-- | Tests header name round-trip.+propHeaderNameRoundTrip :: Property+propHeaderNameRoundTrip = property $ \headerText ->+ let originalTxt = T.pack headerText+ headerName = hdr originalTxt+ roundTrip = TE.decodeUtf8 (fromHeaderName headerName)+ in not (T.null originalTxt) ==> roundTrip === originalTxt++-- | Tests IP extraction consistency.+propIPExtraction :: Property+propIPExtraction = property $ \ipStr ->+ let ip = T.pack ipStr+ req = mkRequestWithXFF ip+ extracted = getClientIPPure req+ expected = T.takeWhile (/= ',') ip+ in not (T.null ip) ==> extracted === expected++-- | Tests rate limiting monotonicity.+propRateLimitingMonotonicity :: Property+propRateLimitingMonotonicity = property $ \limit period ->+ limit > 0 && period > 0 ==> monadicIO $ do+ env <- run $ initConfig (const defaultIPZone)+ let throttle = ThrottleConfig limit period FixedWindow (mkIdentifier 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/Main.hs view
@@ -57,14 +57,13 @@ import Control.Concurrent.STM (atomically, readTVar) import Control.Concurrent (threadDelay)-import Control.Exception (try, SomeException) import Control.Monad (when) import Data.Maybe (isJust) import Data.IORef (readIORef) import Data.Cache (purgeExpired)-import Test.Tasty (TestTree, defaultMain, testGroup)+import Test.Tasty (TestTree, defaultMain, testGroup, after, DependencyType(..)) import Test.Tasty.HUnit (testCase, assertEqual, assertFailure)-import qualified Data.Map.Strict as Map+import qualified Data.HashMap.Strict as HashMap import Data.Text (Text) import qualified Data.Text as Text import qualified Data.Text.Encoding as TE@@ -181,11 +180,7 @@ -- | The main entry point for running the test suite. -- It aggregates all tests and runs them using Tasty. main :: IO ()-main = do- result <- try @SomeException $ defaultMain mainTests- case result of- Left ex -> putStrLn $ "Test suite failed with exception: " ++ show ex- Right () -> putStrLn "Test suite completed successfully"+main = defaultMain $ after AllSucceed "All tests passed" mainTests -- | The root 'TestTree' that combines all test groups from the library. mainTests :: TestTree@@ -418,12 +413,12 @@ threadDelay 1_500_000 -- 1.5 seconds to ensure expiration -- Debug cache state cachesMap <- readIORef (envZoneCachesMap envWithThrottle)- zoneCaches <- case Map.lookup testIPZoneA cachesMap of+ zoneCaches <- case HashMap.lookup testIPZoneA cachesMap of Just caches -> return caches Nothing -> assertFailure "Zone caches not found" >> return undefined let cache :: Cache (InMemoryStore 'FixedWindow) cache = zscCounterCache zoneCaches- key = makeCacheKey FixedWindow testIPZoneA (unsafePerformIO $ RequestUtils.getClientIP req)+ key = makeCacheKey "reset-throttle" FixedWindow testIPZoneA (unsafePerformIO $ RequestUtils.getClientIP req) mVal <- readCache cache key :: IO (Maybe Int) when (isJust mVal) $ putStrLn $ "Cache value before third request: " ++ show mVal -- Try manually deleting if it exists@@ -468,7 +463,7 @@ b3 <- instrument envWithThrottle reqB1 assertEqual "Zone B first request allowed" False b3 --- | Verifies that IPs not matching a specific zone fall back to the default zone and are tracked together.+-- | Verifies that IPs not matching a specific zone fall back to the default zone and are tracked independently by their IP. testIPZoneDefaultFallback :: IO () testIPZoneDefaultFallback = do env <- initConfig mainTestGetRequestIPZone@@ -485,7 +480,7 @@ b1 <- instrument envWithThrottle reqDefault assertEqual "Default zone first request allowed" False b1 b2 <- instrument envWithThrottle reqGeneric- assertEqual "Generic IP default zone first request allowed" False b2+ assertEqual "Generic IP in default zone should also be allowed (independent counter)" False b2 -- | Verifies that 'cacheResetAll' clears the state for all defined IP zones. testIPZoneCacheResetAll :: IO ()@@ -524,10 +519,11 @@ store <- createInMemoryStore @'FixedWindow let cache :: Cache (InMemoryStore 'FixedWindow) cache = newCache FixedWindow store+ throttleName = "test-throttle" ipZone = "zoneX" userKey = "userX"- v1 <- incStoreWithZone cache ipZone userKey 10 :: IO Int- v2 <- incStoreWithZone cache ipZone userKey 10 :: IO Int+ v1 <- incStoreWithZone cache throttleName ipZone userKey 10 :: IO Int+ v2 <- incStoreWithZone cache throttleName ipZone userKey 10 :: IO Int assertEqual "First increment (wrapper)" 1 v1 assertEqual "Second increment (wrapper)" 2 v2 @@ -549,11 +545,12 @@ store <- createInMemoryStore @'FixedWindow let cache :: Cache (InMemoryStore 'FixedWindow) cache = newCache FixedWindow store+ throttleName = "test-throttle" ipZone = "zoneZ" userKey = "userZ" val = 42 :: Int- writeCacheWithZone cache ipZone userKey val 10- mVal <- readCacheWithZone cache ipZone userKey :: IO (Maybe Int)+ writeCacheWithZone cache throttleName ipZone userKey val 10+ mVal <- readCacheWithZone cache throttleName ipZone userKey :: IO (Maybe Int) assertEqual "Read after write (wrapper)" (Just val) mVal -- | Tests the lower-level 'writeCache' and 'readCache' functions with a manually constructed key.
test/TinyLRUTests.hs view
@@ -47,7 +47,6 @@ , deleteCacheWithZone ) import Control.Monad-import Control.Concurrent (threadDelay) import qualified StmContainers.Map as Map import System.Random (randomRIO) import Data.Maybe (isJust)@@ -267,19 +266,22 @@ assertEqual "Tail is Nothing after reset" "Nothing" tailStr atomically $ checkListConsistency cache - , testCase "Integration with Cache" $ do+ , testCase "Integration with Cache: Zone-Based Operations" $ do cache <- newTinyLRUCache- now <- getTime Monotonic- allowed1 <- allowRequest cache "key1" 2 5- assertBool "First request allowed" allowed1- allowed2 <- allowRequest cache "key1" 2 5- assertBool "Second request allowed" allowed2- allowed3 <- allowRequest cache "key1" 2 5- assertBool "Third request denied" (not allowed3)- threadDelay 6000000 -- Wait 6 seconds (> period of 5)- now' <- getTime Monotonic- allowed4 <- allowRequest cache "key1" 2 5- assertBool "Request allowed after expiration" allowed4+ -- writeCacheWithZone :: Cache store -> Text -> Text -> Text -> v -> Int -> IO ()+ -- cache throttle ipZone userKey value expiresIn+ writeCacheWithZone cache "throttle1" "192.168.1.1" "user1" (42 :: Int) 60+ + -- readCacheWithZone :: Cache store -> Text -> Text -> Text -> IO (Maybe v)+ -- cache throttle ipZone userKey+ result1 <- readCacheWithZone cache "throttle1" "192.168.1.1" "user1"+ assertEqual "Zone-based read returns Just 42" (Just 42) result1+ + -- deleteCacheWithZone :: Cache store -> Text -> Text -> Text -> IO ()+ -- cache throttle ipZone userKey+ deleteCacheWithZone cache "throttle1" "192.168.1.1" "user1"+ result2 <- readCacheWithZone cache "throttle1" "192.168.1.1" "user1"+ assertEqual "Zone-based delete removes key" Nothing result2 atomically $ checkWrappedListConsistency cache , testCase "Concurrent Access" $ do@@ -371,11 +373,11 @@ , testCase "Integration with Cache: Zone-Based Operations" $ do cache <- newTinyLRUCache- writeCacheWithZone cache "192.168.1.1" "user1" (42 :: Int) 60- result1 <- readCacheWithZone cache "192.168.1.1" "user1"+ writeCacheWithZone cache "throttle1" "192.168.1.1" "user1" (42 :: Int) 60+ result1 <- readCacheWithZone cache "throttle1" "192.168.1.1" "user1" assertEqual "Zone-based read returns Just 42" (Just 42) result1- deleteCacheWithZone cache "192.168.1.1" "user1"- result2 <- readCacheWithZone cache "192.168.1.1" "user1"+ deleteCacheWithZone cache "throttle1" "192.168.1.1" "user1"+ result2 <- readCacheWithZone cache "throttle1" "192.168.1.1" "user1" assertEqual "Zone-based delete removes key" Nothing result2 atomically $ checkWrappedListConsistency cache