9e7d395 · skrepka · git.iko.soy

git.iko.soy/skrepka

Improve server.knot: auth-tier rate limiting, session counter, concurrent prune guard, drain flush, voluntary logout, message status endpoint, federation ack hardening, verify() atomicity, stats rate-limiting, forwardQueuedMessages optimization, pruneExpired lock splitting, and isBadServerName optimization

9e7d395iko6 months ago

server.knot+223−62

compare with master

unifiedsplit

@@ -80,6 +80,10 @@ maxFederationRateEntries = 100

80maxGlobalRequestRate = 1000 -- max total requests per window

81globalRateWindowMs : Int<Ms>

82globalRateWindowMs = 1000

83+-- Auth-tier global rate limiting (prevents auth flooding from starving client operations)

84+maxAuthGlobalRate = 200 -- max auth requests (challenge + verify + revoke) per window

85+authGlobalRateWindowMs : Int<Ms>

86+authGlobalRateWindowMs = 1000

83-- Session and gossip timing

87-- Session and gossip timing

84sessionExpiry : Int<Ms>

88sessionExpiry : Int<Ms>

85sessionExpiry = 3600000 -- 1 hour

89sessionExpiry = 3600000 -- 1 hour

@@ -402,6 +406,16 @@ type FederationInFlightEntry = {server: Text, count: Int, firstStartedAt: Timest

402-- Last time federation latency metrics were reset (for windowed averages)

406-- Last time federation latency metrics were reset (for windowed averages)

403*federationLatencyResetAt : Timestamp

407*federationLatencyResetAt : Timestamp

404

408

409+-- Total session count across all buckets (avoids O(16) recount on monitoring paths)

410+*totalSessionCount : Int

411+

412+-- Auth-tier global rate limiting state

413+*authGlobalRateCount : Int

414+*authGlobalRateWindowStart : Timestamp

415+

416+-- Max concurrent background prune forks (prevents resource exhaustion from stuck recovery)

417+*activePruneForks : Int

418+

405-- ============================================

419-- ============================================

406-- Routes

420-- Routes

407-- ============================================

421-- ============================================

@@ -421,6 +435,10 @@ route AuthApi where

421 POST {targetToken: Text} /revoke headers {authorization: Text, clientIp: Text}

435 POST {targetToken: Text} /revoke headers {authorization: Text, clientIp: Text}

422 -> {ok: Bool, authorized: Bool}

436 -> {ok: Bool, authorized: Bool}

423 = RevokeSession

437 = RevokeSession

438+ -- Voluntary disconnect: removes session and gossips offline if last session for this key

439+ POST {} /logout headers {authorization: Text, clientIp: Text}

440+ -> {ok: Bool, authorized: Bool}

441+ = Logout

424

442

425-- Client-facing API (requires session token in Authorization header)

443-- Client-facing API (requires session token in Authorization header)

426route ClientApi where

444route ClientApi where

@@ -433,6 +451,10 @@ route ClientApi where

433 POST {to: PubkeyHex, encryptedBlob: BlobHex, timestamp: Timestamp, messageId: OptionalMessageId} / headers {authorization: Text, clientIp: Text}

451 POST {to: PubkeyHex, encryptedBlob: BlobHex, timestamp: Timestamp, messageId: OptionalMessageId} / headers {authorization: Text, clientIp: Text}

434 -> {status: MessageStatus, messageId: OptionalMessageId}

452 -> {status: MessageStatus, messageId: OptionalMessageId}

435 = SendMessage

453 = SendMessage

454+ -- Check delivery status of a sent message (caller must know recipient to prove ownership)

455+ POST {to: PubkeyHex, messageId: MessageIdHex} /status headers {authorization: Text, clientIp: Text}

456+ -> {status: Text, authorized: Bool}

457+ = MessageStatus

436 -- Look up which server a pubkey is present on (for federation routing)

458 -- Look up which server a pubkey is present on (for federation routing)

437 GET /lookup/{pubkey: PubkeyHex} headers {authorization: Text, clientIp: Text}

459 GET /lookup/{pubkey: PubkeyHex} headers {authorization: Text, clientIp: Text}

438 -> {server: AnyServerName, authorized: Bool}

460 -> {server: AnyServerName, authorized: Bool}

@@ -677,10 +699,12 @@ hasDotBeforePort = \s ->

677-- IPv6 shorthand, octal/hex IP tricks, bare IPs without dots, and .arpa domains.

699-- IPv6 shorthand, octal/hex IP tricks, bare IPs without dots, and .arpa domains.

678isBadServerName : Text -> Bool

700isBadServerName : Text -> Bool

679isBadServerName = \sn ->

701isBadServerName = \sn ->

680− sn == ""

702+ let fc = take 1 sn in

681− || take 1 sn == "."

703+ fc == ""

682− || take 1 sn == ":"

704+ || fc == "."

705+ || fc == ":"

683 || containsBadUrlChar sn

706 || containsBadUrlChar sn

707+ || not (hasDotBeforePort sn)

684 || matchesHostnameOrSubdomain sn "localhost"

708 || matchesHostnameOrSubdomain sn "localhost"

685 || matchesHostname sn "localhost.localdomain"

709 || matchesHostname sn "localhost.localdomain"

686 || matchesHostnameOrSubdomain sn "local"

710 || matchesHostnameOrSubdomain sn "local"

@@ -688,7 +712,7 @@ isBadServerName = \sn ->

688 || matchesHostnameOrSubdomain sn "metadata"

712 || matchesHostnameOrSubdomain sn "metadata"

689 || take 4 sn == "127." || take 3 sn == "10."

713 || take 4 sn == "127." || take 3 sn == "10."

690 || take 8 sn == "192.168." || take 8 sn == "169.254."

714 || take 8 sn == "192.168." || take 8 sn == "169.254."

691− || take 2 sn == "0." || take 1 sn == "["

715+ || take 2 sn == "0." || fc == "["

692 || is172Private sn

716 || is172Private sn

693 || take 7 sn == "198.18." || take 7 sn == "198.19."

717 || take 7 sn == "198.18." || take 7 sn == "198.19."

694 || take 8 sn == "192.0.0."

718 || take 8 sn == "192.0.0."

@@ -701,9 +725,8 @@ isBadServerName = \sn ->

701 || containsMultipleColons sn

725 || containsMultipleColons sn

702 || hasOctalOctet sn

726 || hasOctalOctet sn

703 || hasHexPrefix sn

727 || hasHexPrefix sn

704− || (isDigitChar (take 1 sn) && isAllDigitsOrDotsUntilColonOrEnd sn)

728+ || (isDigitChar fc && isAllDigitsOrDotsUntilColonOrEnd sn)

705 || isHighOctetPrefix sn

729 || isHighOctetPrefix sn

706− || not (hasDotBeforePort sn)

707 || containsSuffix ".arpa" sn

730 || containsSuffix ".arpa" sn

708 || containsSuffix ".onion" sn

731 || containsSuffix ".onion" sn

709 || containsSuffix ".i2p" sn

732 || containsSuffix ".i2p" sn

@@ -1302,6 +1325,8 @@ pruneExpired = \t -> do

1302 challenges <- *challenges

1325 challenges <- *challenges

1303 set *challenges = do x <- challenges; where x.expiresAt > t; yield x

1326 set *challenges = do x <- challenges; where x.expiresAt > t; yield x

1304 transformAllSessionBuckets (\bucket -> do x <- bucket; where x.expiresAt > t; yield x)

1327 transformAllSessionBuckets (\bucket -> do x <- bucket; where x.expiresAt > t; yield x)

1328+ sessionCount <- countAllSessions

1329+ set *totalSessionCount = sessionCount

1305 -- Prune messages across all 16 buckets. Each bucket prune is atomic (short lock).

1330 -- Prune messages across all 16 buckets. Each bucket prune is atomic (short lock).

1306 -- Recipient counts are rebuilt incrementally per-bucket — since messages are sharded

1331 -- Recipient counts are rebuilt incrementally per-bucket — since messages are sharded

1307 -- by recipient pubkey, each bucket's counts are disjoint, so simple union merges them.

1332 -- by recipient pubkey, each bucket's counts are disjoint, so simple union merges them.

@@ -1342,22 +1367,29 @@ pruneExpired = \t -> do

1342 set *ackOverflow = do a <- ackOverflow; where a.ackedAt > retainCutoff; yield a

1367 set *ackOverflow = do a <- ackOverflow; where a.ackedAt > retainCutoff; yield a

1343 deliveredIds <- *deliveredIds

1368 deliveredIds <- *deliveredIds

1344 set *deliveredIds = do d <- deliveredIds; where d.deliveredAt > retainCutoff; yield d

1369 set *deliveredIds = do d <- deliveredIds; where d.deliveredAt > retainCutoff; yield d

1345− -- Prune rate limiters and ephemeral state

1370+ -- Prune rate limiters (split into smaller atomic blocks to reduce lock contention)

1346 atomic do

1371 atomic do

1347 sendRates <- *sendRates

1372 sendRates <- *sendRates

1348 set *sendRates = do r <- sendRates; where t - r.windowStart < sendRateWindowMs; yield r

1373 set *sendRates = do r <- sendRates; where t - r.windowStart < sendRateWindowMs; yield r

1349 challengeRates <- *challengeRates

1374 challengeRates <- *challengeRates

1350 set *challengeRates = do r <- challengeRates; where t - r.windowStart < challengeRateWindowMs; yield r

1375 set *challengeRates = do r <- challengeRates; where t - r.windowStart < challengeRateWindowMs; yield r

1376+ recvRates <- *recvRates

1377+ set *recvRates = do r <- recvRates; where t - r.windowStart < recvRateWindowMs; yield r

1378+ -- Reset auth-tier global rate window

1379+ authWS <- *authGlobalRateWindowStart

1380+ when (t - authWS >= authGlobalRateWindowMs) (do

1381+ set *authGlobalRateWindowStart = t

1382+ set *authGlobalRateCount = 0)

1383+ atomic do

1351 federationRates <- *federationRates

1384 federationRates <- *federationRates

1352 set *federationRates = do r <- federationRates; where t - r.windowStart < federationRateWindowMs; yield r

1385 set *federationRates = do r <- federationRates; where t - r.windowStart < federationRateWindowMs; yield r

1353− activePolls <- *activePolls

1354− set *activePolls = do p <- activePolls; where t - p.startedAt < pollTimeoutMs * 3; yield p

1355 ipRates <- *ipRates

1386 ipRates <- *ipRates

1356 set *ipRates = do r <- ipRates; where t - r.windowStart < ipRateWindowMs; yield r

1387 set *ipRates = do r <- ipRates; where t - r.windowStart < ipRateWindowMs; yield r

1357 lookupRates <- *lookupRates

1388 lookupRates <- *lookupRates

1358 set *lookupRates = do r <- lookupRates; where t - r.windowStart < lookupRateWindowMs; yield r

1389 set *lookupRates = do r <- lookupRates; where t - r.windowStart < lookupRateWindowMs; yield r

1359− recvRates <- *recvRates

1390+ atomic do

1360− set *recvRates = do r <- recvRates; where t - r.windowStart < recvRateWindowMs; yield r

1391+ activePolls <- *activePolls

1392+ set *activePolls = do p <- activePolls; where t - p.startedAt < pollTimeoutMs * 3; yield p

1361 recentGossip <- *recentGossip

1393 recentGossip <- *recentGossip

1362 set *recentGossip = do r <- recentGossip; where t - r.seenAt < gossipDedupWindow; yield r

1394 set *recentGossip = do r <- recentGossip; where t - r.seenAt < gossipDedupWindow; yield r

1363 -- Clean up stale in-flight entries: remove zeroed entries and reset leaked slots

1395 -- Clean up stale in-flight entries: remove zeroed entries and reset leaked slots

@@ -1445,6 +1477,21 @@ checkGlobalRate = \t -> atomic do

1445 set *globalRateCount = rateCount + 1

1477 set *globalRateCount = rateCount + 1

1446 yield (False {})

1478 yield (False {})

1447

1479

1480+-- Auth-tier global rate limiting (shared budget for challenge + verify + revoke)

1481+checkAuthGlobalRate = \t -> atomic do

1482+ windowStart <- *authGlobalRateWindowStart

1483+ rateCount <- *authGlobalRateCount

1484+ if t - windowStart >= authGlobalRateWindowMs

1485+ then do

1486+ set *authGlobalRateWindowStart = t

1487+ set *authGlobalRateCount = 1

1488+ yield (False {})

1489+ else if rateCount >= maxAuthGlobalRate

1490+ then yield (True {})

1491+ else do

1492+ set *authGlobalRateCount = rateCount + 1

1493+ yield (False {})

1494+

1448-- IP-based rate limiting; skips if clientIp header is empty (no reverse proxy)

1495-- IP-based rate limiting; skips if clientIp header is empty (no reverse proxy)

1449checkIpRate = \ip t ->

1496checkIpRate = \ip t ->

1450 if ip == "" then yield (False {})

1497 if ip == "" then yield (False {})

@@ -1587,12 +1634,14 @@ retryPendingForwards = \pending -> do

1587-- When a remote key comes online, create forward entries for queued messages and send them

1634-- When a remote key comes online, create forward entries for queued messages and send them

1588forwardQueuedMessages = \targets t -> do

1635forwardQueuedMessages = \targets t -> do

1589 newForwards <- atomic do

1636 newForwards <- atomic do

1590− allMessages <- readAllMessages

1591 forwards <- *forwards

1637 forwards <- *forwards

1592 presence <- *presence

1638 presence <- *presence

1593 let localOnline = localOnlineKeys presence t

1639 let localOnline = localOnlineKeys presence t

1594 let remoteTargets = remoteOnlyTargets targets localOnline

1640 let remoteTargets = remoteOnlyTargets targets localOnline

1641+ -- Read only the message buckets containing target pubkeys (not all 16)

1642+ let neededBuckets = uniqueBucketsFor remoteTargets (\target -> target.pubkey)

1643+ relevantMessages <- readSelectedBuckets neededBuckets

1595− let candidates = forwardCandidates remoteTargets allMessages forwards t

1644+ let candidates = forwardCandidates remoteTargets relevantMessages forwards t

1596 let capped = takeRelation (maxForwardsEntries - count forwards) candidates

1645 let capped = takeRelation (maxForwardsEntries - count forwards) candidates

1597 let updatedForwards = union forwards capped

1646 let updatedForwards = union forwards capped

1598 set *forwards = updatedForwards

1647 set *forwards = updatedForwards

@@ -2068,6 +2117,9 @@ handleChallenge = \pubkey clientIp respond -> do

2068 when (preAuth == "ip_rate_limited") (println ("challenge: IP rate limited for " ++ clientIp))

2117 when (preAuth == "ip_rate_limited") (println ("challenge: IP rate limited for " ++ clientIp))

2069 yield (respond {challenge: "", ok: False {}, expiresAt: 0})

2118 yield (respond {challenge: "", ok: False {}, expiresAt: 0})

2070 else do

2119 else do

2120+ authLimited <- checkAuthGlobalRate t

2121+ if authLimited then yield (respond {challenge: "", ok: False {}, expiresAt: 0})

2122+ else do

2071 rawChallenge <- genToken

2123 rawChallenge <- genToken

2072 let challenge = serverHostN ++ ":" ++ rawChallenge

2124 let challenge = serverHostN ++ ":" ++ rawChallenge

2073 result <- atomic do

2125 result <- atomic do

@@ -2121,11 +2173,13 @@ handleVerifyAuth = \pubkey challenge signature revokeOthers clientIp respond ->

2121 if authRateLimited then do

2173 if authRateLimited then do

2122 println ("verify: rate limited (" ++ show maxChallengeRatePerKey ++ "/s per key), rejecting for " ++ pubkey)

2174 println ("verify: rate limited (" ++ show maxChallengeRatePerKey ++ "/s per key), rejecting for " ++ pubkey)

2123 yield (respond {token: "", expiresAt: 0})

2175 yield (respond {token: "", expiresAt: 0})

2176+ else do

2177+ authGlobalLimited <- checkAuthGlobalRate t

2178+ if authGlobalLimited then yield (respond {token: "", expiresAt: 0})

2124− else verifyAuthInner pubkey challenge signature revokeOthers clientIp respond t

2179+ else verifyAuthInner pubkey challenge signature revokeOthers clientIp respond t

2125

2180

2126verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t -> do

2181verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t -> do

2127 let expiry = t + sessionExpiry

2182 let expiry = t + sessionExpiry

2128− let sigValid = verify (bytesFromHex pubkey) (textToBytes challenge) (bytesFromHex signature)

2129 validation <- atomic do

2183 validation <- atomic do

2130 challenges <- *challenges

2184 challenges <- *challenges

2131 let hasMatch = challengeMatch challenges pubkey challenge clientIp t

2185 let hasMatch = challengeMatch challenges pubkey challenge clientIp t

@@ -2134,11 +2188,17 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->

2134 -- by submitting a bad signature from the same IP. Rate limiter still constrains

2188 -- by submitting a bad signature from the same IP. Rate limiter still constrains

2135 -- brute-force attempts; not consuming on IP mismatch prevents remote attackers

2189 -- brute-force attempts; not consuming on IP mismatch prevents remote attackers

2136 -- from burning challenges.

2190 -- from burning challenges.

2191+ -- verify() is inside the atomic block to prevent a TOCTOU gap between challenge

2192+ -- validation and consumption. It's pure (no IO) so safe in STM.

2137− if not hasMatch || not sigValid

2193+ if not hasMatch

2138 then yield {valid: False {}}

2194 then yield {valid: False {}}

2139 else do

2195 else do

2196+ let sigValid = verify (bytesFromHex pubkey) (textToBytes challenge) (bytesFromHex signature)

2197+ if not sigValid

2198+ then yield {valid: False {}}

2199+ else do

2140− set *challenges = challengesExcluding challenges pubkey challenge

2200+ set *challenges = challengesExcluding challenges pubkey challenge

2141− yield {valid: True {}}

2201+ yield {valid: True {}}

2142 if validation.valid

2202 if validation.valid

2143 then do

2203 then do

2144 token <- genToken

2204 token <- genToken

@@ -2159,6 +2219,8 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->

2159 let nb = tokenBucket token

2219 let nb = tokenBucket token

2160 currentBucket <- readSessionBucket nb

2220 currentBucket <- readSessionBucket nb

2161 writeSessionBucket nb (union currentBucket [newSession])

2221 writeSessionBucket nb (union currentBucket [newSession])

2222+ sc <- *totalSessionCount

2223+ set *totalSessionCount = sc + 1

2162 let presenceWithout = presenceWithoutLocalKey presence pubkey

2224 let presenceWithout = presenceWithoutLocalKey presence pubkey

2163 let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)

