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
, configListenPort :: Word
, configListenHost :: T.Text
, configMaxConnectionsPerClient :: Word
}
deriving (Show)
@ -29,6 +30,7 @@ instance FromYAML ServerConfig where
<*> m .: "currency"
<*> m .: "listen_port"
<*> m .:? "listen_host" .!= "127.0.0.1"
<*> m .:? "max_connections_per_client" .!= 100
newtype Options = Options
{ optConfigLocation :: T.Text

View File

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

View File

@ -5,4 +5,5 @@ db_user: "mateamt"
db_passwd: "mateamt"
listen_port: 8000
#listen_host: "127.0.0.1"
#max_connections_per_client: 100
currency: "€"

View File

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