wai-middleware-throttle (empty) → 0.1.0.0

raw patch · 9 files changed

+459/−0 lines, 9 filesdep +Globdep +HUnitdep +basesetup-changed

Dependencies added: Glob, HUnit, base, blaze-builder, bytestring, doctest, hlint, hspec, http-types, network, process, regex-compat, stm, text, token-bucket, transformers, unordered-containers, wai, wai-extra, wai-middleware-throttle

Files

• LICENSE +30/−0
• Network/Wai/Middleware/Throttle.hs +172/−0
• Setup.hs +2/−0
• test/DocTest.hs +21/−0
• test/HLint.hs +31/−0
• test/Haddock.hs +49/−0
• test/Spec.hs +1/−0
• test/WaiMiddlewareThrottleSpec.hs +64/−0
• wai-middleware-throttle.cabal +89/−0

@@ -0,0 +1,30 @@+Copyright (c) 2015, Christopher Reichert++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++    * Redistributions of source code must retain the above copyright+      notice, this list of conditions and the following disclaimer.++    * Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++    * Neither the name of Christopher Reichert nor the names of other+      contributors may be used to endorse or promote products derived+      from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

@@ -0,0 +1,172 @@+-----------------------------------------------------------------------------+-- |+-- Module      : Network.Wai.Middleware.Throttle+-- Description : WAI Request Throttling Middleware+-- Copyright   : (c) 2015 Christopher Reichert+-- License     : BSD3+-- Maintainer  : Christopher Reichert <creichert07@gmail.com>+-- Stability   : experimental+-- Portability : POSIX+--+-- Uses a <https://en.wikipedia.org/wiki/Token_bucket Token Bucket>+-- algorithm (from the token-bucket package) to throttle WAI Requests.+--+--+-- == Example+--+-- @+-- main = do+--   st <- initThrottler+--   let payload  = "{ \"api\", \"return data\" }"+--       app = throttle defaultThrottleSettings st+--               $ \_ f -> f (responseLBS status200 [] payload)+--   Warp.run 3000 app+-- @++{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RecordWildCards   #-}++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+    , initThrottler++      -- | Throttle settings and configuration+    , ThrottleSettings(..)+    , defaultThrottleSettings+    ) where++import           Control.Applicative+import           Control.Concurrent.STM+import           Control.Concurrent.TokenBucket+import           Control.Monad                  (liftM)+import qualified Data.ByteString.Char8          as BS+import qualified Data.HashMap.Strict            as H+import qualified Data.Text                      as T+import           GHC.Word                       (Word64)+import qualified Network.HTTP.Types.Status      as Http+import           Network.Socket+import           Network.Wai+++newtype WaiThrottle = WT (TVar ThrottleState)+++-- | A 'HashMap' mapping the remote IP address to a 'TokenBucket'+data ThrottleState = ThrottleState !(H.HashMap T.Text 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)++      -- Zone name+      --+      -- TODO use list of zones (rules)+      -- , throttleZone      :: !T.Text++      -- | Rate+    , throttleRate  :: !Integer  -- requests / second++      -- Maximum size of the address cache in MB (similar to nginx)+      --+      -- You can store approximately 160,000 addresses in 10MB with+      -- \$binary_remote_addr.+      -- , throttleCacheSize :: !Integer++      -- | Burst rate+    , throttleBurst :: !Integer+    }+++initThrottler :: IO WaiThrottle+initThrottler = liftM WT $ newTVarIO $ ThrottleState H.empty+++-- | Default settings to throttle requests.+defaultThrottleSettings :: ThrottleSettings+defaultThrottleSettings+    = ThrottleSettings {+        isThrottled         = return . const True+        -- , throttleZone        = "" -- empty zone+      , throttleRate        = 1  -- req / sec+        -- , throttleCacheSize   = 10 -- 10M address cache+      , throttleBurst       = 1  -- 5 concurrent requests+      , onThrottled         = onThrottled'+      }+  where+    bshow = BS.pack . show+    -- remaining = bshow (if 5000 - c < 0+    --                      then 0+    --                      else 5000 - c)+    onThrottled' rt =+      responseLBS+        Http.status429+        [ ("Content-Type", "application/json")+          -- , ("X-RateLimit-Limit", "5000")+          -- , ("X-RateLimit-Remaining", remaining)+        , ("X-RateLimit-Reset",+             bshow (fromIntegral rt / 1000000.0 :: Double))+        ]+        -- match YesodAuth error message renderer+        "{\"message\":\"Too many requests.\"}"+++-- | WAI Request Throttling Middleware.+--+-- Uses a 'Request's 'remoteHost' function to resolve the+-- remote IP address.+throttle :: ThrottleSettings+            -> WaiThrottle+            -> 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 throttleReq+                   else return 0++    if remaining /= 0+        then respond $ onThrottled remaining+        else app req respond+  where+    throttleReq = do++      let SockAddrInet _ host = remoteHost req+      remoteAddr     <- T.pack <$> inet_ntoa host+      throttleState  <- atomically $ readTVar tmap+      (tst, success) <- throttleReq' remoteAddr throttleState++      -- write the throttle state back+      atomically $ writeTVar tmap (ThrottleState tst)+      return success++    throttleReq' remoteAddr (ThrottleState m) = do++      let toInvRate r = round (1e6 / r)+          invRate     = toInvRate (fromInteger throttleRate :: Double)+          burst       = fromInteger throttleBurst++      bucket    <- maybe newTokenBucket return $ H.lookup remoteAddr m+      remaining <- tokenBucketTryAlloc1 bucket burst invRate++      return (H.insert remoteAddr bucket m, remaining)

@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain

