wai-middleware-throttle 0.2.1.0 → 0.2.2.0

raw patch · 2 files changed

+83/−47 lines, 2 filesdep +QuickCheckdep +bytestring-builderdep +mtldep ~token-bucketPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependencies added: QuickCheck, bytestring-builder, mtl, text

Dependency ranges changed: token-bucket

API changes (from Hackage documentation)

Files

• Network/Wai/Middleware/Throttle.hs +64/−33
• wai-middleware-throttle.cabal +19/−14

@@ -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 <creichert07@gmail.com> -- 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

@@ -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