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,7 +2,7 @@ -- | -- Module : Network.Wai.Middleware.Throttle -- Description : WAI Request Throttling Middleware --- Copyright : (c) 2015 Christopher Reichert +-- Copyright : (c) 2015-2017 Christopher Reichert -- License : BSD3 -- Maintainer : Christopher Reichert -- Stability : experimental @@ -23,10 +23,6 @@ -- Warp.run 3000 app -- @ -{-# LANGUAGE CPP #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE RecordWildCards #-} - module Network.Wai.Middleware.Throttle ( -- | Wai Request Throttling Middleware @@ -36,22 +32,29 @@ -- -- Essentially, a TVar with a HashMap for indexing -- remote IP address - , WaiThrottle - , initThrottler + , WaiThrottle, CustomWaiThrottle + , initThrottler, initCustomThrottler -- | Throttle settings and configuration , ThrottleSettings(..) , defaultThrottleSettings + + , RequestHashable(..) ) where -import Control.Applicative ((<$>)) +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 @@ -61,9 +64,13 @@ #define MIN_VERSION_network(a,b,c) 1 #endif -newtype WaiThrottle = WT (TVar ThrottleState) +-- | 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 @@ -93,7 +100,7 @@ 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 = ThrottleState !(IM.IntMap [(Address,TokenBucket)]) +data ThrottleState a = ThrottleState !(IM.IntMap [(a,TokenBucket)]) -- | Settings which control various behaviors in the middleware. @@ -105,9 +112,14 @@ -- | 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 + -- 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 @@ -117,19 +129,25 @@ } +-- |Initialize an IP address throttler initThrottler :: IO WaiThrottle -initThrottler = liftM WT $ newTVarIO $ ThrottleState IM.empty +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 -- req / throttlePeriod - , throttlePeriod = 10^6 -- microseconds - , throttleBurst = 1 -- concurrent requests + , throttleRate = 1 :: Integer -- req / throttlePeriod + , throttlePeriod = 1000000 :: Integer -- microseconds + , throttleBurst = 1 :: Integer -- concurrent requests , onThrottled = onThrottled' + , onRequestError = onRequestError' } where onThrottled' _ = @@ -138,14 +156,27 @@ [ ("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 :: ThrottleSettings - -> WaiThrottle +throttle :: RequestHashable a + => ThrottleSettings + -> CustomWaiThrottle a -> Application -> Application throttle ThrottleSettings{..} (WT tmap) app req respond = do @@ -155,36 +186,36 @@ -- seconds remaining (if the request failed), 0 otherwise. remaining <- if reqIsThrottled - then throttleReq - else return 0 + then runExceptT throttleReq + else return $ Right 0 - if remaining /= 0 - then respond $ onThrottled remaining - else app req respond + case remaining of + Left err -> respond $ onRequestError err + Right 0 -> app req respond + Right n -> respond $ onThrottled n where throttleReq = do - let remoteAddr = Address . remoteHost $ req - throttleState <- atomically $ readTVar tmap - (tst, success) <- throttleReq' remoteAddr throttleState + k <- requestToKey req + throttleState <- liftIO . atomically $ readTVar tmap + (tst, success) <- liftIO $ throttleReq' k throttleState -- write the throttle state back - atomically $ writeTVar tmap (ThrottleState tst) + liftIO . atomically $ writeTVar tmap (ThrottleState tst) return success - throttleReq' remoteAddr (ThrottleState m) = do + throttleReq' k (ThrottleState m) = do let toInvRate r = round (period / r) - period = (fromInteger throttlePeriod :: Double) + period = (fromInteger throttlePeriod :: Double) invRate = toInvRate (fromInteger throttleRate :: Double) burst = fromInteger throttleBurst - bucket <- maybe newTokenBucket return $ addressToBucket remoteAddr m + bucket <- maybe newTokenBucket return $ join $ lookup k <$> IM.lookup (hash k) m remaining <- tokenBucketTryAlloc1 bucket burst invRate - return (insertBucket remoteAddr bucket m, remaining) + return (insertBucket k bucket m, remaining) - addressToBucket remoteAddr m = join (lookup remoteAddr <$> IM.lookup (hash remoteAddr) m) - insertBucket remoteAddr bucket m = + insertBucket k bucket m = let col = unionBy ((==) `on` fst) - in IM.insertWith col (hash remoteAddr) [(remoteAddr, bucket)] m + in IM.insertWith col (hash k) [(k, bucket)] m diff --git a/wai-middleware-throttle.cabal b/wai-middleware-throttle.cabal --- a/wai-middleware-throttle.cabal +++ b/wai-middleware-throttle.cabal @@ -1,11 +1,10 @@ - name: wai-middleware-throttle -version: 0.2.1.0 +version: 0.2.2.0 license: BSD3 license-file: LICENSE author: Christopher Reichert maintainer: creichert07@gmail.com -copyright: (c) 2015, Christopher Reichert +copyright: (c) 2015-2017, Christopher Reichert category: Web build-type: Simple cabal-version: >=1.10 @@ -26,18 +25,23 @@ ghc-options: -Wall -fno-warn-unused-do-bind default-language: Haskell2010 build-depends: base >= 4.6 && < 5.0 - , wai >= 3.0 + , bytestring + , bytestring-builder , containers - , transformers - , network >= 2.1 + , hashable >= 1.2 , http-types + , mtl + , network >= 2.1 , stm - , token-bucket - , hashable >= 1.2 + , text + , token-bucket >= 0.1.0.1 + , transformers + , wai >= 3.0 - other-extensions: CPP - OverloadedStrings - RecordWildCards + default-extensions: CPP + FlexibleContexts + OverloadedStrings + RecordWildCards if impl(ghc < 7.8) build-depends: @@ -52,11 +56,12 @@ default-language: Haskell2010 build-depends: base >= 4.6 && < 5.0 , bytestring - , wai-middleware-throttle - , wai - , wai-extra , hspec >= 1.3 , http-types , HUnit + , QuickCheck , stm , transformers + , wai + , wai-extra + , wai-middleware-throttle