wai-throttler (empty) → 0.1.0.0

raw patch · 4 files changed

+130/−0 lines, 4 filesdep +basedep +bytestringdep +containerssetup-changed

Dependencies added: base, bytestring, containers, http-types, time, wai

Files

• LICENSE +20/−0
• Network/Wai/Middleware/Throttler.hs +80/−0
• Setup.hs +2/−0
• wai-throttler.cabal +28/−0

@@ -0,0 +1,20 @@+Copyright (c) 2014 Maxim Kulkin++Permission is hereby granted, free of charge, to any person obtaining a copy+of this software and associated documentation files (the "Software"), to deal+in the Software without restriction, including without limitation the rights+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell+copies of the Software, and to permit persons to whom the Software is+furnished to do so, subject to the following conditions:++The above copyright notice and this permission notice shall be included in+all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN+THE SOFTWARE.+

@@ -0,0 +1,80 @@+{-# LANGUAGE OverloadedStrings #-}++module Network.Wai.Middleware.Throttler+  ( ThrottleCache(..)+  , newMemoryThrottleCache+  , throttle+  , throttlePath+  ) where++import Control.Concurrent.MVar (MVar, newMVar, takeMVar, putMVar, readMVar)+import Data.Map (Map)+import qualified Data.Map as Map (empty, insert, lookup)+import Data.Maybe (fromMaybe)+import Data.ByteString as B (ByteString, take, length)+import Data.ByteString.Char8 as B8 (pack)+import Data.Time.Clock (getCurrentTime, UTCTime, NominalDiffTime)+import Data.Time.Clock.POSIX (posixSecondsToUTCTime, utcTimeToPOSIXSeconds)+import Network.Wai (Request, Response, Middleware, rawPathInfo, responseLBS)+import Network.HTTP.Types.Status (status403)++class ThrottleCache c where+  cacheCount :: c -> ByteString -> IO Int+++data MemoryThrottleCache = MemoryThrottleCache Int NominalDiffTime (MVar (Map ByteString (UTCTime, Int)))++newMemoryThrottleCache :: Int -> NominalDiffTime -> IO MemoryThrottleCache+newMemoryThrottleCache limit period = fmap (MemoryThrottleCache limit period) $ newMVar Map.empty++instance ThrottleCache MemoryThrottleCache where+  cacheCount (MemoryThrottleCache limit period v) key = do+    map <- takeMVar v+    (t', c') <- case Map.lookup key map of+                  Nothing -> do+                    now <- getCurrentTime+                    return $ (alignTime now period, 1)+                  Just (t, c) -> do+                    case c + 1 > limit of+                      True -> do+                        now <- getCurrentTime+                        let alignedNow = alignTime now period+                        case alignedNow == t of+                          True -> return $ (t, c + 1)+                          False -> return $ (alignedNow, 1)+                      False -> return $ (t, c + 1)+    putMVar v $ Map.insert key (t', c') map+    return c'++    where alignTime :: UTCTime -> NominalDiffTime -> UTCTime+          alignTime time period = posixSecondsToUTCTime . fromIntegral $ a - (a `mod` b)+            where a = floor . utcTimeToPOSIXSeconds $ time+                  b = floor period+++throttlePath :: ThrottleCache c => ByteString -> c -> Int -> (Request -> Maybe ByteString) -> Middleware+throttlePath path cache limit getKey app req = do+  case pathMatches path req of+    False -> app req+    True -> throttle cache limit getKey app req++  where pathMatches :: ByteString -> Request -> Bool+        pathMatches path request = B.take (B.length path) (rawPathInfo request) == path++throttle :: ThrottleCache c => c -> Int -> (Request -> Maybe ByteString) -> Middleware+throttle cache limit getKey app req = do+  case getKey req of+    Nothing -> app req+    Just key -> do+      count <- cacheCount cache key+      if count > limit then return (throttledResponse count limit) else app req++  where throttledResponse :: Int -> Int -> Response+        throttledResponse count limit =+          responseLBS status403+                      [ ("X-RateLimit-Limit", bs limit)+                      , ("X-RateLimit-Remaining", bs remaining)+                      ]+                      ""+          where bs = B8.pack . show+                remaining = max 0 (limit - count)

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

@@ -0,0 +1,28 @@+-- Initial wai-throttler.cabal generated by cabal init.  For further +-- documentation, see http://haskell.org/cabal/users-guide/++name:                wai-throttler+version:             0.1.0.0+synopsis:            Wai middleware for request throttling+-- description:         +license:             MIT+license-file:        LICENSE+author:              Maxim Kulkin+maintainer:          maxim.kulkin@gmail.com+category:            Web+build-type:          Simple+cabal-version:       >=1.10++source-repository head+  type:                git+  location:            https://github.com/maximkulkin/wai-throttler++library+  exposed-modules:     Network.Wai.Middleware.Throttler+  build-depends:       base >=4.6 && <4.7+                     , bytestring == 0.10.*+                     , containers == 0.5.*+                     , time == 1.4.*+                     , wai == 2.*+                     , http-types == 0.8.*+  default-language:    Haskell2010