diff --git a/src/Network/Wai/RateLimit/Postgres.hs b/src/Network/Wai/RateLimit/Postgres.hs --- a/src/Network/Wai/RateLimit/Postgres.hs +++ b/src/Network/Wai/RateLimit/Postgres.hs @@ -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, diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -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!" diff --git a/wai-rate-limit-postgres.cabal b/wai-rate-limit-postgres.cabal --- a/wai-rate-limit-postgres.cabal +++ b/wai-rate-limit-postgres.cabal @@ -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