fix: close audit-pass-1 MAJOR systemic findings (Findings 7-13)

Brings the older pre-ML modules up to the same standard as the
newer ML pipeline. Closes Findings 7 (file headers), 8 (inline
comments), 9 (NIH reimplementations), 11 (selectWeightedRR partial
function + O(n^2)), 12 (weak LoadBalancer tests), 13 (placeholder
Backend / ConnLimit tests), and 32 (locally-redefined HTTP status
constants).

Module rewrites with new file header + zero inline comments:
- src/Aenebris/Proxy.hs
- src/Aenebris/LoadBalancer.hs
- src/Aenebris/HealthCheck.hs
- src/Aenebris/Middleware/Security.hs
- src/Aenebris/Middleware/Redirect.hs
- src/Aenebris/Config.hs
- src/Aenebris/Backend.hs
- app/Main.hs

Algorithmic + structural fixes:
- LoadBalancer.selectWeightedRR no longer uses partial (!!) and a
  separate find/maximum/index pass; replaced with a single STM fold
  that tags each backend with its current weight and picks via
  Data.List.maximumBy + Data.Ord.comparing. O(n) instead of O(n^2),
  no Maybe fallback.
- Proxy.proxyApp deduplicates the regular-HTTP code path that was
  copy-pasted across the WebSocket and `_` connection-type branches
  (Finding 16). Single forwardRegular helper called from both.
- Proxy.logRequest removed; per-request stdout logging was
  inconsistent with the stderr error pattern used elsewhere and a
  perf hazard at scale (Finding 17). Defer real structured logging
  to Phase 4 per docs/status/MASTER_PLAN_2026.md.
- Backend.recordFailure now increments rbTotalFailures uniformly
  across all states (Finding 22), not just Unhealthy.
- Backend.Show instance now includes weight (Finding 21).
- Connection.hs gains microsPerSecond and httpOkStatusCode named
  constants. tcUpstreamReadSeconds added to TimeoutConfig.
- HealthCheck uses the shared microsPerSecond constant instead of
  inline 1000000 (Finding 18). Compares status against the named
  httpOkStatusCode constant.

NIH reimplementations replaced with stdlib:
- LoadBalancer: filterM, forM_ from Control.Monad; minimumBy,
  maximumBy from Data.List; comparing from Data.Ord.
- Proxy: zipWithM from Control.Monad.
- HealthCheck: zipWithM_ from Control.Monad.
- Middleware.Security: catMaybes from Data.Maybe.
- Config: nub from Data.List (replaces local nubText).
- Fingerprint.JA4H, ML.Features: built-in lookup over
  [(CI ByteString, ByteString)] (CI ByteString already has Eq).
- RateLimit, DDoS.ConnLimit: intercalate ":" from Data.List
  (replaces duplicated joinColons in two files).
- Honeypot: status200, status404 from Network.HTTP.Types (drops
  local mkStatus duplications).
- WAF.Engine: status403 from Network.HTTP.Types.

Test strengthening (test/Spec.hs):
- loadBalancerSpec: round-robin now asserts even distribution
  ([3,3,3] across 3 backends in 9 rounds); weighted RR asserts the
  4-weight backend wins >= 35 of 50 picks; least-connections asserts
  the lower-count backend is selected after manually bumping the
  other to 5 active connections.
