From a4a9e1170a45d06030f803fa3aa7146fa2410f2a Mon Sep 17 00:00:00 2001 From: CarterPerez-dev Date: Wed, 29 Apr 2026 01:04:01 -0400 Subject: [PATCH] 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. --- .../haskell-reverse-proxy/aenebris.cabal | 1 + .../src/Aenebris/ML/Middleware.hs | 170 ++++++++++++++++++ .../haskell-reverse-proxy/test/Spec.hs | 104 +++++++++++ 3 files changed, 275 insertions(+) create mode 100644 PROJECTS/advanced/haskell-reverse-proxy/src/Aenebris/ML/Middleware.hs diff --git a/PROJECTS/advanced/haskell-reverse-proxy/aenebris.cabal b/PROJECTS/advanced/haskell-reverse-proxy/aenebris.cabal index 00aa1c34..a95ad7bd 100644 --- a/PROJECTS/advanced/haskell-reverse-proxy/aenebris.cabal +++ b/PROJECTS/advanced/haskell-reverse-proxy/aenebris.cabal @@ -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 diff --git a/PROJECTS/advanced/haskell-reverse-proxy/src/Aenebris/ML/Middleware.hs b/PROJECTS/advanced/haskell-reverse-proxy/src/Aenebris/ML/Middleware.hs new file mode 100644 index 00000000..1433542f --- /dev/null +++ b/PROJECTS/advanced/haskell-reverse-proxy/src/Aenebris/ML/Middleware.hs @@ -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 = + "\n\ + \\n\ + \Verification Required\n\ + \\n\ + \

Verification Required

\n\ + \

Aenebris ML flagged this request as ambiguous. Please complete a \ + \verification challenge to continue.

\n\ + \\n\ + \\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 ++) diff --git a/PROJECTS/advanced/haskell-reverse-proxy/test/Spec.hs b/PROJECTS/advanced/haskell-reverse-proxy/test/Spec.hs index acf0730b..c9f1ca5e 100644 --- a/PROJECTS/advanced/haskell-reverse-proxy/test/Spec.hs +++ b/PROJECTS/advanced/haskell-reverse-proxy/test/Spec.hs @@ -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