@@ -0,0 +1,21 @@+-----------------------------------------------------------------------------+-- |+-- Module      : DocTest.hs+-- Description : Documentation Testing+-- Copyright   : (c) 2015 Christopher Reichert+-- License     : AllRightsReserved+-- Maintainer  : Christopher Reichert <creichert07@gmail.com>+-- Stability   : testing+-- Portability : POSIX+++module Main (main) where+++import System.FilePath.Glob (glob)+import Test.DocTest         (doctest)+++main :: IO ()+main = glob "Network/**/[A-Z]*.hs"+         >>= doctest

@@ -0,0 +1,31 @@+-----------------------------------------------------------------------------+-- |+-- Module      : HLint+-- Description : Lint Testing Coverage+-- Copyright   : (c) 2015 Christopher Reichert+-- License     : BSD3+-- Maintainer  : Christopher Reichert <creichert07@gmail.com>+-- Stability   : testing+-- Portability : POSIX++module Main (main) where++import           Language.Haskell.HLint (hlint)+import           System.Exit            (exitFailure, exitSuccess)+++main :: IO ()+main = do+    hints <- hlint arguments+    print hints+    if null hints+       then exitSuccess+       else exitFailure+++arguments :: [String]+arguments =+      [+        "Network"+      , "test"+      ]

@@ -0,0 +1,49 @@+-----------------------------------------------------------------------------+-- |+-- Module      : Haddock+-- Description : Haddock Test Coverage+-- Copyright   : (c) 2015 Christopher Reichert+-- License     : BSD3+-- Maintainer  : Christopher Reichert <creichert07@gmail.com>+-- Stability   : testing+-- Portability : POSIX++module Main (main) where++import           Data.List      (genericLength)+import           Data.Maybe     (catMaybes)+import           System.Exit    (exitFailure, exitSuccess)+import           System.Process (readProcess)+import           Text.Regex     (matchRegex, mkRegex)+++main :: IO ()+main = do+  output <- readProcess "cabal"+            [+              "haddock"+            , "-v"+            , "--hoogle"+            , "--html"+            , "--internal"+            ] ""++  if average (match output) >= expected+    then exitSuccess+    else putStr output >> exitFailure+++-- | Expected percentage of documentation to pass+-- the test+expected :: Double+expected = 50.0+++average :: (Fractional a, Real b) => [b] -> a+average xs = realToFrac (sum xs) / genericLength xs+++match :: String -> [Int]+match = fmap read . concat . catMaybes . fmap (matchRegex pattern) . lines+  where+    pattern = mkRegex "^ *([0-9]*)% "

@@ -0,0 +1,1 @@+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}

@@ -0,0 +1,64 @@+------------------------------------------------------------------------+-- |+-- 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++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           Network.Wai.Middleware.Throttle+++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") ]+            }+++-- | Test Hmac Authentication+caseThrottle :: Assertion+caseThrottle = do++  st <- liftIO initThrottler++  statuses <- flip runSession (throttleApp st) $ do+    responses <- mapM (const (request defReq)) [ 1 .. 100 :: Integer ]+    mapM (return . simpleStatus) responses++  let msg = "Verifying some of the requests were throttled"+  assertBool msg  $ elem status429 statuses

@@ -0,0 +1,89 @@++name:                wai-middleware-throttle+version:             0.1.0.0+license:             BSD3+license-file:        LICENSE+author:              Christopher Reichert+maintainer:          creichert07@gmail.com+copyright:           (c) 2015, Christopher Reichert+category:            Web+build-type:          Simple+cabal-version:       >=1.10+synopsis:            WAI Middleware for Request Throttling+description:+  WAI Middleware for request rate limiting and throttling.+  .+  Designed to be configured in the spirit of the NGinx module.+++source-repository head+  type:     git+  location: git://github.com/creichert/wai-middleware-throttle.git+++library+  exposed-modules:     Network.Wai.Middleware.Throttle+  ghc-options:         -Wall -fno-warn-unused-do-bind+  default-language:    Haskell2010+  build-depends:       base                  >= 4.6 && < 5.0+                     , wai                   >= 3.0+                     , text+                     , unordered-containers+                     , bytestring+                     , transformers+                     , network+                     , http-types+                     , stm+                     , token-bucket++  if impl(ghc < 7.8)+    build-depends:+      blaze-builder < 0.4++test-suite spec+  type:                exitcode-stdio-1.0+  main-is:             Spec.hs+  hs-source-dirs:      test/+  other-modules:       WaiMiddlewareThrottleSpec+  ghc-options:         -Wall -threaded+  default-language:    Haskell2010+  build-depends:       base                     >= 4.6 && < 5.0+                     , bytestring+                     , wai-middleware-throttle+                     , wai+                     , wai-extra+                     , hspec                    >= 1.3+                     , http-types+                     , HUnit+                     , unordered-containers+                     , stm+                     , transformers+++test-suite doctest+  default-language:    Haskell2010+  ghc-options:         -Wall -Werror -threaded+  main-is:             test/DocTest.hs+  type:                exitcode-stdio-1.0+  build-depends:       base+                     , doctest == 0.9.*+                     , Glob    == 0.7.*+++test-suite haddock+  main-is:             test/Haddock.hs+  ghc-options:         -Wall -Werror -threaded+  type:                exitcode-stdio-1.0+  default-language:    Haskell2010+  build-depends:       base+                     , process+                     , regex-compat+++test-suite hlint+    main-is:          test/HLint.hs+    ghc-options:      -Wall -Werror -threaded+    type:             exitcode-stdio-1.0+    default-language: Haskell2010+    build-depends:    base+                    , hlint == 1.8.*