2225 let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)

2164 let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]

2226 let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]

@@ -2172,6 +2234,7 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->

2172 let afterRevoke = if revokeOthers then do s <- allSessions; where s.pubkey != pubkey; yield s else allSessions

2234 let afterRevoke = if revokeOthers then do s <- allSessions; where s.pubkey != pubkey; yield s else allSessions

2173 let capped = capPerKeyAndGlobal afterRevoke (\s -> s.pubkey == pubkey) maxSessionsPerKey (\s -> s.expiresAt) t maxSessions

2235 let capped = capPerKeyAndGlobal afterRevoke (\s -> s.pubkey == pubkey) maxSessionsPerKey (\s -> s.expiresAt) t maxSessions

2174 let sessionsWithNew = union capped [newSession]

2236 let sessionsWithNew = union capped [newSession]

2237+ set *totalSessionCount = count sessionsWithNew

2175 let bFor = \b -> do s <- sessionsWithNew; where tokenBucket s.token == b; yield s

2238 let bFor = \b -> do s <- sessionsWithNew; where tokenBucket s.token == b; yield s

2176 writeSessionBucket 0 (bFor 0)

2239 writeSessionBucket 0 (bFor 0)

2177 writeSessionBucket 1 (bFor 1)

2240 writeSessionBucket 1 (bFor 1)

