diff --git a/Network/Wai/Middleware/Throttle.hs b/Network/Wai/Middleware/Throttle.hs --- a/Network/Wai/Middleware/Throttle.hs +++ b/Network/Wai/Middleware/Throttle.hs @@ -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 +-- Maintainer : Christopher Reichert , Daniel Fithian -- 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 ) diff --git a/Network/Wai/Middleware/Throttle/Internal.hs b/Network/Wai/Middleware/Throttle/Internal.hs new file mode 100644 --- /dev/null +++ b/Network/Wai/Middleware/Throttle/Internal.hs @@ -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" diff --git a/test/WaiMiddlewareThrottleSpec.hs b/test/WaiMiddlewareThrottleSpec.hs --- a/test/WaiMiddlewareThrottleSpec.hs +++ b/test/WaiMiddlewareThrottleSpec.hs @@ -1,64 +1,57 @@ ------------------------------------------------------------------------- --- | --- Module : WaiMiddlewareThrottleSpec --- Description : WAI Request Throttling Middleware --- Copyright : (c) 2015 Christopher Reichert --- License : BSD3 --- Maintainer : Christopher Reichert --- 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 diff --git a/wai-middleware-throttle.cabal b/wai-middleware-throttle.cabal --- a/wai-middleware-throttle.cabal +++ b/wai-middleware-throttle.cabal @@ -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