361 lines
11 KiB
Haskell
361 lines
11 KiB
Haskell
{-
|
|
©AngelaMos | 2026
|
|
Geo.hs
|
|
-}
|
|
{-# LANGUAGE BangPatterns #-}
|
|
{-# LANGUAGE DeriveGeneric #-}
|
|
{-# LANGUAGE NumericUnderscores #-}
|
|
{-# LANGUAGE OverloadedStrings #-}
|
|
{-# LANGUAGE RecordWildCards #-}
|
|
|
|
module Aenebris.Geo
|
|
( GeoAction(..)
|
|
, GeoConfig(..)
|
|
, GeoConfigYaml(..)
|
|
, GeoInfo(..)
|
|
, GeoDecision(..)
|
|
, AsnWindow(..)
|
|
, Geo(..)
|
|
, defaultGeoLanguage
|
|
, defaultGeoConcentrationWindowSeconds
|
|
, defaultGeoConcentrationThreshold
|
|
, defaultGeoJailCooldownSeconds
|
|
, defaultGeoSweepIntervalMicros
|
|
, defaultGeoFlaggedAsns
|
|
, geoResponseHeaderName
|
|
, parseGeoAction
|
|
, sockAddrToIP
|
|
, countryBlocked
|
|
, emptyGeoInfo
|
|
, lookupGeo
|
|
, bumpAsnCounter
|
|
, asnConcentrationScore
|
|
, purgeAsnCounters
|
|
, startAsnSweeper
|
|
, decideGeo
|
|
, buildGeoConfig
|
|
, openGeo
|
|
, geoMiddleware
|
|
, renderGeoHeader
|
|
) where
|
|
|
|
import Aenebris.DDoS.IPJail (IPJail, jail)
|
|
import Aenebris.RateLimit (clientIPKey)
|
|
|
|
import Control.Concurrent (threadDelay)
|
|
import Control.Concurrent.Async (Async, async)
|
|
import Control.Concurrent.STM
|
|
( STM
|
|
, TVar
|
|
, atomically
|
|
, modifyTVar'
|
|
, newTVarIO
|
|
, readTVar
|
|
, writeTVar
|
|
)
|
|
import Control.Monad (forever, when)
|
|
import Data.Aeson (FromJSON(..), withObject, (.:?), (.!=))
|
|
import Data.ByteString (ByteString)
|
|
import qualified Data.ByteString.Char8 as BC
|
|
import qualified Data.ByteString.Lazy as LBS
|
|
import Data.GeoIP2 (GeoDB, GeoResult(..), AS(..), findGeoData, openGeoDB)
|
|
import Data.IP (IP(..), fromHostAddress, fromHostAddress6)
|
|
import Data.Map.Strict (Map)
|
|
import qualified Data.Map.Strict as Map
|
|
import Data.Maybe (fromMaybe, isJust)
|
|
import Data.Text (Text)
|
|
import qualified Data.Text as T
|
|
import qualified Data.Text.Encoding as TE
|
|
import Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime)
|
|
import GHC.Generics (Generic)
|
|
import Network.HTTP.Types (HeaderName, status403)
|
|
import Network.Socket (SockAddr(..))
|
|
import Network.Wai
|
|
( Middleware
|
|
, mapResponseHeaders
|
|
, remoteHost
|
|
, responseLBS
|
|
)
|
|
|
|
defaultGeoLanguage :: Text
|
|
defaultGeoLanguage = "en"
|
|
|
|
defaultGeoConcentrationWindowSeconds :: Int
|
|
defaultGeoConcentrationWindowSeconds = 60
|
|
|
|
defaultGeoConcentrationThreshold :: Int
|
|
defaultGeoConcentrationThreshold = 500
|
|
|
|
defaultGeoJailCooldownSeconds :: Int
|
|
defaultGeoJailCooldownSeconds = 600
|
|
|
|
defaultGeoSweepIntervalMicros :: Int
|
|
defaultGeoSweepIntervalMicros = 60_000_000
|
|
|
|
defaultGeoFlaggedAsns :: [Int]
|
|
defaultGeoFlaggedAsns = []
|
|
|
|
geoResponseHeaderName :: HeaderName
|
|
geoResponseHeaderName = "x-aenebris-geo"
|
|
|
|
data GeoAction
|
|
= GeoActionLog
|
|
| GeoActionJail
|
|
deriving (Show, Eq, Generic)
|
|
|
|
parseGeoAction :: Text -> GeoAction
|
|
parseGeoAction t = case T.toLower t of
|
|
"jail" -> GeoActionJail
|
|
"log" -> GeoActionLog
|
|
_ -> GeoActionLog
|
|
|
|
data GeoConfig = GeoConfig
|
|
{ gcCountryDb :: !(Maybe FilePath)
|
|
, gcAsnDb :: !(Maybe FilePath)
|
|
, gcBlockedCountries :: ![Text]
|
|
, gcAllowedCountries :: ![Text]
|
|
, gcFlaggedAsns :: ![Int]
|
|
, gcConcentrationWindowSeconds :: !Int
|
|
, gcConcentrationThreshold :: !Int
|
|
, gcJailCooldownSeconds :: !Int
|
|
, gcAction :: !GeoAction
|
|
, gcAnnotateHeader :: !Bool
|
|
, gcLanguage :: !Text
|
|
} deriving (Show, Eq)
|
|
|
|
data GeoConfigYaml = GeoConfigYaml
|
|
{ gcyEnabled :: !Bool
|
|
, gcyCountryDb :: !(Maybe FilePath)
|
|
, gcyAsnDb :: !(Maybe FilePath)
|
|
, gcyBlockedCountries :: ![Text]
|
|
, gcyAllowedCountries :: ![Text]
|
|
, gcyFlaggedAsns :: ![Int]
|
|
, gcyWindowSeconds :: !(Maybe Int)
|
|
, gcyThreshold :: !(Maybe Int)
|
|
, gcyJailCooldownSeconds :: !(Maybe Int)
|
|
, gcyAction :: !(Maybe Text)
|
|
, gcyAnnotateHeader :: !Bool
|
|
, gcyLanguage :: !(Maybe Text)
|
|
} deriving (Show, Eq, Generic)
|
|
|
|
instance FromJSON GeoConfigYaml where
|
|
parseJSON = withObject "GeoConfigYaml" $ \v -> GeoConfigYaml
|
|
<$> v .:? "enabled" .!= False
|
|
<*> v .:? "country_db"
|
|
<*> v .:? "asn_db"
|
|
<*> v .:? "blocked_countries" .!= []
|
|
<*> v .:? "allowed_countries" .!= []
|
|
<*> v .:? "flagged_asns" .!= []
|
|
<*> v .:? "window_seconds"
|
|
<*> v .:? "threshold"
|
|
<*> v .:? "jail_cooldown_seconds"
|
|
<*> v .:? "action"
|
|
<*> v .:? "annotate_header" .!= True
|
|
<*> v .:? "language"
|
|
|
|
buildGeoConfig :: Maybe GeoConfigYaml -> Maybe GeoConfig
|
|
buildGeoConfig Nothing = Nothing
|
|
buildGeoConfig (Just y)
|
|
| not (gcyEnabled y) = Nothing
|
|
| otherwise = Just GeoConfig
|
|
{ gcCountryDb = gcyCountryDb y
|
|
, gcAsnDb = gcyAsnDb y
|
|
, gcBlockedCountries = map T.toUpper (gcyBlockedCountries y)
|
|
, gcAllowedCountries = map T.toUpper (gcyAllowedCountries y)
|
|
, gcFlaggedAsns = gcyFlaggedAsns y
|
|
, gcConcentrationWindowSeconds =
|
|
fromMaybe defaultGeoConcentrationWindowSeconds (gcyWindowSeconds y)
|
|
, gcConcentrationThreshold =
|
|
fromMaybe defaultGeoConcentrationThreshold (gcyThreshold y)
|
|
, gcJailCooldownSeconds =
|
|
fromMaybe defaultGeoJailCooldownSeconds (gcyJailCooldownSeconds y)
|
|
, gcAction = maybe GeoActionLog parseGeoAction (gcyAction y)
|
|
, gcAnnotateHeader = gcyAnnotateHeader y
|
|
, gcLanguage = fromMaybe defaultGeoLanguage (gcyLanguage y)
|
|
}
|
|
|
|
data GeoInfo = GeoInfo
|
|
{ giCountryISO :: !(Maybe Text)
|
|
, giAsnNumber :: !(Maybe Int)
|
|
, giAsnOrg :: !(Maybe Text)
|
|
, giFlaggedAsn :: !Bool
|
|
} deriving (Show, Eq)
|
|
|
|
emptyGeoInfo :: GeoInfo
|
|
emptyGeoInfo = GeoInfo Nothing Nothing Nothing False
|
|
|
|
data GeoDecision
|
|
= GeoAllow
|
|
| GeoBlockCountry !Text
|
|
| GeoJailAsn !Int !Double
|
|
deriving (Show, Eq)
|
|
|
|
data AsnWindow = AsnWindow
|
|
{ awCount :: !Int
|
|
, awWindowStart :: !POSIXTime
|
|
} deriving (Show, Eq)
|
|
|
|
data Geo = Geo
|
|
{ geoCountryDb :: !(Maybe GeoDB)
|
|
, geoAsnDb :: !(Maybe GeoDB)
|
|
, geoConfig :: !GeoConfig
|
|
, geoAsnCounts :: !(TVar (Map Int AsnWindow))
|
|
}
|
|
|
|
openGeo :: GeoConfig -> IO Geo
|
|
openGeo cfg = do
|
|
countryDb <- traverse openGeoDB (gcCountryDb cfg)
|
|
asnDb <- traverse openGeoDB (gcAsnDb cfg)
|
|
counts <- newTVarIO Map.empty
|
|
pure Geo
|
|
{ geoCountryDb = countryDb
|
|
, geoAsnDb = asnDb
|
|
, geoConfig = cfg
|
|
, geoAsnCounts = counts
|
|
}
|
|
|
|
sockAddrToIP :: SockAddr -> Maybe IP
|
|
sockAddrToIP (SockAddrInet _ ha) = Just (IPv4 (fromHostAddress ha))
|
|
sockAddrToIP (SockAddrInet6 _ _ ha6 _) = Just (IPv6 (fromHostAddress6 ha6))
|
|
sockAddrToIP (SockAddrUnix _) = Nothing
|
|
|
|
countryBlocked :: GeoConfig -> Maybe Text -> Maybe Text
|
|
countryBlocked GeoConfig{..} mIso = case mIso of
|
|
Nothing
|
|
| not (null gcAllowedCountries) -> Just "??"
|
|
| otherwise -> Nothing
|
|
Just iso ->
|
|
let up = T.toUpper iso
|
|
inBlocked = up `elem` gcBlockedCountries
|
|
inAllowed = null gcAllowedCountries || up `elem` gcAllowedCountries
|
|
in if inBlocked || not inAllowed then Just up else Nothing
|
|
|
|
lookupGeo :: Geo -> IP -> IO GeoInfo
|
|
lookupGeo Geo{..} ip = do
|
|
let lang = gcLanguage geoConfig
|
|
countryRes = (\db -> findGeoData db lang ip) <$> geoCountryDb
|
|
asnRes = (\db -> findGeoData db lang ip) <$> geoAsnDb
|
|
country = case countryRes of
|
|
Just (Right r) -> geoCountryISO r
|
|
_ -> Nothing
|
|
(asn, org) = case asnRes of
|
|
Just (Right r) -> case geoAS r of
|
|
Just a -> (Just (asNumber a), Just (asOrganization a))
|
|
Nothing -> (Nothing, Nothing)
|
|
_ -> (Nothing, Nothing)
|
|
flagged = case asn of
|
|
Just n -> n `elem` gcFlaggedAsns geoConfig
|
|
Nothing -> False
|
|
pure GeoInfo
|
|
{ giCountryISO = country
|
|
, giAsnNumber = asn
|
|
, giAsnOrg = org
|
|
, giFlaggedAsn = flagged
|
|
}
|
|
|
|
bumpAsnCounter :: Geo -> Int -> POSIXTime -> STM Int
|
|
bumpAsnCounter Geo{..} n now = do
|
|
let window = fromIntegral (gcConcentrationWindowSeconds geoConfig) :: POSIXTime
|
|
m <- readTVar geoAsnCounts
|
|
let entry = case Map.lookup n m of
|
|
Just w
|
|
| now - awWindowStart w < window ->
|
|
w { awCount = awCount w + 1 }
|
|
_ -> AsnWindow { awCount = 1, awWindowStart = now }
|
|
writeTVar geoAsnCounts $! Map.insert n entry m
|
|
pure (awCount entry)
|
|
|
|
asnConcentrationScore :: Geo -> Int -> Double
|
|
asnConcentrationScore Geo{..} count =
|
|
let threshold = max 1 (gcConcentrationThreshold geoConfig)
|
|
ratio = fromIntegral count / fromIntegral threshold :: Double
|
|
in min 1.0 (max 0.0 ratio)
|
|
|
|
purgeAsnCounters :: Geo -> POSIXTime -> STM Int
|
|
purgeAsnCounters Geo{..} now = do
|
|
let window = fromIntegral (gcConcentrationWindowSeconds geoConfig) :: POSIXTime
|
|
m <- readTVar geoAsnCounts
|
|
let (stale, fresh) = Map.partition (\w -> now - awWindowStart w >= window) m
|
|
modifyTVar' geoAsnCounts (const fresh)
|
|
pure (Map.size stale)
|
|
|
|
startAsnSweeper :: Geo -> IO (Async ())
|
|
startAsnSweeper g = async $ forever $ do
|
|
threadDelay defaultGeoSweepIntervalMicros
|
|
now <- getPOSIXTime
|
|
_ <- atomically (purgeAsnCounters g now)
|
|
pure ()
|
|
|
|
decideGeo :: GeoConfig -> GeoInfo -> Int -> GeoDecision
|
|
decideGeo cfg info count =
|
|
case countryBlocked cfg (giCountryISO info) of
|
|
Just iso -> GeoBlockCountry iso
|
|
Nothing ->
|
|
let threshold = gcConcentrationThreshold cfg
|
|
safeThreshold = max 1 threshold
|
|
score = min 1.0
|
|
$ fromIntegral count / fromIntegral safeThreshold :: Double
|
|
shouldJail = gcAction cfg == GeoActionJail
|
|
&& giFlaggedAsn info
|
|
&& count >= threshold
|
|
in case giAsnNumber info of
|
|
Just n | shouldJail -> GeoJailAsn n score
|
|
_ -> GeoAllow
|
|
|
|
renderGeoHeader :: GeoInfo -> Int -> ByteString
|
|
renderGeoHeader info count =
|
|
let countryPart = case giCountryISO info of
|
|
Just iso -> "country=" <> TE.encodeUtf8 iso
|
|
Nothing -> "country=??"
|
|
asnPart = case giAsnNumber info of
|
|
Just n -> "asn=" <> BC.pack (show n)
|
|
Nothing -> "asn=0"
|
|
flagPart = if giFlaggedAsn info then "flag=1" else "flag=0"
|
|
countPart = "count=" <> BC.pack (show count)
|
|
in BC.intercalate " " [countryPart, asnPart, flagPart, countPart]
|
|
|
|
forbiddenByCountry :: Text -> LBS.ByteString
|
|
forbiddenByCountry iso =
|
|
LBS.fromStrict ("403 Forbidden: country " <> TE.encodeUtf8 iso <> " is restricted")
|
|
|
|
geoMiddleware :: Geo -> Maybe IPJail -> Middleware
|
|
geoMiddleware g@Geo{..} mJail app req respond = do
|
|
let mIp = sockAddrToIP (remoteHost req)
|
|
(info, count) <- case mIp of
|
|
Nothing -> pure (emptyGeoInfo, 0)
|
|
Just ip -> do
|
|
infoRaw <- lookupGeo g ip
|
|
now <- getPOSIXTime
|
|
c <- case giAsnNumber infoRaw of
|
|
Just n -> atomically (bumpAsnCounter g n now)
|
|
Nothing -> pure 0
|
|
pure (infoRaw, c)
|
|
let decision = decideGeo geoConfig info count
|
|
annotate = gcAnnotateHeader geoConfig
|
|
headerBS = renderGeoHeader info count
|
|
withGeoHeader =
|
|
if annotate
|
|
then mapResponseHeaders ((geoResponseHeaderName, headerBS) :)
|
|
else id
|
|
case decision of
|
|
GeoAllow -> app req (respond . withGeoHeader)
|
|
GeoBlockCountry iso -> respond $ withGeoHeader $ responseLBS status403
|
|
[ ("Content-Type", "text/plain; charset=utf-8")
|
|
, ("Cache-Control", "no-store")
|
|
]
|
|
(forbiddenByCountry iso)
|
|
GeoJailAsn asn _score -> do
|
|
now <- getPOSIXTime
|
|
when (isJust mJail) $
|
|
atomically $ case mJail of
|
|
Just j -> jail j (clientIPKey req)
|
|
(fromIntegral (gcJailCooldownSeconds geoConfig))
|
|
("geo:asn:" <> BC.pack (show asn))
|
|
now
|
|
Nothing -> pure ()
|
|
respond $ withGeoHeader $ responseLBS status403
|
|
[ ("Content-Type", "text/plain; charset=utf-8")
|
|
, ("Cache-Control", "no-store")
|
|
]
|
|
"403 Forbidden: source network is temporarily restricted"
|