@@ -2380,6 +2443,39 @@ handleLookupPresence = \pubkey authorization clientIp respond -> do

2380 yield {server: maskedServer, authorized: True {}})

2443 yield {server: maskedServer, authorized: True {}})

2381 yield (respond {server: result.server, authorized: result.authorized})

2444 yield (respond {server: result.server, authorized: result.authorized})

2382

2445

2446+-- POST /messages/status: Check delivery status of a previously sent message.

2447+-- Caller must provide both messageId and recipient pubkey (proves they know the recipient).

2448+-- Returns: "pending" (in mailbox), "forwarding" (federation in progress),

2449+-- "delivered" (acked by recipient), "unknown" (expired or never existed).

2450+handleMessageStatus = \to messageId authorization clientIp respond -> do

2451+ t <- now

2452+ globalRateLimited <- checkGlobalRate t

2453+ if globalRateLimited then yield (respond {status: "", authorized: False {}})

2454+ else

2455+ withSessionAuth authorization t clientIp

2456+ (yield (respond {status: "", authorized: False {}}))

2457+ (\senderKey -> do

2458+ result <- atomic do

2459+ let msgBucket = pubkeyBucket to

2460+ bucket <- readMessageBucket msgBucket

2461+ -- Check if message is still in mailbox (pending delivery)

2462+ let inMailbox = any bucket (\m -> m.messageId == messageId && m.toKey == to)

2463+ if inMailbox then yield "pending"

2464+ else do

2465+ -- Check if being forwarded to a remote server

2466+ forwards <- *forwards

2467+ let inForwards = any forwards (\f -> f.messageId == messageId && f.toKey == to)

2468+ if inForwards then yield "forwarding"

2469+ else do

2470+ -- Check if already delivered (acked or in terminal delivered set)

2471+ ackedIds <- *ackedIds

2472+ deliveredIds <- *deliveredIds

2473+ let inAcked = hasMessage messageId to ackedIds

2474+ let inDelivered = hasMessage messageId to deliveredIds

2475+ if inAcked || inDelivered then yield "delivered"

2476+ else yield "unknown"

2477+ yield (respond {status: result, authorized: True {}}))

