add throttle middleware

This commit is contained in:
nek0 2020-04-19 00:49:30 +02:00
parent 0734707632
commit f1a81acb2b
5 changed files with 40 additions and 66 deletions

View File

@ -16,7 +16,7 @@ data ServerConfig = ServerConfig
, configCurrencySymbol :: T.Text
, configListenPort :: Word
, configListenHost :: T.Text
, configMaxConnectionsPerClient :: Word
-- , configMaxConnectionsPerClient :: Word
, configBlockRegistration :: Bool
}
deriving (Show)
@ -31,7 +31,7 @@ instance FromYAML ServerConfig where
<*> m .: "currency"
<*> m .: "listen_port"
<*> m .:? "listen_host" .!= "127.0.0.1"
<*> m .:? "max_connections_per_client" .!= 10
-- <*> m .:? "max_connections_per_client" .!= 10
<*> m .: "block_registration"
newtype Options = Options

View File

@ -11,17 +11,22 @@ import Servant.Server.Experimental.Auth
import Data.Set as S (empty)
import qualified Data.Map.Lazy as M
import Data.ByteString.Lazy as BL hiding (putStrLn)
import Data.ByteString.Char8 as B8 hiding (putStrLn)
import qualified Data.Text as T
import Data.String
import Data.YAML
import Data.Version (showVersion)
import Data.CaseInsensitive (CI)
import qualified Data.CaseInsensitive as CI
import Data.IP
import Database.PostgreSQL.Simple
import Network.Wai
import Network.Wai.Logger
import Network.Wai.Handler.Warp
import Network.Socket (SockAddr(..))
import Network.Wai.Middleware.Throttle
import Network.Socket (SockAddr(..), defaultPort)
import Control.Monad.Reader
@ -30,6 +35,8 @@ import Control.Concurrent.STM.TVar
import Options.Applicative
import System.Clock (TimeSpec(..))
-- internal imports
import API
@ -39,6 +46,7 @@ import AppTypes
import Types
import Control
import Janitor
-- import Middleware
import Paths_mateamt (version)
@ -62,7 +70,7 @@ main = do
sym
lport
lhost
max_conn_per_client
-- max_conn_per_client
block_registration
) -> do
conn <- connectPostgreSQL (
@ -85,8 +93,8 @@ main = do
withStdoutLogger $ \ilog -> do
let settings = setPort (fromIntegral lport) $
setHost (fromString $ T.unpack lhost) $
setOnOpen (addToTracker tracker max_conn_per_client) $
setOnClose (removeFromTracker tracker) $
-- setOnOpen (addToTracker tracker max_conn_per_client) $
-- setOnClose (removeFromTracker tracker) $
setLogger ilog defaultSettings
initState = ReadState
{ rsConnection = conn
@ -94,7 +102,17 @@ main = do
, rsCurrencySymbol = sym
, rsSoftwareVersion = T.pack (showVersion version)
}
runSettings settings (app block_registration initState)
expirationSpec = TimeSpec 5 0 -- five seconds
th <- initCustomThrottler (defaultThrottleSettings expirationSpec)
(\req ->
let headers = requestHeaders req
in case lookup "x-forwarded-for" headers of
Just addrs ->
let addr = fst (B8.break (== ',') addrs)
in Right $ Address $ toSockAddr (read (B8.unpack addr), defaultPort)
Nothing -> Right $ Address $ remoteHost req
)
runSettings settings (throttle th (app block_registration initState))
where
opts = info (options <**> helper)
( fullDesc
@ -102,56 +120,6 @@ main = do
<> header "mateamt - Your friendly mate distribution office"
)
addToTracker
:: TVar (M.Map (Word, Word, Word, Word) Word)
-> Word
-> SockAddr
-> IO Bool
addToTracker tracker maxconn saddr = do
let laddr = translateAddr saddr
(nmap, accept) <- atomically $ do
tmap <- readTVar tracker
let (nmap, accept) = case tmap M.!? laddr of
Just val ->
if val < maxconn
then
(M.adjust (const $ val + 1) laddr tmap, True)
else
(tmap, False)
Nothing ->
(M.insert laddr 1 tmap, True)
writeTVar tracker nmap
return (nmap, accept)
-- print nmap
return accept
translateAddr
:: (Num a, Num b, Num c, Num d)
=> SockAddr
-> (a, b, c, d)
translateAddr saddr =
case saddr of
SockAddrInet _ addr -> (0, 0, 0, fromIntegral addr)
SockAddrInet6 _ _ (a1, a2, a3, a4) _ ->
(fromIntegral a1, fromIntegral a2, fromIntegral a3, fromIntegral a4)
_ -> error "Not made for sockets and the like"
removeFromTracker
:: TVar (M.Map (Word, Word, Word, Word) Word)
-> SockAddr
-> IO ()
removeFromTracker tracker saddr = do
let laddr = translateAddr saddr
atomically $ do
tmap <- readTVar tracker
let num = tmap M.! laddr
nmap = if num -1 > 0
then
M.adjust (const $ num - 1) laddr tmap
else
M.delete laddr tmap
writeTVar tracker nmap
app :: Bool -> ReadState -> Application
-- app conn = serveWithContext userApi genAuthServerContext (users conn)
app block_registration initState =

View File

@ -6,5 +6,5 @@ db_passwd: "mateamt"
listen_port: 8000
#listen_host: "127.0.0.1"
#max_connections_per_client: 10
currency: ""
currency: "meow"
block_registration: false

View File

@ -42,8 +42,12 @@ executable mateamt
, warp
, wai
, wai-logger
, wai-middleware-throttle
, HsYAML >= 0.2.1.0
, optparse-applicative
, case-insensitive
, iproute
, clock
hs-source-dirs: app
ghc-options: -Wall
default-language: Haskell2010

View File

@ -5,11 +5,12 @@ let
inherit (nixpkgs) pkgs;
f = { mkDerivation, aeson, base, base16-bytestring
, base64-bytestring, bytestring, containers, http-api-data
, HsYAML, http-types, mtl, network, opaleye, optparse-applicative
, postgresql-simple, product-profunctors, profunctors, pureMD5
, random-bytestring, servant, servant-rawm, servant-server, stdenv
, stm, text, time, wai, wai-logger, warp
, base64-bytestring, bytestring, case-insensitive, containers
, HsYAML, http-api-data, http-types, iproute, mtl, network, opaleye
, optparse-applicative, postgresql-simple, product-profunctors
, profunctors, pureMD5, random-bytestring, servant, servant-rawm
, servant-server, stdenv, stm, text, time, wai, wai-logger
, wai-middleware-throttle, warp
}:
mkDerivation {
pname = "mateamt";
@ -24,9 +25,10 @@ let
servant-rawm servant-server stm text time wai wai-logger warp
];
executableHaskellDepends = [
base base16-bytestring bytestring containers HsYAML mtl network
opaleye optparse-applicative postgresql-simple servant
servant-server stm text time wai wai-logger warp
base base16-bytestring bytestring case-insensitive containers
HsYAML iproute mtl network opaleye optparse-applicative
postgresql-simple servant servant-server stm text time wai
wai-logger wai-middleware-throttle warp
];
description = "A whole new matemat";
license = stdenv.lib.licenses.agpl3;