add throttle middleware

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

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) $
+                  -- setOnOpen (addToTracker tracker max_conn_per_client) $
-                  setOnClose (removeFromTracker tracker) $
+                  -- 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 =

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

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

shell.nix

@@ -5,11 +5,12 @@ let
   inherit (nixpkgs) pkgs;
 
   f = { mkDerivation, aeson, base, base16-bytestring
-      , base64-bytestring, bytestring, containers, http-api-data
+      , base64-bytestring, bytestring, case-insensitive, containers
-      , HsYAML, http-types, mtl, network, opaleye, optparse-applicative
+      , HsYAML, http-api-data, http-types, iproute, mtl, network, opaleye
-      , postgresql-simple, product-profunctors, profunctors, pureMD5
+      , optparse-applicative, postgresql-simple, product-profunctors
-      , random-bytestring, servant, servant-rawm, servant-server, stdenv
+      , profunctors, pureMD5, random-bytestring, servant, servant-rawm
-      , stm, text, time, wai, wai-logger, warp
+      , 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
+          base base16-bytestring bytestring case-insensitive containers
-          opaleye optparse-applicative postgresql-simple servant
+          HsYAML iproute mtl network opaleye optparse-applicative
-          servant-server stm text time wai wai-logger warp
+          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;