wai-middleware-throttle 0.2.2.1 → 0.3.0.0
raw patch · 4 files changed
+264/−258 lines, 4 filesdep +cachedep +clockdep +safe-exceptionsdep ~mtl
Dependencies added: cache, clock, safe-exceptions
Dependency ranges changed: mtl
Files
- Network/Wai/Middleware/Throttle.hs +18/−203
- Network/Wai/Middleware/Throttle/Internal.hs +190/−0
- test/WaiMiddlewareThrottleSpec.hs +47/−54
- wai-middleware-throttle.cabal +9/−1
Network/Wai/Middleware/Throttle.hs view
@@ -2,9 +2,9 @@ -- | -- Module : Network.Wai.Middleware.Throttle -- Description : WAI Request Throttling Middleware--- Copyright : (c) 2015-2017 Christopher Reichert+-- Copyright : (c) 2015-2018 Christopher Reichert, 2017-2018 Daniel Fithian -- License : BSD3--- Maintainer : Christopher Reichert <creichert07@gmail.com>+-- Maintainer : Christopher Reichert <creichert07@gmail.com>, Daniel Fithian <daniel.m.fithian@gmail.com> -- Stability : experimental -- Portability : POSIX --@@ -16,207 +16,22 @@ -- -- @ -- main = do--- st <- initThrottler--- let payload = "{ \"api\": \"return data\" }"--- app = throttle defaultThrottleSettings st--- $ \_ f -> f (responseLBS status200 [] payload)+-- let expirationSpec = TimeSpec 5 0 -- five seconds+-- th <- initThrottler (defaultThrottleSettings expirationSpec) (const $ Right 1)+-- let appl = throttle th $ \ _ f -> f $+-- responseLBS ok200 [] "ok" -- Warp.run 3000 app -- @--module Network.Wai.Middleware.Throttle (-- -- | Wai Request Throttling Middleware- throttle-- -- | Wai Throttle middleware state.- --- -- Essentially, a TVar with a HashMap for indexing- -- remote IP address- , WaiThrottle, CustomWaiThrottle- , initThrottler, initCustomThrottler-- -- | Throttle settings and configuration- , ThrottleSettings(..)- , defaultThrottleSettings-- , Address(..)- , RequestHashable(..)- ) where--import Control.Applicative ((<$>), pure)-import Control.Concurrent.STM-import Control.Concurrent.TokenBucket-import Control.Monad (join, liftM)-import Control.Monad.IO.Class (liftIO)-import Control.Monad.Except (ExceptT, runExceptT)-import Data.ByteString.Builder (stringUtf8, toLazyByteString)-import Data.Function (on)-import Data.Hashable (Hashable, hash, hashWithSalt)-import qualified Data.IntMap as IM-import Data.List (unionBy)-import Data.Monoid ((<>))-import Data.Text (Text, unpack)-import GHC.Word (Word64)-import qualified Network.HTTP.Types.Status as Http-import Network.Socket-import Network.Wai--#ifndef MIN_VERSION_network-#define MIN_VERSION_network(a,b,c) 1-#endif---- | A throttle for a type implementing 'RequestHashable'-newtype CustomWaiThrottle a = WT (TVar (ThrottleState a))---- | A synonym for an IP address throttle-type WaiThrottle = CustomWaiThrottle Address---newtype Address = Address SockAddr--instance Hashable Address where- hashWithSalt s (Address (SockAddrInet _ a)) = hashWithSalt s a- hashWithSalt s (Address (SockAddrInet6 _ _ a _)) = hashWithSalt s a- hashWithSalt s (Address (SockAddrUnix a)) = hashWithSalt s a-#if MIN_VERSION_network(2,6,1)- hashWithSalt s (Address (SockAddrCan a)) = hashWithSalt s a-#endif--instance Eq Address where- Address (SockAddrInet _ a) == Address (SockAddrInet _ b) = a == b- Address (SockAddrInet6 _ _ a _) == Address (SockAddrInet6 _ _ b _) = a == b- Address (SockAddrUnix a) == Address (SockAddrUnix b) = a == b-#if MIN_VERSION_network(2,6,1)- Address (SockAddrCan a) == Address (SockAddrCan b) = a == b-#endif- _ == _ = False -- not same constructor so cant be equal--instance Ord Address where- Address (SockAddrInet _ a) <= Address (SockAddrInet _ b) = a <= b- Address (SockAddrInet6 _ _ a _) <= Address (SockAddrInet6 _ _ b _) = a <= b- Address (SockAddrUnix a) <= Address (SockAddrUnix b) = a <= b-#if MIN_VERSION_network(2,6,1)- Address (SockAddrCan a) <= Address (SockAddrCan b) = a <= b-#endif- Address a <= Address b = a <= b -- not same constructor so use builtin ordering---- | A 'HashMap' mapping the remote IP address to a 'TokenBucket'-data ThrottleState a = ThrottleState !(IM.IntMap [(a,TokenBucket)])----- | Settings which control various behaviors in the middleware.-data ThrottleSettings = ThrottleSettings- {- -- | Determines whether the 'Request' is throttled- isThrottled :: !(Request -> IO Bool)-- -- | Function to run when the request is throttled.- --- -- The first argument is a 'Word64' containing the amount- -- of microseconds until the next retry should be attempted.- , onThrottled :: !(Word64 -> Response)-- -- | Function to run when the throttler fails to extract the key.- --- -- The first argument is a 'Text' containing the error message.- , onRequestError :: !(Text -> Response)-- -- | Rate- , throttleRate :: !Integer -- requests / throttlePeriod- , throttlePeriod :: !Integer -- microseconds-- -- | Burst rate- , throttleBurst :: !Integer- }----- |Initialize an IP address throttler-initThrottler :: IO WaiThrottle-initThrottler = initCustomThrottler---- |Initialize a "custom" throttler that implements the 'RequestHashable' class-initCustomThrottler :: IO (CustomWaiThrottle a)-initCustomThrottler = liftM WT $ newTVarIO $ ThrottleState IM.empty----- | Default settings to throttle requests.-defaultThrottleSettings :: ThrottleSettings-defaultThrottleSettings- = ThrottleSettings {- isThrottled = return . const True- , throttleRate = 1 :: Integer -- req / throttlePeriod- , throttlePeriod = 1000000 :: Integer -- microseconds- , throttleBurst = 1 :: Integer -- concurrent requests- , onThrottled = onThrottled'- , onRequestError = onRequestError'- }- where- onThrottled' _ =- responseLBS- Http.status429- [ ("Content-Type", "application/json")- ]- "{\"message\":\"Too many requests.\"}"- onRequestError' reason =- responseLBS- Http.status400- [ ("Content-Type", "application/json")- ]- ("{\"message\":\"" <> toLazyByteString (stringUtf8 $ unpack reason) <> "\"}")---- | A class extracting a hashable key from a 'Request' to store in an in-memory map-class (Eq a, Ord a, Hashable a) => RequestHashable a where- requestToKey :: (Functor m, Monad m) => Request -> ExceptT Text m a--instance RequestHashable Address where- requestToKey = pure . Address . remoteHost---- | WAI Request Throttling Middleware.------ Uses a 'Request's 'remoteHost' function to resolve the--- remote IP address.-throttle :: RequestHashable a- => ThrottleSettings- -> CustomWaiThrottle a- -> Application- -> Application-throttle ThrottleSettings{..} (WT tmap) app req respond = do-- -- determine whether the request needs throttling- reqIsThrottled <- isThrottled req-- -- seconds remaining (if the request failed), 0 otherwise.- remaining <- if reqIsThrottled- then runExceptT throttleReq- else return $ Right 0-- case remaining of- Left err -> respond $ onRequestError err- Right 0 -> app req respond- Right n -> respond $ onThrottled n- where- throttleReq = do-- k <- requestToKey req- throttleState <- liftIO . atomically $ readTVar tmap- (tst, success) <- liftIO $ throttleReq' k throttleState-- -- write the throttle state back- liftIO . atomically $ writeTVar tmap (ThrottleState tst)- return success-- throttleReq' k (ThrottleState m) = do-- let toInvRate r = round (period / r)- period = (fromInteger throttlePeriod :: Double)- invRate = toInvRate (fromInteger throttleRate :: Double)- burst = fromInteger throttleBurst-- bucket <- maybe newTokenBucket return $ join $ lookup k <$> IM.lookup (hash k) m- remaining <- tokenBucketTryAlloc1 bucket burst invRate-- return (insertBucket k bucket m, remaining)+module Network.Wai.Middleware.Throttle+ ( module Network.Wai.Middleware.Throttle.Internal+ ) where - insertBucket k bucket m =- let col = unionBy ((==) `on` fst)- in IM.insertWith col (hash k) [(k, bucket)] m+import Network.Wai.Middleware.Throttle.Internal+ ( Address (..)+ , Throttle (..)+ , ThrottleSettings (..)+ , defaultThrottleSettings+ , extractAddress+ , initCustomThrottler+ , initThrottler+ , throttle )
+ Network/Wai/Middleware/Throttle/Internal.hs view
@@ -0,0 +1,190 @@+module Network.Wai.Middleware.Throttle.Internal where++import Prelude hiding (lookup)++import Control.Concurrent.TokenBucket (TokenBucket, newTokenBucket, tokenBucketTryAlloc1)+import Control.Exception.Safe (onException)+#if MIN_VERSION_cache(0,1,1)+import Control.Monad.STM (STM, atomically)+import Data.Cache (Cache, delete, insert, insertSTM, lookupSTM, newCache)+#else+import Data.Cache (Cache, delete, insert, insert', lookup, newCache)+#endif+import Data.Hashable (Hashable, hashWithSalt)+import GHC.Word (Word64)+import Network.HTTP.Types.Status (status429)+import Network.Socket (SockAddr (SockAddrInet, SockAddrInet6, SockAddrUnix))+#if MIN_VERSION_network(2,6,1)+import Network.Socket (SockAddr (SockAddrCan))+#endif+import Network.Wai (Application, Request, Response, remoteHost, responseLBS)+import System.Clock (Clock (Monotonic), TimeSpec, getTime)++newtype Address = Address SockAddr++instance Hashable Address where+ hashWithSalt s (Address (SockAddrInet _ a)) = hashWithSalt s a+ hashWithSalt s (Address (SockAddrInet6 _ _ a _)) = hashWithSalt s a+ hashWithSalt s (Address (SockAddrUnix a)) = hashWithSalt s a+#if MIN_VERSION_network(2,6,1)+ hashWithSalt s (Address (SockAddrCan a)) = hashWithSalt s a+#endif++instance Eq Address where+ Address (SockAddrInet _ a) == Address (SockAddrInet _ b) = a == b+ Address (SockAddrInet6 _ _ a _) == Address (SockAddrInet6 _ _ b _) = a == b+ Address (SockAddrUnix a) == Address (SockAddrUnix b) = a == b+#if MIN_VERSION_network(2,6,1)+ Address (SockAddrCan a) == Address (SockAddrCan b) = a == b+#endif+ _ == _ = False -- not same constructor so cant be equal++instance Ord Address where+ Address (SockAddrInet _ a) <= Address (SockAddrInet _ b) = a <= b+ Address (SockAddrInet6 _ _ a _) <= Address (SockAddrInet6 _ _ b _) = a <= b+ Address (SockAddrUnix a) <= Address (SockAddrUnix b) = a <= b+#if MIN_VERSION_network(2,6,1)+ Address (SockAddrCan a) <= Address (SockAddrCan b) = a <= b+#endif+ Address a <= Address b = a <= b -- not same constructor so use builtin ordering++extractAddress :: Request -> Either Response Address+extractAddress = Right . Address . remoteHost++data CacheState a+ = CacheStatePresent a+ | CacheStateInitializing++data CacheResult a+ = CacheResultExists a+ | CacheResultEmpty++-- |A throttle for a hashable key type. Initialize using 'initThrottler' with 'defaultThrottleSettings'.+data Throttle a = Throttle+ { throttleSettings :: ThrottleSettings+ -- ^ The throttle settings+ , throttleCache :: Cache a (CacheState TokenBucket)+ -- ^ The cache, initialized in 'initThrottler'+ , throttleGetKey :: Request -> Either Response a+ -- ^ The function to extract a throttle key from a 'Network.Wai.Request'+ }++-- |Throttle settings for controlling token bucket algorithm and cache expiration.+data ThrottleSettings = ThrottleSettings+ { throttleSettingsRate :: Double+ -- ^ Number of requests per throttle period allowed, defaults to 1+ , throttleSettingsPeriod :: Double+ -- ^ Microseconds, defaults to 1 second+ , throttleSettingsBurst :: Word64+ -- ^ Number of concurrent requests allowed - should be greater than rate / period, defaults to 1+ , throttleSettingsCacheExpiration :: TimeSpec+ -- ^ The amount of time before a stale token bucket is purged from the cache+ , throttleSettingsIsThrottled :: Request -> Bool+ -- ^ Whether or not a request is throttled, defaults to true+ , throttleSettingsOnThrottled :: Word64 -> Response+ -- ^ The response when a request is throttled - defaults to a vanilla 429+ }++-- |Default throttle settings.+defaultThrottleSettings :: TimeSpec -> ThrottleSettings+defaultThrottleSettings expirationInterval = ThrottleSettings+ { throttleSettingsRate = 1+ , throttleSettingsPeriod = 1000000+ , throttleSettingsBurst = 1+ , throttleSettingsCacheExpiration = expirationInterval+ , throttleSettingsIsThrottled = const True+ , throttleSettingsOnThrottled = const $+ responseLBS status429 [("Content-Type", "application/json")] "{\"message\":\"Too many requests.\"}"+ }++initThrottler :: ThrottleSettings -> IO (Throttle Address)+initThrottler = flip initCustomThrottler extractAddress++-- |Initialize a throttle using settings and a way to extract the key from the request.+initCustomThrottler :: ThrottleSettings -> (Request -> Either Response a) -> IO (Throttle a)+initCustomThrottler throttleSettings@(ThrottleSettings {..}) throttleGetKey = do+ throttleCache <- newCache $ Just throttleSettingsCacheExpiration+ pure Throttle {..}++-- |Internal use only. Retrieve a token bucket from the cache.+#if MIN_VERSION_cache(0,1,1)+retrieveCache :: (Eq a, Hashable a) => Throttle a -> TimeSpec -> a -> STM (CacheResult TokenBucket)+retrieveCache th time throttleKey = do+ let cache = throttleCache th+ lookupSTM True throttleKey cache time >>= \ case+ Just (CacheStatePresent oldBucket) -> pure $ CacheResultExists oldBucket+ Just CacheStateInitializing -> retrieveCache th time throttleKey+ Nothing -> do+ insertSTM throttleKey CacheStateInitializing cache Nothing+ pure CacheResultEmpty+#else+retrieveCache :: (Eq a, Hashable a) => Throttle a -> TimeSpec -> a -> IO (CacheResult TokenBucket)+retrieveCache th time throttleKey = do+ let cache = throttleCache th+ lookup cache throttleKey >>= \ case+ Just (CacheStatePresent oldBucket) -> pure $ CacheResultExists oldBucket+ Just CacheStateInitializing -> retrieveCache th time throttleKey+ Nothing -> do+ insert' cache Nothing throttleKey CacheStateInitializing+ pure CacheResultEmpty+#endif++-- |Internal use only. Create a token bucket if it wasn't in the cache.+processCacheResult :: (Eq a, Hashable a) => Throttle a -> a -> CacheResult TokenBucket -> IO TokenBucket+processCacheResult th throttleKey cacheResult = case cacheResult of+ CacheResultExists bucket -> pure bucket+ CacheResultEmpty -> do+ let cache = throttleCache th+ initializeBucket = do+ bucket <- newTokenBucket+ insert cache throttleKey (CacheStatePresent bucket)+ pure bucket+ cleanupBucket = delete cache throttleKey+ initializeBucket `onException` cleanupBucket++-- |Internal use only. Retrieve or initialize a token bucket depending on if it was found in the cache.+retrieveOrInitializeBucket :: (Eq a, Hashable a) => Throttle a -> a -> IO TokenBucket+retrieveOrInitializeBucket th throttleKey = do+ now <- getTime Monotonic+#if MIN_VERSION_cache(0,1,1)+ cacheResult <- atomically $ retrieveCache th now throttleKey+#else+ cacheResult <- retrieveCache th now throttleKey+#endif+ processCacheResult th throttleKey cacheResult++-- |Internal use only. Throttle a request by the throttle key.+throttleRequest :: (Eq a, Hashable a) => Throttle a -> a -> IO Word64+throttleRequest th throttleKey = do+ bucket <- retrieveOrInitializeBucket th throttleKey+ let settings = throttleSettings th+ rate = throttleSettingsRate settings+ period = throttleSettingsPeriod settings+ burst = throttleSettingsBurst settings+ tokenBucketTryAlloc1 bucket burst $ round (period / rate)++-- |Run the throttling middleware given a throttle that has been initialized.+throttle :: (Eq a, Hashable a) => Throttle a -> Application -> Application+throttle th app req respond = do+ let settings = throttleSettings th+ getKey = throttleGetKey th+ isThrottled = throttleSettingsIsThrottled settings+ onThrottled = throttleSettingsOnThrottled settings+ case isThrottled req of+ False -> app req respond+ True -> case getKey req of+ Left response -> respond response+ Right throttleKey -> do+ throttleRequest th throttleKey >>= \ case+ 0 -> app req respond+ n -> respond $ onThrottled n++instance Show (CacheState a) where+ show = \ case+ CacheStatePresent _ -> "Present"+ CacheStateInitializing -> "Initializing"++instance Show (CacheResult a) where+ show = \ case+ CacheResultExists _ -> "Exists"+ CacheResultEmpty -> "Empty"
test/WaiMiddlewareThrottleSpec.hs view
@@ -1,64 +1,57 @@---------------------------------------------------------------------------- |--- Module : WaiMiddlewareThrottleSpec--- Description : WAI Request Throttling Middleware--- Copyright : (c) 2015 Christopher Reichert--- License : BSD3--- Maintainer : Christopher Reichert <creichert07@gmail.com>--- Stability : unstable--- Portability : POSIX----{-# LANGUAGE OverloadedStrings #-}--module WaiMiddlewareThrottleSpec (- spec- ) where+module WaiMiddlewareThrottleSpec where -import Control.Monad.IO.Class-import Network.HTTP.Types-import Network.HTTP.Types.Status-import Network.Wai-import Network.Wai.Test-import Test.Hspec-import Test.HUnit hiding (Test)+import Prelude hiding (lookup) -import Network.Wai.Middleware.Throttle+import Control.Concurrent (threadDelay)+import Control.Monad (replicateM, void)+import Data.Cache (lookup)+import Data.Maybe (isJust, isNothing)+import Network.HTTP.Types.Status (status200, status429)+import Network.Wai (defaultRequest, responseLBS)+import Network.Wai.Test (request, runSession, simpleStatus)+import System.Clock (TimeSpec (TimeSpec))+import Test.Hspec (Spec, before, describe, it, shouldSatisfy) +-- the modules being tested+import Network.Wai.Middleware.Throttle+import Network.Wai.Middleware.Throttle.Internal spec :: Spec-spec = describe "Network.Wai.Middleware.Throttle" $- it "throttles requests" caseThrottle----- | Simple Hmac Middleware App------ This app has preloaded api keys to simulate--- some database or service which can access the--- private keys.-throttleApp :: WaiThrottle -> Application-throttleApp st = throttle defaultThrottleSettings st- $ \_ f -> f response- where- payload = "{ \"api\", \"return data\" }"- response = responseLBS status200 [] payload---defReq :: Request-defReq = defaultRequest- { requestMethod = "GET"- , requestHeaders = [ ("Content-Type", "application/json") ]- }+spec = do+ let expirationSpec = TimeSpec 5 0+ runBefore :: IO (Throttle Int)+ runBefore = initCustomThrottler (defaultThrottleSettings expirationSpec) (const $ Right 1)+ before runBefore $+ describe "Network.Wai.Middleware.Throttle" $ do+ describe "Bucket Operations" $ do + it "initializes bucket when missing" $ \ th -> do+ let throttleKey = 1+ cache = throttleCache th+ void $ retrieveOrInitializeBucket th throttleKey+ lookup cache throttleKey >>= \ b -> b `shouldSatisfy` isJust --- | Test Hmac Authentication-caseThrottle :: Assertion-caseThrottle = do+ it "retrieves bucket on subsequent calls" $ \ th -> do+ let throttleKey = 1+ cache = throttleCache th+ void $ retrieveOrInitializeBucket th throttleKey+ void $ retrieveOrInitializeBucket th throttleKey+ lookup cache throttleKey >>= \ b -> b `shouldSatisfy` isJust - st <- liftIO initThrottler+ it "expires buckets" $ \ th -> do+ let throttleKey = 1+ cache = throttleCache th+ void $ retrieveOrInitializeBucket th throttleKey+ lookup cache throttleKey >>= \ b -> b `shouldSatisfy` isJust+ threadDelay 5000000+ lookup cache throttleKey >>= \ b -> b `shouldSatisfy` isNothing - statuses <- flip runSession (throttleApp st) $ do- responses <- mapM (const (request defReq)) [ 1 .. 100 :: Integer ]- mapM (return . simpleStatus) responses+ describe "Throttling Behavior" $ do - let msg = "Verifying some of the requests were throttled"- assertBool msg $ elem status429 statuses+ it "throttles requests" $ \ th -> do+ let appl = throttle th $ \ _ f -> f $+ responseLBS status200 [] "ok"+ statuses <- flip runSession appl $ do+ responses <- replicateM 100 (request defaultRequest)+ pure $ simpleStatus <$> responses+ statuses `shouldSatisfy` elem status429
wai-middleware-throttle.cabal view
@@ -1,5 +1,5 @@ name: wai-middleware-throttle-version: 0.2.2.1+version: 0.3.0.0 license: BSD3 license-file: LICENSE author: Christopher Reichert@@ -22,16 +22,20 @@ library exposed-modules: Network.Wai.Middleware.Throttle+ Network.Wai.Middleware.Throttle.Internal ghc-options: -Wall -fno-warn-unused-do-bind default-language: Haskell2010 build-depends: base >= 4.6 && < 5.0 , bytestring , bytestring-builder+ , cache+ , clock , containers , hashable >= 1.2 , http-types , mtl , network >= 2.1+ , safe-exceptions , stm , text , token-bucket >= 0.1.0.1@@ -40,6 +44,7 @@ default-extensions: CPP FlexibleContexts+ LambdaCase OverloadedStrings RecordWildCards @@ -56,6 +61,8 @@ default-language: Haskell2010 build-depends: base >= 4.6 && < 5.0 , bytestring+ , cache+ , clock , hspec >= 1.3 , http-types , HUnit@@ -65,3 +72,4 @@ , wai , wai-extra , wai-middleware-throttle+ default-extensions: OverloadedStrings