Something went wrong. Try again.
A Haskell Webserver Template
Something went wrong. Try again.
2.8 kB · 64 lines
Haskell
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465module Webserver.Data.Sql.Cache.RateLimit where
----------------------------------------------------------------------------------------------------
import Data.Functor.Contravariant (contramap)import Data.Int (Int32)import qualified Data.Time.Clock as Clockimport Data.Tuple (fst, snd)import GHC.Real (fromIntegral)import qualified Hasql.Decoders as Decodersimport qualified Hasql.Encoders as Encodersimport qualified Hasql.Pool as Poolimport qualified Hasql.Session as Sessionimport qualified Hasql.Statement as Statementimport Webserver.Control.Nil (Nil (..))import qualified Webserver.Control.Sql as Sql
----------------------------------------------------------------------------------------------------
$(Sql.generateTable "cache" "rate_limits")
----------------------------------------------------------------------------------------------------
instance Nil RateLimits where nil = RateLimits {id = nil, key = nil, hits = nil, windowStart = nil}
----------------------------------------------------------------------------------------------------
{- Atomic UPSERT that increments hits within a window or resets if the window has expired. Returns the current hit count after the upsert. Parameters: $1 = key (Text), $2 = window duration (Interval).-}checkRateLimitStatement :: Statement.Statement (Text, Clock.DiffTime) Int32checkRateLimitStatement = Statement.Statement "INSERT INTO cache.rate_limits (key, hits, window_start) \ \VALUES ($1, 1, NOW()) \ \ON CONFLICT (key) DO UPDATE \ \SET hits = CASE \ \ WHEN cache.rate_limits.window_start + $2 < NOW() THEN 1 \ \ ELSE cache.rate_limits.hits + 1 \ \ END, \ \ window_start = CASE \ \ WHEN cache.rate_limits.window_start + $2 < NOW() THEN NOW() \ \ ELSE cache.rate_limits.window_start \ \ END \ \RETURNING hits" ( contramap fst (Encoders.param (Encoders.nonNullable Encoders.text)) <> contramap snd (Encoders.param (Encoders.nonNullable Encoders.interval)) ) (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int4))) True
----------------------------------------------------------------------------------------------------
{- Check if a request identified by the given key is within the rate limit. Returns True if the request is allowed, False if rate limited.-}checkRateLimit :: Pool.Pool -> Text -> Int -> Clock.DiffTime -> IO BoolcheckRateLimit pool key_ maxHits window = do hits_ <- Sql.runSession pool <| Session.statement (key_, window) checkRateLimitStatement return <| fromIntegral hits_ <= maxHits