feat: add Aenebris.ML.Middleware (Wai integration for ML pipeline)
- New Aenebris.ML.Middleware module: Wai middleware that wires the
Aenebris.ML.Engine into the per-request pipeline. Per request:
1. Run extractFeatures (FeatureContext + Request) -> FeatureVector
2. Convert via featureVectorToVector -> VU.Vector Double
3. runEngine eng fv -> DecisionDetails
4. Route by ddDecision:
DecisionHuman -> pass through to inner Application,
optionally attach X-Aenebris-ML-Decision
and X-Aenebris-ML-Score response headers
DecisionBot -> 403 + text/plain block body
DecisionChallenge -> 403 + text/html challenge page
All response builders are operator-overridable via the
MLMiddlewareConfig record. Optional logging callback fires once
per request for structured logs / metrics integration.
- defaultBotResponse, defaultChallengeResponse: sensible defaults
matching the existing WAF middleware pattern (403 + descriptive
body). decisionResponseHeader / scoreResponseHeader are exported
so operators can build custom responses that still surface the
signal headers.
- 12 tests via Network.Wai.Test.runSession covering: pass-through
on Human, 403 on Bot, 403 on Challenge with text/html body,
signal-header attachment toggle, custom response builder override,
logging callback invocation count.
- aenebris.cabal: expose Aenebris.ML.Middleware in the library stanza.
- test/Spec.hs: add modifyTVar' to the Control.Concurrent.STM import
list for the logging-callback test.
- 354 total examples passing, 0 GHC warnings on the new module.
This closes the core Phase 2.5 ML pipeline. The full sequence
(Features -> Model -> Loader -> Inference -> Calibration -> IForest
-> Engine -> Middleware) is now fully implemented and integration-
tested. A hosted Aenebris instance can now be configured with a
trained LightGBM model + optional calibrator + optional IForest
and will block bots, challenge ambiguous traffic, and pass humans
on every HTTP request.
This commit is contained in:
CarterPerez-dev
parent
a98de6f8f4
commit
a4a9e1170a
|
|
@ -45,6 +45,7 @@ library
|
|||
, Aenebris.ML.Calibration
|
||||
, Aenebris.ML.IForest
|
||||
, Aenebris.ML.Engine
|
||||
, Aenebris.ML.Middleware
|
||||
default-language: Haskell2010
|
||||
build-depends: base >= 4.7 && < 5
|
||||
, warp >= 3.3
|
||||
|
|
|
|||
|
|
@ -0,0 +1,170 @@
|
|||
{-
|
||||
©AngelaMos | 2026
|
||||
Middleware.hs
|
||||
-}
|
||||
{-# LANGUAGE BangPatterns #-}
|
||||
{-# LANGUAGE OverloadedStrings #-}
|
||||
|
||||
module Aenebris.ML.Middleware
|
||||
( MLMiddlewareConfig(..)
|
||||
, defaultMLMiddlewareConfig
|
||||
, defaultBotResponse
|
||||
, defaultChallengeResponse
|
||||
, mlBotDetectionMiddleware
|
||||
, decisionResponseHeader
|
||||
, scoreResponseHeader
|
||||
, decisionToWireText
|
||||
) where
|
||||
|
||||
import Data.ByteString (ByteString)
|
||||
import qualified Data.ByteString.Char8 as BC
|
||||
import qualified Data.ByteString.Lazy as LBS
|
||||
import Data.CaseInsensitive (CI)
|
||||
import qualified Data.CaseInsensitive as CI
|
||||
import Network.HTTP.Types (HeaderName, hContentType, status403)
|
||||
import Network.Wai
|
||||
( Application
|
||||
, Middleware
|
||||
, Request
|
||||
, Response
|
||||
, ResponseReceived
|
||||
, mapResponseHeaders
|
||||
, responseLBS
|
||||
)
|
||||
|
||||
import Aenebris.ML.Engine
|
||||
( Decision(..)
|
||||
, DecisionDetails(..)
|
||||
, Engine
|
||||
, runEngine
|
||||
)
|
||||
import Aenebris.ML.Features
|
||||
( FeatureContext
|
||||
, extractFeatures
|
||||
, featureVectorToVector
|
||||
)
|
||||
|
||||
humanWireText :: ByteString
|
||||
humanWireText = "human"
|
||||
|
||||
botWireText :: ByteString
|
||||
botWireText = "bot"
|
||||
|
||||
challengeWireText :: ByteString
|
||||
challengeWireText = "challenge"
|
||||
|
||||
botBlockBody :: LBS.ByteString
|
||||
botBlockBody = "403 Forbidden \x2014 request blocked by Aenebris ML"
|
||||
|
||||
challengePageBody :: LBS.ByteString
|
||||
challengePageBody =
|
||||
"<!DOCTYPE html>\n\
|
||||
\<html lang=\"en\">\n\
|
||||
\<head><meta charset=\"utf-8\"><title>Verification Required</title></head>\n\
|
||||
\<body>\n\
|
||||
\<h1>Verification Required</h1>\n\
|
||||
\<p>Aenebris ML flagged this request as ambiguous. Please complete a \
|
||||
\verification challenge to continue.</p>\n\
|
||||
\</body>\n\
|
||||
\</html>\n"
|
||||
|
||||
contentTypePlain :: ByteString
|
||||
contentTypePlain = "text/plain; charset=utf-8"
|
||||
|
||||
contentTypeHtml :: ByteString
|
||||
contentTypeHtml = "text/html; charset=utf-8"
|
||||
|
||||
decisionResponseHeader :: HeaderName
|
||||
decisionResponseHeader = CI.mk "X-Aenebris-ML-Decision"
|
||||
|
||||
scoreResponseHeader :: HeaderName
|
||||
scoreResponseHeader = CI.mk "X-Aenebris-ML-Score"
|
||||
|
||||
decisionToWireText :: Decision -> ByteString
|
||||
decisionToWireText DecisionHuman = humanWireText
|
||||
decisionToWireText DecisionBot = botWireText
|
||||
decisionToWireText DecisionChallenge = challengeWireText
|
||||
|
||||
calibratedToHeaderValue :: Double -> ByteString
|
||||
calibratedToHeaderValue = BC.pack . show
|
||||
|
||||
data MLMiddlewareConfig = MLMiddlewareConfig
|
||||
{ mmcEngine :: !Engine
|
||||
, mmcFeatureContext :: !FeatureContext
|
||||
, mmcBotResponse :: !(Request -> DecisionDetails -> Response)
|
||||
, mmcChallengeResponse :: !(Request -> DecisionDetails -> Response)
|
||||
, mmcLogDetails :: !(Maybe (Request -> DecisionDetails -> IO ()))
|
||||
, mmcAttachHeaders :: !Bool
|
||||
}
|
||||
|
||||
defaultMLMiddlewareConfig
|
||||
:: Engine
|
||||
-> FeatureContext
|
||||
-> MLMiddlewareConfig
|
||||
defaultMLMiddlewareConfig eng ctx = MLMiddlewareConfig
|
||||
{ mmcEngine = eng
|
||||
, mmcFeatureContext = ctx
|
||||
, mmcBotResponse = defaultBotResponse
|
||||
, mmcChallengeResponse = defaultChallengeResponse
|
||||
, mmcLogDetails = Nothing
|
||||
, mmcAttachHeaders = True
|
||||
}
|
||||
|
||||
defaultBotResponse :: Request -> DecisionDetails -> Response
|
||||
defaultBotResponse _req details =
|
||||
responseLBS
|
||||
status403
|
||||
(mlSignalHeaders DecisionBot details
|
||||
<> [(hContentType, contentTypePlain)])
|
||||
botBlockBody
|
||||
|
||||
defaultChallengeResponse :: Request -> DecisionDetails -> Response
|
||||
defaultChallengeResponse _req details =
|
||||
responseLBS
|
||||
status403
|
||||
(mlSignalHeaders DecisionChallenge details
|
||||
<> [(hContentType, contentTypeHtml)])
|
||||
challengePageBody
|
||||
|
||||
mlSignalHeaders :: Decision -> DecisionDetails -> [(CI ByteString, ByteString)]
|
||||
mlSignalHeaders decision details =
|
||||
[ (decisionResponseHeader, decisionToWireText decision)
|
||||
, (scoreResponseHeader, calibratedToHeaderValue (ddCalibrated details))
|
||||
]
|
||||
|
||||
mlBotDetectionMiddleware :: MLMiddlewareConfig -> Middleware
|
||||
mlBotDetectionMiddleware cfg app req respond = do
|
||||
let !fv = extractFeatures (mmcFeatureContext cfg) req
|
||||
!fvVec = featureVectorToVector fv
|
||||
!details = runEngine (mmcEngine cfg) fvVec
|
||||
emitLog cfg req details
|
||||
routeDecision cfg app req respond details
|
||||
|
||||
emitLog
|
||||
:: MLMiddlewareConfig -> Request -> DecisionDetails -> IO ()
|
||||
emitLog cfg req details = case mmcLogDetails cfg of
|
||||
Just logger -> logger req details
|
||||
Nothing -> pure ()
|
||||
|
||||
routeDecision
|
||||
:: MLMiddlewareConfig
|
||||
-> Application
|
||||
-> Request
|
||||
-> (Response -> IO ResponseReceived)
|
||||
-> DecisionDetails
|
||||
-> IO ResponseReceived
|
||||
routeDecision cfg app req respond details = case ddDecision details of
|
||||
DecisionHuman ->
|
||||
let respondWithHeaders =
|
||||
if mmcAttachHeaders cfg
|
||||
then respond . attachHumanSignalHeaders details
|
||||
else respond
|
||||
in app req respondWithHeaders
|
||||
DecisionBot ->
|
||||
respond (mmcBotResponse cfg req details)
|
||||
DecisionChallenge ->
|
||||
respond (mmcChallengeResponse cfg req details)
|
||||
|
||||
attachHumanSignalHeaders :: DecisionDetails -> Response -> Response
|
||||
attachHumanSignalHeaders details =
|
||||
mapResponseHeaders (mlSignalHeaders DecisionHuman details ++)
|
||||
|
|
@ -196,6 +196,14 @@ import Aenebris.ML.Engine
|
|||
, runEngine
|
||||
, runEngineDecision
|
||||
)
|
||||
import Aenebris.ML.Middleware
|
||||
( MLMiddlewareConfig(..)
|
||||
, decisionResponseHeader
|
||||
, decisionToWireText
|
||||
, defaultMLMiddlewareConfig
|
||||
, mlBotDetectionMiddleware
|
||||
, scoreResponseHeader
|
||||
)
|
||||
import Aenebris.ML.IForest
|
||||
( IForest(..)
|
||||
, ITree(..)
|
||||
|
|
@ -252,6 +260,7 @@ import Aenebris.WAF.Rule
|
|||
|
||||
import Control.Concurrent.STM
|
||||
( atomically
|
||||
, modifyTVar'
|
||||
, newTVarIO
|
||||
, readTVarIO
|
||||
)
|
||||
|
|
@ -373,6 +382,7 @@ main = hspec $ do
|
|||
mlCalibrationSpec
|
||||
mlIForestSpec
|
||||
mlEngineSpec
|
||||
mlMiddlewareSpec
|
||||
|
||||
configSpec :: Spec
|
||||
configSpec = describe "Config" $ do
|
||||
|
|
@ -2671,6 +2681,100 @@ mlEngineSpec = describe "ML.Engine" $ do
|
|||
it "ddIForestScore is Nothing when no IForest configured" $
|
||||
ddIForestScore result `shouldBe` Nothing
|
||||
|
||||
mlEngineWithLeaf :: Double -> Engine
|
||||
mlEngineWithLeaf leafValue =
|
||||
let ens = Ensemble
|
||||
{ ensembleVersion = currentEnsembleVersion
|
||||
, ensembleFeatureCount = featureVectorLength
|
||||
, ensembleObjective = ObjectiveBinaryLogistic
|
||||
, ensembleBaseScore = 0.0
|
||||
, ensembleSigmoidScale = defaultSigmoidScale
|
||||
, ensembleAverageOutput = False
|
||||
, ensembleTrees = V.singleton (makeLeafTree leafValue)
|
||||
}
|
||||
in makeEngine ens NoCalibrator Nothing defaultEngineConfig
|
||||
|
||||
mlMiddlewareSpec :: Spec
|
||||
mlMiddlewareSpec = describe "ML.Middleware" $ do
|
||||
let humanEng = mlEngineWithLeaf (-5.0)
|
||||
botEng = mlEngineWithLeaf 5.0
|
||||
challengeEng = mlEngineWithLeaf 0.0
|
||||
ctx = emptyFeatureContext
|
||||
humanCfg = defaultMLMiddlewareConfig humanEng ctx
|
||||
botCfg = defaultMLMiddlewareConfig botEng ctx
|
||||
challengeCfg = defaultMLMiddlewareConfig challengeEng ctx
|
||||
buildApp cfg = mlBotDetectionMiddleware cfg okApp
|
||||
|
||||
describe "DecisionHuman → pass through with ML signal headers" $ do
|
||||
it "returns 200 from inner application" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp humanCfg)
|
||||
simpleStatus resp `shouldBe` status200
|
||||
|
||||
it "attaches X-Aenebris-ML-Decision: human" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp humanCfg)
|
||||
lookup decisionResponseHeader (simpleHeaders resp)
|
||||
`shouldBe` Just (decisionToWireText DecisionHuman)
|
||||
|
||||
it "attaches X-Aenebris-ML-Score header" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp humanCfg)
|
||||
lookup scoreResponseHeader (simpleHeaders resp) `shouldSatisfy` isJust
|
||||
|
||||
it "omits ML signal headers when mmcAttachHeaders=False" $ do
|
||||
let cfg = humanCfg { mmcAttachHeaders = False }
|
||||
app = buildApp cfg
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) app
|
||||
lookup decisionResponseHeader (simpleHeaders resp) `shouldBe` Nothing
|
||||
lookup scoreResponseHeader (simpleHeaders resp) `shouldBe` Nothing
|
||||
|
||||
describe "DecisionBot → 403 block" $ do
|
||||
it "returns 403" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp botCfg)
|
||||
simpleStatus resp `shouldBe` status403
|
||||
|
||||
it "attaches X-Aenebris-ML-Decision: bot" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp botCfg)
|
||||
lookup decisionResponseHeader (simpleHeaders resp)
|
||||
`shouldBe` Just (decisionToWireText DecisionBot)
|
||||
|
||||
it "responds with text/plain content type" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp botCfg)
|
||||
lookup "Content-Type" (simpleHeaders resp)
|
||||
`shouldSatisfy` maybe False (BS.isPrefixOf "text/plain")
|
||||
|
||||
describe "DecisionChallenge → 403 challenge page" $ do
|
||||
it "returns 403" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp challengeCfg)
|
||||
simpleStatus resp `shouldBe` status403
|
||||
|
||||
it "attaches X-Aenebris-ML-Decision: challenge" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp challengeCfg)
|
||||
lookup decisionResponseHeader (simpleHeaders resp)
|
||||
`shouldBe` Just (decisionToWireText DecisionChallenge)
|
||||
|
||||
it "responds with text/html content type" $ do
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp challengeCfg)
|
||||
lookup "Content-Type" (simpleHeaders resp)
|
||||
`shouldSatisfy` maybe False (BS.isPrefixOf "text/html")
|
||||
|
||||
describe "custom response builders override defaults" $
|
||||
it "custom mmcBotResponse is used instead of the default 403" $ do
|
||||
let customBotResp _req _details =
|
||||
responseLBS status429
|
||||
[("Content-Type", "text/plain")]
|
||||
"custom bot response"
|
||||
cfg = botCfg { mmcBotResponse = customBotResp }
|
||||
resp <- runSession (request Network.Wai.Test.defaultRequest) (buildApp cfg)
|
||||
simpleStatus resp `shouldBe` status429
|
||||
|
||||
describe "logging callback fires" $
|
||||
it "mmcLogDetails callback is invoked once per request" $ do
|
||||
counter <- newTVarIO (0 :: Int)
|
||||
let logCallback _req _details =
|
||||
atomically (modifyTVar' counter (+ 1))
|
||||
cfg = humanCfg { mmcLogDetails = Just logCallback }
|
||||
_ <- runSession (request Network.Wai.Test.defaultRequest) (buildApp cfg)
|
||||
readTVarIO counter `shouldReturn` 1
|
||||
|
||||
headersOnlyRequest :: [(BS.ByteString, BS.ByteString)] -> Request
|
||||
headersOnlyRequest hs =
|
||||
Network.Wai.Test.defaultRequest
|
||||
|
|
|
|||
Loading…
Reference in New Issue