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.Calibration
|
||||||
, Aenebris.ML.IForest
|
, Aenebris.ML.IForest
|
||||||
, Aenebris.ML.Engine
|
, Aenebris.ML.Engine
|
||||||
|
, Aenebris.ML.Middleware
|
||||||
default-language: Haskell2010
|
default-language: Haskell2010
|
||||||
build-depends: base >= 4.7 && < 5
|
build-depends: base >= 4.7 && < 5
|
||||||
, warp >= 3.3
|
, 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
|
, runEngine
|
||||||
, runEngineDecision
|
, runEngineDecision
|
||||||
)
|
)
|
||||||
|
import Aenebris.ML.Middleware
|
||||||
|
( MLMiddlewareConfig(..)
|
||||||
|
, decisionResponseHeader
|
||||||
|
, decisionToWireText
|
||||||
|
, defaultMLMiddlewareConfig
|
||||||
|
, mlBotDetectionMiddleware
|
||||||
|
, scoreResponseHeader
|
||||||
|
)
|
||||||
import Aenebris.ML.IForest
|
import Aenebris.ML.IForest
|
||||||
( IForest(..)
|
( IForest(..)
|
||||||
, ITree(..)
|
, ITree(..)
|
||||||
|
|
@ -252,6 +260,7 @@ import Aenebris.WAF.Rule
|
||||||
|
|
||||||
import Control.Concurrent.STM
|
import Control.Concurrent.STM
|
||||||
( atomically
|
( atomically
|
||||||
|
, modifyTVar'
|
||||||
, newTVarIO
|
, newTVarIO
|
||||||
, readTVarIO
|
, readTVarIO
|
||||||
)
|
)
|
||||||
|
|
@ -373,6 +382,7 @@ main = hspec $ do
|
||||||
mlCalibrationSpec
|
mlCalibrationSpec
|
||||||
mlIForestSpec
|
mlIForestSpec
|
||||||
mlEngineSpec
|
mlEngineSpec
|
||||||
|
mlMiddlewareSpec
|
||||||
|
|
||||||
configSpec :: Spec
|
configSpec :: Spec
|
||||||
configSpec = describe "Config" $ do
|
configSpec = describe "Config" $ do
|
||||||
|
|
@ -2671,6 +2681,100 @@ mlEngineSpec = describe "ML.Engine" $ do
|
||||||
it "ddIForestScore is Nothing when no IForest configured" $
|
it "ddIForestScore is Nothing when no IForest configured" $
|
||||||
ddIForestScore result `shouldBe` Nothing
|
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 :: [(BS.ByteString, BS.ByteString)] -> Request
|
||||||
headersOnlyRequest hs =
|
headersOnlyRequest hs =
|
||||||
Network.Wai.Test.defaultRequest
|
Network.Wai.Test.defaultRequest
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue