on master help sign in
Improve server.knot: message sharding, federation security, observability, and resource management

- Shard messages into 16 buckets by recipient pubkey for O(n/16) hot-path scans - Add federation replay protection (reject stale timestamps) - Fix challenge consumption to allow concurrent device auth - Add per-key rate limit to /auth/revoke - Block RFC 2606 reserved TLDs (.test, .example, .invalid) - Incremental recipient count maintenance during prune - Use totalMessageCount in stats/health (avoid full message list scans) - Reset federation latency metrics every 5 minutes (windowed averages) - Fix in-flight slot leak: clean stale entries - Drain completion tracking for safe server shutdown

commit5ef07c20e20ba35556ded036a601a0a2ac87edcc authoriko <mail@iko.soy> dateFri 24 Apr 2026, 21:19 (6 months ago) parent104bf5e
1 file changed+238−64
server.knot+238−64
server.knot+238−64
@@ -107,6 +107,10 @@ federationFetchTimeoutMs : Int<Ms>
107107federationFetchTimeoutMs = 30000 -- 30 seconds; consider server unresponsive if fetch hasn't returned
108108maxFederationInFlight = 100
109109
110+-- Federation latency metric window
111+federationLatencyWindowMs : Int<Ms>
112+federationLatencyWindowMs = 300000 -- 5 minutes; reset latency metrics after this interval
113+
110114-- Background task intervals
111115backgroundPruneInterval : Int<Ms>
112116backgroundPruneInterval = 30000 -- 30 seconds; also retries failed forwards/acks
@@ -274,7 +278,24 @@ type FederationInFlightEntry = {server: Text, count: Int, firstStartedAt: Timest
274278*sessionBucket14 : [Session]
275279*sessionBucket15 : [Session]
276280
277−*messages : [Message]
281+-- Messages sharded into 16 buckets by first hex char of recipient pubkey (toKey).
282+-- Hot-path operations (poll, send, ack) scan only one bucket (~1/16 of messages).
283+*messageBucket0 : [Message]
284+*messageBucket1 : [Message]
285+*messageBucket2 : [Message]
286+*messageBucket3 : [Message]
287+*messageBucket4 : [Message]
288+*messageBucket5 : [Message]
289+*messageBucket6 : [Message]
290+*messageBucket7 : [Message]
291+*messageBucket8 : [Message]
292+*messageBucket9 : [Message]
293+*messageBucket10 : [Message]
294+*messageBucket11 : [Message]
295+*messageBucket12 : [Message]
296+*messageBucket13 : [Message]
297+*messageBucket14 : [Message]
298+*messageBucket15 : [Message]
278299
279300*totalMessageCount : Int
280301
@@ -374,6 +395,12 @@ type FederationInFlightEntry = {server: Text, count: Int, firstStartedAt: Timest
374395-- Server instance generation (epoch timestamp, for cursor staleness detection)
375396*serverGeneration : Int
376397
398+-- Timestamp when drain federation ops completed (0 if in progress or not draining)
399+*drainCompletedAt : Timestamp
400+
401+-- Last time federation latency metrics were reset (for windowed averages)
402+*federationLatencyResetAt : Timestamp
403+
377404-- ============================================
378405-- Routes
379406-- ============================================
@@ -444,7 +471,7 @@ route HealthApi where
444471-- Returns counts of in-memory collections (requires federation or maintenance secret)
445472route StatsApi where
446473 GET /stats headers {authorization: Text}
447− -> {sessions: Int, sessionsMax: Int, messages: Int, messagesMax: Int, presence: Int, presenceMax: Int, forwards: Int, forwardsMax: Int, challenges: Int, challengesMax: Int, ackedIds: Int, ackedIdsMax: Int, ackOverflow: Int, ackOverflowMax: Int, deliveredIds: Int, deliveredIdsMax: Int, forwardFailures: Int, activePolls: Int, activePollsMax: Int, sendRates: Int, challengeRates: Int, federationRates: Int, globalRateCount: Int, ipRates: Int, lookupRates: Int, recvRates: Int, pruneCount: Int, heartbeatCount: Int, uptime: Int, lastPruneAt: Timestamp, lastHeartbeatAt: Timestamp, draining: Bool, federationEnabled: Bool, recentGossip: Int, totalMessageCount: Int, recipientCounts: Int, federationInFlight: Int, messagesPct: Int, presencePct: Int, forwardsPct: Int, sessionsPct: Int, federationAvgLatencyMs: Int, federationMaxLatencyMs: Int, federationCallCount: Int, ok: Bool}
474+ -> {sessions: Int, sessionsMax: Int, messages: Int, messagesMax: Int, presence: Int, presenceMax: Int, forwards: Int, forwardsMax: Int, challenges: Int, challengesMax: Int, ackedIds: Int, ackedIdsMax: Int, ackOverflow: Int, ackOverflowMax: Int, deliveredIds: Int, deliveredIdsMax: Int, forwardFailures: Int, activePolls: Int, activePollsMax: Int, sendRates: Int, challengeRates: Int, federationRates: Int, globalRateCount: Int, ipRates: Int, lookupRates: Int, recvRates: Int, pruneCount: Int, heartbeatCount: Int, uptime: Int, lastPruneAt: Timestamp, lastHeartbeatAt: Timestamp, draining: Bool, federationEnabled: Bool, recentGossip: Int, totalMessageCount: Int, recipientCounts: Int, federationInFlight: Int, messagesPct: Int, presencePct: Int, forwardsPct: Int, sessionsPct: Int, federationAvgLatencyMs: Int, federationMaxLatencyMs: Int, federationCallCount: Int, drainCompletedAt: Timestamp, ok: Bool}
448475 = Stats
449476 GET /stats/peers headers {authorization: Text}
450477 -> {peers: [PeerStatus], ok: Bool}
@@ -683,6 +710,9 @@ isBadServerName = \sn ->
683710 || containsSuffix ".home" sn
684711 || containsSuffix ".lan" sn
685712 || containsSuffix ".intranet" sn
713+ || containsSuffix ".test" sn
714+ || containsSuffix ".example" sn
715+ || containsSuffix ".invalid" sn
686716
687717-- Look up a non-expired session matching the given token (constant-time comparison).
688718-- Caller passes the appropriate session bucket (sharded by token prefix).
@@ -841,6 +871,100 @@ transformAllSessionBuckets = \f -> do
841871 b14 <- *sessionBucket14; set *sessionBucket14 = f b14
842872 b15 <- *sessionBucket15; set *sessionBucket15 = f b15
843873
874+-- ============================================
875+-- Message Bucket Helpers
876+-- ============================================
877+-- Messages are sharded into 16 buckets by the first hex char of the recipient pubkey.
878+-- This reduces poll/send/ack scans from O(n) to O(n/16) on the hot path.
879+
880+readMessageBucket : Int -> IO [Message]
881+readMessageBucket = \b ->
882+ if b == 0 then do s <- *messageBucket0; yield s
883+ else if b == 1 then do s <- *messageBucket1; yield s
884+ else if b == 2 then do s <- *messageBucket2; yield s
885+ else if b == 3 then do s <- *messageBucket3; yield s
886+ else if b == 4 then do s <- *messageBucket4; yield s
887+ else if b == 5 then do s <- *messageBucket5; yield s
888+ else if b == 6 then do s <- *messageBucket6; yield s
889+ else if b == 7 then do s <- *messageBucket7; yield s
890+ else if b == 8 then do s <- *messageBucket8; yield s
891+ else if b == 9 then do s <- *messageBucket9; yield s
892+ else if b == 10 then do s <- *messageBucket10; yield s
893+ else if b == 11 then do s <- *messageBucket11; yield s
894+ else if b == 12 then do s <- *messageBucket12; yield s
895+ else if b == 13 then do s <- *messageBucket13; yield s
896+ else if b == 14 then do s <- *messageBucket14; yield s
897+ else do s <- *messageBucket15; yield s
898+
899+writeMessageBucket : Int -> [Message] -> IO ()
900+writeMessageBucket = \b msgs ->
901+ if b == 0 then set *messageBucket0 = msgs
902+ else if b == 1 then set *messageBucket1 = msgs
903+ else if b == 2 then set *messageBucket2 = msgs
904+ else if b == 3 then set *messageBucket3 = msgs
905+ else if b == 4 then set *messageBucket4 = msgs
906+ else if b == 5 then set *messageBucket5 = msgs
907+ else if b == 6 then set *messageBucket6 = msgs
908+ else if b == 7 then set *messageBucket7 = msgs
909+ else if b == 8 then set *messageBucket8 = msgs
910+ else if b == 9 then set *messageBucket9 = msgs
911+ else if b == 10 then set *messageBucket10 = msgs
912+ else if b == 11 then set *messageBucket11 = msgs
913+ else if b == 12 then set *messageBucket12 = msgs
914+ else if b == 13 then set *messageBucket13 = msgs
915+ else if b == 14 then set *messageBucket14 = msgs
916+ else set *messageBucket15 = msgs
917+
918+-- Read all 16 message buckets (for cold-path operations: federation ack, forward candidates, stats)
919+readAllMessages : IO [Message]
920+readAllMessages = do
921+ m0 <- *messageBucket0
922+ m1 <- *messageBucket1
923+ m2 <- *messageBucket2
924+ m3 <- *messageBucket3
925+ m4 <- *messageBucket4
926+ m5 <- *messageBucket5
927+ m6 <- *messageBucket6
928+ m7 <- *messageBucket7
929+ m8 <- *messageBucket8
930+ m9 <- *messageBucket9
931+ m10 <- *messageBucket10
932+ m11 <- *messageBucket11
933+ m12 <- *messageBucket12
934+ m13 <- *messageBucket13
935+ m14 <- *messageBucket14
936+ m15 <- *messageBucket15
937+ yield (union m0 (union m1 (union m2 (union m3 (union m4 (union m5 (union m6 (union m7 (union m8 (union m9 (union m10 (union m11 (union m12 (union m13 (union m14 m15)))))))))))))))
938+
939+-- Apply a filter/transform to every message bucket
940+transformAllMessageBuckets : ([Message] -> [Message]) -> IO ()
941+transformAllMessageBuckets = \f -> do
942+ b0 <- *messageBucket0; set *messageBucket0 = f b0
943+ b1 <- *messageBucket1; set *messageBucket1 = f b1
944+ b2 <- *messageBucket2; set *messageBucket2 = f b2
945+ b3 <- *messageBucket3; set *messageBucket3 = f b3
946+ b4 <- *messageBucket4; set *messageBucket4 = f b4
947+ b5 <- *messageBucket5; set *messageBucket5 = f b5
948+ b6 <- *messageBucket6; set *messageBucket6 = f b6
949+ b7 <- *messageBucket7; set *messageBucket7 = f b7
950+ b8 <- *messageBucket8; set *messageBucket8 = f b8
951+ b9 <- *messageBucket9; set *messageBucket9 = f b9
952+ b10 <- *messageBucket10; set *messageBucket10 = f b10
953+ b11 <- *messageBucket11; set *messageBucket11 = f b11
954+ b12 <- *messageBucket12; set *messageBucket12 = f b12
955+ b13 <- *messageBucket13; set *messageBucket13 = f b13
956+ b14 <- *messageBucket14; set *messageBucket14 = f b14
957+ b15 <- *messageBucket15; set *messageBucket15 = f b15
958+
959+-- Prune a single message bucket and return kept count + per-recipient counts.
960+-- Since messages are sharded by recipient pubkey, each bucket's recipient counts
961+-- are disjoint from other buckets — simple union merges them.
962+pruneMessageBucketCounting = \b messageCutoff -> atomic do
963+ bucket <- readMessageBucket b
964+ let kept = do m <- bucket; where m.receivedAt > messageCutoff; yield m
965+ writeMessageBucket b kept
966+ yield {count: count kept, recipients: recomputeRecipientCounts kept}
967+
844968-- ============================================
845969-- Recipient Count Helpers
846970-- ============================================
@@ -1140,16 +1264,31 @@ pruneExpired = \t -> do
11401264 challenges <- *challenges
11411265 set *challenges = do x <- challenges; where x.expiresAt > t; yield x
11421266 transformAllSessionBuckets (\bucket -> do x <- bucket; where x.expiresAt > t; yield x)
1143− -- Prune messages (largest collection, isolated for minimal lock time)
1144− prunedMessages <- atomic do
1145− let messageCutoff = t - unackedMessageMaxAge
1146− messages <- *messages
1147− let prunedMessages = do m <- messages; where m.receivedAt > messageCutoff; yield m
1148− set *messages = prunedMessages
1149− set *totalMessageCount = count prunedMessages
1150− yield prunedMessages
1151− -- Recompute recipient counts from pruned messages (outside main message lock)
1152− atomic do set *recipientMsgCounts = recomputeRecipientCounts prunedMessages
1267+ -- Prune messages across all 16 buckets. Each bucket prune is atomic (short lock).
1268+ -- Recipient counts are rebuilt incrementally per-bucket — since messages are sharded
1269+ -- by recipient pubkey, each bucket's counts are disjoint, so simple union merges them.
1270+ let messageCutoff = t - unackedMessageMaxAge
1271+ r0 <- pruneMessageBucketCounting 0 messageCutoff
1272+ r1 <- pruneMessageBucketCounting 1 messageCutoff
1273+ r2 <- pruneMessageBucketCounting 2 messageCutoff
1274+ r3 <- pruneMessageBucketCounting 3 messageCutoff
1275+ r4 <- pruneMessageBucketCounting 4 messageCutoff
1276+ r5 <- pruneMessageBucketCounting 5 messageCutoff
1277+ r6 <- pruneMessageBucketCounting 6 messageCutoff
1278+ r7 <- pruneMessageBucketCounting 7 messageCutoff
1279+ r8 <- pruneMessageBucketCounting 8 messageCutoff
1280+ r9 <- pruneMessageBucketCounting 9 messageCutoff
1281+ r10 <- pruneMessageBucketCounting 10 messageCutoff
1282+ r11 <- pruneMessageBucketCounting 11 messageCutoff
1283+ r12 <- pruneMessageBucketCounting 12 messageCutoff
1284+ r13 <- pruneMessageBucketCounting 13 messageCutoff
1285+ r14 <- pruneMessageBucketCounting 14 messageCutoff
1286+ r15 <- pruneMessageBucketCounting 15 messageCutoff
1287+ let prunedTotal = r0.count + r1.count + r2.count + r3.count + r4.count + r5.count + r6.count + r7.count + r8.count + r9.count + r10.count + r11.count + r12.count + r13.count + r14.count + r15.count
1288+ let prunedRecipients = union r0.recipients (union r1.recipients (union r2.recipients (union r3.recipients (union r4.recipients (union r5.recipients (union r6.recipients (union r7.recipients (union r8.recipients (union r9.recipients (union r10.recipients (union r11.recipients (union r12.recipients (union r13.recipients (union r14.recipients r15.recipients))))))))))))))
1289+ atomic do
1290+ set *totalMessageCount = prunedTotal
1291+ set *recipientMsgCounts = prunedRecipients
11531292 -- Prune federation state
11541293 atomic do
11551294 let retainCutoff = t - messageMaxAge
@@ -1183,9 +1322,11 @@ pruneExpired = \t -> do
11831322 set *recvRates = do r <- recvRates; where t - r.windowStart < recvRateWindowMs; yield r
11841323 recentGossip <- *recentGossip
11851324 set *recentGossip = do r <- recentGossip; where t - r.seenAt < gossipDedupWindow; yield r
1186− -- Clean up stale in-flight entries (hung fetches that never completed)
1325+ -- Clean up stale in-flight entries: remove zeroed entries and reset leaked slots
1326+ -- where the forked task died without calling clearFederationInFlight.
1327+ -- Uses 3x timeout threshold to allow the skip-hung-server logic to take effect first.
11871328 inFlight <- *federationInFlight
1188− set *federationInFlight = do r <- inFlight; where r.count > 0; yield r
1329+ set *federationInFlight = do r <- inFlight; where r.count > 0 && (r.firstStartedAt == 0 || t - r.firstStartedAt < federationFetchTimeoutMs * 3); yield r
11891330 yield presenceResult
11901331
11911332
@@ -1407,12 +1548,12 @@ retryPendingForwards = \pending -> do
14071548-- When a remote key comes online, create forward entries for queued messages and send them
14081549forwardQueuedMessages = \targets t -> do
14091550 newForwards <- atomic do
1410− messages <- *messages
1551+ allMessages <- readAllMessages
14111552 forwards <- *forwards
14121553 presence <- *presence
14131554 let localOnline = localOnlineKeys presence t
14141555 let remoteTargets = remoteOnlyTargets targets localOnline
1415− let candidates = forwardCandidates remoteTargets messages forwards t
1556+ let candidates = forwardCandidates remoteTargets allMessages forwards t
14161557 let capped = takeRelation (maxForwardsEntries - count forwards) candidates
14171558 let updatedForwards = union forwards capped
14181559 set *forwards = updatedForwards
@@ -1549,6 +1690,15 @@ pruneAndRetry = \t -> do
15491690 let refreshPeers = if shouldRefresh then collectPeerServers presence (\p -> p.server) else []
15501691 when shouldRefresh (set *lastPresenceRefresh = t)
15511692 yield {extendedActive: toExtend, nearExpiryActive: notExtended, refreshPeers: refreshPeers, shouldRefresh: shouldRefresh}
1693+ -- Reset federation latency metrics periodically for windowed averages (prevents
1694+ -- lifetime-cumulative average from masking current conditions)
1695+ atomic do
1696+ resetAt <- *federationLatencyResetAt
1697+ when (t - resetAt >= federationLatencyWindowMs) (do
1698+ set *federationLatencyTotal = 0
1699+ set *federationLatencyCount = 0
1700+ set *federationLatencyMax = 0
1701+ set *federationLatencyResetAt = t)
15521702 let result = {pendingForwards: pendingForwards, pendingAcks: ackState.pendingAcks, pendingOverflow: ackState.pendingOverflow, extendedActive: refreshResult.extendedActive, nearExpiryActive: refreshResult.nearExpiryActive, refreshPeers: refreshResult.refreshPeers, shouldRefresh: refreshResult.shouldRefresh}
15531703 -- Fork federation I/O so hung peers don't block next prune cycle
15541704 fork (do
@@ -1621,12 +1771,13 @@ readMessageSignal = \pubkey ->
16211771-- and create ack entries for federated messages (to notify origin servers).
16221772-- Returns pending acks and any evicted entries that exceeded ackedIds capacity.
16231773ackPolledMessages = \pubkey cursor t -> do
1624− messages <- *messages
1625− let toAck = do m <- messages; where m.toKey == pubkey && m.receivedAt <= cursor; yield m
1774+ let b = pubkeyBucket pubkey
1775+ bucket <- readMessageBucket b
1776+ let toAck = do m <- bucket; where m.toKey == pubkey && m.receivedAt <= cursor; yield m
16261777 let ackCount = count toAck
16271778 if ackCount > 0
16281779 then do
1629− set *messages = diff messages toAck
1780+ writeMessageBucket b (diff bucket toAck)
16301781 totalCount <- *totalMessageCount
16311782 set *totalMessageCount = totalCount - ackCount
16321783 recipientCounts <- *recipientMsgCounts
@@ -1644,8 +1795,8 @@ ackPolledMessages = \pubkey cursor t -> do
16441795 else yield {pendingAcks: [], evictedEntries: []}
16451796
16461797pollNewEvents = \pubkey clampedCursor t -> do
1647− messages <- *messages
1648− let limitedMsgs = takeRelation maxPollEvents (sortBy (\m -> m.receivedAt) (pollableMessages messages pubkey clampedCursor t))
1798+ bucket <- readMessageBucket (pubkeyBucket pubkey)
1799+ let limitedMsgs = takeRelation maxPollEvents (sortBy (\m -> m.receivedAt) (pollableMessages bucket pubkey clampedCursor t))
16491800 let evts = do
16501801 m <- limitedMsgs
16511802 yield {id: m.messageId, eventType: "message",
@@ -1704,8 +1855,8 @@ pollOnce = \token cursor clientIp -> do
17041855 pollResult <- do
17051856 -- Phase 1: check for existing messages (atomic read)
17061857 initialResult <- atomic do
1707− messages <- *messages
1708− if count (pollableMessages messages s.pubkey s.clampedCursor t) > 0
1858+ bucket <- readMessageBucket (pubkeyBucket s.pubkey)
1859+ if count (pollableMessages bucket s.pubkey s.clampedCursor t) > 0
17091860 then do
17101861 r <- pollNewEvents s.pubkey s.clampedCursor t
17111862 yield (Just {value: {events: r.events, cursor: r.cursor}})
@@ -1923,9 +2074,10 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
19232074 validation <- atomic do
19242075 challenges <- *challenges
19252076 let hasMatch = challengeMatch challenges pubkey challenge clientIp t
1926− -- Only consume challenge on IP match (rate limiter handles brute-force from same IP;
1927− -- not consuming on IP mismatch prevents remote attackers from burning challenges)
1928− when hasMatch (set *challenges = do c <- challenges; where c.pubkey != pubkey; yield c)
2077+ -- Consume only the matched challenge on IP match (rate limiter handles brute-force from same IP;
2078+ -- not consuming on IP mismatch prevents remote attackers from burning challenges).
2079+ -- Other challenges for the same pubkey are preserved (concurrent auth from multiple devices).
2080+ when hasMatch (set *challenges = challengesExcluding challenges pubkey challenge)
19292081 if not hasMatch || not sigValid
19302082 then yield {valid: False {}}
19312083 else yield {valid: True {}}
@@ -2020,11 +2172,12 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization clientIp
20202172 -- Reject self-sends (sender == recipient wastes a mailbox slot)
20212173 if senderKey == to then yield {send: SendFailed {status: MsRejected {}}, senderRateLimited: False {}, globalCapReached: False {}}
20222174 else do
2023− messages <- *messages
2175+ let msgBucket = pubkeyBucket to
2176+ bucket <- readMessageBucket msgBucket
20242177 ackedIds <- *ackedIds
20252178 deliveredIds <- *deliveredIds
2026− -- Dedup check (only when messageId is non-empty; skips O(n) scan otherwise)
2027− let isDup = checkMessageDup messages messageId to || (messageId != "" && (hasMessage messageId to ackedIds || hasMessage messageId to deliveredIds))
2179+ -- Dedup check (only when messageId is non-empty; scans only the recipient's bucket)
2180+ let isDup = checkMessageDup bucket messageId to || (messageId != "" && (hasMessage messageId to ackedIds || hasMessage messageId to deliveredIds))
20282181 if isDup
20292182 then yield {send: SendDuplicate {}, senderRateLimited: False {}, globalCapReached: False {}}
20302183 else do
@@ -2060,9 +2213,9 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization clientIp
20602213 seqCounter <- *seqCounter
20612214 let seqTs = nextSeqTs t seqCounter
20622215 set *seqCounter = seqTs
2063− set *messages = union messages
2216+ writeMessageBucket msgBucket (union bucket
20642217 [{messageId: effectiveMsgId, toKey: to,
2065− encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: "", receivedAt: seqTs}]
2218+ encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: "", receivedAt: seqTs}])
20662219 set *totalMessageCount = totalMsgCount + 1
20672220 set *recipientMsgCounts = adjustRecipientCount recipientCounts to 1
20682221 signalMessageForKey to seqTs
@@ -2150,6 +2303,16 @@ handleRevoke = \targetToken authorization clientIp respond -> do
21502303 withSessionAuth authorization t clientIp
21512304 (yield (respond {ok: False {}, authorized: False {}}))
21522305 (\callerPubkey -> do
2306+ revokeLimited <- atomic do
2307+ challengeRates <- *challengeRates
2308+ let rateResult = checkRateWindow challengeRates (\r -> r.pubkey == callerPubkey)
2309+ (\ts -> {pubkey: callerPubkey, windowStart: ts, count: 1})
2310+ (\r -> {pubkey: r.pubkey, windowStart: r.windowStart, count: r.count + 1})
2311+ maxChallengeRatePerKey challengeRateWindowMs maxChallengeRateEntries t
2312+ set *challengeRates = rateResult.entries
2313+ yield rateResult.limited
2314+ if revokeLimited then yield (respond {ok: False {}, authorized: True {}})
2315+ else do
21532316 let targetBucket = tokenBucket targetToken
21542317 removed <- atomic do
21552318 sessions <- readSessionBucket targetBucket
@@ -2210,12 +2373,19 @@ handleRecvGossip = \events hopsRemaining fromServer authorization respond -> do
22102373handleForwardMessage = \toKey encryptedBlob timestamp originServer messageId fromServer authorization respond -> do
22112374 t <- now
22122375 withFederationAuth authorization t respond fromServer (not (isLocalServer originServer) && fromServer == originServer) (do
2376+ -- Replay protection: reject forwards with timestamps older than messageMaxAge.
2377+ -- Legitimate queued messages have timestamps within this window. Replayed captures
2378+ -- become stale once dedup state (messages + ackedIds + deliveredIds) is pruned (~30 days).
2379+ if t - timestamp > messageMaxAge
2380+ then yield (respond {ok: False {}, error: "stale_timestamp"})
2381+ else do
22132382 let clampedTs = clampTimestamp timestamp t
22142383 result <- atomic do
2215− messages <- *messages
2384+ let msgBucket = pubkeyBucket toKey
2385+ bucket <- readMessageBucket msgBucket
22162386 ackedIds <- *ackedIds
22172387 deliveredIds <- *deliveredIds
2218− let isDup = checkMessageDup messages messageId toKey || hasMessage messageId toKey ackedIds || hasMessage messageId toKey deliveredIds
2388+ let isDup = checkMessageDup bucket messageId toKey || hasMessage messageId toKey ackedIds || hasMessage messageId toKey deliveredIds
22192389 if isDup
22202390 then yield {forward: ForwardDuplicate {}, globalCapReached: False {}, recvRateLimited: False {}}
22212391 else do
@@ -2240,10 +2410,9 @@ handleForwardMessage = \toKey encryptedBlob timestamp originServer messageId fro
22402410 seqCounter <- *seqCounter
22412411 let seqTs = nextSeqTs t seqCounter
22422412 set *seqCounter = seqTs
2243− let updatedMessages = union messages
2413+ writeMessageBucket msgBucket (union bucket
22442414 [{messageId: messageId, toKey: toKey,
2245− encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: originServer, receivedAt: seqTs}]
2246− set *messages = updatedMessages
2415+ encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: originServer, receivedAt: seqTs}])
22472416 set *totalMessageCount = totalMsgCount + 1
22482417 set *recipientMsgCounts = adjustRecipientCount recipientCounts toKey 1
22492418 signalMessageForKey toKey seqTs
@@ -2272,7 +2441,8 @@ handleFederationAck = \messageIds originServer fromServer authorization respond
22722441 let ackedForwards = do f <- forwards; where f.toServer == fromServer && any dedupedIds (\mid -> f.messageId == mid); yield f
22732442 let updatedForwards = forwardsExcluding forwards fromServer dedupedIds
22742443 set *forwards = updatedForwards
2275− messages <- *messages
2444+ -- Read all message buckets (cold path: federation ack handles at most maxMessageIds entries)
2445+ allMessages <- readAllMessages
22762446 presence <- *presence
22772447 let removable = do
22782448 f <- ackedForwards
@@ -2280,20 +2450,18 @@ handleFederationAck = \messageIds originServer fromServer authorization respond
22802450 where none presence (\p -> p.pubkey == f.toKey && isLocalServer p.server && p.expiresAt > t)
22812451 yield f
22822452 -- Only remove messages stored long enough for local clients to have polled them (grace period)
2283− let pastGrace = do f <- removable; where any messages (\m -> m.messageId == f.messageId && m.toKey == f.toKey && t - m.receivedAt >= ackGracePeriod); yield f
2284− -- Single-pass fold: filter messages and count removals together
2285− let msgResult = fold (\acc m ->
2286− if none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey)
2287− then {kept: union acc.kept [m], removed: acc.removed}
2288− else {kept: acc.kept, removed: acc.removed + 1}
2289− ) {kept: [], removed: 0} messages
2290− set *messages = msgResult.kept
2291− totalMsgCount <- *totalMessageCount
2292− set *totalMessageCount = totalMsgCount - msgResult.removed
2293− -- Adjust per-recipient counts for removed messages
2294− recipientCounts <- *recipientMsgCounts
2295− let adjustedCounts = fold (\acc f -> adjustRecipientCount acc f.toKey (0 - 1)) recipientCounts pastGrace
2296− set *recipientMsgCounts = adjustedCounts
2453+ let pastGrace = do f <- removable; where any allMessages (\m -> m.messageId == f.messageId && m.toKey == f.toKey && t - m.receivedAt >= ackGracePeriod); yield f
2454+ let removed = count pastGrace
2455+ -- Remove matched messages from their buckets (pastGrace is small, so scanning all 16 buckets is fast)
2456+ when (removed > 0) (do
2457+ transformAllMessageBuckets (\bucket ->
2458+ do m <- bucket; where none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey); yield m)
2459+ totalMsgCount <- *totalMessageCount
2460+ set *totalMessageCount = totalMsgCount - removed
2461+ -- Adjust per-recipient counts for removed messages
2462+ recipientCounts <- *recipientMsgCounts
2463+ let adjustedCounts = fold (\acc f -> adjustRecipientCount acc f.toKey (0 - 1)) recipientCounts pastGrace
2464+ set *recipientMsgCounts = adjustedCounts)
22972465 yield (respond {ok: True {}, error: ""}))
22982466
22992467-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).
@@ -2317,10 +2485,14 @@ handleDrain = \authorization respond -> do
23172485 let pendingForwardCount = count forwards
23182486 let pendingAckCount = count ackedIds
23192487 -- Fork gossip so drain responds immediately (hung peers can't block the response).
2320− -- Operator should allow a few seconds after response before killing the process.
2321− fork (when (federationEnabled && drainedCount > 0 && count drainResult.peerServers > 0) (do
2322− let gossipEvents = do p <- drainResult.localEntries; yield {eventType: GossipOffline {}, pubkey: p.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}
2323− sendBatchedGossip drainResult.peerServers gossipEvents gossipOriginHops))
2488+ -- Operator can poll /stats for drainCompletedAt > 0 to know when it's safe to stop.
2489+ atomic do set *drainCompletedAt = 0
2490+ fork (do
2491+ when (federationEnabled && drainedCount > 0 && count drainResult.peerServers > 0) (do
2492+ let gossipEvents = do p <- drainResult.localEntries; yield {eventType: GossipOffline {}, pubkey: p.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}
2493+ sendBatchedGossip drainResult.peerServers gossipEvents gossipOriginHops)
2494+ t2 <- now
2495+ atomic do set *drainCompletedAt = t2)
23242496 println ("drain: broadcast offline for " ++ show drainedCount ++ " local keys")
23252497 when (pendingForwardCount > 0)
23262498 (println ("drain: warning: " ++ show pendingForwardCount ++ " pending forwards will not be retried"))
@@ -2362,12 +2534,12 @@ handleHealth = \respond -> do
23622534 forwards <- *forwards
23632535 sessionCount <- countAllSessions
23642536 activePolls <- *activePolls
2365− messages <- *messages
2537+ tmc <- *totalMessageCount
23662538 lastFedSuccess <- *lastFederationSuccessAt
23672539 forwardFailures <- *forwardFailures
23682540 let warmup = t - sa < healthPruneMaxAge
23692541 -- Capacity pressure: flag when any capped collection exceeds 90% full
2370− let capacityOk = count messages * 10 <= maxMessages * 9
2542+ let capacityOk = tmc * 10 <= maxMessages * 9
23712543 && count presence * 10 <= maxPresenceEntries * 9
23722544 && count forwards * 10 <= maxForwardsEntries * 9
23732545 && sessionCount * 10 <= maxSessions * 9
@@ -2386,7 +2558,7 @@ handleHealth = \respond -> do
23862558 capacityHealthy: result.capacityOk,
23872559 federationHealthy: result.federationOk})
23882560
2389−zeroStats = {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, 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, ok: False {}}
2561+zeroStats = {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, 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 {}}
23902562
23912563handleStats = \authorization respond -> do
23922564 let authorized = verifyMaintenanceAuth authorization
@@ -2394,9 +2566,10 @@ handleStats = \authorization respond -> do
23942566 then do
23952567 t <- now
23962568 -- Read all refs atomically for a consistent snapshot, count outside to minimize lock time
2569+ -- Read all refs atomically for a consistent snapshot. Messages use totalMessageCount
2570+ -- (maintained counter) instead of scanning all 16 message buckets.
23972571 snapshot <- atomic do
23982572 sc <- countAllSessions
2399− m <- *messages
24002573 p <- *presence
24012574 f <- *forwards
24022575 c <- *challenges
@@ -2425,14 +2598,15 @@ handleStats = \authorization respond -> do
24252598 flt <- *federationLatencyTotal
24262599 flc <- *federationLatencyCount
24272600 flm <- *federationLatencyMax
2428− yield {sc: sc, m: m, p: p, f: f, c: c, a: a, ao: ao, d: d, ff: ff, ap: ap, sr: sr, cr: cr, fr: fr, grc: grc, ir: ir, lr: lr, rr: rr, pc: pc, hc: hc, sa: sa, lp: lp, lh: lh, dr: dr, rg: rg, tmc: tmc, rc: rc, fi: fi, flt: flt, flc: flc, flm: flm}
2429− let result = {sessions: snapshot.sc, messages: count snapshot.m, presence: count snapshot.p, forwards: count snapshot.f, challenges: count snapshot.c, ackedIds: count snapshot.a, ackOverflow: count snapshot.ao, deliveredIds: count snapshot.d, forwardFailures: count snapshot.ff, activePolls: count snapshot.ap, sendRates: count snapshot.sr, challengeRates: count snapshot.cr, federationRates: count snapshot.fr, globalRateCount: snapshot.grc, ipRates: count snapshot.ir, lookupRates: count snapshot.lr, recvRates: count snapshot.rr, pruneCount: snapshot.pc, heartbeatCount: snapshot.hc, uptime: t - snapshot.sa, lastPruneAt: snapshot.lp, lastHeartbeatAt: snapshot.lh, draining: snapshot.dr, recentGossip: count snapshot.rg, totalMessageCount: snapshot.tmc, recipientCounts: count snapshot.rc, federationInFlight: count snapshot.fi, federationLatencyTotal: snapshot.flt, federationLatencyCount: snapshot.flc, federationLatencyMax: snapshot.flm}
2601+ dca <- *drainCompletedAt
2602+ yield {sc: sc, p: p, f: f, c: c, a: a, ao: ao, d: d, ff: ff, ap: ap, sr: sr, cr: cr, fr: fr, grc: grc, ir: ir, lr: lr, rr: rr, pc: pc, hc: hc, sa: sa, lp: lp, lh: lh, dr: dr, rg: rg, tmc: tmc, rc: rc, fi: fi, flt: flt, flc: flc, flm: flm, dca: dca}
2603+ let result = {sessions: snapshot.sc, messages: snapshot.tmc, presence: count snapshot.p, forwards: count snapshot.f, challenges: count snapshot.c, ackedIds: count snapshot.a, ackOverflow: count snapshot.ao, deliveredIds: count snapshot.d, forwardFailures: count snapshot.ff, activePolls: count snapshot.ap, sendRates: count snapshot.sr, challengeRates: count snapshot.cr, federationRates: count snapshot.fr, globalRateCount: snapshot.grc, ipRates: count snapshot.ir, lookupRates: count snapshot.lr, recvRates: count snapshot.rr, pruneCount: snapshot.pc, heartbeatCount: snapshot.hc, uptime: t - snapshot.sa, lastPruneAt: snapshot.lp, lastHeartbeatAt: snapshot.lh, draining: snapshot.dr, recentGossip: count snapshot.rg, totalMessageCount: snapshot.tmc, recipientCounts: count snapshot.rc, federationInFlight: count snapshot.fi, federationLatencyTotal: snapshot.flt, federationLatencyCount: snapshot.flc, federationLatencyMax: snapshot.flm, drainCompletedAt: snapshot.dca}
24302604 let msgPct = if maxMessages > 0 then result.messages * 100 / maxMessages else 0
24312605 let presPct = if maxPresenceEntries > 0 then result.presence * 100 / maxPresenceEntries else 0
24322606 let fwdPct = if maxForwardsEntries > 0 then result.forwards * 100 / maxForwardsEntries else 0
24332607 let sessPct = if maxSessions > 0 then result.sessions * 100 / maxSessions else 0
24342608 let fedAvg = if result.federationLatencyCount > 0 then result.federationLatencyTotal / result.federationLatencyCount else 0
2435− yield (respond {sessions: result.sessions, sessionsMax: maxSessions, messages: result.messages, messagesMax: maxMessages, presence: result.presence, presenceMax: maxPresenceEntries, forwards: result.forwards, forwardsMax: maxForwardsEntries, challenges: result.challenges, challengesMax: maxChallenges, ackedIds: result.ackedIds, ackedIdsMax: maxAckedIds, ackOverflow: result.ackOverflow, ackOverflowMax: maxAckOverflow, deliveredIds: result.deliveredIds, deliveredIdsMax: maxDeliveredIds, forwardFailures: result.forwardFailures, activePolls: result.activePolls, activePollsMax: maxActivePolls, sendRates: result.sendRates, challengeRates: result.challengeRates, federationRates: result.federationRates, globalRateCount: result.globalRateCount, ipRates: result.ipRates, lookupRates: result.lookupRates, recvRates: result.recvRates, pruneCount: result.pruneCount, heartbeatCount: result.heartbeatCount, uptime: result.uptime, lastPruneAt: result.lastPruneAt, lastHeartbeatAt: result.lastHeartbeatAt, draining: result.draining, federationEnabled: federationEnabled, recentGossip: result.recentGossip, totalMessageCount: result.totalMessageCount, recipientCounts: result.recipientCounts, federationInFlight: result.federationInFlight, messagesPct: msgPct, presencePct: presPct, forwardsPct: fwdPct, sessionsPct: sessPct, federationAvgLatencyMs: fedAvg, federationMaxLatencyMs: result.federationLatencyMax, federationCallCount: result.federationLatencyCount, ok: True {}})
2609+ yield (respond {sessions: result.sessions, sessionsMax: maxSessions, messages: result.messages, messagesMax: maxMessages, presence: result.presence, presenceMax: maxPresenceEntries, forwards: result.forwards, forwardsMax: maxForwardsEntries, challenges: result.challenges, challengesMax: maxChallenges, ackedIds: result.ackedIds, ackedIdsMax: maxAckedIds, ackOverflow: result.ackOverflow, ackOverflowMax: maxAckOverflow, deliveredIds: result.deliveredIds, deliveredIdsMax: maxDeliveredIds, forwardFailures: result.forwardFailures, activePolls: result.activePolls, activePollsMax: maxActivePolls, sendRates: result.sendRates, challengeRates: result.challengeRates, federationRates: result.federationRates, globalRateCount: result.globalRateCount, ipRates: result.ipRates, lookupRates: result.lookupRates, recvRates: result.recvRates, pruneCount: result.pruneCount, heartbeatCount: result.heartbeatCount, uptime: result.uptime, lastPruneAt: result.lastPruneAt, lastHeartbeatAt: result.lastHeartbeatAt, draining: result.draining, federationEnabled: federationEnabled, recentGossip: result.recentGossip, totalMessageCount: result.totalMessageCount, recipientCounts: result.recipientCounts, federationInFlight: result.federationInFlight, messagesPct: msgPct, presencePct: presPct, forwardsPct: fwdPct, sessionsPct: sessPct, federationAvgLatencyMs: fedAvg, federationMaxLatencyMs: result.federationLatencyMax, federationCallCount: result.federationLatencyCount, drainCompletedAt: result.drainCompletedAt, ok: True {}})
24362610 else yield (respond zeroStats)
24372611
24382612handleStatsPeers = \authorization respond -> do
@@ -2451,12 +2625,12 @@ handleStatsKey = \pubkey authorization respond -> do
24512625 then do
24522626 t <- now
24532627 result <- atomic do
2454− messages <- *messages
2628+ keyBucket <- readMessageBucket (pubkeyBucket pubkey)
24552629 presence <- *presence
24562630 allSessions <- readAllSessions
24572631 forwards <- *forwards
24582632 activePolls <- *activePolls
2459− let keyMessages = count (do m <- messages; where m.toKey == pubkey; yield m)
2633+ let keyMessages = count (do m <- keyBucket; where m.toKey == pubkey; yield m)
24602634 let keyPresence = do p <- presence; where p.pubkey == pubkey && p.expiresAt > t; yield {server: p.server, expiresAt: p.expiresAt}
24612635 let keySessions = count (do s <- allSessions; where s.pubkey == pubkey; yield s)
24622636 let keyForwards = count (do f <- forwards; where f.toKey == pubkey; yield f)