module Main (main) where
import Control.Concurrent
( ThreadId,
forkIO,
killThread,
threadDelay,
)
import Control.Exception (bracket, throwIO)
import Control.Monad (replicateM, when)
import Data.ByteString (ByteString)
import Data.Either (isRight)
import Data.Pool (Pool, createPool)
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8')
import Data.Text.Encoding.Error (UnicodeException)
import qualified Data.Time.Units as U
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 (postgresBackend)
import qualified Network.Wai.RateLimit.Strategy as R
import System.Environment.Blank (getEnv)
import Test.Tasty (defaultMain)
import Test.Tasty.HUnit (assertFailure, testCaseSteps)
-- | 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_1"
let app _ respond = respond $ Wai.responseLBS status200 [] "Ok!"
s = U.fromMicroseconds $ 1_000_000 * seconds
strategy = R.fixedWindow pgBackend s 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 =
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!"
step "Allow excessive requests, slow down and then be allowed"
bracket
(mkTestWaiApp 1 1 "key4" >>= launchApp)
tearDownApp
$ \_ -> do
mgr <- newManager defaultManagerSettings
rs <- replicateM 3 $ mkReq mgr
when (rs /= [200, 429, 429]) $
assertFailure "Unexpected result!"
threadDelay 1_000_000
rs2 <- replicateM 1 $ mkReq mgr
when (rs2 /= [200]) $
assertFailure "Unexpected result!"
step "Allow keys that are invalid as unicode strings"
let invalidUtf8String = "\187" :: ByteString
when (isRight (decodeUtf8' invalidUtf8String :: Either UnicodeException Text)) $
assertFailure "`invalidUtf8String` appears to have a valid UTF8 string"
bracket
(mkTestWaiApp 1 3 invalidUtf8String >>= launchApp)
tearDownApp
$ \_ -> do
mgr <- newManager defaultManagerSettings
rs <- replicateM 3 $ mkReq mgr
when (rs /= [200, 200, 200]) $
assertFailure $
"Unexpected result: " ++ show rs