diff --git a/app/AppTypes/Configuration.hs b/app/AppTypes/Configuration.hs index e39e5d5..e06dd01 100644 --- a/app/AppTypes/Configuration.hs +++ b/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 diff --git a/app/Main.hs b/app/Main.hs index 484bd1d..8d7060f 100644 --- a/app/Main.hs +++ b/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 = diff --git a/config.yaml b/config.yaml index 33f77f7..bb5ece7 100644 --- a/config.yaml +++ b/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: "€" diff --git a/mateamt.cabal b/mateamt.cabal index 7c275c3..a7c48af 100644 --- a/mateamt.cabal +++ b/mateamt.cabal @@ -36,6 +36,7 @@ executable mateamt , opaleye , postgresql-simple , stm + , network , servant , servant-server , time