module Network.Wai.RateLimit.Strategy (
Strategy(..),
fixedWindow,
slidingWindow
) where
import Control.Monad
import Data.Time.Units
import Network.Wai
import Network.Wai.RateLimit.Backend
newtype Strategy = MkStrategy {
Strategy -> Request -> IO Bool
strategyOnRequest :: Request -> IO Bool
}
windowStrategy
:: Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> (Integer -> Bool)
-> Request
-> IO Bool
windowStrategy :: Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> (Integer -> Bool)
-> Request
-> IO Bool
windowStrategy MkBackend{key -> IO Integer
key -> Integer -> IO Integer
key -> Integer -> IO ()
backendExpireIn :: forall key. Backend key -> key -> Integer -> IO ()
backendIncAndGetUsage :: forall key. Backend key -> key -> Integer -> IO Integer
backendGetUsage :: forall key. Backend key -> key -> IO Integer
backendExpireIn :: key -> Integer -> IO ()
backendIncAndGetUsage :: key -> Integer -> IO Integer
backendGetUsage :: key -> IO Integer
..} Second
seconds Integer
capacity Request -> IO key
getKey Integer -> Bool
cond Request
req = do
key
key <- Request -> IO key
getKey Request
req
Integer
used <- key -> Integer -> IO Integer
backendIncAndGetUsage key
key Integer
1
if Integer
used Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
capacity
then do
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Integer -> Bool
cond Integer
used) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ key -> Integer -> IO ()
backendExpireIn key
key (Second -> Integer
forall a. Integral a => a -> Integer
toInteger Second
seconds)
Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
else Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
fixedWindow
:: Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> Strategy
fixedWindow :: Backend key -> Second -> Integer -> (Request -> IO key) -> Strategy
fixedWindow Backend key
backend Second
seconds Integer
capacity Request -> IO key
getKey = MkStrategy :: (Request -> IO Bool) -> Strategy
MkStrategy{
strategyOnRequest :: Request -> IO Bool
strategyOnRequest =
let cond :: Integer -> Bool
cond = Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
(==) Integer
1
in Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> (Integer -> Bool)
-> Request
-> IO Bool
forall key.
Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> (Integer -> Bool)
-> Request
-> IO Bool
windowStrategy Backend key
backend Second
seconds Integer
capacity Request -> IO key
getKey Integer -> Bool
cond
}
slidingWindow
:: Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> Strategy
slidingWindow :: Backend key -> Second -> Integer -> (Request -> IO key) -> Strategy
slidingWindow Backend key
backend Second
seconds Integer
capacity Request -> IO key
getKey = MkStrategy :: (Request -> IO Bool) -> Strategy
MkStrategy{
strategyOnRequest :: Request -> IO Bool
strategyOnRequest =
let cond :: b -> Bool
cond = Bool -> b -> Bool
forall a b. a -> b -> a
const Bool
True
in Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> (Integer -> Bool)
-> Request
-> IO Bool
forall key.
Backend key
-> Second
-> Integer
-> (Request -> IO key)
-> (Integer -> Bool)
-> Request
-> IO Bool
windowStrategy Backend key
backend Second
seconds Integer
capacity Request -> IO key
getKey Integer -> Bool
forall b. b -> Bool
cond
}