add tracker to stop DDOSing clients

app/AppTypes/Configuration.hs

@@ -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

app/Main.hs

@@ -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 =

config.yaml

@@ -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: "€"

mateamt.cabal

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