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

View File

@ -11,17 +11,22 @@ import Servant.Server.Experimental.Auth
import Data.Set as S (empty) import Data.Set as S (empty)
import qualified Data.Map.Lazy as M import qualified Data.Map.Lazy as M
import Data.ByteString.Lazy as BL hiding (putStrLn) import Data.ByteString.Lazy as BL hiding (putStrLn)
import Data.ByteString.Char8 as B8 hiding (putStrLn)
import qualified Data.Text as T import qualified Data.Text as T
import Data.String import Data.String
import Data.YAML import Data.YAML
import Data.Version (showVersion) import Data.Version (showVersion)
import Data.CaseInsensitive (CI)
import qualified Data.CaseInsensitive as CI
import Data.IP
import Database.PostgreSQL.Simple import Database.PostgreSQL.Simple
import Network.Wai import Network.Wai
import Network.Wai.Logger import Network.Wai.Logger
import Network.Wai.Handler.Warp import Network.Wai.Handler.Warp
import Network.Socket (SockAddr(..)) import Network.Wai.Middleware.Throttle
import Network.Socket (SockAddr(..), defaultPort)
import Control.Monad.Reader import Control.Monad.Reader
@ -30,6 +35,8 @@ import Control.Concurrent.STM.TVar
import Options.Applicative import Options.Applicative
import System.Clock (TimeSpec(..))
-- internal imports -- internal imports
import API import API
@ -39,6 +46,7 @@ import AppTypes
import Types import Types
import Control import Control
import Janitor import Janitor
-- import Middleware
import Paths_mateamt (version) import Paths_mateamt (version)
@ -62,7 +70,7 @@ main = do
sym sym
lport lport
lhost lhost
max_conn_per_client -- max_conn_per_client
block_registration block_registration
) -> do ) -> do
conn <- connectPostgreSQL ( conn <- connectPostgreSQL (
@ -85,8 +93,8 @@ main = do
withStdoutLogger $ \ilog -> do withStdoutLogger $ \ilog -> do
let settings = setPort (fromIntegral lport) $ let settings = setPort (fromIntegral lport) $
setHost (fromString $ T.unpack lhost) $ setHost (fromString $ T.unpack lhost) $
setOnOpen (addToTracker tracker max_conn_per_client) $ -- setOnOpen (addToTracker tracker max_conn_per_client) $
setOnClose (removeFromTracker tracker) $ -- setOnClose (removeFromTracker tracker) $
setLogger ilog defaultSettings setLogger ilog defaultSettings
initState = ReadState initState = ReadState
{ rsConnection = conn { rsConnection = conn
@ -94,7 +102,17 @@ main = do
, rsCurrencySymbol = sym , rsCurrencySymbol = sym
, rsSoftwareVersion = T.pack (showVersion version) , 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 where
opts = info (options <**> helper) opts = info (options <**> helper)
( fullDesc ( fullDesc
@ -102,56 +120,6 @@ main = do
<> header "mateamt - Your friendly mate distribution office" <> 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 :: Bool -> ReadState -> Application
-- app conn = serveWithContext userApi genAuthServerContext (users conn) -- app conn = serveWithContext userApi genAuthServerContext (users conn)
app block_registration initState = app block_registration initState =

View File

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

View File

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

View File

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