diff --git a/app/AppTypes/Configuration.hs b/app/AppTypes/Configuration.hs index 7924d64..3c708c7 100644 --- a/app/AppTypes/Configuration.hs +++ b/app/AppTypes/Configuration.hs @@ -16,7 +16,7 @@ data ServerConfig = ServerConfig , configCurrencySymbol :: T.Text , configListenPort :: Word , configListenHost :: T.Text - , configMaxConnectionsPerClient :: Word + -- , configMaxConnectionsPerClient :: Word , configBlockRegistration :: Bool } deriving (Show) @@ -31,7 +31,7 @@ instance FromYAML ServerConfig where <*> m .: "currency" <*> m .: "listen_port" <*> m .:? "listen_host" .!= "127.0.0.1" - <*> m .:? "max_connections_per_client" .!= 10 + -- <*> m .:? "max_connections_per_client" .!= 10 <*> m .: "block_registration" newtype Options = Options diff --git a/app/Main.hs b/app/Main.hs index 8f2db84..7462a4b 100644 --- a/app/Main.hs +++ b/app/Main.hs @@ -11,17 +11,22 @@ 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 Data.ByteString.Char8 as B8 hiding (putStrLn) import qualified Data.Text as T import Data.String import Data.YAML import Data.Version (showVersion) +import Data.CaseInsensitive (CI) +import qualified Data.CaseInsensitive as CI +import Data.IP import Database.PostgreSQL.Simple import Network.Wai import Network.Wai.Logger import Network.Wai.Handler.Warp -import Network.Socket (SockAddr(..)) +import Network.Wai.Middleware.Throttle +import Network.Socket (SockAddr(..), defaultPort) import Control.Monad.Reader @@ -30,6 +35,8 @@ import Control.Concurrent.STM.TVar import Options.Applicative +import System.Clock (TimeSpec(..)) + -- internal imports import API @@ -39,6 +46,7 @@ import AppTypes import Types import Control import Janitor +-- import Middleware import Paths_mateamt (version) @@ -62,7 +70,7 @@ main = do sym lport lhost - max_conn_per_client + -- max_conn_per_client block_registration ) -> do conn <- connectPostgreSQL ( @@ -85,8 +93,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) $ + -- setOnOpen (addToTracker tracker max_conn_per_client) $ + -- setOnClose (removeFromTracker tracker) $ setLogger ilog defaultSettings initState = ReadState { rsConnection = conn @@ -94,7 +102,17 @@ main = do , rsCurrencySymbol = sym , 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 opts = info (options <**> helper) ( fullDesc @@ -102,56 +120,6 @@ 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 :: Bool -> ReadState -> Application -- app conn = serveWithContext userApi genAuthServerContext (users conn) app block_registration initState = diff --git a/config.yaml b/config.yaml index 6beb350..62e3be6 100644 --- a/config.yaml +++ b/config.yaml @@ -6,5 +6,5 @@ db_passwd: "mateamt" listen_port: 8000 #listen_host: "127.0.0.1" #max_connections_per_client: 10 -currency: "€" +currency: "meow" block_registration: false diff --git a/mateamt.cabal b/mateamt.cabal index 28351d1..33fc215 100644 --- a/mateamt.cabal +++ b/mateamt.cabal @@ -42,8 +42,12 @@ executable mateamt , warp , wai , wai-logger + , wai-middleware-throttle , HsYAML >= 0.2.1.0 , optparse-applicative + , case-insensitive + , iproute + , clock hs-source-dirs: app ghc-options: -Wall default-language: Haskell2010 diff --git a/shell.nix b/shell.nix index 43ccdc5..f46a366 100644 --- a/shell.nix +++ b/shell.nix @@ -5,11 +5,12 @@ let inherit (nixpkgs) pkgs; f = { mkDerivation, aeson, base, base16-bytestring - , base64-bytestring, bytestring, containers, http-api-data - , HsYAML, http-types, mtl, network, opaleye, optparse-applicative - , postgresql-simple, product-profunctors, profunctors, pureMD5 - , random-bytestring, servant, servant-rawm, servant-server, stdenv - , stm, text, time, wai, wai-logger, warp + , base64-bytestring, bytestring, case-insensitive, containers + , HsYAML, http-api-data, http-types, iproute, mtl, network, opaleye + , optparse-applicative, postgresql-simple, product-profunctors + , profunctors, pureMD5, random-bytestring, servant, servant-rawm + , servant-server, stdenv, stm, text, time, wai, wai-logger + , wai-middleware-throttle, warp }: mkDerivation { pname = "mateamt"; @@ -24,9 +25,10 @@ let servant-rawm servant-server stm text time wai wai-logger warp ]; executableHaskellDepends = [ - base base16-bytestring bytestring containers HsYAML mtl network - opaleye optparse-applicative postgresql-simple servant - servant-server stm text time wai wai-logger warp + base base16-bytestring bytestring case-insensitive containers + HsYAML iproute mtl network opaleye optparse-applicative + postgresql-simple servant servant-server stm text time wai + wai-logger wai-middleware-throttle warp ]; description = "A whole new matemat"; license = stdenv.lib.licenses.agpl3;