89 lines
2.2 KiB
Haskell
89 lines
2.2 KiB
Haskell
{-
|
|
©AngelaMos | 2026
|
|
ConnLimit.hs
|
|
-}
|
|
{-# LANGUAGE OverloadedStrings #-}
|
|
{-# LANGUAGE RecordWildCards #-}
|
|
|
|
module Aenebris.DDoS.ConnLimit
|
|
( ConnLimiter
|
|
, ConnLimitConfig(..)
|
|
, defaultConnLimitConfig
|
|
, newConnLimiter
|
|
, tryAcquire
|
|
, release
|
|
, currentCount
|
|
, connLimitOnOpen
|
|
, connLimitOnClose
|
|
, defaultPerIPLimit
|
|
, ipBytesFromSockAddr
|
|
) where
|
|
|
|
import Control.Concurrent.STM
|
|
( STM
|
|
, TVar
|
|
, atomically
|
|
, modifyTVar'
|
|
, newTVarIO
|
|
, readTVar
|
|
, writeTVar
|
|
)
|
|
import Aenebris.Net.IP (sockAddrToIPBytes)
|
|
import Data.ByteString (ByteString)
|
|
import Data.Map.Strict (Map)
|
|
import qualified Data.Map.Strict as Map
|
|
import Network.Socket (SockAddr)
|
|
|
|
defaultPerIPLimit :: Int
|
|
defaultPerIPLimit = 16
|
|
|
|
data ConnLimitConfig = ConnLimitConfig
|
|
{ clcPerIPLimit :: !Int
|
|
} deriving (Show, Eq)
|
|
|
|
defaultConnLimitConfig :: ConnLimitConfig
|
|
defaultConnLimitConfig = ConnLimitConfig { clcPerIPLimit = defaultPerIPLimit }
|
|
|
|
data ConnLimiter = ConnLimiter
|
|
{ clCounts :: TVar (Map ByteString Int)
|
|
, clConfig :: !ConnLimitConfig
|
|
}
|
|
|
|
newConnLimiter :: ConnLimitConfig -> IO ConnLimiter
|
|
newConnLimiter cfg = do
|
|
tv <- newTVarIO Map.empty
|
|
pure ConnLimiter { clCounts = tv, clConfig = cfg }
|
|
|
|
tryAcquire :: ConnLimiter -> ByteString -> STM Bool
|
|
tryAcquire ConnLimiter{..} ip = do
|
|
m <- readTVar clCounts
|
|
let current = Map.findWithDefault 0 ip m
|
|
limit = clcPerIPLimit clConfig
|
|
if current >= limit
|
|
then pure False
|
|
else do
|
|
writeTVar clCounts (Map.insert ip (current + 1) m)
|
|
pure True
|
|
|
|
release :: ConnLimiter -> ByteString -> STM ()
|
|
release ConnLimiter{..} ip =
|
|
modifyTVar' clCounts (Map.update decrement ip)
|
|
where
|
|
decrement n
|
|
| n <= 1 = Nothing
|
|
| otherwise = Just (n - 1)
|
|
|
|
currentCount :: ConnLimiter -> ByteString -> STM Int
|
|
currentCount ConnLimiter{..} ip = do
|
|
m <- readTVar clCounts
|
|
pure (Map.findWithDefault 0 ip m)
|
|
|
|
connLimitOnOpen :: ConnLimiter -> SockAddr -> IO Bool
|
|
connLimitOnOpen cl sa = atomically (tryAcquire cl (sockAddrToIPBytes sa))
|
|
|
|
connLimitOnClose :: ConnLimiter -> SockAddr -> IO ()
|
|
connLimitOnClose cl sa = atomically (release cl (sockAddrToIPBytes sa))
|
|
|
|
ipBytesFromSockAddr :: SockAddr -> ByteString
|
|
ipBytesFromSockAddr = sockAddrToIPBytes
|