2478+

2383-- POST /auth/revoke: Revoke a specific session token.

2479-- POST /auth/revoke: Revoke a specific session token.

2384-- Caller authenticates via Authorization header, then specifies the target token to remove.

2480-- Caller authenticates via Authorization header, then specifies the target token to remove.

2385-- Only removes the target if it belongs to the same pubkey as the caller (prevents cross-user revocation).

2481-- Only removes the target if it belongs to the same pubkey as the caller (prevents cross-user revocation).

@@ -2408,11 +2504,59 @@ handleRevoke = \targetToken authorization clientIp respond -> do

2408 if hasTarget

2504 if hasTarget

2409 then do

2505 then do

2410 writeSessionBucket targetBucket (do s <- sessions; where not (safeCompare s.token targetToken && s.pubkey == callerPubkey); yield s)

2506 writeSessionBucket targetBucket (do s <- sessions; where not (safeCompare s.token targetToken && s.pubkey == callerPubkey); yield s)

2507+ sc <- *totalSessionCount

2508+ set *totalSessionCount = max 0 (sc - 1)

2411 yield (True {})

2509 yield (True {})

2412 else yield (False {})

2510 else yield (False {})

2413 when removed (println ("revoke: session revoked for " ++ callerPubkey))

2511 when removed (println ("revoke: session revoked for " ++ callerPubkey))