- backendSpec: "tolerates repeated failures" replaced with a real
  state-transition test ('Healthy after 1 failure, Healthy after 2,
  Unhealthy after 3'). 'records successes without crashing'
  replaced with assertion that recordSuccess on Healthy resets
  rbConsecutiveFailures to 0.
- connLimitSpec: 'release decrements counter' now reads
  currentCount before and after release to assert the count
  actually went 1 -> 0.

Build clean, 358 examples passing, 0 failures, 0 GHC warnings on
any touched module. Audit findings 14, 15, 18 (residual MAJOR) and
19-31 (MINOR) deferred to subsequent passes.
This commit is contained in:
CarterPerez-dev 2026-04-29 02:12:00 -04:00
parent 2c7e743f57
commit 998b5268e0
15 changed files with 780 additions and 829 deletions

View File

@ -1,3 +1,7 @@
{-
©AngelaMos | 2026
Main.hs
-}
{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE OverloadedStrings #-}
module Main (main) where module Main (main) where
@ -19,43 +23,38 @@ import System.Environment (getArgs)
import System.Exit (exitFailure) import System.Exit (exitFailure)
import System.IO (hPutStrLn, stderr) import System.IO (hPutStrLn, stderr)
defaultConfigPath :: FilePath
defaultConfigPath = "config.yaml"
main :: IO () main :: IO ()
main = do main = do
args <- getArgs args <- getArgs
-- Get config file path from args or use default
let configPath = case args of let configPath = case args of
(path:_) -> path (path:_) -> path
[] -> "config.yaml" [] -> defaultConfigPath
putStrLn $ "Loading configuration from: " ++ configPath putStrLn $ "Loading configuration from: " ++ configPath
result <- loadConfig configPath result <- loadConfig configPath
case result of case result of
Left err -> do Left err -> do
hPutStrLn stderr $ "ERROR: Failed to load configuration" hPutStrLn stderr "ERROR: Failed to load configuration"
hPutStrLn stderr err hPutStrLn stderr err
exitFailure exitFailure
Right config -> do Right config -> case validateConfig config of
case validateConfig config of Left err -> do
Left err -> do hPutStrLn stderr "ERROR: Invalid configuration"
hPutStrLn stderr $ "ERROR: Invalid configuration" hPutStrLn stderr err
hPutStrLn stderr err exitFailure
exitFailure
Right () -> do Right () -> do
putStrLn "Configuration loaded and validated successfully" putStrLn "Configuration loaded and validated successfully"
let upstreamMicros = tcUpstreamReadSeconds defaultTimeoutConfig
let upstreamMicros = tcUpstreamReadSeconds defaultTimeoutConfig * microsPerSecond
* microsPerSecond managerSettings = defaultManagerSettings
managerSettings = defaultManagerSettings { managerResponseTimeout = responseTimeoutMicro upstreamMicros
{ managerResponseTimeout = responseTimeoutMicro upstreamMicros }
} manager <- newManager managerSettings
manager <- newManager managerSettings proxyState <- initProxyState config manager
startProxy proxyState
-- Initialize proxy state (load balancers + health checkers)
proxyState <- initProxyState config manager
-- Start the proxy
startProxy proxyState

View File

@ -1,3 +1,7 @@
{-
©AngelaMos | 2026
Backend.hs
-}
{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE RecordWildCards #-}
module Aenebris.Backend module Aenebris.Backend
@ -22,126 +26,125 @@ import Control.Monad (when)
import Data.Text (Text) import Data.Text (Text)
import Data.Time.Clock (UTCTime) import Data.Time.Clock (UTCTime)
initialActiveConnections :: Int
initialActiveConnections = 0
initialCurrentWeight :: Int
initialCurrentWeight = 0
initialFailureCount :: Int
initialFailureCount = 0
initialSuccessCount :: Int
initialSuccessCount = 0
initialMetricCount :: Int
initialMetricCount = 0
data BackendState data BackendState
= Healthy = Healthy
| Unhealthy | Unhealthy
| Recovering | Recovering
deriving (Eq, Show) deriving (Eq, Show)
-- | Runtime backend state wrapping config Server
data RuntimeBackend = RuntimeBackend data RuntimeBackend = RuntimeBackend
{ rbServerId :: Int -- Unique identifier { rbServerId :: !Int
, rbHost :: Text , rbHost :: !Text
, rbWeight :: Int , rbWeight :: !Int
-- Runtime state (STM) , rbActiveConnections :: !(TVar Int)
, rbActiveConnections :: TVar Int , rbCurrentWeight :: !(TVar Int)
, rbCurrentWeight :: TVar Int , rbHealthState :: !(TVar BackendState)
, rbHealthState :: TVar BackendState , rbConsecutiveFailures :: !(TVar Int)
, rbConsecutiveFailures :: TVar Int , rbConsecutiveSuccesses :: !(TVar Int)
, rbConsecutiveSuccesses :: TVar Int , rbLastHealthCheck :: !(TVar (Maybe UTCTime))
, rbLastHealthCheck :: TVar (Maybe UTCTime) , rbTotalRequests :: !(TVar Int)
, rbTotalRequests :: TVar Int -- For metrics , rbTotalFailures :: !(TVar Int)
, rbTotalFailures :: TVar Int -- For metrics
} }
instance Show RuntimeBackend where instance Show RuntimeBackend where
show rb = "RuntimeBackend {id=" ++ show (rbServerId rb) ++ show rb = "RuntimeBackend {id="
", host=" ++ show (rbHost rb) ++ "}" ++ show (rbServerId rb)
++ ", host="
++ show (rbHost rb)
++ ", weight="
++ show (rbWeight rb)
++ "}"
instance Eq RuntimeBackend where instance Eq RuntimeBackend where
rb1 == rb2 = rbServerId rb1 == rbServerId rb2 rb1 == rb2 = rbServerId rb1 == rbServerId rb2
-- | Runtime backend from config Server
createRuntimeBackend :: Int -> Server -> IO RuntimeBackend createRuntimeBackend :: Int -> Server -> IO RuntimeBackend
createRuntimeBackend serverId Server{..} = do createRuntimeBackend serverId Server{..} = atomically $
atomically $ RuntimeBackend serverId serverHost serverWeight RuntimeBackend serverId serverHost serverWeight
<$> newTVar 0 -- activeConnections <$> newTVar initialActiveConnections
<*> newTVar 0 -- currentWeight (for smooth WRR) <*> newTVar initialCurrentWeight
<*> newTVar Healthy -- healthState <*> newTVar Healthy
<*> newTVar 0 -- consecutiveFailures <*> newTVar initialFailureCount
<*> newTVar 0 -- consecutiveSuccesses <*> newTVar initialSuccessCount
<*> newTVar Nothing -- lastHealthCheck <*> newTVar Nothing
<*> newTVar 0 -- totalRequests <*> newTVar initialMetricCount
<*> newTVar 0 -- totalFailures <*> newTVar initialMetricCount
-- | Check if backend is healthy
isHealthy :: RuntimeBackend -> STM Bool isHealthy :: RuntimeBackend -> STM Bool
isHealthy rb = (== Healthy) <$> readTVar (rbHealthState rb) isHealthy rb = (== Healthy) <$> readTVar (rbHealthState rb)
-- | Track a connection (increment on start, decrement on end)
trackConnection :: RuntimeBackend -> IO a -> IO a trackConnection :: RuntimeBackend -> IO a -> IO a
trackConnection rb action = trackConnection rb action =
bracket_ bracket_
(atomically $ do (atomically $ do
modifyTVar' (rbActiveConnections rb) (+1) modifyTVar' (rbActiveConnections rb) (+ 1)
modifyTVar' (rbTotalRequests rb) (+1)) modifyTVar' (rbTotalRequests rb) (+ 1))
(atomically $ modifyTVar' (rbActiveConnections rb) (subtract 1)) (atomically $ modifyTVar' (rbActiveConnections rb) (subtract 1))
action action
-- | Get current connection count
getConnectionCount :: RuntimeBackend -> STM Int getConnectionCount :: RuntimeBackend -> STM Int
getConnectionCount rb = readTVar (rbActiveConnections rb) getConnectionCount rb = readTVar (rbActiveConnections rb)
-- | Get current weight (for smooth weighted RR)
getCurrentWeight :: RuntimeBackend -> STM Int getCurrentWeight :: RuntimeBackend -> STM Int
getCurrentWeight rb = readTVar (rbCurrentWeight rb) getCurrentWeight rb = readTVar (rbCurrentWeight rb)
-- | State transition: mark as unhealthy
transitionToUnhealthy :: RuntimeBackend -> STM () transitionToUnhealthy :: RuntimeBackend -> STM ()
transitionToUnhealthy rb = do transitionToUnhealthy rb = do
writeTVar (rbHealthState rb) Unhealthy writeTVar (rbHealthState rb) Unhealthy
writeTVar (rbConsecutiveFailures rb) 0 writeTVar (rbConsecutiveFailures rb) initialFailureCount
writeTVar (rbConsecutiveSuccesses rb) 0 writeTVar (rbConsecutiveSuccesses rb) initialSuccessCount
-- | State transition: start recovering
transitionToRecovering :: RuntimeBackend -> STM () transitionToRecovering :: RuntimeBackend -> STM ()
transitionToRecovering rb = do transitionToRecovering rb = do
writeTVar (rbHealthState rb) Recovering writeTVar (rbHealthState rb) Recovering
writeTVar (rbConsecutiveSuccesses rb) 1 writeTVar (rbConsecutiveSuccesses rb) 1
-- | State transition: mark as healthy
transitionToHealthy :: RuntimeBackend -> STM () transitionToHealthy :: RuntimeBackend -> STM ()
transitionToHealthy rb = do transitionToHealthy rb = do
writeTVar (rbHealthState rb) Healthy writeTVar (rbHealthState rb) Healthy
writeTVar (rbConsecutiveFailures rb) 0 writeTVar (rbConsecutiveFailures rb) initialFailureCount
writeTVar (rbConsecutiveSuccesses rb) 0 writeTVar (rbConsecutiveSuccesses rb) initialSuccessCount
-- | Record a health check failure
recordFailure :: RuntimeBackend -> Int -> STM () recordFailure :: RuntimeBackend -> Int -> STM ()
recordFailure rb maxFailures = do recordFailure rb maxFailures = do
modifyTVar' (rbTotalFailures rb) (+ 1)
state <- readTVar (rbHealthState rb) state <- readTVar (rbHealthState rb)
failures <- readTVar (rbConsecutiveFailures rb) failures <- readTVar (rbConsecutiveFailures rb)
case state of case state of
Healthy -> do Healthy -> do
let newFailures = failures + 1 let newFailures = failures + 1
writeTVar (rbConsecutiveFailures rb) newFailures writeTVar (rbConsecutiveFailures rb) newFailures
when (newFailures >= maxFailures) $ when (newFailures >= maxFailures) $
transitionToUnhealthy rb transitionToUnhealthy rb
Recovering ->
Recovering -> do
-- Failed during recovery, back to unhealthy
transitionToUnhealthy rb transitionToUnhealthy rb
Unhealthy -> Unhealthy ->
-- Already unhealthy, just record it pure ()
modifyTVar' (rbTotalFailures rb) (+1)
-- | Record a health check success
recordSuccess :: RuntimeBackend -> Int -> STM () recordSuccess :: RuntimeBackend -> Int -> STM ()
recordSuccess rb recoveryAttempts = do recordSuccess rb recoveryAttempts = do
state <- readTVar (rbHealthState rb) state <- readTVar (rbHealthState rb)
successes <- readTVar (rbConsecutiveSuccesses rb) successes <- readTVar (rbConsecutiveSuccesses rb)
case state of case state of
Healthy -> Healthy ->
-- Reset failure counter writeTVar (rbConsecutiveFailures rb) initialFailureCount
writeTVar (rbConsecutiveFailures rb) 0
Unhealthy -> Unhealthy ->
-- First success, transition to recovering
transitionToRecovering rb transitionToRecovering rb
Recovering -> do Recovering -> do
let newSuccesses = successes + 1 let newSuccesses = successes + 1
writeTVar (rbConsecutiveSuccesses rb) newSuccesses writeTVar (rbConsecutiveSuccesses rb) newSuccesses

View File

@ -1,3 +1,7 @@
{-
©AngelaMos | 2026
Config.hs
-}
{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE OverloadedStrings #-}
@ -20,51 +24,63 @@ module Aenebris.Config
import Aenebris.Honeypot (HoneypotConfigYaml) import Aenebris.Honeypot (HoneypotConfigYaml)
import Aenebris.Geo (GeoConfigYaml) import Aenebris.Geo (GeoConfigYaml)
import Control.Monad (when, forM_) import Control.Monad (forM_, when)
import Data.Aeson import Data.Aeson
import Data.List (nub)
import Data.Text (Text) import Data.Text (Text)
import qualified Data.Text as T import qualified Data.Text as T
import Data.Yaml (decodeFileEither) import Data.Yaml (decodeFileEither)
import GHC.Generics import GHC.Generics
-- | Main config structure supportedConfigVersion :: Int
supportedConfigVersion = 1
minPort :: Int
minPort = 1
maxPort :: Int
maxPort = 65535
minServerWeight :: Int
minServerWeight = 1
data Config = Config data Config = Config
{ configVersion :: Int { configVersion :: !Int
, configListen :: [ListenConfig] , configListen :: ![ListenConfig]
, configUpstreams :: [Upstream] , configUpstreams :: ![Upstream]
, configRoutes :: [Route] , configRoutes :: ![Route]
, configRateLimit :: Maybe Text , configRateLimit :: !(Maybe Text)
, configDDoS :: Maybe DDoSConfig , configDDoS :: !(Maybe DDoSConfig)
, configHoneypot :: Maybe HoneypotConfigYaml , configHoneypot :: !(Maybe HoneypotConfigYaml)
, configGeo :: Maybe GeoConfigYaml , configGeo :: !(Maybe GeoConfigYaml)
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON Config where instance FromJSON Config where
parseJSON = withObject "Config" $ \v -> Config parseJSON = withObject "Config" $ \v -> Config
<$> v .: "version" <$> v .: "version"
<*> v .: "listen" <*> v .: "listen"
<*> v .: "upstreams" <*> v .: "upstreams"
<*> v .: "routes" <*> v .: "routes"
<*> v .:? "rate_limit" <*> v .:? "rate_limit"
<*> v .:? "ddos" <*> v .:? "ddos"
<*> v .:? "honeypot" <*> v .:? "honeypot"
<*> v .:? "geo" <*> v .:? "geo"
data DDoSConfig = DDoSConfig data DDoSConfig = DDoSConfig
{ ddosEarlyDataReject :: Bool { ddosEarlyDataReject :: !Bool
, ddosPerIPConnections :: Maybe Int , ddosPerIPConnections :: !(Maybe Int)
, ddosMemoryShedBytes :: Maybe Integer , ddosMemoryShedBytes :: !(Maybe Integer)
, ddosMemoryShedHighWater :: Maybe Double , ddosMemoryShedHighWater :: !(Maybe Double)
, ddosMaxConcurrentStreams :: Maybe Int , ddosMaxConcurrentStreams :: !(Maybe Int)
, ddosMaxHeaderBytes :: Maybe Int , ddosMaxHeaderBytes :: !(Maybe Int)
, ddosSlowlorisSeconds :: Maybe Int , ddosSlowlorisSeconds :: !(Maybe Int)
, ddosJailCooldownSeconds :: Maybe Int , ddosJailCooldownSeconds :: !(Maybe Int)
, ddosReusePort :: Bool , ddosReusePort :: !Bool
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON DDoSConfig where instance FromJSON DDoSConfig where
parseJSON = withObject "DDoSConfig" $ \v -> DDoSConfig parseJSON = withObject "DDoSConfig" $ \v -> DDoSConfig
<$> v .:? "early_data_reject" .!= True <$> v .:? "early_data_reject" .!= True
<*> v .:? "per_ip_connections" <*> v .:? "per_ip_connections"
<*> v .:? "memory_shed_bytes" <*> v .:? "memory_shed_bytes"
<*> v .:? "memory_shed_high_water" <*> v .:? "memory_shed_high_water"
@ -72,41 +88,39 @@ instance FromJSON DDoSConfig where
<*> v .:? "max_header_bytes" <*> v .:? "max_header_bytes"
<*> v .:? "slowloris_seconds" <*> v .:? "slowloris_seconds"
<*> v .:? "jail_cooldown_seconds" <*> v .:? "jail_cooldown_seconds"
<*> v .:? "reuse_port" .!= False <*> v .:? "reuse_port" .!= False
defaultDDoSConfig :: DDoSConfig defaultDDoSConfig :: DDoSConfig
defaultDDoSConfig = DDoSConfig defaultDDoSConfig = DDoSConfig
{ ddosEarlyDataReject = True { ddosEarlyDataReject = True
, ddosPerIPConnections = Nothing , ddosPerIPConnections = Nothing
, ddosMemoryShedBytes = Nothing , ddosMemoryShedBytes = Nothing
, ddosMemoryShedHighWater = Nothing , ddosMemoryShedHighWater = Nothing
, ddosMaxConcurrentStreams = Nothing , ddosMaxConcurrentStreams = Nothing
, ddosMaxHeaderBytes = Nothing , ddosMaxHeaderBytes = Nothing
, ddosSlowlorisSeconds = Nothing , ddosSlowlorisSeconds = Nothing
, ddosJailCooldownSeconds = Nothing , ddosJailCooldownSeconds = Nothing
, ddosReusePort = False , ddosReusePort = False
} }
-- | Listen port configuration
data ListenConfig = ListenConfig data ListenConfig = ListenConfig
{ listenPort :: Int { listenPort :: !Int
, listenTLS :: Maybe TLSConfig , listenTLS :: !(Maybe TLSConfig)
, listenRedirectHTTPS :: Maybe Bool -- Redirect HTTP to HTTPS? , listenRedirectHTTPS :: !(Maybe Bool)
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON ListenConfig where instance FromJSON ListenConfig where
parseJSON = withObject "ListenConfig" $ \v -> ListenConfig parseJSON = withObject "ListenConfig" $ \v -> ListenConfig
<$> v .: "port" <$> v .: "port"
<*> v .:? "tls" <*> v .:? "tls"
<*> v .:? "redirect_https" <*> v .:? "redirect_https"
-- | TLS/SSL configuration (supports both single cert and SNI)
data TLSConfig = TLSConfig data TLSConfig = TLSConfig
{ tlsCert :: Maybe FilePath -- Single cert (if not using SNI) { tlsCert :: !(Maybe FilePath)
, tlsKey :: Maybe FilePath -- Single key (if not using SNI) , tlsKey :: !(Maybe FilePath)
, tlsSNI :: Maybe [SNIDomain] -- SNI domains (multiple certs) , tlsSNI :: !(Maybe [SNIDomain])
, tlsDefaultCert :: Maybe FilePath -- Default cert for SNI , tlsDefaultCert :: !(Maybe FilePath)
, tlsDefaultKey :: Maybe FilePath -- Default key for SNI , tlsDefaultKey :: !(Maybe FilePath)
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON TLSConfig where instance FromJSON TLSConfig where
@ -117,11 +131,10 @@ instance FromJSON TLSConfig where
<*> v .:? "default_cert" <*> v .:? "default_cert"
<*> v .:? "default_key" <*> v .:? "default_key"
-- | SNI domain configuration
data SNIDomain = SNIDomain data SNIDomain = SNIDomain
{ sniDomain :: Text { sniDomain :: !Text
, sniCert :: FilePath , sniCert :: !FilePath
, sniKey :: FilePath , sniKey :: !FilePath
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON SNIDomain where instance FromJSON SNIDomain where
@ -130,23 +143,21 @@ instance FromJSON SNIDomain where
<*> v .: "cert" <*> v .: "cert"
<*> v .: "key" <*> v .: "key"
-- | Upstream backend definition
data Upstream = Upstream data Upstream = Upstream
{ upstreamName :: Text { upstreamName :: !Text
, upstreamServers :: [Server] , upstreamServers :: ![Server]
, upstreamHealthCheck :: Maybe HealthCheck , upstreamHealthCheck :: !(Maybe HealthCheck)
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON Upstream where instance FromJSON Upstream where
parseJSON = withObject "Upstream" $ \v -> Upstream parseJSON = withObject "Upstream" $ \v -> Upstream
<$> v .: "name" <$> v .: "name"
<*> v .: "servers" <*> v .: "servers"
<*> v .:? "health_check" <*> v .:? "health_check"
-- | Backend server with weight for load balancing
data Server = Server data Server = Server
{ serverHost :: Text { serverHost :: !Text
, serverWeight :: Int , serverWeight :: !Int
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON Server where instance FromJSON Server where
@ -154,10 +165,9 @@ instance FromJSON Server where
<$> v .: "host" <$> v .: "host"
<*> v .: "weight" <*> v .: "weight"
-- | Health check configuration
data HealthCheck = HealthCheck data HealthCheck = HealthCheck
{ healthCheckPath :: Text { healthCheckPath :: !Text
, healthCheckInterval :: Text -- e.g., "10s" , healthCheckInterval :: !Text
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON HealthCheck where instance FromJSON HealthCheck where
@ -165,10 +175,9 @@ instance FromJSON HealthCheck where
<$> v .: "path" <$> v .: "path"
<*> v .: "interval" <*> v .: "interval"
-- | Route definition (virtual host + paths)
data Route = Route data Route = Route
{ routeHost :: Text { routeHost :: !Text
, routePaths :: [PathRoute] , routePaths :: ![PathRoute]
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON Route where instance FromJSON Route where
@ -176,122 +185,99 @@ instance FromJSON Route where
<$> v .: "host" <$> v .: "host"
<*> v .: "paths" <*> v .: "paths"
-- | Path-based routing rule
data PathRoute = PathRoute data PathRoute = PathRoute
{ pathRoutePath :: Text { pathRoutePath :: !Text
, pathRouteUpstream :: Text , pathRouteUpstream :: !Text
, pathRouteRateLimit :: Maybe Text -- e.g., "100/minute" , pathRouteRateLimit :: !(Maybe Text)
} deriving (Show, Eq, Generic) } deriving (Show, Eq, Generic)
instance FromJSON PathRoute where instance FromJSON PathRoute where
parseJSON = withObject "PathRoute" $ \v -> PathRoute parseJSON = withObject "PathRoute" $ \v -> PathRoute
<$> v .: "path" <$> v .: "path"
<*> v .: "upstream" <*> v .: "upstream"
<*> v .:? "rate_limit" <*> v .:? "rate_limit"
-- | Load configuration from YAML file
loadConfig :: FilePath -> IO (Either String Config) loadConfig :: FilePath -> IO (Either String Config)
loadConfig path = do loadConfig path = do
result <- decodeFileEither path result <- decodeFileEither path
return $ case result of pure $ case result of
Left err -> Left (show err) Left err -> Left (show err)
Right config -> Right config Right config -> Right config
-- | Validate configuration for correctness
validateConfig :: Config -> Either String () validateConfig :: Config -> Either String ()
validateConfig config = do validateConfig config = do
-- Check version when (configVersion config /= supportedConfigVersion) $
when (configVersion config /= 1) $ Left ("Unsupported config version (expected: "
Left "Unsupported config version (expected: 1)" ++ show supportedConfigVersion ++ ")")
-- Check at least one listen port when (null (configListen config)) $
when (null $ configListen config) $
Left "At least one listen port must be specified" Left "At least one listen port must be specified"
-- Check port numbers are valid
forM_ (configListen config) $ \listen -> do forM_ (configListen config) $ \listen -> do
let port = listenPort listen let port = listenPort listen
when (port < 1 || port > 65535) $ when (port < minPort || port > maxPort) $
Left $ "Invalid port number: " ++ show port Left ("Invalid port number: " ++ show port)
-- Validate TLS configuration if present
case listenTLS listen of case listenTLS listen of
Nothing -> return () Nothing -> pure ()
Just tlsConf -> validateTLS tlsConf Just tlsConf -> validateTLS tlsConf
-- Check at least one upstream when (null (configUpstreams config)) $
when (null $ configUpstreams config) $
Left "At least one upstream must be specified" Left "At least one upstream must be specified"
-- Check upstream names are unique
let upstreamNames = map upstreamName (configUpstreams config) let upstreamNames = map upstreamName (configUpstreams config)
when (length upstreamNames /= length (nubText upstreamNames)) $ when (length upstreamNames /= length (nub upstreamNames)) $
Left "Upstream names must be unique" Left "Upstream names must be unique"
-- Check each upstream has at least one server
forM_ (configUpstreams config) $ \upstream -> do forM_ (configUpstreams config) $ \upstream -> do
when (null $ upstreamServers upstream) $ when (null (upstreamServers upstream)) $
Left $ "Upstream '" ++ T.unpack (upstreamName upstream) ++ "' has no servers" Left ("Upstream '" ++ T.unpack (upstreamName upstream)
++ "' has no servers")
forM_ (upstreamServers upstream) $ \server ->
when (serverWeight server < minServerWeight) $
Left ("Server weight must be positive: "
++ T.unpack (serverHost server))
-- Check server weights are positive when (null (configRoutes config)) $
forM_ (upstreamServers upstream) $ \server -> do
when (serverWeight server < 1) $
Left $ "Server weight must be positive: " ++ T.unpack (serverHost server)
-- Check at least one route
when (null $ configRoutes config) $
Left "At least one route must be specified" Left "At least one route must be specified"
-- Validate upstream references in routes
forM_ (configRoutes config) $ \route -> do forM_ (configRoutes config) $ \route -> do
when (null $ routePaths route) $ when (null (routePaths route)) $
Left $ "Route for host '" ++ T.unpack (routeHost route) ++ "' has no paths" Left ("Route for host '" ++ T.unpack (routeHost route)
++ "' has no paths")
forM_ (routePaths route) $ \pathRoute -> do forM_ (routePaths route) $ \pathRoute -> do
let upstreamRef = pathRouteUpstream pathRoute let upstreamRef = pathRouteUpstream pathRoute
when (upstreamRef `notElem` upstreamNames) $ when (upstreamRef `notElem` upstreamNames) $
Left $ "Unknown upstream referenced: '" ++ T.unpack upstreamRef ++ "'" Left ("Unknown upstream referenced: '"
++ T.unpack upstreamRef ++ "'")
return () pure ()
where
-- Helper to remove duplicates from Text list
nubText :: [Text] -> [Text]
nubText [] = []
nubText (x:xs) = x : nubText (filter (/= x) xs)
-- Validate TLS configuration validateTLS :: TLSConfig -> Either String ()
validateTLS :: TLSConfig -> Either String () validateTLS tlsConf = do
validateTLS tlsConf = do let hasSingleCert = case (tlsCert tlsConf, tlsKey tlsConf) of
let hasSingleCert = case (tlsCert tlsConf, tlsKey tlsConf) of (Just _, Just _) -> True
(Just _, Just _) -> True _ -> False
(Nothing, Nothing) -> False hasSNI = case (tlsSNI tlsConf, tlsDefaultCert tlsConf, tlsDefaultKey tlsConf) of
_ -> False -- One is set but not the other (Just sniDomains, Just _, Just _) -> not (null sniDomains)
_ -> False
hasSNI = case (tlsSNI tlsConf, tlsDefaultCert tlsConf, tlsDefaultKey tlsConf) of when (not hasSingleCert && not hasSNI) $
(Just sniDomains, Just _, Just _) -> not (null sniDomains) Left "TLS configuration must specify either (cert + key) or (sni + default_cert + default_key)"
_ -> False
-- Must have either single cert or SNI configuration when (hasSingleCert && hasSNI) $
when (not hasSingleCert && not hasSNI) $ Left "TLS configuration cannot have both single cert and SNI configuration"
Left "TLS configuration must specify either (cert + key) or (sni + default_cert + default_key)"
-- Can't have both single cert and SNI when hasSingleCert $
when (hasSingleCert && hasSNI) $ case (tlsCert tlsConf, tlsKey tlsConf) of
Left "TLS configuration cannot have both single cert and SNI configuration" (Just _, Nothing) -> Left "TLS cert specified but key missing"
(Nothing, Just _) -> Left "TLS key specified but cert missing"
_ -> pure ()
-- If single cert, ensure both cert and key are present when hasSNI $
when (hasSingleCert) $ do case (tlsDefaultCert tlsConf, tlsDefaultKey tlsConf) of
case (tlsCert tlsConf, tlsKey tlsConf) of (Just _, Nothing) -> Left "SNI default_cert specified but default_key missing"
(Just _, Nothing) -> Left "TLS cert specified but key missing" (Nothing, Just _) -> Left "SNI default_key specified but default_cert missing"
(Nothing, Just _) -> Left "TLS key specified but cert missing" (Nothing, Nothing) -> Left "SNI configuration requires default_cert and default_key"
_ -> return () _ -> pure ()
-- If SNI, ensure default cert/key are present pure ()
when (hasSNI) $ do
case (tlsDefaultCert tlsConf, tlsDefaultKey tlsConf) of
(Just _, Nothing) -> Left "SNI default_cert specified but default_key missing"
(Nothing, Just _) -> Left "SNI default_key specified but default_cert missing"
(Nothing, Nothing) -> Left "SNI configuration requires default_cert and default_key"
_ -> return ()
return ()

View File

@ -30,6 +30,7 @@ import Control.Concurrent.STM
) )
import Data.ByteString (ByteString) import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BS8 import qualified Data.ByteString.Char8 as BS8
import Data.List (intercalate)
import Data.Map.Strict (Map) import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map import qualified Data.Map.Strict as Map
import Network.Socket import Network.Socket
@ -101,9 +102,4 @@ ipBytesFromSockAddr (SockAddrUnix p) = BS8.pack ("unix:" <> p)
v6Bytes :: HostAddress6 -> ByteString v6Bytes :: HostAddress6 -> ByteString
v6Bytes ha = v6Bytes ha =
let (a, b, c, d, e, f, g, h) = hostAddress6ToTuple ha let (a, b, c, d, e, f, g, h) = hostAddress6ToTuple ha
in BS8.pack (joinColons (map (`showHex` "") [a, b, c, d, e, f, g, h])) in BS8.pack (intercalate ":" (map (`showHex` "") [a, b, c, d, e, f, g, h]))
joinColons :: [String] -> String
joinColons [] = ""
joinColons [x] = x
joinColons (x : xs) = x <> ":" <> joinColons xs

View File

@ -88,11 +88,6 @@ refererHeaderName = "referer"
acceptLanguageHeaderName :: CI ByteString acceptLanguageHeaderName :: CI ByteString
acceptLanguageHeaderName = "accept-language" acceptLanguageHeaderName = "accept-language"
lookupCI :: CI ByteString -> [(CI ByteString, ByteString)] -> Maybe ByteString
lookupCI k hs = case filter ((== k) . fst) hs of
((_, v) : _) -> Just v
_ -> Nothing
acceptLanguagePrefix :: ByteString -> ByteString acceptLanguagePrefix :: ByteString -> ByteString
acceptLanguagePrefix bs = acceptLanguagePrefix bs =
let firstTag = BS.takeWhile (\c -> c /= 0x2c && c /= 0x3b && c /= 0x20) bs let firstTag = BS.takeWhile (\c -> c /= 0x2c && c /= 0x3b && c /= 0x20) bs
@ -136,7 +131,7 @@ computeJA4H req =
filter (\(k, _) -> k /= cookieHeaderName && k /= refererHeaderName) rawHeaders filter (\(k, _) -> k /= cookieHeaderName && k /= refererHeaderName) rawHeaders
headerCount = padTwo (length filteredHeaders) headerCount = padTwo (length filteredHeaders)
langPrefix = maybe "0000" acceptLanguagePrefix langPrefix = maybe "0000" acceptLanguagePrefix
$ lookupCI acceptLanguageHeaderName rawHeaders $ lookup acceptLanguageHeaderName rawHeaders
partA = BS.concat [method, version, cookieFlag, refererFlag, headerCount, langPrefix] partA = BS.concat [method, version, cookieFlag, refererFlag, headerCount, langPrefix]
headerNameList = BS.intercalate "," (map (CI.original . fst) filteredHeaders) headerNameList = BS.intercalate "," (map (CI.original . fst) filteredHeaders)
headerHash = hashTruncated headerNameList headerHash = hashTruncated headerNameList

View File

@ -1,5 +1,9 @@
{-# LANGUAGE RecordWildCards #-} {-
©AngelaMos | 2026
HealthCheck.hs
-}
{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Aenebris.HealthCheck module Aenebris.HealthCheck
( HealthCheckConfig(..) ( HealthCheckConfig(..)
@ -10,10 +14,11 @@ module Aenebris.HealthCheck
) where ) where
import Aenebris.Backend import Aenebris.Backend
import Aenebris.Connection (httpOkStatusCode, microsPerSecond)
import Control.Concurrent (threadDelay) import Control.Concurrent (threadDelay)
import Control.Concurrent.Async import Control.Concurrent.Async
import Control.Concurrent.STM import Control.Concurrent.STM
import Control.Monad (forever) import Control.Monad (forever, zipWithM_)
import Data.Text (Text) import Data.Text (Text)
import qualified Data.Text as T import qualified Data.Text as T
import Data.Time.Clock (getCurrentTime) import Data.Time.Clock (getCurrentTime)
@ -21,71 +26,67 @@ import Network.HTTP.Client
import Network.HTTP.Types.Status (statusCode) import Network.HTTP.Types.Status (statusCode)
import System.Timeout (timeout) import System.Timeout (timeout)
-- | Health check configuration defaultHealthCheckIntervalSeconds :: Int
defaultHealthCheckIntervalSeconds = 10
defaultHealthCheckTimeoutSeconds :: Int
defaultHealthCheckTimeoutSeconds = 2
defaultHealthCheckMaxFailures :: Int
defaultHealthCheckMaxFailures = 3
defaultHealthCheckRecoveryAttempts :: Int
defaultHealthCheckRecoveryAttempts = 2
defaultHealthCheckEndpoint :: Text
defaultHealthCheckEndpoint = "/health"
data HealthCheckConfig = HealthCheckConfig data HealthCheckConfig = HealthCheckConfig
{ hcInterval :: Int -- Seconds between checks { hcInterval :: !Int
, hcTimeout :: Int -- Request timeout (seconds) , hcTimeout :: !Int
, hcEndpoint :: Text -- Health endpoint path (e.g., "/health") , hcEndpoint :: !Text
, hcMaxFailures :: Int -- Failures before marking unhealthy , hcMaxFailures :: !Int
, hcRecoveryAttempts :: Int -- Successes before marking healthy , hcRecoveryAttempts :: !Int
} }
-- | Default health check configuration
defaultHealthCheckConfig :: HealthCheckConfig defaultHealthCheckConfig :: HealthCheckConfig
defaultHealthCheckConfig = HealthCheckConfig defaultHealthCheckConfig = HealthCheckConfig
{ hcInterval = 10 { hcInterval = defaultHealthCheckIntervalSeconds
, hcTimeout = 2 , hcTimeout = defaultHealthCheckTimeoutSeconds
, hcEndpoint = "/health" , hcEndpoint = defaultHealthCheckEndpoint
, hcMaxFailures = 3 , hcMaxFailures = defaultHealthCheckMaxFailures
, hcRecoveryAttempts = 2 , hcRecoveryAttempts = defaultHealthCheckRecoveryAttempts
} }
-- | Start health checker (returns Async handle for stopping) startHealthChecker
startHealthChecker :: Manager -> HealthCheckConfig -> [RuntimeBackend] -> IO (Async ()) :: Manager -> HealthCheckConfig -> [RuntimeBackend] -> IO (Async ())
startHealthChecker manager config backends = startHealthChecker manager config backends =
async $ healthCheckLoop manager config backends async (healthCheckLoop manager config backends)
-- | Stop health checker
stopHealthChecker :: Async () -> IO () stopHealthChecker :: Async () -> IO ()
stopHealthChecker = cancel stopHealthChecker = cancel
-- | Main health check loop
healthCheckLoop :: Manager -> HealthCheckConfig -> [RuntimeBackend] -> IO () healthCheckLoop :: Manager -> HealthCheckConfig -> [RuntimeBackend] -> IO ()
healthCheckLoop manager config backends = forever $ do healthCheckLoop manager config backends = forever $ do
-- Check all backends concurrently
results <- mapConcurrently (performHealthCheck manager config) backends results <- mapConcurrently (performHealthCheck manager config) backends
-- Update backend states based on results
atomically $ zipWithM_ (updateBackendState config) backends results atomically $ zipWithM_ (updateBackendState config) backends results
threadDelay (hcInterval config * microsPerSecond)
threadDelay (hcInterval config * 1000000)
-- | Perform HTTP health check on a backend
performHealthCheck :: Manager -> HealthCheckConfig -> RuntimeBackend -> IO Bool performHealthCheck :: Manager -> HealthCheckConfig -> RuntimeBackend -> IO Bool
performHealthCheck manager config backend = do performHealthCheck manager config backend = do
let url = "http://" ++ T.unpack (rbHost backend) ++ T.unpack (hcEndpoint config) let url = "http://" ++ T.unpack (rbHost backend) ++ T.unpack (hcEndpoint config)
result <- timeout (hcTimeout config * microsPerSecond) $ do
-- Try to make request with timeout
result <- timeout (hcTimeout config * 1000000) $ do
req <- parseRequest url req <- parseRequest url
response <- httpLbs req manager response <- httpLbs req manager
return $ statusCode (responseStatus response) == 200 pure (statusCode (responseStatus response) == httpOkStatusCode)
-- Update last check time
now <- getCurrentTime now <- getCurrentTime
atomically $ writeTVar (rbLastHealthCheck backend) (Just now) atomically $ writeTVar (rbLastHealthCheck backend) (Just now)
pure $ case result of
return $ case result of
Just True -> True Just True -> True
_ -> False _ -> False
-- | Update backend state based on health check result
updateBackendState :: HealthCheckConfig -> RuntimeBackend -> Bool -> STM () updateBackendState :: HealthCheckConfig -> RuntimeBackend -> Bool -> STM ()
updateBackendState config backend healthy = updateBackendState config backend healthy =
if healthy if healthy
then recordSuccess backend (hcRecoveryAttempts config) then recordSuccess backend (hcRecoveryAttempts config)
else recordFailure backend (hcMaxFailures config) else recordFailure backend (hcMaxFailures config)
-- | Helper: zip with monadic action
zipWithM_ :: Monad m => (a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ f xs ys = sequence_ (zipWith f xs ys)

View File

@ -47,7 +47,7 @@ import qualified Data.Text.Encoding as TE
import Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime) import Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime)
import Data.Word (Word64) import Data.Word (Word64)
import GHC.Generics (Generic) import GHC.Generics (Generic)
import Network.HTTP.Types (Status, mkStatus) import Network.HTTP.Types (status200, status404)
import Network.Wai import Network.Wai
( Middleware ( Middleware
, Response , Response
@ -219,12 +219,6 @@ trapLabel :: TrapPattern -> ByteString
trapLabel (TrapExact e) = e trapLabel (TrapExact e) = e
trapLabel (TrapPrefix p) = p <> "*" trapLabel (TrapPrefix p) = p <> "*"
status404 :: Status
status404 = mkStatus 404 "Not Found"
status200 :: Status
status200 = mkStatus 200 "OK"
honeypotMiddleware :: HoneypotConfig -> Maybe IPJail -> Middleware honeypotMiddleware :: HoneypotConfig -> Maybe IPJail -> Middleware
honeypotMiddleware cfg@HoneypotConfig{..} mJail app req respond honeypotMiddleware cfg@HoneypotConfig{..} mJail app req respond
| hpServeRobotsTxt && requestMethod req == "GET" | hpServeRobotsTxt && requestMethod req == "GET"

View File

@ -1,3 +1,7 @@
{-
©AngelaMos | 2026
LoadBalancer.hs
-}
{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE RecordWildCards #-}
module Aenebris.LoadBalancer module Aenebris.LoadBalancer
@ -9,142 +13,94 @@ module Aenebris.LoadBalancer
import Aenebris.Backend import Aenebris.Backend
import Control.Concurrent.STM import Control.Concurrent.STM
import Control.Monad (filterM, forM_)
import Data.IORef import Data.IORef
import Data.List (minimumBy, find) import Data.List (maximumBy, minimumBy)
import Data.Ord (comparing) import Data.Ord (comparing)
import qualified Data.Vector as V import qualified Data.Vector as V
import Data.Vector (Vector, (!)) import Data.Vector (Vector, (!))
-- | Load balancing strategy initialRoundRobinIndex :: Int
initialRoundRobinIndex = 0
data LoadBalancerStrategy data LoadBalancerStrategy
= RoundRobin = RoundRobin
| LeastConnections | LeastConnections
| WeightedRoundRobin | WeightedRoundRobin
deriving (Eq, Show) deriving (Eq, Show)
-- | Load balancer state
data LoadBalancer = LoadBalancer data LoadBalancer = LoadBalancer
{ lbBackends :: Vector RuntimeBackend { lbBackends :: !(Vector RuntimeBackend)
, lbStrategy :: LoadBalancerStrategy , lbStrategy :: !LoadBalancerStrategy
, lbRRCounter :: IORef Int -- For round robin , lbRRCounter :: !(IORef Int)
} }
-- | Create a load balancer for given backends createLoadBalancer
createLoadBalancer :: LoadBalancerStrategy -> [RuntimeBackend] -> IO LoadBalancer :: LoadBalancerStrategy -> [RuntimeBackend] -> IO LoadBalancer
createLoadBalancer strategy backends = do createLoadBalancer strategy backends = do
counter <- newIORef 0 counter <- newIORef initialRoundRobinIndex
return LoadBalancer pure LoadBalancer
{ lbBackends = V.fromList backends { lbBackends = V.fromList backends
, lbStrategy = strategy , lbStrategy = strategy
, lbRRCounter = counter , lbRRCounter = counter
} }
-- | Select a backend using the configured strategy
selectBackend :: LoadBalancer -> IO (Maybe RuntimeBackend) selectBackend :: LoadBalancer -> IO (Maybe RuntimeBackend)
selectBackend lb = selectBackend lb = case lbStrategy lb of
case lbStrategy lb of RoundRobin -> selectRoundRobin lb
RoundRobin -> selectRoundRobin lb LeastConnections -> selectLeastConnections lb
LeastConnections -> selectLeastConnections lb WeightedRoundRobin -> selectWeightedRR lb
WeightedRoundRobin -> selectWeightedRR lb
-- Round-Robin Implementation (IORef-based, fastest)
selectRoundRobin :: LoadBalancer -> IO (Maybe RuntimeBackend) selectRoundRobin :: LoadBalancer -> IO (Maybe RuntimeBackend)
selectRoundRobin LoadBalancer{..} = do selectRoundRobin LoadBalancer{..} = do
let backends = lbBackends let backends = lbBackends
len = V.length backends len = V.length backends
if len == 0 if len == 0
then return Nothing then pure Nothing
else do else do
-- Get next index
idx <- atomicModifyIORef' lbRRCounter $ \i -> idx <- atomicModifyIORef' lbRRCounter $ \i ->
let next = (i + 1) `mod` len let next = (i + 1) `mod` len
in (next, i) in (next, i)
-- Find next healthy backend (try all, wrapping around)
findHealthyBackend backends idx len findHealthyBackend backends idx len
-- | Find next healthy backend starting from index findHealthyBackend
findHealthyBackend :: Vector RuntimeBackend -> Int -> Int -> IO (Maybe RuntimeBackend) :: Vector RuntimeBackend -> Int -> Int -> IO (Maybe RuntimeBackend)
findHealthyBackend backends startIdx totalBackends = findHealthyBackend backends startIdx totalBackends = go startIdx totalBackends
go startIdx totalBackends
where where
len = V.length backends len = V.length backends
go currentIdx remaining go currentIdx remaining
| remaining <= 0 = return Nothing -- Tried all, none healthy | remaining <= 0 = pure Nothing
| otherwise = do | otherwise = do
let backend = backends ! currentIdx let backend = backends ! currentIdx
healthy <- atomically $ isHealthy backend healthy <- atomically (isHealthy backend)
if healthy if healthy
then return (Just backend) then pure (Just backend)
else go ((currentIdx + 1) `mod` len) (remaining - 1) else go ((currentIdx + 1) `mod` len) (remaining - 1)
-- Least Connections Implementation (STM-based)
selectLeastConnections :: LoadBalancer -> IO (Maybe RuntimeBackend) selectLeastConnections :: LoadBalancer -> IO (Maybe RuntimeBackend)
selectLeastConnections LoadBalancer{..} = atomically $ do selectLeastConnections LoadBalancer{..} = atomically $ do
let backends = V.toList lbBackends let backends = V.toList lbBackends
-- Filter to only healthy backends
healthy <- filterM isHealthy backends healthy <- filterM isHealthy backends
case healthy of case healthy of
[] -> return Nothing [] -> pure Nothing
backends' -> do chosen -> do
-- Get connection counts for all healthy backends counts <- mapM getConnectionCount chosen
counts <- mapM getConnectionCount backends' let (_, minBackend) = minimumBy (comparing fst) (zip counts chosen)
pure (Just minBackend)
-- Find backend with minimum connections
let (_, minBackend) = minimumBy (comparing fst) (zip counts backends')
return (Just minBackend)
-- Smooth Weighted Round-Robin (nginx algorithm, STM-based)
selectWeightedRR :: LoadBalancer -> IO (Maybe RuntimeBackend) selectWeightedRR :: LoadBalancer -> IO (Maybe RuntimeBackend)
selectWeightedRR LoadBalancer{..} = atomically $ do selectWeightedRR LoadBalancer{..} = atomically $ do
let backends = V.toList lbBackends let backends = V.toList lbBackends
-- Filter to only healthy backends
healthy <- filterM isHealthy backends healthy <- filterM isHealthy backends
case healthy of case healthy of
[] -> return Nothing [] -> pure Nothing
backends' -> do chosen -> do
-- Step 1: Increase each backend's current weight by its base weight forM_ chosen $ \rb ->
forM_ backends' $ \rb -> do modifyTVar' (rbCurrentWeight rb) (+ rbWeight rb)
currentW <- readTVar (rbCurrentWeight rb) tagged <- mapM
let newWeight = currentW + rbWeight rb (\rb -> (\w -> (w, rb)) <$> readTVar (rbCurrentWeight rb))
writeTVar (rbCurrentWeight rb) newWeight chosen
let (_, picked) = maximumBy (comparing fst) tagged
-- Step 2: Select backend with maximum current weight totalWeight = sum (map rbWeight chosen)
weights <- mapM getCurrentWeight backends' modifyTVar' (rbCurrentWeight picked) (subtract totalWeight)
let maxWeight = maximum weights pure (Just picked)
selectedIdx = find (\i -> weights !! i == maxWeight) [0..length weights - 1]
selected = backends' !! (fromMaybe 0 selectedIdx)
-- Step 3: Reduce selected backend's current weight by total weight
let totalWeight = sum (map rbWeight backends')
currentW <- readTVar (rbCurrentWeight selected)
writeTVar (rbCurrentWeight selected) (currentW - totalWeight)
return (Just selected)
-- Helper: STM filter
filterM :: Monad m => (a -> m Bool) -> [a] -> m [a]
filterM _ [] = return []
filterM p (x:xs) = do
b <- p x
rest <- filterM p xs
return $ if b then x : rest else rest
-- Helper: fromMaybe
fromMaybe :: a -> Maybe a -> a
fromMaybe def Nothing = def
fromMaybe _ (Just x) = x
-- Helper: forM_
forM_ :: Monad m => [a] -> (a -> m b) -> m ()
forM_ xs f = sequence_ (map f xs)

View File

@ -494,17 +494,17 @@ extractFeatures FeatureContext{..} req =
let !headers = requestHeaders req let !headers = requestHeaders req
!path = rawPathInfo req !path = rawPathInfo req
!method = requestMethod req !method = requestMethod req
!mAcceptLang = lookupCI acceptLanguageHeader headers !mAcceptLang = lookup acceptLanguageHeader headers
!mUserAgent = lookupCI userAgentHeader headers !mUserAgent = lookup userAgentHeader headers
!mAcceptEnc = lookupCI acceptEncodingHeader headers !mAcceptEnc = lookup acceptEncodingHeader headers
!mReferer = lookupCI refererHeader headers !mReferer = lookup refererHeader headers
!mCookie = lookupCI cookieHeader headers !mCookie = lookup cookieHeader headers
!mSecChUa = lookupCI secChUaHeader headers !mSecChUa = lookup secChUaHeader headers
!mSecChPlat = lookupCI secChUaPlatformHeader headers !mSecChPlat = lookup secChUaPlatformHeader headers
!mSecFetchSite = lookupCI secFetchSiteHeader headers !mSecFetchSite = lookup secFetchSiteHeader headers
!mSecFetchMode = lookupCI secFetchModeHeader headers !mSecFetchMode = lookup secFetchModeHeader headers
!mSecFetchDest = lookupCI secFetchDestHeader headers !mSecFetchDest = lookup secFetchDestHeader headers
!mAccept = lookupCI acceptHeader headers !mAccept = lookup acceptHeader headers
!uaBytes = fromMaybe BS.empty mUserAgent !uaBytes = fromMaybe BS.empty mUserAgent
!uaLen = BS.length uaBytes !uaLen = BS.length uaBytes
!depth = pathDepth path !depth = pathDepth path
@ -552,8 +552,3 @@ extractFeatures FeatureContext{..} req =
, fHeaderOrderCanonical = boolToDouble , fHeaderOrderCanonical = boolToDouble
(headerOrderIsCanonicalBrowser headers) (headerOrderIsCanonicalBrowser headers)
} }
lookupCI :: CI ByteString -> [(CI ByteString, ByteString)] -> Maybe ByteString
lookupCI k hs = case filter ((== k) . fst) hs of
((_, v) : _) -> Just v
_ -> Nothing

View File

@ -1,3 +1,7 @@
{-
©AngelaMos | 2026
Redirect.hs
-}
{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE OverloadedStrings #-}
module Aenebris.Middleware.Redirect module Aenebris.Middleware.Redirect
@ -7,38 +11,44 @@ module Aenebris.Middleware.Redirect
import qualified Data.ByteString.Char8 as BS import qualified Data.ByteString.Char8 as BS
import Data.Maybe (fromMaybe) import Data.Maybe (fromMaybe)
import Network.HTTP.Types (status301, hLocation) import Network.HTTP.Types (hLocation, status301)
import Network.Wai (Middleware, responseLBS, requestHeaderHost, rawPathInfo, rawQueryString, isSecure) import Network.Wai
( Middleware
, isSecure
, rawPathInfo
, rawQueryString
, requestHeaderHost
, responseLBS
)
defaultHostFallback :: BS.ByteString
defaultHostFallback = "localhost"
httpsScheme :: BS.ByteString
httpsScheme = "https://"
standardHttpsPort :: Int
standardHttpsPort = 443
redirectBody :: BS.ByteString
redirectBody = "Redirecting to HTTPS"
-- | Redirect HTTP requests to HTTPS (assumes HTTPS is on port 443)
httpsRedirect :: Middleware httpsRedirect :: Middleware
httpsRedirect = httpsRedirectWithPort Nothing httpsRedirect = httpsRedirectWithPort Nothing
-- | Redirect HTTP requests to HTTPS with optional custom port
-- If port is Nothing, assumes 443 (standard HTTPS port, no port in URL)
-- If port is Just n, includes :n in the redirect URL
httpsRedirectWithPort :: Maybe Int -> Middleware httpsRedirectWithPort :: Maybe Int -> Middleware
httpsRedirectWithPort httpsPort app req respond httpsRedirectWithPort httpsPort app req respond
| isSecure req = app req respond -- Already HTTPS, pass through | isSecure req = app req respond
| otherwise = do | otherwise = do
-- Get host from Host header let hostHeader = fromMaybe defaultHostFallback (requestHeaderHost req)
let hostHeader = fromMaybe "localhost" $ requestHeaderHost req
-- Build HTTPS URL with optional port
host = case httpsPort of host = case httpsPort of
Nothing -> hostHeader -- Standard 443, don't include port Nothing -> hostHeader
Just 443 -> hostHeader -- Standard 443, don't include port Just port | port == standardHttpsPort -> hostHeader
Just port -> hostHeader <> ":" <> BS.pack (show port) Just port -> hostHeader <> ":" <> BS.pack (show port)
path = rawPathInfo req
-- Get path and query string (already encoded in rawPathInfo) query = rawQueryString req
path = rawPathInfo req redirectUrl = httpsScheme <> host <> path <> query
query = rawQueryString req
-- Build full redirect URL
redirectUrl = "https://" <> host <> path <> query
-- Send 301 permanent redirect
respond $ responseLBS respond $ responseLBS
status301 status301
[(hLocation, redirectUrl)] [(hLocation, redirectUrl)]
"Redirecting to HTTPS" (BS.fromStrict redirectBody)

View File

@ -1,3 +1,7 @@
{-
©AngelaMos | 2026
Security.hs
-}
{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE OverloadedStrings #-}
module Aenebris.Middleware.Security module Aenebris.Middleware.Security
@ -11,125 +15,101 @@ module Aenebris.Middleware.Security
import Data.ByteString (ByteString) import Data.ByteString (ByteString)
import qualified Data.CaseInsensitive as CI import qualified Data.CaseInsensitive as CI
import Data.Maybe (catMaybes)
import Network.HTTP.Types (Header, ResponseHeaders) import Network.HTTP.Types (Header, ResponseHeaders)
import Network.Wai (Middleware, mapResponseHeaders) import Network.Wai (Middleware, mapResponseHeaders)
-- | Security level presets
data SecurityLevel data SecurityLevel
= Testing -- Short HSTS, permissive CSP, for development = Testing
| Production -- Balanced security for production | Production
| Strict -- Maximum security, strict CSP, HSTS preload | Strict
deriving (Show, Eq) deriving (Show, Eq)
-- | Security configuration
data SecurityConfig = SecurityConfig data SecurityConfig = SecurityConfig
{ scHSTS :: Maybe ByteString -- Strict-Transport-Security header { scHSTS :: !(Maybe ByteString)
, scCSP :: Maybe ByteString -- Content-Security-Policy header , scCSP :: !(Maybe ByteString)
, scFrameOptions :: Maybe ByteString -- X-Frame-Options header , scFrameOptions :: !(Maybe ByteString)
, scContentTypeOptions :: Bool -- X-Content-Type-Options: nosniff , scContentTypeOptions :: !Bool
, scReferrerPolicy :: Maybe ByteString -- Referrer-Policy header , scReferrerPolicy :: !(Maybe ByteString)
, scPermissionsPolicy :: Maybe ByteString -- Permissions-Policy header , scPermissionsPolicy :: !(Maybe ByteString)
, scXSSProtection :: Maybe ByteString -- X-XSS-Protection (legacy, but some crawlers check) , scXSSProtection :: !(Maybe ByteString)
, scExpectCT :: Maybe ByteString -- Expect-CT (transitional) , scExpectCT :: !(Maybe ByteString)
, scServerHeader :: Maybe ByteString -- Server header (hide or customize) , scServerHeader :: !(Maybe ByteString)
, scRemovePoweredBy :: Bool -- Remove X-Powered-By headers , scRemovePoweredBy :: !Bool
} deriving (Show, Eq) } deriving (Show, Eq)
-- | Testing/development security configuration
-- Use short HSTS for easy testing, permissive CSP
testingSecurityConfig :: SecurityConfig testingSecurityConfig :: SecurityConfig
testingSecurityConfig = SecurityConfig testingSecurityConfig = SecurityConfig
{ scHSTS = Just "max-age=300" -- 5 minutes for testing { scHSTS = Just "max-age=300"
, scCSP = Just "default-src 'self' 'unsafe-inline' 'unsafe-eval'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'" , scCSP = Just "default-src 'self' 'unsafe-inline' 'unsafe-eval'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'"
, scFrameOptions = Just "SAMEORIGIN" , scFrameOptions = Just "SAMEORIGIN"
, scContentTypeOptions = True , scContentTypeOptions = True
, scReferrerPolicy = Just "strict-origin-when-cross-origin" , scReferrerPolicy = Just "strict-origin-when-cross-origin"
, scPermissionsPolicy = Just "geolocation=(), microphone=(), camera=()" , scPermissionsPolicy = Just "geolocation=(), microphone=(), camera=()"
, scXSSProtection = Just "1; mode=block" , scXSSProtection = Just "1; mode=block"
, scExpectCT = Nothing , scExpectCT = Nothing
, scServerHeader = Just "Aenebris/0.1.0" , scServerHeader = Just "Aenebris/0.1.0"
, scRemovePoweredBy = True , scRemovePoweredBy = True
} }
-- | Production security configuration
-- Balanced security, 1-month HSTS
defaultSecurityConfig :: SecurityConfig defaultSecurityConfig :: SecurityConfig
defaultSecurityConfig = SecurityConfig defaultSecurityConfig = SecurityConfig
{ scHSTS = Just "max-age=2592000; includeSubDomains" -- 30 days { scHSTS = Just "max-age=2592000; includeSubDomains"
, scCSP = Just "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data: https:; font-src 'self' data:; connect-src 'self'; frame-ancestors 'none'" , scCSP = Just "default-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data: https:; font-src 'self' data:; connect-src 'self'; frame-ancestors 'none'"
, scFrameOptions = Just "DENY" , scFrameOptions = Just "DENY"
, scContentTypeOptions = True , scContentTypeOptions = True
, scReferrerPolicy = Just "strict-origin-when-cross-origin" , scReferrerPolicy = Just "strict-origin-when-cross-origin"
, scPermissionsPolicy = Just "geolocation=(), microphone=(), camera=(), payment=(), usb=(), magnetometer=(), gyroscope=(), accelerometer=()" , scPermissionsPolicy = Just "geolocation=(), microphone=(), camera=(), payment=(), usb=(), magnetometer=(), gyroscope=(), accelerometer=()"
, scXSSProtection = Just "1; mode=block" , scXSSProtection = Just "1; mode=block"
, scExpectCT = Just "max-age=86400, enforce" , scExpectCT = Just "max-age=86400, enforce"
, scServerHeader = Just "Aenebris" -- Don't reveal version in production , scServerHeader = Just "Aenebris"
, scRemovePoweredBy = True , scRemovePoweredBy = True
} }
-- | Strict security configuration for maximum protection
-- 2-year HSTS with preload, very restrictive CSP
strictSecurityConfig :: SecurityConfig strictSecurityConfig :: SecurityConfig
strictSecurityConfig = SecurityConfig strictSecurityConfig = SecurityConfig
{ scHSTS = Just "max-age=63072000; includeSubDomains; preload" -- 2 years + preload { scHSTS = Just "max-age=63072000; includeSubDomains; preload"
, scCSP = Just "default-src 'none'; script-src 'self'; style-src 'self'; img-src 'self' data:; font-src 'self'; connect-src 'self'; frame-ancestors 'none'; base-uri 'self'; form-action 'self'; upgrade-insecure-requests" , scCSP = Just "default-src 'none'; script-src 'self'; style-src 'self'; img-src 'self' data:; font-src 'self'; connect-src 'self'; frame-ancestors 'none'; base-uri 'self'; form-action 'self'; upgrade-insecure-requests"
, scFrameOptions = Just "DENY" , scFrameOptions = Just "DENY"
, scContentTypeOptions = True , scContentTypeOptions = True
, scReferrerPolicy = Just "no-referrer" -- Strictest, no referrer leakage , scReferrerPolicy = Just "no-referrer"
, scPermissionsPolicy = Just "geolocation=(), microphone=(), camera=(), payment=(), usb=(), magnetometer=(), gyroscope=(), accelerometer=(), bluetooth=(), display-capture=(), document-domain=()" , scPermissionsPolicy = Just "geolocation=(), microphone=(), camera=(), payment=(), usb=(), magnetometer=(), gyroscope=(), accelerometer=(), bluetooth=(), display-capture=(), document-domain=()"
, scXSSProtection = Just "1; mode=block" , scXSSProtection = Just "1; mode=block"
, scExpectCT = Just "max-age=86400, enforce" , scExpectCT = Just "max-age=86400, enforce"
, scServerHeader = Nothing -- Hide completely , scServerHeader = Nothing
, scRemovePoweredBy = True , scRemovePoweredBy = True
} }
-- | Middleware that adds security headers to all responses
addSecurityHeaders :: SecurityConfig -> Middleware addSecurityHeaders :: SecurityConfig -> Middleware
addSecurityHeaders config app req respond = addSecurityHeaders config app req respond =
app req $ \res -> app req $ \res ->
respond $ mapResponseHeaders (addHeaders config) res respond $ mapResponseHeaders (addHeaders config) res
-- | Add security headers to response headers
addHeaders :: SecurityConfig -> ResponseHeaders -> ResponseHeaders addHeaders :: SecurityConfig -> ResponseHeaders -> ResponseHeaders
addHeaders config headers = addHeaders config headers =
let let cleaned = if scRemovePoweredBy config
-- Remove headers we want to control then filter (not . isPoweredBy) headers
cleaned = if scRemovePoweredBy config else headers
then filter (not . isPoweredBy) headers newHeaders = catMaybes
else headers [ fmap (\v -> ("Strict-Transport-Security", v)) (scHSTS config)
, fmap (\v -> ("Content-Security-Policy", v)) (scCSP config)
-- Build new security headers , fmap (\v -> ("X-Frame-Options", v)) (scFrameOptions config)
newHeaders = catMaybes , if scContentTypeOptions config
[ fmap (\v -> ("Strict-Transport-Security", v)) (scHSTS config) then Just ("X-Content-Type-Options", "nosniff")
, fmap (\v -> ("Content-Security-Policy", v)) (scCSP config) else Nothing
, fmap (\v -> ("X-Frame-Options", v)) (scFrameOptions config) , fmap (\v -> ("Referrer-Policy", v)) (scReferrerPolicy config)
, if scContentTypeOptions config , fmap (\v -> ("Permissions-Policy", v)) (scPermissionsPolicy config)
then Just ("X-Content-Type-Options", "nosniff") , fmap (\v -> ("X-XSS-Protection", v)) (scXSSProtection config)
else Nothing , fmap (\v -> ("Expect-CT", v)) (scExpectCT config)
, fmap (\v -> ("Referrer-Policy", v)) (scReferrerPolicy config) ]
, fmap (\v -> ("Permissions-Policy", v)) (scPermissionsPolicy config) serverHeader = case scServerHeader config of
, fmap (\v -> ("X-XSS-Protection", v)) (scXSSProtection config) Just v -> [("Server", v)]
, fmap (\v -> ("Expect-CT", v)) (scExpectCT config) Nothing -> []
] withoutServer = filter (not . isServerHeader) cleaned
-- Handle Server header specially
serverHeader = case scServerHeader config of
Just v -> [("Server", v)]
Nothing -> [] -- Remove Server header completely
-- Remove existing Server header if we're replacing it
withoutServer = filter (not . isServerHeader) cleaned
in withoutServer ++ newHeaders ++ serverHeader in withoutServer ++ newHeaders ++ serverHeader
-- | Check if header is X-Powered-By
isPoweredBy :: Header -> Bool isPoweredBy :: Header -> Bool
isPoweredBy (name, _) = CI.mk name == CI.mk "X-Powered-By" isPoweredBy (name, _) = CI.mk name == CI.mk "X-Powered-By"
-- | Check if header is Server
isServerHeader :: Header -> Bool isServerHeader :: Header -> Bool
isServerHeader (name, _) = CI.mk name == CI.mk "Server" isServerHeader (name, _) = CI.mk name == CI.mk "Server"
-- | catMaybes implementation (since we're not importing Data.Maybe)
catMaybes :: [Maybe a] -> [a]
catMaybes = foldr (\mx xs -> case mx of Just x -> x:xs; Nothing -> xs) []

View File

@ -1,6 +1,10 @@
{-
©AngelaMos | 2026
Proxy.hs
-}
{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Aenebris.Proxy module Aenebris.Proxy
( ProxyState(..) ( ProxyState(..)
@ -25,7 +29,12 @@ import Aenebris.TLS
import Aenebris.Tunnel import Aenebris.Tunnel
import Aenebris.Middleware.Security import Aenebris.Middleware.Security
import Aenebris.Middleware.Redirect import Aenebris.Middleware.Redirect
import Aenebris.RateLimit (RateLimiter, createRateLimiter, parseRateSpec, rateLimitMiddleware) import Aenebris.RateLimit
( RateLimiter
, createRateLimiter
, parseRateSpec
, rateLimitMiddleware
)
import Aenebris.DDoS.EarlyData (earlyDataGuard) import Aenebris.DDoS.EarlyData (earlyDataGuard)
import Aenebris.DDoS.MemoryShed import Aenebris.DDoS.MemoryShed
( MemoryShed ( MemoryShed
@ -72,7 +81,13 @@ import Aenebris.Geo
) )
import Control.Concurrent.STM (TVar, newTVarIO) import Control.Concurrent.STM (TVar, newTVarIO)
import Control.Concurrent.Async (Async, async, waitAnyCancel) import Control.Concurrent.Async (Async, async, waitAnyCancel)
import Control.Exception (try, SomeException) import Control.Exception (SomeException, try)
import Control.Monad (unless, zipWithM)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import Data.ByteString.Builder (byteString)
import qualified Data.ByteString.Lazy as LBS
import Data.Function ((&)) import Data.Function ((&))
import Data.List (sortBy) import Data.List (sortBy)
import Data.Map.Strict (Map) import Data.Map.Strict (Map)
@ -82,12 +97,16 @@ import Data.Ord (comparing)
import Data.Text (Text) import Data.Text (Text)
import qualified Data.Text as T import qualified Data.Text as T
import qualified Data.Text.Encoding as TE import qualified Data.Text.Encoding as TE
import Network.HTTP.Client (Manager, withResponse, parseRequest, RequestBody(..), brRead) import Network.HTTP.Client
( Manager
, RequestBody(..)
, brRead
, parseRequest
, withResponse
)
import qualified Network.HTTP.Client as HTTP import qualified Network.HTTP.Client as HTTP
import Network.HTTP.Types import Network.HTTP.Types
import Network.Wai import Network.Wai
import Data.ByteString.Builder (byteString)
import Control.Monad (unless)
import Network.Wai.Handler.Warp import Network.Wai.Handler.Warp
( Settings ( Settings
, defaultSettings , defaultSettings
@ -102,33 +121,77 @@ import Network.Wai.Handler.WarpTLS (runTLS)
import System.Exit (exitFailure) import System.Exit (exitFailure)
import System.IO (hPutStrLn, stderr) import System.IO (hPutStrLn, stderr)
import System.Timeout (timeout) import System.Timeout (timeout)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import qualified Data.ByteString.Lazy as LBS
-- | Proxy runtime state memoryShedPollIntervalMicros :: Int
memoryShedPollIntervalMicros = microsPerSecond
contentTypePlain :: ByteString
contentTypePlain = "text/plain"
bodyNotFound :: LBS.ByteString
bodyNotFound = "Not Found: No route configured for this host/path"
bodyUpstreamMisconfigured :: LBS.ByteString
bodyUpstreamMisconfigured = "Internal Server Error: Upstream configuration error"
bodyNoHealthyBackends :: LBS.ByteString
bodyNoHealthyBackends = "Service Unavailable: No healthy backends available"
bodyBadGateway :: LBS.ByteString
bodyBadGateway = "Bad Gateway: Could not connect to backend server"
bodyGatewayTimeout :: LBS.ByteString
bodyGatewayTimeout = "504 Gateway Timeout: upstream did not respond in time"
bodyWebSocketUpgradeFailed :: LBS.ByteString
bodyWebSocketUpgradeFailed = "WebSocket upgrade failed"
hopByHopRequestHeaders :: [HeaderName]
hopByHopRequestHeaders =
[ "Connection"
, "Keep-Alive"
, "Proxy-Authenticate"
, "Proxy-Authorization"
, "TE"
, "Trailers"
, "Transfer-Encoding"
, "Upgrade"
]
hopByHopResponseHeaders :: [HeaderName]
hopByHopResponseHeaders =
[ "Transfer-Encoding"
, "Connection"
, "Keep-Alive"
]
eventStreamContentType :: ByteString
eventStreamContentType = "text/event-stream"
chunkedEncoding :: ByteString
chunkedEncoding = "chunked"
httpScheme :: String
httpScheme = "http://"
data ProxyState = ProxyState data ProxyState = ProxyState
{ psConfig :: Config { psConfig :: !Config
, psLoadBalancers :: Map Text LoadBalancer -- upstream name -> load balancer , psLoadBalancers :: !(Map Text LoadBalancer)
, psHealthCheckers :: [Async ()] , psHealthCheckers :: ![Async ()]
, psManager :: Manager , psManager :: !Manager
, psRateLimiter :: Maybe RateLimiter , psRateLimiter :: !(Maybe RateLimiter)
, psMemoryShed :: Maybe MemoryShed , psMemoryShed :: !(Maybe MemoryShed)
, psIPJail :: Maybe IPJail , psIPJail :: !(Maybe IPJail)
, psConnLimiter :: Maybe ConnLimiter , psConnLimiter :: !(Maybe ConnLimiter)
, psWafRuleSet :: TVar RuleSet , psWafRuleSet :: !(TVar RuleSet)
, psGeo :: Maybe Geo , psGeo :: !(Maybe Geo)
} }
-- | Initialize proxy state from config
initProxyState :: Config -> Manager -> IO ProxyState initProxyState :: Config -> Manager -> IO ProxyState
initProxyState config manager = do initProxyState config manager = do
-- Create load balancers for each upstream
lbs <- mapM createUpstreamLoadBalancer (configUpstreams config) lbs <- mapM createUpstreamLoadBalancer (configUpstreams config)
let lbMap = Map.fromList (zip (map upstreamName $ configUpstreams config) lbs) let lbMap = Map.fromList (zip (map upstreamName (configUpstreams config)) lbs)
-- Start health checkers for all upstreams
checkers <- mapM startUpstreamHealthChecker (configUpstreams config) checkers <- mapM startUpstreamHealthChecker (configUpstreams config)
rateLimiter <- case configRateLimit config >>= parseRateSpec of rateLimiter <- case configRateLimit config >>= parseRateSpec of
@ -142,8 +205,9 @@ initProxyState config manager = do
ms <- newMemoryShed ms <- newMemoryShed
let cfg = MemoryShedConfig let cfg = MemoryShedConfig
{ mscHeapBudgetBytes = fromInteger budgetBytes { mscHeapBudgetBytes = fromInteger budgetBytes
, mscHighWaterFraction = fromMaybe defaultHighWaterFraction (ddos >>= ddosMemoryShedHighWater) , mscHighWaterFraction = fromMaybe defaultHighWaterFraction
, mscPollIntervalMicros = 1000000 (ddos >>= ddosMemoryShedHighWater)
, mscPollIntervalMicros = memoryShedPollIntervalMicros
} }
_ <- startMemoryShedPoller cfg ms _ <- startMemoryShedPoller cfg ms
pure (Just ms) pure (Just ms)
@ -170,57 +234,45 @@ initProxyState config manager = do
Nothing -> pure Nothing Nothing -> pure Nothing
return ProxyState return ProxyState
{ psConfig = config { psConfig = config
, psLoadBalancers = lbMap , psLoadBalancers = lbMap
, psHealthCheckers = checkers , psHealthCheckers = checkers
, psManager = manager , psManager = manager
, psRateLimiter = rateLimiter , psRateLimiter = rateLimiter
, psMemoryShed = memShed , psMemoryShed = memShed
, psIPJail = ipJail , psIPJail = ipJail
, psConnLimiter = connLimiter , psConnLimiter = connLimiter
, psWafRuleSet = wafVar , psWafRuleSet = wafVar
, psGeo = geoHandle , psGeo = geoHandle
} }
where where
-- Create load balancer for an upstream
createUpstreamLoadBalancer :: Upstream -> IO LoadBalancer createUpstreamLoadBalancer :: Upstream -> IO LoadBalancer
createUpstreamLoadBalancer upstream = do createUpstreamLoadBalancer upstream = do
-- Convert Config Servers to RuntimeBackends
backends <- zipWithM createRuntimeBackend [0..] (upstreamServers upstream) backends <- zipWithM createRuntimeBackend [0..] (upstreamServers upstream)
-- Determine strategy (for now, use weighted if weights differ, else round-robin)
let weights = map serverWeight (upstreamServers upstream) let weights = map serverWeight (upstreamServers upstream)
strategy = case weights of strategy = case weights of
[] -> RoundRobin -- No backends, shouldn't happen but be safe [] -> RoundRobin
(w:ws) -> if all (== w) ws (w:ws)
then RoundRobin | all (== w) ws -> RoundRobin
else WeightedRoundRobin | otherwise -> WeightedRoundRobin
createLoadBalancer strategy backends createLoadBalancer strategy backends
-- Start health checker for an upstream
startUpstreamHealthChecker :: Upstream -> IO (Async ()) startUpstreamHealthChecker :: Upstream -> IO (Async ())
startUpstreamHealthChecker upstream = do startUpstreamHealthChecker upstream = do
backends <- zipWithM createRuntimeBackend [0..] (upstreamServers upstream) backends <- zipWithM createRuntimeBackend [0..] (upstreamServers upstream)
-- Use health check config from upstream, or defaults
let hcConfig = case upstreamHealthCheck upstream of let hcConfig = case upstreamHealthCheck upstream of
Just hc -> defaultHealthCheckConfig Just hc -> defaultHealthCheckConfig
{ hcInterval = 10 -- TODO: parse interval from config { hcEndpoint = healthCheckPath hc
, hcEndpoint = healthCheckPath hc
} }
Nothing -> defaultHealthCheckConfig Nothing -> defaultHealthCheckConfig
startHealthChecker manager hcConfig backends startHealthChecker manager hcConfig backends
-- | Start the proxy server with given configuration
-- Supports multiple ports with HTTP and HTTPS (including SNI)
startProxy :: ProxyState -> IO () startProxy :: ProxyState -> IO ()
startProxy ProxyState{..} = do startProxy ProxyState{..} = do
putStrLn $ "Starting Ᾰenebris reverse proxy" putStrLn "Starting Aenebris reverse proxy"
putStrLn $ "Loaded " ++ show (length $ configUpstreams psConfig) ++ " upstream(s)" putStrLn $ "Loaded " ++ show (length (configUpstreams psConfig)) ++ " upstream(s)"
putStrLn $ "Loaded " ++ show (length $ configRoutes psConfig) ++ " route(s)" putStrLn $ "Loaded " ++ show (length (configRoutes psConfig)) ++ " route(s)"
putStrLn $ "Health checking enabled for all upstreams" putStrLn "Health checking enabled for all upstreams"
case configListen psConfig of case configListen psConfig of
[] -> do [] -> do
@ -228,28 +280,37 @@ startProxy ProxyState{..} = do
exitFailure exitFailure
listenConfigs -> do listenConfigs -> do
case psRateLimiter of case psRateLimiter of
Just _ -> putStrLn "Rate limiting enabled" Just _ -> putStrLn "Rate limiting enabled"
Nothing -> pure () Nothing -> pure ()
putStrLn "WAF enabled (Phase 1: paranoia level 2, default rule pack)" putStrLn "WAF enabled (Phase 1: paranoia level 2, default rule pack)"
case buildHoneypotConfig (configHoneypot psConfig) of case buildHoneypotConfig (configHoneypot psConfig) of
Just hp -> putStrLn $ "Honeypot enabled (" ++ show (length (hpPatterns hp)) Just hp -> putStrLn $
++ " trap patterns, action=" ++ show (hpAction hp) ++ ")" "Honeypot enabled ("
++ show (length (hpPatterns hp))
++ " trap patterns, action="
++ show (hpAction hp)
++ ")"
Nothing -> pure () Nothing -> pure ()
case psGeo of case psGeo of
Just g -> Just g ->
let gc = geoConfig g let gc = geoConfig g
parts = [ "country_db=" ++ maybe "off" (const "on") (gcCountryDb gc) parts =
, "asn_db=" ++ maybe "off" (const "on") (gcAsnDb gc) [ "country_db=" ++ maybe "off" (const "on") (gcCountryDb gc)
, "blocked=" ++ show (length (gcBlockedCountries gc)) , "asn_db=" ++ maybe "off" (const "on") (gcAsnDb gc)
, "flagged_asns=" ++ show (length (gcFlaggedAsns gc)) , "blocked=" ++ show (length (gcBlockedCountries gc))
] , "flagged_asns=" ++ show (length (gcFlaggedAsns gc))
]
in putStrLn $ "Geo/ASN enabled (" ++ unwords parts ++ ")" in putStrLn $ "Geo/ASN enabled (" ++ unwords parts ++ ")"
Nothing -> pure () Nothing -> pure ()
servers <- mapM (launchServer psConfig psLoadBalancers psManager psRateLimiter psMemoryShed psIPJail psConnLimiter psWafRuleSet psGeo) listenConfigs
servers <- mapM
(launchServer psConfig psLoadBalancers psManager
psRateLimiter psMemoryShed psIPJail
psConnLimiter psWafRuleSet psGeo)
listenConfigs
_ <- waitAnyCancel servers _ <- waitAnyCancel servers
putStrLn "All servers stopped" putStrLn "All servers stopped"
launchServer launchServer
@ -264,265 +325,253 @@ launchServer
-> Maybe Geo -> Maybe Geo
-> ListenConfig -> ListenConfig
-> IO (Async ()) -> IO (Async ())
launchServer config loadBalancers manager mRateLimiter mMemShed mIPJail mConnLim wafVar mGeo listenConfig = async $ do launchServer config loadBalancers manager mRateLimiter mMemShed mIPJail mConnLim wafVar mGeo listenConfig =
let port = listenPort listenConfig async $ do
shouldRedirect = fromMaybe False (listenRedirectHTTPS listenConfig) let port = listenPort listenConfig
ddosCfg = fromMaybe defaultDDoSConfig (configDDoS config) shouldRedirect = fromMaybe False (listenRedirectHTTPS listenConfig)
ddosCfg = fromMaybe defaultDDoSConfig (configDDoS config)
baseApp = proxyApp config loadBalancers manager baseApp = proxyApp config loadBalancers manager
fingerprintedApp = ja4hMiddleware baseApp
wafApp = wafMiddleware wafVar fingerprintedApp
securedApp = addSecurityHeaders defaultSecurityConfig wafApp
fingerprintedApp = ja4hMiddleware baseApp earlyDataApp = if ddosEarlyDataReject ddosCfg
then earlyDataGuard securedApp
else securedApp
wafApp = wafMiddleware wafVar fingerprintedApp mHoneypotCfg = buildHoneypotConfig (configHoneypot config)
honeypotApp = case mHoneypotCfg of
Just hp -> honeypotMiddleware hp mIPJail earlyDataApp
Nothing -> earlyDataApp
securedApp = addSecurityHeaders defaultSecurityConfig wafApp geoApp = case mGeo of
Just g -> geoMiddleware g mIPJail honeypotApp
Nothing -> honeypotApp
earlyDataApp = if ddosEarlyDataReject ddosCfg jailedApp = case mIPJail of
then earlyDataGuard securedApp Just j -> ipJailMiddleware j geoApp
else securedApp Nothing -> geoApp
mHoneypotCfg = buildHoneypotConfig (configHoneypot config) shedApp = case mMemShed of
Just ms -> memoryShedMiddleware ms jailedApp
Nothing -> jailedApp
honeypotApp = case mHoneypotCfg of limitedApp = case mRateLimiter of
Just hp -> honeypotMiddleware hp mIPJail earlyDataApp Just rl -> rateLimitMiddleware rl shedApp
Nothing -> earlyDataApp Nothing -> shedApp
geoApp = case mGeo of warpSettings = applyDDoSSettings ddosCfg mConnLim
Just g -> geoMiddleware g mIPJail honeypotApp (defaultSettings & setPort port)
Nothing -> honeypotApp
jailedApp = case mIPJail of case listenTLS listenConfig of
Just j -> ipJailMiddleware j geoApp Nothing -> do
Nothing -> geoApp let app = if shouldRedirect
then httpsRedirect limitedApp
else limitedApp
putStrLn $ "* HTTP server listening on :" ++ show port
if shouldRedirect
then putStrLn " Redirecting all traffic to HTTPS"
else pure ()
runSettings warpSettings app
shedApp = case mMemShed of Just tlsConfig -> do
Just ms -> memoryShedMiddleware ms jailedApp let isSNI = case tlsSNI tlsConfig of
Nothing -> jailedApp Just domains -> not (null domains)
Nothing -> False
limitedApp = case mRateLimiter of if isSNI
Just rl -> rateLimitMiddleware rl shedApp then launchHTTPSWithSNI port tlsConfig limitedApp
Nothing -> shedApp else launchHTTPS port tlsConfig limitedApp
warpSettings = applyDDoSSettings ddosCfg mConnLim (defaultSettings & setPort port)
case listenTLS listenConfig of
Nothing -> do
let app = if shouldRedirect
then httpsRedirect limitedApp
else limitedApp
putStrLn $ "✓ HTTP server listening on :" ++ show port
if shouldRedirect
then putStrLn $ " └─ Redirecting all traffic to HTTPS"
else return ()
runSettings warpSettings app
Just tlsConfig -> do
let isSNI = case tlsSNI tlsConfig of
Just domains -> not (null domains)
Nothing -> False
if isSNI
then launchHTTPSWithSNI port tlsConfig limitedApp
else launchHTTPS port tlsConfig limitedApp
applyDDoSSettings :: DDoSConfig -> Maybe ConnLimiter -> Settings -> Settings applyDDoSSettings :: DDoSConfig -> Maybe ConnLimiter -> Settings -> Settings
applyDDoSSettings ddos mConnLim s0 = applyDDoSSettings ddos mConnLim s0 =
let s1 = case ddosMaxHeaderBytes ddos of let s1Inner = case ddosSlowlorisSeconds ddos of
Just n -> setMaxTotalHeaderLength n s1Inner Just n -> setTimeout n s0
Nothing -> s1Inner
s1Inner = case ddosSlowlorisSeconds ddos of
Just n -> setTimeout n s0
Nothing -> s0 Nothing -> s0
s1 = case ddosMaxHeaderBytes ddos of
Just n -> setMaxTotalHeaderLength n s1Inner
Nothing -> s1Inner
s2 = case mConnLim of s2 = case mConnLim of
Just cl -> setOnClose (connLimitOnClose cl) (setOnOpen (connLimitOnOpen cl) s1) Just cl -> setOnClose (connLimitOnClose cl)
(setOnOpen (connLimitOnOpen cl) s1)
Nothing -> s1 Nothing -> s1
in s2 in s2
-- | Launch HTTPS server with single certificate
launchHTTPS :: Int -> TLSConfig -> Application -> IO () launchHTTPS :: Int -> TLSConfig -> Application -> IO ()
launchHTTPS port tlsConfig app = do launchHTTPS port tlsConfig app =
case (tlsCert tlsConfig, tlsKey tlsConfig) of case (tlsCert tlsConfig, tlsKey tlsConfig) of
(Just certFile, Just keyFile) -> do (Just certFile, Just keyFile) -> do
-- Load TLS settings
tlsResult <- createTLSSettings certFile keyFile tlsResult <- createTLSSettings certFile keyFile
case tlsResult of case tlsResult of
Left err -> do Left err -> do
hPutStrLn stderr "ERROR: Failed to load TLS certificate" hPutStrLn stderr "ERROR: Failed to load TLS certificate"
hPutStrLn stderr $ " " ++ show err hPutStrLn stderr $ " " ++ show err
exitFailure exitFailure
Right tlsSettings -> do Right tlsSettings -> do
let warpSettings = defaultSettings & setPort port let warpSettings = defaultSettings & setPort port
putStrLn $ " HTTPS server listening on :" ++ show port putStrLn $ "* HTTPS server listening on :" ++ show port
putStrLn $ " ├─ Certificate: " ++ certFile putStrLn $ " Certificate: " ++ certFile
putStrLn $ " ├─ TLS 1.2 + TLS 1.3 enabled" putStrLn " TLS 1.2 + TLS 1.3 enabled"
putStrLn $ " ├─ HTTP/2 enabled (ALPN)" putStrLn " HTTP/2 enabled (ALPN)"
putStrLn $ " └─ Strong cipher suites enforced" putStrLn " Strong cipher suites enforced"
runTLS tlsSettings warpSettings app runTLS tlsSettings warpSettings app
_ -> do _ -> do
hPutStrLn stderr "ERROR: TLS configuration requires both cert and key" hPutStrLn stderr "ERROR: TLS configuration requires both cert and key"
exitFailure exitFailure
-- | Launch HTTPS server with SNI support (multiple certificates)
launchHTTPSWithSNI :: Int -> TLSConfig -> Application -> IO () launchHTTPSWithSNI :: Int -> TLSConfig -> Application -> IO ()
launchHTTPSWithSNI port tlsConfig app = do launchHTTPSWithSNI port tlsConfig app =
case (tlsSNI tlsConfig, tlsDefaultCert tlsConfig, tlsDefaultKey tlsConfig) of case (tlsSNI tlsConfig, tlsDefaultCert tlsConfig, tlsDefaultKey tlsConfig) of
(Just sniDomains, Just defaultCert, Just defaultKey) -> do (Just sniDomains, Just defaultCert, Just defaultKey) -> do
-- Convert SNIDomain list to the format expected by createSNISettings
let domainList = [(sniDomain d, sniCert d, sniKey d) | d <- sniDomains] let domainList = [(sniDomain d, sniCert d, sniKey d) | d <- sniDomains]
-- Load SNI TLS settings
tlsResult <- createSNISettings domainList defaultCert defaultKey tlsResult <- createSNISettings domainList defaultCert defaultKey
case tlsResult of case tlsResult of
Left err -> do Left err -> do
hPutStrLn stderr "ERROR: Failed to load SNI certificates" hPutStrLn stderr "ERROR: Failed to load SNI certificates"
hPutStrLn stderr $ " " ++ show err hPutStrLn stderr $ " " ++ show err
exitFailure exitFailure
Right tlsSettings -> do Right tlsSettings -> do
let warpSettings = defaultSettings & setPort port let warpSettings = defaultSettings & setPort port
putStrLn $ "✓ HTTPS server with SNI listening on :" ++ show port putStrLn $ "* HTTPS server with SNI listening on :" ++ show port
putStrLn $ " ├─ SNI domains: " ++ show (length sniDomains) ++ " configured" putStrLn $ " SNI domains: " ++ show (length sniDomains) ++ " configured"
mapM_ (\d -> putStrLn $ " │ • " ++ T.unpack (sniDomain d) ++ " -> " ++ sniCert d) sniDomains mapM_
putStrLn $ " ├─ Default certificate: " ++ defaultCert (\d -> putStrLn $
putStrLn $ " ├─ TLS 1.2 + TLS 1.3 enabled" " " ++ T.unpack (sniDomain d) ++ " -> " ++ sniCert d)
putStrLn $ " ├─ HTTP/2 enabled (ALPN)" sniDomains
putStrLn $ " └─ Strong cipher suites enforced" putStrLn $ " Default certificate: " ++ defaultCert
putStrLn " TLS 1.2 + TLS 1.3 enabled"
putStrLn " HTTP/2 enabled (ALPN)"
putStrLn " Strong cipher suites enforced"
runTLS tlsSettings warpSettings app runTLS tlsSettings warpSettings app
_ -> do _ -> do
hPutStrLn stderr hPutStrLn stderr "ERROR: SNI requires sni, default_cert, and default_key"
"ERROR: SNI requires sni, default_cert, and default_key"
exitFailure exitFailure
-- | Main proxy application (WAI)
proxyApp :: Config -> Map Text LoadBalancer -> Manager -> Application proxyApp :: Config -> Map Text LoadBalancer -> Manager -> Application
proxyApp config loadBalancers manager req respond = do proxyApp config loadBalancers manager req respond = do
logRequest req let hostHeader = lookup "Host" (requestHeaders req)
let hostHeader = lookup "Host" (requestHeaders req)
requestPath = rawPathInfo req requestPath = rawPathInfo req
headers = requestHeaders req headers = requestHeaders req
connType = detectConnectionType headers connType = detectConnectionType headers
case selectRoute config hostHeader requestPath of case selectRoute config hostHeader requestPath of
Nothing -> do Nothing -> do
hPutStrLn stderr $ "ERROR: No route found for request" hPutStrLn stderr "ERROR: No route found for request"
respond $ responseLBS respond $ responseLBS
status404 status404
[("Content-Type", "text/plain")] [(hContentType, contentTypePlain)]
"Not Found: No route configured for this host/path" bodyNotFound
Just (upstreamName, _pathRoute) -> do Just (upstreamName, _pathRoute) ->
case Map.lookup upstreamName loadBalancers of case Map.lookup upstreamName loadBalancers of
Nothing -> do Nothing -> do
hPutStrLn stderr $ "ERROR: Load balancer not found: " ++ T.unpack upstreamName hPutStrLn stderr $
"ERROR: Load balancer not found: " ++ T.unpack upstreamName
respond $ responseLBS respond $ responseLBS
status500 status500
[("Content-Type", "text/plain")] [(hContentType, contentTypePlain)]
"Internal Server Error: Upstream configuration error" bodyUpstreamMisconfigured
Just loadBalancer -> do Just loadBalancer -> do
mBackend <- selectBackend loadBalancer mBackend <- selectBackend loadBalancer
case mBackend of case mBackend of
Nothing -> do Nothing -> do
hPutStrLn stderr $ "ERROR: No healthy backends available" hPutStrLn stderr "ERROR: No healthy backends available"
respond $ responseLBS respond $ responseLBS
status503 status503
[("Content-Type", "text/plain")] [(hContentType, contentTypePlain)]
"Service Unavailable: No healthy backends available" bodyNoHealthyBackends
Just backend -> do Just backend -> case connType of
case connType of WebSocket -> do
WebSocket -> do hPutStrLn stderr "[WS] WebSocket upgrade detected"
hPutStrLn stderr $ "[WS] WebSocket upgrade detected" handleWebSocketUpgrade req respond backend
handleWebSocketUpgrade req respond backend _ ->
forwardRegular manager backend req respond
RegularHttp -> do forwardRegular
result <- try $ trackConnection backend $ :: Manager
forwardRequest manager req (rbHost backend) respond -> RuntimeBackend
-> Request
-> (Response -> IO ResponseReceived)
-> IO ResponseReceived
forwardRegular manager backend req respond = do
result <- try $ trackConnection backend $
forwardRequest manager req (rbHost backend) respond
case result of
Left (err :: SomeException) -> do
hPutStrLn stderr $ "ERROR: " ++ show err
respond $ responseLBS
status502
[(hContentType, contentTypePlain)]
bodyBadGateway
Right responseReceived ->
pure responseReceived
case result of handleWebSocketUpgrade
Left (err :: SomeException) -> do :: Request
hPutStrLn stderr $ "ERROR: " ++ show err -> (Response -> IO ResponseReceived)
respond $ responseLBS -> RuntimeBackend
status502 -> IO ResponseReceived
[("Content-Type", "text/plain")]
"Bad Gateway: Could not connect to backend server"
Right responseReceived ->
return responseReceived
_ -> do
result <- try $ trackConnection backend $
forwardRequest manager req (rbHost backend) respond
case result of
Left (err :: SomeException) -> do
hPutStrLn stderr $ "ERROR: " ++ show err
respond $ responseLBS
status502
[("Content-Type", "text/plain")]
"Bad Gateway: Could not connect to backend server"
Right responseReceived ->
return responseReceived
handleWebSocketUpgrade :: Request -> (Response -> IO ResponseReceived) -> RuntimeBackend -> IO ResponseReceived
handleWebSocketUpgrade req respond backend = do handleWebSocketUpgrade req respond backend = do
let backendHost = rbHost backend let backendHost = rbHost backend
backupResponse = responseLBS backupResponse = responseLBS
status502 status502
[("Content-Type", "text/plain")] [(hContentType, contentTypePlain)]
"WebSocket upgrade failed" bodyWebSocketUpgradeFailed
respond $ responseRaw (wsHandler req backendHost) backupResponse respond $ responseRaw (wsHandler req backendHost) backupResponse
wsHandler :: Request -> Text -> IO ByteString -> (ByteString -> IO ()) -> IO () wsHandler
:: Request
-> Text
-> IO ByteString
-> (ByteString -> IO ())
-> IO ()
wsHandler req backendHost recv send = do wsHandler req backendHost recv send = do
hPutStrLn stderr $ "[WS] Starting WebSocket tunnel to " ++ T.unpack backendHost hPutStrLn stderr $ "[WS] Starting WebSocket tunnel to " ++ T.unpack backendHost
tunnelWebSocket req backendHost send recv tunnelWebSocket req backendHost send recv
-- | Select a route based on Host header and path selectRoute
selectRoute :: Config -> Maybe BS.ByteString -> BS.ByteString -> Maybe (Text, PathRoute) :: Config
selectRoute config hostHeader requestPath = -> Maybe BS.ByteString
case hostHeader of -> BS.ByteString
Nothing -> Nothing -- No Host header, can't route -> Maybe (Text, PathRoute)
Just host -> do selectRoute config hostHeader requestPath = case hostHeader of
-- Find route matching this host Nothing -> Nothing
let hostText = TE.decodeUtf8 host Just host -> do
matchingRoutes = filter (\r -> routeHost r == hostText) (configRoutes config) let hostText = TE.decodeUtf8 host
matchingRoutes = filter (\r -> routeHost r == hostText)
(configRoutes config)
route <- listToMaybe matchingRoutes
let requestPathText = TE.decodeUtf8 requestPath
sortedPaths = sortBy
(comparing (negate . T.length . pathRoutePath))
(routePaths route)
matchingPaths = filter
(\p -> pathMatches (pathRoutePath p) requestPathText)
sortedPaths
pathRoute <- listToMaybe matchingPaths
return (pathRouteUpstream pathRoute, pathRoute)
-- Find first matching path within the route
route <- listToMaybe matchingRoutes
let requestPathText = TE.decodeUtf8 requestPath
-- Sort paths by length (longest first) so more specific paths match first
sortedPaths = sortBy (comparing (negate . T.length . pathRoutePath)) (routePaths route)
matchingPaths = filter (\p -> pathMatches (pathRoutePath p) requestPathText) sortedPaths
pathRoute <- listToMaybe matchingPaths
return (pathRouteUpstream pathRoute, pathRoute)
-- | Check if a path pattern matches a request path
pathMatches :: Text -> Text -> Bool pathMatches :: Text -> Text -> Bool
pathMatches pattern requestPath = pathMatches pattern requestPath =
pattern == "/" || T.isPrefixOf pattern requestPath pattern == "/" || T.isPrefixOf pattern requestPath
-- | Select an upstream for a request (exported for testing) selectUpstream
selectUpstream :: Config -> Maybe BS.ByteString -> BS.ByteString -> Maybe Text :: Config -> Maybe BS.ByteString -> BS.ByteString -> Maybe Text
selectUpstream config hostHeader requestPath = selectUpstream config hostHeader requestPath =
fmap fst $ selectRoute config hostHeader requestPath fst <$> selectRoute config hostHeader requestPath
-- | Forward request to backend server with streaming support forwardRequest
forwardRequest :: Manager -> Request -> Text -> (Response -> IO ResponseReceived) -> IO ResponseReceived :: Manager
-> Request
-> Text
-> (Response -> IO ResponseReceived)
-> IO ResponseReceived
forwardRequest manager clientReq backendHost respond = do forwardRequest manager clientReq backendHost respond = do
let backendUrl = "http://" ++ T.unpack backendHost ++ let backendUrl = httpScheme ++ T.unpack backendHost
BS8.unpack (rawPathInfo clientReq) ++ ++ BS8.unpack (rawPathInfo clientReq)
BS8.unpack (rawQueryString clientReq) ++ BS8.unpack (rawQueryString clientReq)
initReq <- parseRequest backendUrl initReq <- parseRequest backendUrl
@ -536,56 +585,53 @@ forwardRequest manager clientReq backendHost respond = do
backendReq = initReq backendReq = initReq
{ HTTP.method = requestMethod clientReq { HTTP.method = requestMethod clientReq
, HTTP.requestHeaders = filterHeaders (requestHeaders clientReq) , HTTP.requestHeaders = filterRequestHeaders (requestHeaders clientReq)
, HTTP.requestBody = streamingBody , HTTP.requestBody = streamingBody
} }
let upstreamMicros = tcUpstreamReadSeconds defaultTimeoutConfig upstreamMicros = tcUpstreamReadSeconds defaultTimeoutConfig
* microsPerSecond * microsPerSecond
mResult <- timeout upstreamMicros $ mResult <- timeout upstreamMicros $
withResponse backendReq manager $ \backendResponse -> do withResponse backendReq manager $ \backendResponse -> do
let status = HTTP.responseStatus backendResponse let status = HTTP.responseStatus backendResponse
headers = HTTP.responseHeaders backendResponse headers = HTTP.responseHeaders backendResponse
bodyReader = HTTP.responseBody backendResponse bodyReader = HTTP.responseBody backendResponse
if shouldStreamResponse headers if shouldStreamResponse headers
then do then do
hPutStrLn stderr "[STREAM] Streaming response detected" hPutStrLn stderr "[STREAM] Streaming response detected"
respond $ responseStream status (filterResponseHeaders headers) $ \write flush -> do respond $ responseStream status (filterResponseHeaders headers) $
let loop = do \write flush -> do
chunk <- brRead bodyReader let loop = do
unless (BS.null chunk) $ do chunk <- brRead bodyReader
write (byteString chunk) unless (BS.null chunk) $ do
flush write (byteString chunk)
loop flush
loop loop
loop
else do else do
body <- readFullBody bodyReader body <- readFullBody bodyReader
respond $ responseLBS status (filterResponseHeaders headers) body respond $ responseLBS status (filterResponseHeaders headers) body
case mResult of case mResult of
Just rr -> return rr Just rr -> pure rr
Nothing -> respond $ responseLBS Nothing -> respond $ responseLBS
status504 status504
[("Content-Type", "text/plain")] [(hContentType, contentTypePlain)]
"504 Gateway Timeout: upstream did not respond in time" bodyGatewayTimeout
shouldStreamResponse :: [(HeaderName, BS.ByteString)] -> Bool shouldStreamResponse :: [(HeaderName, BS.ByteString)] -> Bool
shouldStreamResponse headers = shouldStreamResponse headers = isSSE || isChunkedWithoutLength
isSSE || isChunkedWithoutLength
where where
isSSE = case lookup "Content-Type" headers of isSSE = case lookup "Content-Type" headers of
Just ct -> "text/event-stream" `BS.isInfixOf` ct Just ct -> eventStreamContentType `BS.isInfixOf` ct
Nothing -> False Nothing -> False
isChunkedWithoutLength = hasChunkedEncoding && not hasContentLength
isChunkedWithoutLength =
hasChunkedEncoding && not hasContentLength
hasChunkedEncoding = case lookup "Transfer-Encoding" headers of hasChunkedEncoding = case lookup "Transfer-Encoding" headers of
Just te -> "chunked" `BS.isInfixOf` te Just te -> chunkedEncoding `BS.isInfixOf` te
Nothing -> False Nothing -> False
hasContentLength = case lookup "Content-Length" headers of hasContentLength = case lookup "Content-Length" headers of
Just _ -> True Just _ -> True
Nothing -> False Nothing -> False
readFullBody :: HTTP.BodyReader -> IO LBS.ByteString readFullBody :: HTTP.BodyReader -> IO LBS.ByteString
@ -594,45 +640,17 @@ readFullBody bodyReader = LBS.fromChunks <$> go
go = do go = do
chunk <- brRead bodyReader chunk <- brRead bodyReader
if BS.null chunk if BS.null chunk
then return [] then pure []
else do else do
rest <- go rest <- go
return (chunk : rest) pure (chunk : rest)
filterResponseHeaders :: [(HeaderName, BS.ByteString)] -> [(HeaderName, BS.ByteString)] filterResponseHeaders
filterResponseHeaders = filter (\(name, _) -> name `notElem` hopByHopHeaders) :: [(HeaderName, BS.ByteString)] -> [(HeaderName, BS.ByteString)]
where filterResponseHeaders =
hopByHopHeaders = filter (\(name, _) -> name `notElem` hopByHopResponseHeaders)
[ "Transfer-Encoding"
, "Connection"
, "Keep-Alive"
]
-- | Filter headers for regular HTTP (remove hop-by-hop headers) filterRequestHeaders
filterHeaders :: [(HeaderName, BS.ByteString)] -> [(HeaderName, BS.ByteString)] :: [(HeaderName, BS.ByteString)] -> [(HeaderName, BS.ByteString)]
filterHeaders headers = filter (\(name, _) -> name `notElem` hopByHopHeaders) headers filterRequestHeaders =
where filter (\(name, _) -> name `notElem` hopByHopRequestHeaders)
hopByHopHeaders =
[ "Connection"
, "Keep-Alive"
, "Proxy-Authenticate"
, "Proxy-Authorization"
, "TE"
, "Trailers"
, "Transfer-Encoding"
, "Upgrade"
]
-- | Log incoming request
logRequest :: Request -> IO ()
logRequest req = do
let method' = BS8.unpack (requestMethod req)
path = BS8.unpack (rawPathInfo req)
query = BS8.unpack (rawQueryString req)
host = fromMaybe "unknown" $ lookup "Host" (requestHeaders req)
putStrLn $ "[→] " ++ method' ++ " " ++ path ++ query ++ " (Host: " ++ BS8.unpack host ++ ")"
-- Helper: zipWithM
zipWithM :: Monad m => (a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM f xs ys = sequence (zipWith f xs ys)

View File

@ -31,6 +31,7 @@ import Control.Concurrent.STM
) )
import Data.ByteString (ByteString) import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BS8 import qualified Data.ByteString.Char8 as BS8
import Data.List (intercalate)
import Data.Map.Strict (Map) import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map import qualified Data.Map.Strict as Map
import Data.Text (Text) import Data.Text (Text)
@ -175,12 +176,7 @@ clientIPKey req = case remoteHost req of
v6Bytes ha = v6Bytes ha =
let (a, b, c, d, e, f, g, h) = hostAddress6ToTuple ha let (a, b, c, d, e, f, g, h) = hostAddress6ToTuple ha
parts = [a, b, c, d, e, f, g, h] parts = [a, b, c, d, e, f, g, h]
in BS8.pack (joinColons (map (`showHex` "") parts)) in BS8.pack (intercalate ":" (map (`showHex` "") parts))
joinColons :: [String] -> String
joinColons [] = ""
joinColons [x] = x
joinColons (x : xs) = x <> ":" <> joinColons xs
pathClassKey :: Request -> ByteString pathClassKey :: Request -> ByteString
pathClassKey req = pathClassKey req =

View File

@ -39,7 +39,7 @@ import Data.CaseInsensitive (CI)
import qualified Data.CaseInsensitive as CI import qualified Data.CaseInsensitive as CI
import Data.Char (toLower) import Data.Char (toLower)
import Data.Word (Word32) import Data.Word (Word32)
import Network.HTTP.Types (Status, mkStatus) import Network.HTTP.Types (status403)
import Network.HTTP.Types.URI (urlDecode) import Network.HTTP.Types.URI (urlDecode)
import Network.Wai import Network.Wai
( Middleware ( Middleware
@ -159,9 +159,6 @@ decisionFromScore hasBlock score threshold matches
wafResponseHeader :: CI ByteString wafResponseHeader :: CI ByteString
wafResponseHeader = "x-aenebris-waf" wafResponseHeader = "x-aenebris-waf"
status403 :: Status
status403 = mkStatus 403 "Forbidden"
wafMiddleware :: TVar RuleSet -> Middleware wafMiddleware :: TVar RuleSet -> Middleware
wafMiddleware rsVar app req respond = do wafMiddleware rsVar app req respond = do
rs <- readTVarIO rsVar rs <- readTVarIO rsVar

View File

@ -10,10 +10,13 @@ import Aenebris.Backend
( createRuntimeBackend ( createRuntimeBackend
, getConnectionCount , getConnectionCount
, isHealthy , isHealthy
, rbActiveConnections
, rbConsecutiveFailures
, rbServerId , rbServerId
, rbWeight , rbWeight
, recordFailure , recordFailure
, recordSuccess , recordSuccess
, trackConnection
, transitionToHealthy , transitionToHealthy
, transitionToRecovering , transitionToRecovering
, transitionToUnhealthy , transitionToUnhealthy
@ -30,7 +33,8 @@ import Aenebris.Config
, validateConfig , validateConfig
) )
import Aenebris.DDoS.ConnLimit import Aenebris.DDoS.ConnLimit
( defaultConnLimitConfig ( currentCount
, defaultConnLimitConfig
, defaultPerIPLimit , defaultPerIPLimit
, ipBytesFromSockAddr , ipBytesFromSockAddr
, newConnLimiter , newConnLimiter
@ -262,6 +266,7 @@ import Control.Concurrent.STM
( atomically ( atomically
, modifyTVar' , modifyTVar'
, newTVarIO , newTVarIO
, readTVar
, readTVarIO , readTVarIO
) )
import qualified Data.ByteString as BS import qualified Data.ByteString as BS
@ -477,31 +482,46 @@ loadBalancerSpec = describe "LoadBalancer" $ do
lb <- createLoadBalancer RoundRobin [] lb <- createLoadBalancer RoundRobin []
selectBackend lb `shouldReturn` Nothing selectBackend lb `shouldReturn` Nothing
it "selects from backend pool with round robin" $ do it "round robin distributes evenly across the pool" $ do
bks <- mapM (\(i, h) -> createRuntimeBackend i (Server h 1)) bks <- mapM (\(i, h) -> createRuntimeBackend i (Server h 1))
[(0, "host-a:80"), (1, "host-b:80"), (2, "host-c:80")] [(0, "host-a:80"), (1, "host-b:80"), (2, "host-c:80")]
lb <- createLoadBalancer RoundRobin bks lb <- createLoadBalancer RoundRobin bks
let getName = fmap (fmap rbServerId) (selectBackend lb) let totalRounds = 9 :: Int
a <- getName selections <- mapM
b <- getName (\_ -> fmap (fmap rbServerId) (selectBackend lb))
c <- getName [1 .. totalRounds]
isJust a `shouldBe` True let counts =
isJust b `shouldBe` True [ length (filter (== Just sid) selections)
isJust c `shouldBe` True | sid <- [0, 1, 2]
]
counts `shouldBe` [3, 3, 3]
it "selects backend with weighted round robin" $ do it "weighted round robin selects proportionally" $ do
bks <- mapM (\(i, h, w) -> createRuntimeBackend i (Server h w)) bks <- mapM (\(i, h, w) -> createRuntimeBackend i (Server h w))
[(0, "host-a:80", 1), (1, "host-b:80", 4)] [(0, "host-a:80", 1), (1, "host-b:80", 4)]
lb <- createLoadBalancer WeightedRoundRobin bks lb <- createLoadBalancer WeightedRoundRobin bks
selected <- selectBackend lb let totalRounds = 50 :: Int
isJust selected `shouldBe` True selections <- mapM
(\_ -> fmap (fmap rbServerId) (selectBackend lb))
[1 .. totalRounds]
let countA = length (filter (== Just 0) selections)
countB = length (filter (== Just 1) selections)
countB `shouldSatisfy` (>= 35)
countA `shouldSatisfy` (<= 15)
it "selects least connections backend" $ do it "least connections picks the backend with fewest active connections" $ do
bks <- mapM (\(i, h) -> createRuntimeBackend i (Server h 1)) bks <- mapM (\(i, h) -> createRuntimeBackend i (Server h 1))
[(0, "host-a:80"), (1, "host-b:80")] [(0, "host-a:80"), (1, "host-b:80")]
lb <- createLoadBalancer LeastConnections bks case bks of
selected <- selectBackend lb [a, _b] -> do
isJust selected `shouldBe` True atomically $
modifyTVar'
(rbActiveConnections a)
(+ 5)
lb <- createLoadBalancer LeastConnections bks
selected <- selectBackend lb
fmap rbServerId selected `shouldBe` Just 1
_ -> expectationFailure "expected exactly two backends"
backendSpec :: Spec backendSpec :: Spec
backendSpec = describe "Backend" $ do backendSpec = describe "Backend" $ do
@ -526,17 +546,21 @@ backendSpec = describe "Backend" $ do
bk <- createRuntimeBackend 0 (Server "host:80" 1) bk <- createRuntimeBackend 0 (Server "host:80" 1)
atomically (getConnectionCount bk) `shouldReturn` 0 atomically (getConnectionCount bk) `shouldReturn` 0
it "tolerates repeated failures" $ do it "transitions to Unhealthy after maxFailures consecutive failures" $ do
bk <- createRuntimeBackend 0 (Server "host:80" 10) bk <- createRuntimeBackend 0 (Server "host:80" 10)
atomically $ recordFailure bk 3 atomically (recordFailure bk 3)
atomically $ recordFailure bk 3 atomically (isHealthy bk) `shouldReturn` True
atomically $ recordFailure bk 3 atomically (recordFailure bk 3)
rbWeight bk `shouldBe` 10 atomically (isHealthy bk) `shouldReturn` True
atomically (recordFailure bk 3)
atomically (isHealthy bk) `shouldReturn` False
it "records successes without crashing" $ do it "recordSuccess on Healthy resets the failure counter" $ do
bk <- createRuntimeBackend 0 (Server "host:80" 5) bk <- createRuntimeBackend 0 (Server "host:80" 5)
atomically $ recordSuccess bk 5 atomically (recordFailure bk 5)
pure () atomically (readTVar (rbConsecutiveFailures bk)) `shouldReturn` 1
atomically (recordSuccess bk 5)
atomically (readTVar (rbConsecutiveFailures bk)) `shouldReturn` 0
securitySpec :: Spec securitySpec :: Spec
securitySpec = describe "Security headers" $ do securitySpec = describe "Security headers" $ do
@ -752,11 +776,12 @@ connLimitSpec = describe "ConnLimit" $ do
res <- atomically (tryAcquire cl "9.9.9.9") res <- atomically (tryAcquire cl "9.9.9.9")
res `shouldBe` False res `shouldBe` False
it "release decrements counter" $ do it "release decrements counter back to 0" $ do
cl <- newConnLimiter defaultConnLimitConfig cl <- newConnLimiter defaultConnLimitConfig
_ <- atomically (tryAcquire cl "1.2.3.4") _ <- atomically (tryAcquire cl "1.2.3.4")
atomically (currentCount cl "1.2.3.4") `shouldReturn` 1
atomically (release cl "1.2.3.4") atomically (release cl "1.2.3.4")
pure () atomically (currentCount cl "1.2.3.4") `shouldReturn` 0
ja4hSpec :: Spec ja4hSpec :: Spec
ja4hSpec = describe "JA4H fingerprint" $ do ja4hSpec = describe "JA4H fingerprint" $ do