add tracker to stop DDOSing clients

This commit is contained in:
nek0 2019-10-31 23:58:50 +01:00
parent 173c4f4ee7
commit 99bf138a92
4 changed files with 61 additions and 0 deletions

View file

@ -16,6 +16,7 @@ data ServerConfig = ServerConfig
, configCurrencySymbol :: T.Text , configCurrencySymbol :: T.Text
, configListenPort :: Word , configListenPort :: Word
, configListenHost :: T.Text , configListenHost :: T.Text
, configMaxConnectionsPerClient :: Word
} }
deriving (Show) deriving (Show)
@ -29,6 +30,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" .!= 100
newtype Options = Options newtype Options = Options
{ optConfigLocation :: T.Text { optConfigLocation :: T.Text

View file

@ -9,6 +9,7 @@ import Servant
import Servant.Server.Experimental.Auth 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 Data.ByteString.Lazy as BL hiding (putStrLn) import Data.ByteString.Lazy as BL hiding (putStrLn)
import qualified Data.Text as T import qualified Data.Text as T
import Data.String import Data.String
@ -19,9 +20,11 @@ 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 Control.Monad.Reader import Control.Monad.Reader
import Control.Concurrent.STM (atomically)
import Control.Concurrent.STM.TVar import Control.Concurrent.STM.TVar
import Options.Applicative import Options.Applicative
@ -56,6 +59,7 @@ main = do
sym sym
lport lport
lhost lhost
max_conn_per_client
) -> do ) -> do
conn <- connectPostgreSQL ( conn <- connectPostgreSQL (
"host='" <> fromString (T.unpack db_host) <> "' " <> "host='" <> fromString (T.unpack db_host) <> "' " <>
@ -65,6 +69,7 @@ main = do
"password='" <> fromString (T.unpack db_passwd) <> "'" "password='" <> fromString (T.unpack db_passwd) <> "'"
) )
store <- newTVarIO S.empty store <- newTVarIO S.empty
tracker <- newTVarIO M.empty
void $ execute_ conn initAvatar void $ execute_ conn initAvatar
void $ execute_ conn initUser void $ execute_ conn initUser
void $ execute_ conn initProduct void $ execute_ conn initProduct
@ -76,6 +81,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) $
setOnClose (removeFromTracker tracker) $
setLogger ilog defaultSettings setLogger ilog defaultSettings
initState = ReadState initState = ReadState
{ rsConnection = conn { rsConnection = conn
@ -89,6 +96,56 @@ 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 :: ReadState -> Application app :: ReadState -> Application
-- app conn = serveWithContext userApi genAuthServerContext (users conn) -- app conn = serveWithContext userApi genAuthServerContext (users conn)
app initState = app initState =

View file

@ -5,4 +5,5 @@ db_user: "mateamt"
db_passwd: "mateamt" 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: 100
currency: "€" currency: "€"

View file

@ -36,6 +36,7 @@ executable mateamt
, opaleye , opaleye
, postgresql-simple , postgresql-simple
, stm , stm
, network
, servant , servant
, servant-server , servant-server
, time , time