diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/Network/Wai/Middleware/Throttle.hs b/Network/Wai/Middleware/Throttle.hs new file mode 100644 --- /dev/null +++ b/Network/Wai/Middleware/Throttle.hs @@ -0,0 +1,172 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Network.Wai.Middleware.Throttle +-- Description : WAI Request Throttling Middleware +-- Copyright : (c) 2015 Christopher Reichert +-- License : BSD3 +-- Maintainer : Christopher Reichert +-- Stability : experimental +-- Portability : POSIX +-- +-- Uses a +-- 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) diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/test/DocTest.hs b/test/DocTest.hs new file mode 100644 --- /dev/null +++ b/test/DocTest.hs @@ -0,0 +1,21 @@ +----------------------------------------------------------------------------- +-- | +-- Module : DocTest.hs +-- Description : Documentation Testing +-- Copyright : (c) 2015 Christopher Reichert +-- License : AllRightsReserved +-- Maintainer : Christopher Reichert +-- 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 diff --git a/test/HLint.hs b/test/HLint.hs new file mode 100644 --- /dev/null +++ b/test/HLint.hs @@ -0,0 +1,31 @@ +----------------------------------------------------------------------------- +-- | +-- Module : HLint +-- Description : Lint Testing Coverage +-- Copyright : (c) 2015 Christopher Reichert +-- License : BSD3 +-- Maintainer : Christopher Reichert +-- 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" + ] diff --git a/test/Haddock.hs b/test/Haddock.hs new file mode 100644 --- /dev/null +++ b/test/Haddock.hs @@ -0,0 +1,49 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Haddock +-- Description : Haddock Test Coverage +-- Copyright : (c) 2015 Christopher Reichert +-- License : BSD3 +-- Maintainer : Christopher Reichert +-- 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]*)% " diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-} diff --git a/test/WaiMiddlewareThrottleSpec.hs b/test/WaiMiddlewareThrottleSpec.hs new file mode 100644 --- /dev/null +++ b/test/WaiMiddlewareThrottleSpec.hs @@ -0,0 +1,64 @@ +------------------------------------------------------------------------ +-- | +-- 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 + +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 diff --git a/wai-middleware-throttle.cabal b/wai-middleware-throttle.cabal new file mode 100644 --- /dev/null +++ b/wai-middleware-throttle.cabal @@ -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.*