2414 yield (respond {ok: removed, authorized: True {}}))

2512 yield (respond {ok: removed, authorized: True {}}))

2415

2513

2514+-- POST /auth/logout: Voluntary disconnect. Removes the caller's session and, if it was

2515+-- the last active session for this pubkey, clears local presence and gossips offline.

2516+-- This gives faster offline detection than waiting for session expiry (1 hour).

2517+handleLogout = \authorization clientIp respond -> do

2518+ t <- now

2519+ globalRateLimited <- checkGlobalRate t

2520+ if globalRateLimited then yield (respond {ok: False {}, authorized: False {}})

2521+ else

2522+ case extractToken authorization of

2523+ Nothing {} -> yield (respond {ok: False {}, authorized: False {}})

2524+ Just {value: token} -> do

2525+ let bucket = tokenBucket token

2526+ result <- atomic do

2527+ sessions <- readSessionBucket bucket

2528+ case authenticateSession sessions token t of

2529+ Unauthorized {} -> yield (Nothing {})

2530+ Authorized {pubkey, ip} ->

2531+ if clientIp != "" && ip != "" && clientIp != ip

2532+ then yield (Nothing {})

2533+ else do

2534+ -- Remove the caller's session

2535+ writeSessionBucket bucket (do s <- sessions; where not (safeCompare s.token token); yield s)

2536+ sc <- *totalSessionCount

2537+ set *totalSessionCount = max 0 (sc - 1)

2538+ -- Check if this was the last session for this key

2539+ allSessions <- readAllSessions

2540+ let remaining = count (do s <- allSessions; where s.pubkey == pubkey; yield s)

2541+ if remaining == 0

2542+ then do

2543+ -- Remove local presence immediately

2544+ presence <- *presence

2545+ set *presence = presenceWithoutLocalKey presence pubkey

2546+ yield (Just {value: {pubkey: pubkey, gossipOffline: True {}}})

2547+ else yield (Just {value: {pubkey: pubkey, gossipOffline: False {}}})

2548+ case result of

2549+ Nothing {} -> yield (respond {ok: False {}, authorized: False {}})

2550+ Just {value: r} -> do

2551+ when r.gossipOffline (do

2552+ println ("logout: last session for " ++ r.pubkey ++ ", gossiping offline")

2553+ when federationEnabled (do

2554+ presence <- *presence

2555+ let peerServers = collectPeerServers presence (\p -> p.server)

2556+ when (count peerServers > 0)

2557+ (fork (sendBatchedGossip peerServers [{eventType: GossipOffline {}, pubkey: r.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}] gossipOriginHops))))

2558+ yield (respond {ok: True {}, authorized: True {}})

2559+

2416-- POST /federation/gossip: Receive presence gossip from a peer server.

2560-- POST /federation/gossip: Receive presence gossip from a peer server.

2417-- Validates events, merges into presence table (last-write-wins), and re-gossips

2561-- Validates events, merges into presence table (last-write-wins), and re-gossips

2418-- to other known peers with decremented hop count. Triggers message forwarding

2562-- to other known peers with decremented hop count. Triggers message forwarding

@@ -2532,33 +2676,12 @@ handleFederationAck = \messageIds originServer fromServer authorization respond

2532 t <- now

2676 t <- now

2533 withFederationAuth authorization t respond fromServer (isLocalServer originServer && not (isLocalServer fromServer)) (do

2677 withFederationAuth authorization t respond fromServer (isLocalServer originServer && not (isLocalServer fromServer)) (do

2534 let dedupedIds = dedupMessageIds messageIds

2678 let dedupedIds = dedupMessageIds messageIds

2679+ -- Only clear forward entries; origin messages are retained until the recipient

2680+ -- polls from this server or they age out via unackedMessageMaxAge. This prevents

2681+ -- a compromised peer from deleting messages by sending unverified acks.

2535 atomic do

2682 atomic do

2536 forwards <- *forwards

2683 forwards <- *forwards

2537− let ackedForwards = do f <- forwards; where f.toServer == fromServer && any dedupedIds (\mid -> f.messageId == mid); yield f

2538− let updatedForwards = forwardsExcluding forwards fromServer dedupedIds

2684+ set *forwards = forwardsExcluding forwards fromServer dedupedIds

2539− set *forwards = updatedForwards

2540− -- Read only message buckets containing acked recipients (not all 16)

2541− let neededBuckets = uniqueBucketsFor ackedForwards (\f -> f.toKey)

2542− relevantMessages <- readSelectedBuckets neededBuckets

2543− presence <- *presence

2544− let removable = do

2545− f <- ackedForwards

2546− where none updatedForwards (\rf -> rf.messageId == f.messageId && rf.toKey == f.toKey)

2547− where none presence (\p -> p.pubkey == f.toKey && isLocalServer p.server && p.expiresAt > t)

2548− yield f

2549− -- Only remove messages stored long enough for local clients to have polled them (grace period)

2550− let pastGrace = do f <- removable; where any relevantMessages (\m -> m.messageId == f.messageId && m.toKey == f.toKey && t - m.receivedAt >= ackGracePeriod); yield f

2551− let removed = count pastGrace

2552− -- Remove matched messages from only the affected buckets

2553− when (removed > 0) (do

2554− transformSelectedBuckets neededBuckets (\bucket ->

2555− do m <- bucket; where none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey); yield m)

2556− totalMsgCount <- *totalMessageCount

2557− set *totalMessageCount = totalMsgCount - removed

2558− -- Batch-adjust per-recipient counts (single pass instead of per-entry fold)

2559− recipientCounts <- *recipientMsgCounts

2560− let removedCounts = countByRecipient pastGrace

2561− set *recipientMsgCounts = batchDecrementRecipientCounts recipientCounts removedCounts)

2562 yield (respond {ok: True {}, error: ""}))

