wai-rate-limit-postgres 0.0.0.0 → 0.1.0.0
raw patch · 3 files changed
+156/−21 lines, 3 filesdep +http-clientdep +http-typesdep +postgresql-simple-urldep ~basedep ~postgresql-simpledep ~relude
Dependencies added: http-client, http-types, postgresql-simple-url, tasty, tasty-hunit, wai, warp
Dependency ranges changed: base, postgresql-simple, relude, resource-pool, text, wai-rate-limit
Files
- src/Network/Wai/RateLimit/Postgres.hs +48/−13
- test/Spec.hs +92/−2
- wai-rate-limit-postgres.cabal +16/−6
src/Network/Wai/RateLimit/Postgres.hs view
@@ -10,6 +10,7 @@ ) where +import Control.Concurrent (forkIO, threadDelay) import Control.Exception (Handler (..), catches, throwIO, try) import Data.Pool (Pool, withResource) import qualified Data.Text as T@@ -47,8 +48,7 @@ tableName, "(key VARCHAR PRIMARY KEY,", "usage INT8 NOT NULL,",- "start_time TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,",- "expires_in INTERVAL NOT NULL DEFAULT '1 week'::INTERVAL)"+ "expires_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP + '1 week'::INTERVAL)" ] sqlHandlers :: [Handler a]@@ -80,12 +80,13 @@ [ "SELECT usage FROM", tableName, "WHERE key = ?",- "AND start_time + expires_in < CURRENT_TIMESTAMP"+ "AND expires_at > CURRENT_TIMESTAMP" ] pgBackendIncAndGetUsage :: Pool PG.Connection -> Text -> ByteString -> Integer -> IO (Either PGBackendError Integer) pgBackendIncAndGetUsage p tableName key usage = withResource p $ \c -> do res <- try $ PG.query c incAndGetQuery (key, usage) `catches` sqlHandlers+ print res return $ do rows <- res case rows of@@ -101,18 +102,15 @@ tableName, "(key, usage) VALUES (?, ?)", "ON CONFLICT (key) DO UPDATE SET",- "usage = CASE WHEN start_time + expires_in < CURRENT_TIMESTAMP",- "THEN usage + EXCLUDED.usage",- "ELSE EXCLUDED.usage END,",- "start_time = CASE WHEN start_time + expires_in < CURRENT_TIMESTAMP",- "THEN start_time",- "ELSE CURRENT_TIMESTAMP END",+ "usage = CASE WHEN rate_limiter.expires_at > CURRENT_TIMESTAMP THEN rate_limiter.usage + EXCLUDED.usage ELSE EXCLUDED.usage END,",+ "expires_at = CASE WHEN rate_limiter.expires_at > CURRENT_TIMESTAMP THEN rate_limiter.expires_at ELSE CURRENT_TIMESTAMP + '1 week'::INTERVAL END", "RETURNING usage" ] pgBackendExpireIn :: Pool PG.Connection -> Text -> ByteString -> Integer -> IO (Either PGBackendError ()) pgBackendExpireIn p tableName key seconds = withResource p $ \c -> do- res <- try $ PG.execute c expireInQuery (key, seconds) `catches` sqlHandlers+ putText "Called expire in!"+ res <- try $ PG.execute c expireInQuery (seconds, key) `catches` sqlHandlers return $ do count <- res if count /= 1@@ -126,16 +124,53 @@ " " [ "UPDATE", tableName,- "SET start_time = CURRENT_TIMESTAMP,",- "expires_in = ?",+ "SET expires_at = CURRENT_TIMESTAMP + '? second'::interval", "WHERE key = ?" ] +pgBackendCleanup :: Pool PG.Connection -> Text -> IO ()+pgBackendCleanup p tableName = void $+ forkIO $ do+ forever $ do+ res <- withResource p $ \c -> do+ tryDBErr $ PG.execute_ c removeExpired `catches` sqlHandlers+ case res of+ Left _ -> threadDelay d10s+ Right n -> delay n+ where+ d10s = 10_000_000+ d1s = 1_000_000+ d100ms = 100_000++ -- Try to ensure we cleanup as fast as garbage is created.+ delay n+ | n == 5000 = threadDelay d100ms+ | n > 0 = threadDelay d1s+ | otherwise = threadDelay d10s++ tryDBErr :: IO a -> IO (Either PGBackendError a)+ tryDBErr a = try a++ removeExpired =+ fromString $+ toString $+ T.intercalate+ " "+ [ "DELETE FROM",+ tableName,+ "WHERE key IN (SELECT key FROM",+ tableName,+ "WHERE expires_at < CURRENT_TIMESTAMP LIMIT 5000)"+ ]+ -- | Initialize a postgres backend for rate-limiting. Takes a connection pool--- and table name to use for storage.+-- and table name to use for storage. The table will be created if it does not+-- exist. A thread is also launched to periodically clean up expired rows from+-- the table. postgresBackend :: Pool PG.Connection -> Text -> IO (Backend ByteString PGBackendError) postgresBackend p tableName = do initPostgresBackend p tableName+ pgBackendCleanup p tableName return $ MkBackend { backendGetUsage = pgBackendGetUsage p tableName,
test/Spec.hs view
@@ -1,6 +1,96 @@ module Main (main) where --- import Network.Wai.RateLimit.Postgres+import Control.Concurrent+import Control.Exception+import Data.Pool+import qualified Database.PostgreSQL.Simple as PG+import Database.PostgreSQL.Simple.URL (parseDatabaseUrl)+import Network.HTTP.Client (Manager, defaultManagerSettings, httpLbs, newManager, parseRequest, responseStatus)+import Network.HTTP.Types (status200, statusCode)+import qualified Network.Wai as Wai+import qualified Network.Wai.Handler.Warp as Warp+import qualified Network.Wai.RateLimit as R+import Network.Wai.RateLimit.Postgres+import qualified Network.Wai.RateLimit.Strategy as R+import System.Environment.Blank+import System.IO.Error+import Test.Tasty+import Test.Tasty.HUnit +-- | Example value = "postgres://postgres:postgres@localhost:5432/postgres"+testPGEnv :: String+testPGEnv = "PG_DB_URI"++mkConnPool :: IO (Pool PG.Connection)+mkConnPool = do+ connURI <- getEnv testPGEnv >>= maybe (throwIO $ userError "PG_DB_URI not set") return+ connectInfo <- maybe (throwIO $ userError "invalid uri") return $ parseDatabaseUrl connURI+ createPool+ (PG.connect connectInfo)+ PG.close+ 1+ 5+ 10++-- | The 'key' argument specified the PGBackend key to be used for every+-- request made to this test app.+mkTestWaiApp :: Integer -> Integer -> ByteString -> IO Wai.Application+mkTestWaiApp seconds limit key = do+ pool <- mkConnPool+ pgBackend <- postgresBackend pool "rate_limiter"+ let app _ respond = respond $ Wai.responseLBS status200 [] "Ok!"+ strategy = R.fixedWindow pgBackend seconds limit (const $ return key)+ middleware = R.rateLimiting strategy+ return $ middleware app++mkReq :: Manager -> IO Int+mkReq manager = do+ req <- parseRequest "http://localhost:11222"+ resp <- httpLbs req manager+ return $ statusCode $ responseStatus resp++launchApp :: Wai.Application -> IO ThreadId+launchApp app = forkIO $ Warp.run 11222 app++tearDownApp :: ThreadId -> IO ()+tearDownApp = killThread+ main :: IO ()-main = putStrLn "Tests not implemented!"+main =+ defaultMain $+ testCaseSteps+ "Rate Limiting Tests"+ $ \step ->+ do+ step "Limit excessive requests (1)"++ bracket+ (mkTestWaiApp 1 2 "key1" >>= launchApp)+ tearDownApp+ $ \_ -> do+ mgr <- newManager defaultManagerSettings+ rs <- replicateM 3 $ mkReq mgr+ when (rs /= [200, 200, 429]) $ do+ assertFailure "Not ratelimited!"++ step "Limit excessive requests (2)"++ bracket+ (mkTestWaiApp 1 3 "key2" >>= launchApp)+ tearDownApp+ $ \_ -> do+ mgr <- newManager defaultManagerSettings+ rs <- replicateM 10 $ mkReq mgr+ when (rs /= replicate 3 200 ++ replicate 7 429) $+ assertFailure "Unexpected result!"++ step "Allow non-excessive requests"++ bracket+ (mkTestWaiApp 1 3 "key3" >>= launchApp)+ tearDownApp+ $ \_ -> do+ mgr <- newManager defaultManagerSettings+ rs <- replicateM 3 $ mkReq mgr+ when (rs /= [200, 200, 200]) $+ assertFailure "Unexpected result!"
wai-rate-limit-postgres.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: wai-rate-limit-postgres-version: 0.0.0.0+version: 0.1.0.0 category: Security, Web, Network synopsis: See README for more info description:@@ -24,11 +24,11 @@ common common-options build-depends: base >= 4.14.3.0 && < 4.16- , relude- , wai-rate-limit- , postgresql-simple- , text- , resource-pool+ , relude ^>= 1.0+ , wai-rate-limit ^>= 0.1.0.0+ , postgresql-simple ^>= 0.6+ , text ^>= 1.2.4+ , resource-pool ^>= 0.2.3.2 mixins: base hiding (Prelude) , relude (Relude as Prelude)@@ -58,6 +58,7 @@ , QuasiQuotes , OverloadedStrings , LambdaCase+ , NumericUnderscores library import: common-options@@ -70,6 +71,15 @@ hs-source-dirs: test main-is: Spec.hs build-depends: wai-rate-limit-postgres+ , postgresql-simple-url+ , resource-pool+ , tasty+ , tasty-hunit+ , http-client+ , http-types+ , warp+ , wai-rate-limit+ , wai ghc-options: -threaded -rtsopts -with-rtsopts=-N