2685 yield (respond {ok: True {}, error: ""}))

2563

2686

2564-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).

2687-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).

@@ -2579,14 +2702,20 @@ handleDrain = \authorization respond -> do

2579 let drainedCount = count drainResult.localEntries

2702 let drainedCount = count drainResult.localEntries

2580 forwards <- *forwards

2703 forwards <- *forwards

2581 ackedIds <- *ackedIds

2704 ackedIds <- *ackedIds

2705+ let pendingAcks = do a <- ackedIds; where not (isLocalServer a.originServer); yield a

2582 let pendingForwardCount = count forwards

2706 let pendingForwardCount = count forwards

2583− let pendingAckCount = count ackedIds

2707+ let pendingAckCount = count pendingAcks

2584− -- Fork gossip so drain responds immediately (hung peers can't block the response).

2708+ -- Fork best-effort flush of pending federation ops, then gossip, then wait for in-flight.

2585− -- After gossip, wait for any in-flight federation requests (from prior forks) to

2586− -- complete before signaling drainCompletedAt. Operator polls /stats for drainCompletedAt > 0

2709+ -- Operator polls /stats for drainCompletedAt > 0 to know when safe to stop.

2587− -- to know when it's truly safe to stop the process.

2588 atomic do set *drainCompletedAt = 0

2710 atomic do set *drainCompletedAt = 0

2589 fork (do

2711 fork (do

2712+ -- Best-effort flush of pending forwards and acks before stopping

2713+ when (federationEnabled && pendingForwardCount > 0) (do

2714+ println ("drain: flushing " ++ show pendingForwardCount ++ " pending forwards")

2715+ retryPendingForwards forwards)

2716+ when (federationEnabled && pendingAckCount > 0) (do

2717+ println ("drain: flushing " ++ show pendingAckCount ++ " pending acks")

2718+ sendAcksAndClean pendingAcks)

2590 when (federationEnabled && drainedCount > 0 && count drainResult.peerServers > 0) (do

2719 when (federationEnabled && drainedCount > 0 && count drainResult.peerServers > 0) (do

2591 let gossipEvents = do p <- drainResult.localEntries; yield {eventType: GossipOffline {}, pubkey: p.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}

2720 let gossipEvents = do p <- drainResult.localEntries; yield {eventType: GossipOffline {}, pubkey: p.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}

2592 sendBatchedGossip drainResult.peerServers gossipEvents gossipOriginHops)

2721 sendBatchedGossip drainResult.peerServers gossipEvents gossipOriginHops)

@@ -2596,9 +2725,9 @@ handleDrain = \authorization respond -> do

2596 atomic do set *drainCompletedAt = t2)

2725 atomic do set *drainCompletedAt = t2)

2597 println ("drain: broadcast offline for " ++ show drainedCount ++ " local keys")

2726 println ("drain: broadcast offline for " ++ show drainedCount ++ " local keys")

2598 when (pendingForwardCount > 0)

2727 when (pendingForwardCount > 0)

2599− (println ("drain: warning: " ++ show pendingForwardCount ++ " pending forwards will not be retried"))

2728+ (println ("drain: flushing " ++ show pendingForwardCount ++ " pending forwards"))

2600 when (pendingAckCount > 0)

2729 when (pendingAckCount > 0)

2601− (println ("drain: warning: " ++ show pendingAckCount ++ " pending acks will not be sent"))

2730+ (println ("drain: flushing " ++ show pendingAckCount ++ " pending acks"))

2602 yield (respond {ok: True {}, drained: drainedCount})

2731 yield (respond {ok: True {}, drained: drainedCount})

2603

2732

2604-- POST /maintenance/undrain: Resume accepting requests after a drain.

2733-- POST /maintenance/undrain: Resume accepting requests after a drain.

@@ -2634,7 +2763,7 @@ handleHealth = \respond -> do

2634 sa <- *startedAt

2763 sa <- *startedAt

2635 presence <- *presence

2764 presence <- *presence

2636 forwards <- *forwards

2765 forwards <- *forwards

2637− sessionCount <- countAllSessions

2766+ sessionCount <- *totalSessionCount

2638 activePolls <- *activePolls

2767 activePolls <- *activePolls

2639 tmc <- *totalMessageCount

2768 tmc <- *totalMessageCount

2640 lastFedSuccess <- *lastFederationSuccessAt

2769 lastFedSuccess <- *lastFederationSuccessAt

@@ -2665,14 +2794,17 @@ handleHealth = \respond -> do

2665zeroStats = {sessions: 0, sessionsMax: 0, messages: 0, messagesMax: 0, presence: 0, presenceMax: 0, forwards: 0, forwardsMax: 0, challenges: 0, challengesMax: 0, ackedIds: 0, ackedIdsMax: 0, ackOverflow: 0, ackOverflowMax: 0, deliveredIds: 0, deliveredIdsMax: 0, forwardFailures: 0, activePolls: 0, activePollsMax: 0, sendRates: 0, challengeRates: 0, federationRates: 0, globalRateCount: 0, ipRates: 0, lookupRates: 0, recvRates: 0, pruneCount: 0, pruneRunning: 0, heartbeatCount: 0, uptime: 0, lastPruneAt: 0, lastHeartbeatAt: 0, draining: False {}, federationEnabled: False {}, recentGossip: 0, totalMessageCount: 0, recipientCounts: 0, federationInFlight: 0, messagesPct: 0, presencePct: 0, forwardsPct: 0, sessionsPct: 0, federationAvgLatencyMs: 0, federationMaxLatencyMs: 0, federationCallCount: 0, drainCompletedAt: 0, ok: False {}}

2794zeroStats = {sessions: 0, sessionsMax: 0, messages: 0, messagesMax: 0, presence: 0, presenceMax: 0, forwards: 0, forwardsMax: 0, challenges: 0, challengesMax: 0, ackedIds: 0, ackedIdsMax: 0, ackOverflow: 0, ackOverflowMax: 0, deliveredIds: 0, deliveredIdsMax: 0, forwardFailures: 0, activePolls: 0, activePollsMax: 0, sendRates: 0, challengeRates: 0, federationRates: 0, globalRateCount: 0, ipRates: 0, lookupRates: 0, recvRates: 0, pruneCount: 0, pruneRunning: 0, heartbeatCount: 0, uptime: 0, lastPruneAt: 0, lastHeartbeatAt: 0, draining: False {}, federationEnabled: False {}, recentGossip: 0, totalMessageCount: 0, recipientCounts: 0, federationInFlight: 0, messagesPct: 0, presencePct: 0, forwardsPct: 0, sessionsPct: 0, federationAvgLatencyMs: 0, federationMaxLatencyMs: 0, federationCallCount: 0, drainCompletedAt: 0, ok: False {}}

2666

2795

2667handleStats = \authorization respond -> do

2796handleStats = \authorization respond -> do

2797+ t <- now

2798+ globalRateLimited <- checkGlobalRate t

2799+ if globalRateLimited then yield (respond zeroStats)

2800+ else do

2668 let authorized = verifyMaintenanceAuth authorization

2801 let authorized = verifyMaintenanceAuth authorization

2669 if authorized

2802 if authorized

2670 then do

2803 then do

2671− t <- now

2672 -- Split into two atomic reads: core capacity state (consistent) and counters/timing

2804 -- Split into two atomic reads: core capacity state (consistent) and counters/timing

2673 -- (slight inconsistency acceptable for monitoring). This narrows the lock window.

2805 -- (slight inconsistency acceptable for monitoring). This narrows the lock window.

2674 core <- atomic do

2806 core <- atomic do

2675− sc <- countAllSessions

2807+ sc <- *totalSessionCount

2676 p <- *presence

2808 p <- *presence

2677 f <- *forwards

2809 f <- *forwards

2678 c <- *challenges

2810 c <- *challenges

@@ -2716,20 +2848,26 @@ handleStats = \authorization respond -> do

2716 else yield (respond zeroStats)

2848 else yield (respond zeroStats)

2717

2849

2718handleStatsPeers = \authorization respond -> do

2850handleStatsPeers = \authorization respond -> do

2851+ t <- now

2852+ globalRateLimited <- checkGlobalRate t

2853+ if globalRateLimited then yield (respond {peers: [], ok: False {}})

2854+ else do

2719 let authorized = verifyMaintenanceAuth authorization

2855 let authorized = verifyMaintenanceAuth authorization

2720 if authorized

2856 if authorized

2721 then do

2857 then do

2722− t <- now

2723 failures <- *forwardFailures

2858 failures <- *forwardFailures

2724 let peers = do f <- failures; yield {server: f.server, failures: f.failures, lastFailedAt: f.failedAt, firstFailedAt: f.firstFailedAt, backoffMs: retryBackoff f.failures, retriable: shouldRetryServer f.server t failures}

2859 let peers = do f <- failures; yield {server: f.server, failures: f.failures, lastFailedAt: f.failedAt, firstFailedAt: f.firstFailedAt, backoffMs: retryBackoff f.failures, retriable: shouldRetryServer f.server t failures}

2725 yield (respond {peers: peers, ok: True {}})

2860 yield (respond {peers: peers, ok: True {}})

2726 else yield (respond {peers: [], ok: False {}})

2861 else yield (respond {peers: [], ok: False {}})

2727

2862

2728handleStatsKey = \pubkey authorization respond -> do

2863handleStatsKey = \pubkey authorization respond -> do

2864+ t <- now

2865+ globalRateLimited <- checkGlobalRate t

2866+ if globalRateLimited then yield (respond {messages: 0, presence: [], sessions: 0, forwards: 0, activePolls: 0, ok: False {}})

2867+ else do

2729 let authorized = verifyMaintenanceAuth authorization

2868 let authorized = verifyMaintenanceAuth authorization

2730 if authorized

2869 if authorized

2731 then do

2870 then do

2732− t <- now

2733 result <- atomic do

2871 result <- atomic do

2734 keyBucket <- readMessageBucket (pubkeyBucket pubkey)

2872 keyBucket <- readMessageBucket (pubkeyBucket pubkey)

2735 presence <- *presence

2873 presence <- *presence

@@ -2749,8 +2887,10 @@ serve = \req -> case req of

2749 ReqChallenge {pubkey, clientIp, respond} -> handleChallenge pubkey clientIp respond

2887 ReqChallenge {pubkey, clientIp, respond} -> handleChallenge pubkey clientIp respond

2750 VerifyAuth {pubkey, challenge, signature, revokeOthers, clientIp, respond} -> handleVerifyAuth pubkey challenge signature revokeOthers clientIp respond

2888 VerifyAuth {pubkey, challenge, signature, revokeOthers, clientIp, respond} -> handleVerifyAuth pubkey challenge signature revokeOthers clientIp respond

2751 RevokeSession {targetToken, authorization, clientIp, respond} -> handleRevoke targetToken authorization clientIp respond

2889 RevokeSession {targetToken, authorization, clientIp, respond} -> handleRevoke targetToken authorization clientIp respond

2890+ Logout {authorization, clientIp, respond} -> handleLogout authorization clientIp respond

2752 Poll {cursor, generation, watchKeys, authorization, clientIp, respond} -> handlePoll cursor generation watchKeys authorization clientIp respond

2891 Poll {cursor, generation, watchKeys, authorization, clientIp, respond} -> handlePoll cursor generation watchKeys authorization clientIp respond

2753 SendMessage {to, encryptedBlob, timestamp, messageId, authorization, clientIp, respond} -> handleSendMessage to encryptedBlob timestamp messageId authorization clientIp respond

2892 SendMessage {to, encryptedBlob, timestamp, messageId, authorization, clientIp, respond} -> handleSendMessage to encryptedBlob timestamp messageId authorization clientIp respond

2893+ MessageStatus {to, messageId, authorization, clientIp, respond} -> handleMessageStatus to messageId authorization clientIp respond

2754 LookupPresence {pubkey, authorization, clientIp, respond} -> handleLookupPresence pubkey authorization clientIp respond

2894 LookupPresence {pubkey, authorization, clientIp, respond} -> handleLookupPresence pubkey authorization clientIp respond

2755 RecvGossip {events, hopsRemaining, fromServer, authorization, respond} -> handleRecvGossip events hopsRemaining fromServer authorization respond

2895 RecvGossip {events, hopsRemaining, fromServer, authorization, respond} -> handleRecvGossip events hopsRemaining fromServer authorization respond

2756 ForwardMessage {toKey, encryptedBlob, timestamp, originServer, messageId, fromServer, authorization, respond} -> handleForwardMessage toKey encryptedBlob timestamp originServer messageId fromServer authorization respond

2896 ForwardMessage {toKey, encryptedBlob, timestamp, originServer, messageId, fromServer, authorization, respond} -> handleForwardMessage toKey encryptedBlob timestamp originServer messageId fromServer authorization respond

@@ -2801,13 +2941,24 @@ backgroundPrune = do

2801 when status.recovered

2941 when status.recovered

2802 (println ("backgroundPrune: previous cycle stuck for " ++ show (t - status.stuckSince) ++ "ms, forcing recovery"))

2942 (println ("backgroundPrune: previous cycle stuck for " ++ show (t - status.stuckSince) ++ "ms, forcing recovery"))

2803 if status.canRun

2943 if status.canRun

2944+ then do

2945+ -- Guard against unbounded concurrent prune forks from repeated stuck recovery

2946+ forks <- *activePruneForks

2947+ if forks >= 2

2948+ then println "backgroundPrune: skipping (max concurrent prune forks reached)"

2949+ else do

2950+ atomic do

2951+ af <- *activePruneForks

2952+ set *activePruneForks = af + 1

2804− then fork (do

2953+ fork (do

2805− pruneAndRetry t

2954+ pruneAndRetry t

2806− -- Only clear running flag if this is still the latest generation

2955+ -- Only clear running flag if this is still the latest generation

2807− -- (prevents a stuck old fork from clearing the flag after a recovery cycle started)

2956+ -- (prevents a stuck old fork from clearing the flag after a recovery cycle started)

2808− atomic do

2957+ atomic do

2809− currentGen <- *pruneGeneration

2958+ currentGen <- *pruneGeneration

2810− when (currentGen == status.gen) (set *pruneRunning = 0))

2959+ when (currentGen == status.gen) (set *pruneRunning = 0)

2960+ af <- *activePruneForks

2961+ set *activePruneForks = max 0 (af - 1))

2811 else println "backgroundPrune: skipping (previous cycle still running)"

2962 else println "backgroundPrune: skipping (previous cycle still running)"

2812 backgroundPrune

2963 backgroundPrune

2813

2964

@@ -2851,6 +3002,16 @@ main = do

2851 println (" federation: " ++ if federationEnabled then "enabled" else "disabled")

3002 println (" federation: " ++ if federationEnabled then "enabled" else "disabled")

2852 println (" prune interval: " ++ show backgroundPruneInterval ++ "ms")

3003 println (" prune interval: " ++ show backgroundPruneInterval ++ "ms")

2853 println (" poll timeout: " ++ show pollTimeoutMs ++ "ms")

3004 println (" poll timeout: " ++ show pollTimeoutMs ++ "ms")

3005+ println (" max messages: " ++ show maxMessages ++ " (per recipient: " ++ show maxMessagesPerRecipient ++ ")")

3006+ println (" max sessions: " ++ show maxSessions ++ " (per key: " ++ show maxSessionsPerKey ++ ")")

3007+ println (" max presence: " ++ show maxPresenceEntries)

3008+ println (" max forwards: " ++ show maxForwardsEntries)

3009+ println (" session expiry: " ++ show sessionExpiry ++ "ms")

3010+ println (" message max age: " ++ show messageMaxAge ++ "ms")

3011+ println (" global rate limit: " ++ show maxGlobalRequestRate ++ "/" ++ show globalRateWindowMs ++ "ms")

3012+ println (" auth rate limit: " ++ show maxAuthGlobalRate ++ "/" ++ show authGlobalRateWindowMs ++ "ms")

3013+ when (count federationPeerSecrets > 0)

3014+ (println (" per-peer secrets: " ++ show (count federationPeerSecrets) ++ " configured"))

2854 fork backgroundPrune

3015 fork backgroundPrune

2855 fork pollHeartbeat

3016 fork pollHeartbeat

2856 listen 8080 serve

3017 listen 8080 serve