@@ -18,10 +18,11 @@ serverHost = "localhost:8080"
1818federationSecret = "" -- must be set to enable federation
1919maintenanceSecret = "" -- set to enable /maintenance/prune independently of federation
2020federationScheme = "https://" -- scheme for federation HTTP calls
21−-- NOTE: Outbound federation HTTP calls have no timeout. A slow/unresponsive peer can hold
22−-- a concurrency slot (maxFederationConcurrency) indefinitely. Add a fetch timeout if the
23−-- Knot runtime supports it; until then, exponential backoff on failed fetches mitigates
24−-- but does not fully prevent resource exhaustion from hung connections.
21+-- NOTE: Outbound federation HTTP calls have no runtime-level timeout. The fetchFederation
22+-- wrapper tracks in-flight requests per server and skips servers with outstanding requests
23+-- older than federationFetchTimeoutMs (see shouldRetryServerEx). This prevents a single
24+-- slow peer from consuming all concurrency slots, though truly hung connections still hold
25+-- one slot until the forked task terminates.
2526-- Per-peer federation secrets (overrides global federationSecret for specific peers)
2627-- Format: [{server: "peer.example.com", secret: "peer-specific-secret"}]
2728federationPeerSecrets : [{server: Text, secret: Text}]
@@ -101,6 +102,10 @@ retryBackoffMax = 480000 -- 8 minutes
101102maxForwardFailures = 1000 -- max tracked server failure records
102103maxServerDowntime : Int<Ms>
103104maxServerDowntime = 86400000 -- 24 hours; stop retrying a server after this duration of consecutive failures
105+-- Federation in-flight tracking: skip servers with outstanding requests older than this threshold
106+federationFetchTimeoutMs : Int<Ms>
107+federationFetchTimeoutMs = 30000 -- 30 seconds; consider server unresponsive if fetch hasn't returned
108+maxFederationInFlight = 100
104109
105110-- Background task intervals
106111backgroundPruneInterval : Int<Ms>
@@ -208,7 +213,7 @@ type Timestamp = Int<Ms>
208213-- Pending challenge for Ed25519 challenge-response auth
209214type Challenge = {pubkey: PubkeyHex, challenge: Text, expiresAt: Timestamp, ip: Text}
210215-- Authenticated session; token is a random hex string
211−type Session = {token: Text, pubkey: PubkeyHex, expiresAt: Timestamp}
216+type Session = {token: Text, pubkey: PubkeyHex, expiresAt: Timestamp, ip: Text}
212217-- Stored message; originServer is "" for locally-originated messages, set to sender server for federated forwards
213218type Message = {messageId: MessageIdHex, toKey: PubkeyHex, encryptedBlob: BlobHex, timestamp: Timestamp, originServer: Text, receivedAt: Timestamp}
214219-- Presence: tracks that a pubkey is online at a given server, with expiry
@@ -229,7 +234,7 @@ type ChallengeRateEntry = {pubkey: PubkeyHex, windowStart: Timestamp, count: Int
229234type FederationRateEntry = {server: Text, windowStart: Timestamp, count: Int}
230235-- Tracks active long-poll connections to enforce per-key concurrency limits
231236type ActivePoll = {pubkey: PubkeyHex, startedAt: Timestamp}
232−data AuthResult = Authorized {pubkey: PubkeyHex} | Unauthorized {}
237+data AuthResult = Authorized {pubkey: PubkeyHex, ip: Text} | Unauthorized {}
233238data SendResult = SendFailed {status: MessageStatus} | SendDone {delivery: DeliveryStatus, remoteServers: [ServerName]} | SendDuplicate {}
234239data ForwardResult = ForwardDuplicate {} | ForwardRejected {} | ForwardStored {}
235240
@@ -239,6 +244,8 @@ type LookupRateEntry = {pubkey: PubkeyHex, windowStart: Timestamp, count: Int}
239244type RecentGossipEntry = {pubkey: PubkeyHex, server: Text, timestamp: Timestamp, seenAt: Timestamp}
240245type PeerStatus = {server: Text, failures: Int, lastFailedAt: Timestamp, firstFailedAt: Timestamp, backoffMs: Int, retriable: Bool}
241246type KeyPresenceEntry = {server: Text, expiresAt: Timestamp}
247+type RecipientCount = {pubkey: PubkeyHex, count: Int}
248+type FederationInFlightEntry = {server: Text, count: Int, firstStartedAt: Timestamp}
242249
243250-- ============================================
244251-- Persisted State
@@ -248,7 +255,24 @@ type KeyPresenceEntry = {server: Text, expiresAt: Timestamp}
248255
249256*challenges : [Challenge]
250257
251−*sessions : [Session]
258+-- Sessions sharded into 16 buckets by first hex char of token.
259+-- Hot-path auth scans only one bucket (~1/16 of sessions).
260+*sessionBucket0 : [Session]
261+*sessionBucket1 : [Session]
262+*sessionBucket2 : [Session]
263+*sessionBucket3 : [Session]
264+*sessionBucket4 : [Session]
265+*sessionBucket5 : [Session]
266+*sessionBucket6 : [Session]
267+*sessionBucket7 : [Session]
268+*sessionBucket8 : [Session]
269+*sessionBucket9 : [Session]
270+*sessionBucket10 : [Session]
271+*sessionBucket11 : [Session]
272+*sessionBucket12 : [Session]
273+*sessionBucket13 : [Session]
274+*sessionBucket14 : [Session]
275+*sessionBucket15 : [Session]
252276
253277*messages : [Message]
254278
@@ -336,6 +360,12 @@ type KeyPresenceEntry = {server: Text, expiresAt: Timestamp}
336360
337361*lastFederationSuccessAt : Timestamp
338362
363+-- Per-recipient message counts (avoids O(n) scan of all messages on send/forward)
364+*recipientMsgCounts : [RecipientCount]
365+
366+-- Tracks outstanding federation fetches per server (detects hung connections)
367+*federationInFlight : [FederationInFlightEntry]
368+
339369-- ============================================
340370-- Routes
341371-- ============================================
@@ -356,16 +386,16 @@ route AuthApi where
356386-- Client-facing API (requires session token in Authorization header)
357387route ClientApi where
358388 -- Long-poll for new messages; cursor tracks last-seen receivedAt timestamp
359− POST {cursor: Timestamp} /poll headers {authorization: Text}
389+ POST {cursor: Timestamp} /poll headers {authorization: Text, clientIp: Text}
360390 -> {events: [PollEvent], cursor: Timestamp, authorized: Bool, rateLimited: Bool, draining: Bool}
361391 = Poll
362392 /messages
363393 -- Send an E2E encrypted message to a recipient pubkey
364− POST {to: PubkeyHex, encryptedBlob: BlobHex, timestamp: Timestamp, messageId: OptionalMessageId} / headers {authorization: Text}
394+ POST {to: PubkeyHex, encryptedBlob: BlobHex, timestamp: Timestamp, messageId: OptionalMessageId} / headers {authorization: Text, clientIp: Text}
365395 -> {status: MessageStatus, messageId: OptionalMessageId}
366396 = SendMessage
367397 -- Look up which server a pubkey is present on (for federation routing)
368− GET /lookup/{pubkey: PubkeyHex} headers {authorization: Text}
398+ GET /lookup/{pubkey: PubkeyHex} headers {authorization: Text, clientIp: Text}
369399 -> {server: AnyServerName, authorized: Bool}
370400 = LookupPresence
371401
@@ -400,7 +430,7 @@ route HealthApi where
400430-- Returns counts of in-memory collections (requires federation or maintenance secret)
401431route StatsApi where
402432 GET /stats headers {authorization: Text}
403− -> {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, ok: Bool}
433+ -> {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, ok: Bool}
404434 = Stats
405435 GET /stats/peers headers {authorization: Text}
406436 -> {peers: [PeerStatus], ok: Bool}
@@ -633,16 +663,16 @@ isBadServerName = \sn ->
633663 || isHighOctetPrefix sn
634664 || not (hasDotBeforePort sn)
635665 || containsSuffix ".arpa" sn
666+ || containsSuffix ".onion" sn
667+ || containsSuffix ".i2p" sn
636668
637−-- Look up a non-expired session matching the given token (constant-time comparison)
638−-- Pre-filter by token prefix to skip ~15/16 of sessions before expensive safeCompare.
639−-- Prefix match is not constant-time but only leaks approximate session count (already in /stats).
669+-- Look up a non-expired session matching the given token (constant-time comparison).
670+-- Caller passes the appropriate session bucket (sharded by token prefix).
640671authenticateSession : [Session] -> Text -> Timestamp -> AuthResult
641672authenticateSession = \sessions token t ->
642− let prefix = take 1 token in
643− case headOf (do s <- sessions; where take 1 s.token == prefix && s.expiresAt > t && safeCompare s.token token; yield s) of
673+ case headOf (do s <- sessions; where s.expiresAt > t && safeCompare s.token token; yield s) of
644674 Nothing {} -> Unauthorized {}
645− Just {value: s} -> Authorized {pubkey: s.pubkey}
675+ Just {value: s} -> Authorized {pubkey: s.pubkey, ip: s.ip}
646676
647677-- Extract bearer token from "Authorization: Bearer <token>" header
648678extractToken = \authHeader ->
@@ -684,6 +714,181 @@ federationAuthForServer = \server ->
684714 Just {value: p} -> "Bearer " ++ p.secret
685715 Nothing {} -> federationAuth
686716
717+-- ============================================
718+-- Session Bucket Helpers
719+-- ============================================
720+-- Sessions are sharded into 16 buckets by the first hex char of the token.
721+-- This reduces auth scan from O(n) to O(n/16) on the hot path.
722+
723+tokenBucket : Text -> Int
724+tokenBucket = \token -> clamp 0 15 (hexCharIndex (take 1 token))
725+
726+readSessionBucket : Int -> IO [Session]
727+readSessionBucket = \b ->
728+ if b == 0 then do s <- *sessionBucket0; yield s
729+ else if b == 1 then do s <- *sessionBucket1; yield s
730+ else if b == 2 then do s <- *sessionBucket2; yield s
731+ else if b == 3 then do s <- *sessionBucket3; yield s
732+ else if b == 4 then do s <- *sessionBucket4; yield s
733+ else if b == 5 then do s <- *sessionBucket5; yield s
734+ else if b == 6 then do s <- *sessionBucket6; yield s
735+ else if b == 7 then do s <- *sessionBucket7; yield s
736+ else if b == 8 then do s <- *sessionBucket8; yield s
737+ else if b == 9 then do s <- *sessionBucket9; yield s
738+ else if b == 10 then do s <- *sessionBucket10; yield s
739+ else if b == 11 then do s <- *sessionBucket11; yield s
740+ else if b == 12 then do s <- *sessionBucket12; yield s
741+ else if b == 13 then do s <- *sessionBucket13; yield s
742+ else if b == 14 then do s <- *sessionBucket14; yield s
743+ else do s <- *sessionBucket15; yield s
744+
745+writeSessionBucket : Int -> [Session] -> IO ()
746+writeSessionBucket = \b sessions ->
747+ if b == 0 then set *sessionBucket0 = sessions
748+ else if b == 1 then set *sessionBucket1 = sessions
749+ else if b == 2 then set *sessionBucket2 = sessions
750+ else if b == 3 then set *sessionBucket3 = sessions
751+ else if b == 4 then set *sessionBucket4 = sessions
752+ else if b == 5 then set *sessionBucket5 = sessions
753+ else if b == 6 then set *sessionBucket6 = sessions
754+ else if b == 7 then set *sessionBucket7 = sessions
755+ else if b == 8 then set *sessionBucket8 = sessions
756+ else if b == 9 then set *sessionBucket9 = sessions
757+ else if b == 10 then set *sessionBucket10 = sessions
758+ else if b == 11 then set *sessionBucket11 = sessions
759+ else if b == 12 then set *sessionBucket12 = sessions
760+ else if b == 13 then set *sessionBucket13 = sessions
761+ else if b == 14 then set *sessionBucket14 = sessions
762+ else set *sessionBucket15 = sessions
763+
764+-- Read all 16 session buckets (for per-key operations, prune, stats)
765+readAllSessions : IO [Session]
766+readAllSessions = do
767+ s0 <- *sessionBucket0
768+ s1 <- *sessionBucket1
769+ s2 <- *sessionBucket2
770+ s3 <- *sessionBucket3
771+ s4 <- *sessionBucket4
772+ s5 <- *sessionBucket5
773+ s6 <- *sessionBucket6
774+ s7 <- *sessionBucket7
775+ s8 <- *sessionBucket8
776+ s9 <- *sessionBucket9
777+ s10 <- *sessionBucket10
778+ s11 <- *sessionBucket11
779+ s12 <- *sessionBucket12
780+ s13 <- *sessionBucket13
781+ s14 <- *sessionBucket14
782+ s15 <- *sessionBucket15
783+ yield (union s0 (union s1 (union s2 (union s3 (union s4 (union s5 (union s6 (union s7 (union s8 (union s9 (union s10 (union s11 (union s12 (union s13 (union s14 s15)))))))))))))))
784+
785+-- Count sessions across all 16 buckets without materializing the full list
786+countAllSessions : IO Int
787+countAllSessions = do
788+ s0 <- *sessionBucket0
789+ s1 <- *sessionBucket1
790+ s2 <- *sessionBucket2
791+ s3 <- *sessionBucket3
792+ s4 <- *sessionBucket4
793+ s5 <- *sessionBucket5
794+ s6 <- *sessionBucket6
795+ s7 <- *sessionBucket7
796+ s8 <- *sessionBucket8
797+ s9 <- *sessionBucket9
798+ s10 <- *sessionBucket10
799+ s11 <- *sessionBucket11
800+ s12 <- *sessionBucket12
801+ s13 <- *sessionBucket13
802+ s14 <- *sessionBucket14
803+ s15 <- *sessionBucket15
804+ yield (count s0 + count s1 + count s2 + count s3 + count s4 + count s5 + count s6 + count s7 + count s8 + count s9 + count s10 + count s11 + count s12 + count s13 + count s14 + count s15)
805+
806+-- Apply a filter/transform to every session bucket
807+transformAllSessionBuckets : ([Session] -> [Session]) -> IO ()
808+transformAllSessionBuckets = \f -> do
809+ b0 <- *sessionBucket0; set *sessionBucket0 = f b0
810+ b1 <- *sessionBucket1; set *sessionBucket1 = f b1
811+ b2 <- *sessionBucket2; set *sessionBucket2 = f b2
812+ b3 <- *sessionBucket3; set *sessionBucket3 = f b3
813+ b4 <- *sessionBucket4; set *sessionBucket4 = f b4
814+ b5 <- *sessionBucket5; set *sessionBucket5 = f b5
815+ b6 <- *sessionBucket6; set *sessionBucket6 = f b6
816+ b7 <- *sessionBucket7; set *sessionBucket7 = f b7
817+ b8 <- *sessionBucket8; set *sessionBucket8 = f b8
818+ b9 <- *sessionBucket9; set *sessionBucket9 = f b9
819+ b10 <- *sessionBucket10; set *sessionBucket10 = f b10
820+ b11 <- *sessionBucket11; set *sessionBucket11 = f b11
821+ b12 <- *sessionBucket12; set *sessionBucket12 = f b12
822+ b13 <- *sessionBucket13; set *sessionBucket13 = f b13
823+ b14 <- *sessionBucket14; set *sessionBucket14 = f b14
824+ b15 <- *sessionBucket15; set *sessionBucket15 = f b15
825+
826+-- ============================================
827+-- Recipient Count Helpers
828+-- ============================================
829+-- Per-recipient message counts avoid O(n) scan of all messages on send/forward.
830+
831+getRecipientCount : PubkeyHex -> [RecipientCount] -> Int
832+getRecipientCount = \pubkey counts ->
833+ fold (\acc c -> if c.pubkey == pubkey then c.count else acc) 0 counts
834+
835+adjustRecipientCount : [RecipientCount] -> PubkeyHex -> Int -> [RecipientCount]
836+adjustRecipientCount = \counts pubkey delta ->
837+ if delta == 0 then counts
838+ else if any counts (\c -> c.pubkey == pubkey)
839+ then do c <- counts; yield (if c.pubkey == pubkey then {pubkey: c.pubkey, count: max 0 (c.count + delta)} else c)
840+ else if delta > 0 then union counts [{pubkey: pubkey, count: delta}]
841+ else counts
842+
843+recomputeRecipientCounts : [Message] -> [RecipientCount]
844+recomputeRecipientCounts = \messages ->
845+ fold (\acc m ->
846+ if any acc (\c -> c.pubkey == m.toKey)
847+ then do c <- acc; yield (if c.pubkey == m.toKey then {pubkey: c.pubkey, count: c.count + 1} else c)
848+ else union acc [{pubkey: m.toKey, count: 1}]
849+ ) [] messages
850+
851+-- Check message dedup only (without counting — count comes from recipientMsgCounts)
852+checkMessageDup : [Message] -> Text -> PubkeyHex -> Bool
853+checkMessageDup = \messages messageId toKey ->
854+ messageId != "" && any messages (\m -> m.messageId == messageId && m.toKey == toKey)
855+
856+-- ============================================
857+-- Federation In-Flight Tracking
858+-- ============================================
859+-- Tracks outstanding federation fetches per server. Servers with hung requests
860+-- (older than federationFetchTimeoutMs) are skipped to prevent slot exhaustion.
861+
862+markFederationInFlight = \server -> do
863+ t <- now
864+ atomic do
865+ inFlight <- *federationInFlight
866+ let existing = headOf (do r <- inFlight; where r.server == server; yield r)
867+ case existing of
868+ Nothing {} -> set *federationInFlight =
869+ let updated = union inFlight [{server: server, count: 1, firstStartedAt: t}] in
870+ if count updated > maxFederationInFlight then evictOneOldest updated (\r -> r.firstStartedAt) else updated
871+ Just {} -> set *federationInFlight = do r <- inFlight; yield (if r.server == server then {server: r.server, count: r.count + 1, firstStartedAt: r.firstStartedAt} else r)
872+
873+clearFederationInFlight = \server -> atomic do
874+ inFlight <- *federationInFlight
875+ set *federationInFlight = do r <- inFlight; yield (if r.server == server then {server: r.server, count: max 0 (r.count - 1), firstStartedAt: if r.count <= 1 then 0 else r.firstStartedAt} else r)
876+
877+hasHungFederationRequest : Text -> Timestamp -> [FederationInFlightEntry] -> Bool
878+hasHungFederationRequest = \server t inFlight ->
879+ any inFlight (\r -> r.server == server && r.count > 0 && r.firstStartedAt > 0 && t - r.firstStartedAt >= federationFetchTimeoutMs)
880+
881+-- Wrapper around fetch for federation calls: tracks in-flight state, logs slow responses, records failures
882+fetchFederation = \server request -> do
883+ markFederationInFlight server
884+ startTime <- now
885+ result <- fetch (federationUrl server) request
886+ clearFederationInFlight server
887+ endTime <- now
888+ when (endTime - startTime > federationFetchTimeoutMs)
889+ (println ("federation: slow response from " ++ server ++ " (" ++ show (endTime - startTime) ++ "ms)"))
890+ yield result
891+
687892clampTtl : Int<S> -> Int<S>
688893clampTtl = clamp 1 maxPresenceTtl
689894
@@ -714,22 +919,24 @@ dedupMessageIds = \messageIds ->
714919
715920-- Constant-time string comparison to prevent timing side-channels.
716921-- Always iterates maxTokenLen characters regardless of input length;
717−-- rejects empty strings to avoid vacuous equality.
718−-- Loop terminates strictly on counter (n starts at max of both lengths and maxTokenLen,
719−-- so both strings are exhausted by the time n reaches 0). Avoiding s1/s2 emptiness
720−-- checks in the branch prevents a potential timing side-channel from branch prediction.
721−safeCompareLoop : Text -> Text -> Int -> Int -> Int
722−safeCompareLoop = \s1 s2 n result ->
922+-- rejects empty strings. Tracks both character mismatches and length
923+-- mismatch (one string ended before the other) within the loop itself,
924+-- avoiding variable-time count(chars) calls that would leak string length.
925+safeCompareLoop : Text -> Text -> Int -> Int -> Int -> Int
926+safeCompareLoop = \s1 s2 n diff lenDiff ->
723927 if n <= 0
724− then result
725− else safeCompareLoop (drop 1 s1) (drop 1 s2) (n - 1) (result + (if take 1 s1 == take 1 s2 then 0 else 1))
928+ then diff + lenDiff
929+ else
930+ let c1 = take 1 s1 in
931+ let c2 = take 1 s2 in
932+ let newDiff = diff + (if c1 == c2 then 0 else 1) in
933+ let endMismatch = if c1 == "" then (if c2 == "" then 0 else 1) else (if c2 == "" then 1 else 0) in
934+ safeCompareLoop (drop 1 s1) (drop 1 s2) (n - 1) newDiff (lenDiff + endMismatch)
726935
727936safeCompare : Text -> Text -> Bool
728937safeCompare = \a b ->
729− let lenA = count (chars a) in
730− let lenB = count (chars b) in
731− let lenDiff = if lenA == lenB && lenA > 0 then 0 else 1 in
732− safeCompareLoop a b maxTokenLen 0 + lenDiff == 0
938+ if take 1 a == "" || take 1 b == "" then False {}
939+ else safeCompareLoop a b maxTokenLen 0 0 == 0
733940
734941-- Generic sliding-window rate limiter. Returns whether the request is rate-limited
735942-- and the updated entries list. Callers provide match/create/update functions for their entry type.
@@ -798,7 +1005,8 @@ removeOneForKey = \items pubkey ->
7981005 let others = do x <- items; where x.pubkey != pubkey; yield x in
7991006 union others (evictOneOldest forKey (\x -> x.startedAt))
8001007
801−-- Extract unique remote (non-local) server names from a list of items
1008+-- Extract unique remote (non-local) server names from a list of items.
1009+-- O(n*k) where k = unique servers; acceptable since k is typically small (<20).
8021010collectPeerServers : [a] -> (a -> Text) -> [Text]
8031011collectPeerServers = \items getServer ->
8041012 fold (\acc item ->
@@ -852,14 +1060,6 @@ capGossipPresence = \args ->
8521060 (do e <- reclaimedEntries; yield (pubkeyServer e)),
8531061 withinCapacity: False {}}
8541062
855−-- Single-pass dedup + per-recipient count over messages (avoids two O(n) scans in send/forward)
856−checkMessageStore : [Message] -> Text -> PubkeyHex -> Timestamp -> {isDup: Bool, recipientCount: Int}
857−checkMessageStore = \messages messageId toKey t ->
858− fold (\acc m ->
859− {isDup: acc.isDup || (messageId != "" && m.messageId == messageId && m.toKey == toKey),
860− recipientCount: if m.toKey == toKey && m.receivedAt >= t - messageMaxAge then acc.recipientCount + 1 else acc.recipientCount}
861− ) {isDup: False {}, recipientCount: 0} messages
862−
8631063hasMessage : MessageIdHex -> PubkeyHex -> [{messageId: MessageIdHex, toKey: PubkeyHex | a}] -> Bool
8641064hasMessage = \mid toKey items -> any items (\x -> x.messageId == mid && x.toKey == toKey)
8651065
@@ -913,15 +1113,17 @@ pruneExpired = \t -> do
9131113 atomic do
9141114 challenges <- *challenges
9151115 set *challenges = do x <- challenges; where x.expiresAt > t; yield x
916− sessions <- *sessions
917− set *sessions = do x <- sessions; where x.expiresAt > t; yield x
1116+ transformAllSessionBuckets (\bucket -> do x <- bucket; where x.expiresAt > t; yield x)
9181117 -- Prune messages (largest collection, isolated for minimal lock time)
919− atomic do
1118+ prunedMessages <- atomic do
9201119 let messageCutoff = t - unackedMessageMaxAge
9211120 messages <- *messages
9221121 let prunedMessages = do m <- messages; where m.receivedAt > messageCutoff; yield m
9231122 set *messages = prunedMessages
9241123 set *totalMessageCount = count prunedMessages
1124+ yield prunedMessages
1125+ -- Recompute recipient counts from pruned messages (outside main message lock)
1126+ atomic do set *recipientMsgCounts = recomputeRecipientCounts prunedMessages
9251127 -- Prune federation state
9261128 atomic do
9271129 let retainCutoff = t - messageMaxAge
@@ -955,6 +1157,9 @@ pruneExpired = \t -> do
9551157 set *recvRates = do r <- recvRates; where t - r.windowStart < recvRateWindowMs; yield r
9561158 recentGossip <- *recentGossip
9571159 set *recentGossip = do r <- recentGossip; where t - r.seenAt < gossipDedupWindow; yield r
1160+ -- Clean up stale in-flight entries (hung fetches that never completed)
1161+ inFlight <- *federationInFlight
1162+ set *federationInFlight = do r <- inFlight; where r.count > 0; yield r
9581163 yield presenceResult
9591164
9601165
@@ -974,7 +1179,7 @@ forEachBatch = \items batchSize process ->
9741179sendBatchedGossipToServer : ServerName -> [GossipEvent] -> Int -> IO {network} {}
9751180sendBatchedGossipToServer = \server allEvents hops ->
9761181 forEachBatch allEvents maxGossipEvents (\batch -> do
977− result <- fetch (federationUrl server) (RecvGossip {
1182+ result <- fetchFederation server (RecvGossip {
9781183 authorization: federationAuthForServer server,
9791184 events: batch, hopsRemaining: hops, fromServer: serverHostN})
9801185 case result of
@@ -1080,14 +1285,18 @@ withFederationAuth = \authorization t respond fromServer precondition action ->
10801285 clearFailure fromServer
10811286 action
10821287
1083−withSessionAuth = \authorization t onUnauth onAuth -> do
1084− sessions <- *sessions
1288+withSessionAuth = \authorization t clientIp onUnauth onAuth -> do
10851289 case extractToken authorization of
10861290 Nothing {} -> onUnauth
1087− Just {value: token} ->
1291+ Just {value: token} -> do
1292+ let bucket = tokenBucket token
1293+ sessions <- readSessionBucket bucket
10881294 case authenticateSession sessions token t of
10891295 Unauthorized {} -> onUnauth
1090− Authorized {pubkey} -> onAuth pubkey
1296+ Authorized {pubkey, ip} -> do
1297+ when (clientIp != "" && ip != "" && clientIp != ip)
1298+ (println ("session: IP mismatch for " ++ pubkey ++ " (session: " ++ ip ++ ", request: " ++ clientIp ++ ")"))
1299+ onAuth pubkey
10911300
10921301removeForward = \mid srv -> atomic do
10931302 forwards <- *forwards
@@ -1107,6 +1316,11 @@ shouldRetryServer = \server t failures ->
11071316 Nothing {} -> True {}
11081317 Just {value: f} -> t - f.firstFailedAt < maxServerDowntime && t - f.failedAt >= retryBackoff f.failures
11091318
1319+-- Extended retry check: also skip servers with hung in-flight requests
1320+shouldRetryServerEx : Text -> Timestamp -> [ForwardFailure] -> [FederationInFlightEntry] -> Bool
1321+shouldRetryServerEx = \server t failures inFlight ->
1322+ shouldRetryServer server t failures && not (hasHungFederationRequest server t inFlight)
1323+
11101324recordFailure = \server -> do
11111325 t <- now
11121326 atomic do
@@ -1132,7 +1346,7 @@ retryForwardsForServer = \entries remaining ->
11321346 Nothing {} -> yield {}
11331347 Just {value: f} -> do
11341348 let rest = do x <- entries; where x.messageId != f.messageId; yield x
1135− result <- fetch (federationUrl f.toServer) (ForwardMessage {
1349+ result <- fetchFederation f.toServer (ForwardMessage {
11361350 toKey: f.toKey, encryptedBlob: f.encryptedBlob,
11371351 timestamp: f.timestamp, originServer: serverHostN, messageId: f.messageId,
11381352 fromServer: serverHostN, authorization: federationAuthForServer f.toServer})
@@ -1152,7 +1366,8 @@ retryPendingForwards = \pending -> do
11521366 let servers = collectPeerServers pending (\x -> x.toServer)
11531367 t <- now
11541368 failures <- *forwardFailures
1155− let eligibleServers = do s <- servers; where shouldRetryServer s t failures; yield s
1369+ inFlight <- *federationInFlight
1370+ let eligibleServers = do s <- servers; where shouldRetryServerEx s t failures inFlight; yield s
11561371 sendToEach eligibleServers (\server ->
11571372 retryForwardsForServer (forwardsForServer pending server) maxForwardRetriesPerServer)
11581373 yield {}
@@ -1172,12 +1387,13 @@ forwardQueuedMessages = \targets t -> do
11721387 yield capped
11731388 let servers = collectPeerServers newForwards (\x -> x.toServer)
11741389 failures <- *forwardFailures
1175− let eligibleServers = do s <- servers; where shouldRetryServer s t failures; yield s
1390+ inFlight <- *federationInFlight
1391+ let eligibleServers = do s <- servers; where shouldRetryServerEx s t failures inFlight; yield s
11761392 sendToEach eligibleServers (\server ->
11771393 retryForwardsForServer (forwardsForServer newForwards server) maxForwardRetriesPerServer)
11781394
11791395fetchFederationAck = \origin messageIds ->
1180− fetch (federationUrl origin) (FederationAck {
1396+ fetchFederation origin (FederationAck {
11811397 authorization: federationAuthForServer origin,
11821398 messageIds: messageIds,
11831399 originServer: origin, fromServer: serverHostN})
@@ -1203,7 +1419,8 @@ sendAcksAndClean = \pendingAcks -> do
12031419 let servers = collectPeerServers pendingAcks (\a -> a.originServer)
12041420 t <- now
12051421 failures <- *forwardFailures
1206− let eligibleServers = do s <- servers; where shouldRetryServer s t failures; yield s
1422+ inFlight <- *federationInFlight
1423+ let eligibleServers = do s <- servers; where shouldRetryServerEx s t failures inFlight; yield s
12071424 sendToEach eligibleServers (\server ->
12081425 sendBatchedAcksToServer (acksForServer pendingAcks server) server)
12091426
@@ -1227,7 +1444,7 @@ retryEvictedAcks = \evictedEntries pendingAcks ->
12271444 if excess <= 0
12281445 then do set *ackedIds = combined; yield []
12291446 else do
1230− let eviction = evictSmallest combined excess (\a -> a.ackedAt)
1447+ let eviction = evictSmallest combined (max 0 excess) (\a -> a.ackedAt)
12311448 set *ackedIds = eviction.kept
12321449 yield eviction.evicted
12331450 when (count overflow > 0) (do
@@ -1238,14 +1455,14 @@ retryEvictedAcks = \evictedEntries pendingAcks ->
12381455 if overflowExcess <= 0
12391456 then do set *ackOverflow = updated; yield 0
12401457 else do
1241− let overflowEviction = evictSmallest updated overflowExcess (\a -> a.ackedAt)
1458+ let overflowEviction = evictSmallest updated (max 0 overflowExcess) (\a -> a.ackedAt)
12421459 set *ackOverflow = overflowEviction.kept
12431460 deliveredIds <- *deliveredIds
12441461 let newDelivered = do a <- overflowEviction.evicted; yield {messageId: a.messageId, toKey: a.toKey, deliveredAt: a.ackedAt}
12451462 let combinedDelivered = union deliveredIds newDelivered
12461463 let deliveredExcess = count combinedDelivered - maxDeliveredIds
12471464 set *deliveredIds = if deliveredExcess <= 0 then combinedDelivered
1248− else (evictSmallest combinedDelivered deliveredExcess (\d -> d.deliveredAt)).kept
1465+ else (evictSmallest combinedDelivered (max 0 deliveredExcess) (\d -> d.deliveredAt)).kept
12491466 yield overflowExcess
12501467 when (dropped > 0)
12511468 (println ("ackMessages: overflow at capacity (" ++ show maxAckOverflow ++ "), " ++ show dropped ++ " ack entries moved to delivered set for " ++ server)))
@@ -1256,7 +1473,8 @@ retryOverflowAcks = \overflow -> do
12561473 let servers = collectPeerServers overflow (\a -> a.originServer)
12571474 t <- now
12581475 failures <- *forwardFailures
1259− let eligibleServers = do s <- servers; where shouldRetryServer s t failures; yield s
1476+ inFlight <- *federationInFlight
1477+ let eligibleServers = do s <- servers; where shouldRetryServerEx s t failures inFlight; yield s
12601478 sendToEach eligibleServers (\server ->
12611479 forEachBatch (acksForServer overflow server) maxMessageIds (\batch -> do
12621480 result <- fetchFederationAck server (do a <- batch; yield a.messageId)
@@ -1285,11 +1503,11 @@ pruneAndRetry = \t -> do
12851503 lastRefresh <- *lastPresenceRefresh
12861504 let shouldRefresh = federationEnabled && t - lastRefresh >= presenceRefreshInterval
12871505 presence <- *presence
1288− sessions <- *sessions
1506+ allSessions <- readAllSessions
12891507 -- Extend presence for keys with active sessions nearing expiry (sliding window)
12901508 let localNearExpiry = if shouldRefresh then (do p <- presence; where isLocalServer p.server && p.expiresAt > t && p.expiresAt - t < presenceRefreshThreshold; yield p) else []
1291− let toExtend = do p <- localNearExpiry; where any sessions (\s -> s.pubkey == p.pubkey && s.expiresAt > t); yield p
1292− let notExtended = do p <- localNearExpiry; where none sessions (\s -> s.pubkey == p.pubkey && s.expiresAt > t); yield p
1509+ let toExtend = do p <- localNearExpiry; where any allSessions (\s -> s.pubkey == p.pubkey && s.expiresAt > t); yield p
1510+ let notExtended = do p <- localNearExpiry; where none allSessions (\s -> s.pubkey == p.pubkey && s.expiresAt > t); yield p
12931511 when (count toExtend > 0) (
12941512 set *presence = do pp <- presence; yield (if any toExtend (\e -> e.pubkey == pp.pubkey && isLocalServer pp.server) then {pubkey: pp.pubkey, server: pp.server, expiresAt: t + sessionExpiry, timestamp: t} else pp))
12951513 let refreshPeers = if shouldRefresh then collectPeerServers presence (\p -> p.server) else []
@@ -1368,18 +1586,21 @@ readMessageSignal = \pubkey ->
13681586ackPolledMessages = \pubkey cursor t -> do
13691587 messages <- *messages
13701588 let toAck = do m <- messages; where m.toKey == pubkey && m.receivedAt <= cursor; yield m
1371− if count toAck > 0
1589+ let ackCount = count toAck
1590+ if ackCount > 0
13721591 then do
13731592 set *messages = diff messages toAck
13741593 totalCount <- *totalMessageCount
1375− set *totalMessageCount = totalCount - count toAck
1594+ set *totalMessageCount = totalCount - ackCount
1595+ recipientCounts <- *recipientMsgCounts
1596+ set *recipientMsgCounts = adjustRecipientCount recipientCounts pubkey (0 - ackCount)
13761597 let ackedMsgIds = do m <- toAck; yield m.messageId
13771598 forwards <- *forwards
13781599 set *forwards = do f <- forwards; where none ackedMsgIds (\mid -> f.messageId == mid && f.toKey == pubkey); yield f
13791600 let ackedFederated = federatedAckEntries toAck t
13801601 ackedIds <- *ackedIds
13811602 let combinedAckedIds = union ackedIds ackedFederated
1382− let excessCount = count combinedAckedIds - maxAckedIds
1603+ let excessCount = max 0 (count combinedAckedIds - maxAckedIds)
13831604 let eviction = evictSmallest combinedAckedIds excessCount (\a -> a.ackedAt)
13841605 set *ackedIds = eviction.kept
13851606 yield {pendingAcks: ackedFederated, evictedEntries: eviction.evicted}
@@ -1403,15 +1624,18 @@ pollNewEvents = \pubkey clampedCursor t -> do
14031624-- 3. Check for existing messages; if none, block on signal/heartbeat (retry-based wait)
14041625-- 4. Re-authenticate after waking (session may have expired during wait)
14051626-- 5. Return new events and release the poll slot
1406−pollOnce = \token cursor -> do
1627+pollOnce = \token cursor clientIp -> do
14071628 t <- now
1629+ let bucket = tokenBucket token
14081630 setup <- atomic do
1409− sessions <- *sessions
1631+ sessions <- readSessionBucket bucket
14101632 case authenticateSession sessions token t of
14111633 Unauthorized {} -> yield (Nothing {})
1412− Authorized {pubkey} -> do
1634+ Authorized {pubkey, ip} -> do
1635+ when (clientIp != "" && ip != "" && clientIp != ip)
1636+ (println ("poll: session IP mismatch for " ++ pubkey ++ " (session: " ++ ip ++ ", request: " ++ clientIp ++ ")"))
14131637 -- Refresh session expiry (sliding window keeps active sessions alive)
1414− set *sessions = do s <- sessions; yield (if s.token == token then {token: s.token, pubkey: s.pubkey, expiresAt: t + sessionExpiry} else s)
1638+ writeSessionBucket bucket (do s <- sessions; yield (if safeCompare s.token token then {token: s.token, pubkey: s.pubkey, expiresAt: t + sessionExpiry, ip: s.ip} else s))
14151639 seqCounter <- *seqCounter
14161640 let clampedCursor = clamp 0 (max seqCounter t) cursor
14171641 ackResult <- if clampedCursor > 0
@@ -1460,14 +1684,15 @@ pollOnce = \token cursor -> do
14601684 -- Phase 3: re-authenticate and read messages after waking
14611685 t2 <- now
14621686 result <- atomic do
1463− sessions2 <- *sessions
1687+ sessions2 <- readSessionBucket bucket
14641688 case authenticateSession sessions2 token t2 of
14651689 Unauthorized {} -> yield {authorized: False {}, events: [], cursor: s.clampedCursor}
14661690 Authorized {} -> do
14671691 r <- pollNewEvents s.pubkey s.clampedCursor t2
14681692 yield {authorized: True {}, events: r.events, cursor: r.cursor}
14691693 yield {authorized: result.authorized, events: result.events, cursor: result.cursor}
1470− -- Release poll slot
1694+ -- Release poll slot. Safety net: if this is never reached (e.g., panic between
1695+ -- acquire and release), pruneExpired cleans stale entries after pollTimeoutMs * 3.
14711696 atomic do
14721697 activePolls <- *activePolls
14731698 set *activePolls = removeOneForKey activePolls s.pubkey
@@ -1667,34 +1892,58 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
16671892 if validation.valid
16681893 then do
16691894 token <- genToken
1895+ let newBucket = tokenBucket token
16701896 authResult <- atomic do
16711897 isDraining <- *draining
16721898 if isDraining then yield (Nothing {})
16731899 else do
1674− sessions <- *sessions
1900+ -- Read all session buckets for per-key enforcement (revokeOthers, capPerKey)
1901+ allSessions <- readAllSessions
16751902 presence <- *presence
1676− let afterRevoke = if revokeOthers then do s <- sessions; where s.pubkey != pubkey; yield s else sessions
1903+ let revokedCount = if revokeOthers then count (do s <- allSessions; where s.pubkey == pubkey; yield s) else 0
1904+ let afterRevoke = if revokeOthers then do s <- allSessions; where s.pubkey != pubkey; yield s else allSessions
16771905 let capped = capPerKeyAndGlobal afterRevoke (\s -> s.pubkey == pubkey) maxSessionsPerKey (\s -> s.expiresAt) t maxSessions
1678− set *sessions = union capped [{token: token, pubkey: pubkey, expiresAt: expiry}]
1906+ let newSession = {token: token, pubkey: pubkey, expiresAt: expiry, ip: clientIp}
1907+ -- Redistribute capped sessions + new session across 16 buckets
1908+ let sessionsWithNew = union capped [newSession]
1909+ let bFor = \b -> do s <- sessionsWithNew; where tokenBucket s.token == b; yield s
1910+ writeSessionBucket 0 (bFor 0)
1911+ writeSessionBucket 1 (bFor 1)
1912+ writeSessionBucket 2 (bFor 2)
1913+ writeSessionBucket 3 (bFor 3)
1914+ writeSessionBucket 4 (bFor 4)
1915+ writeSessionBucket 5 (bFor 5)
1916+ writeSessionBucket 6 (bFor 6)
1917+ writeSessionBucket 7 (bFor 7)
1918+ writeSessionBucket 8 (bFor 8)
1919+ writeSessionBucket 9 (bFor 9)
1920+ writeSessionBucket 10 (bFor 10)
1921+ writeSessionBucket 11 (bFor 11)
1922+ writeSessionBucket 12 (bFor 12)
1923+ writeSessionBucket 13 (bFor 13)
1924+ writeSessionBucket 14 (bFor 14)
1925+ writeSessionBucket 15 (bFor 15)
16791926 let presenceWithout = presenceWithoutLocalKey presence pubkey
16801927 let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)
16811928 let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]
16821929 set *presence = updatedPresence
1683− yield (Just {value: if federationEnabled
1930+ yield (Just {value: {peerServers: if federationEnabled
16841931 then collectPeerServers updatedPresence (\p -> p.server)
1685− else []})
1932+ else [], revokedCount: revokedCount}})
16861933 case authResult of
16871934 Nothing {} -> yield (respond {token: "", expiresAt: 0})
1688− Just {value: peerServers} -> do
1935+ Just {value: result} -> do
1936+ when (revokeOthers && result.revokedCount > 0)
1937+ (println ("verify: revokeOthers revoked " ++ show result.revokedCount ++ " sessions for " ++ pubkey))
16891938 fork (when federationEnabled
1690− (sendBatchedGossip peerServers [{eventType: GossipOnline {}, pubkey: pubkey, server: serverHostN, timestamp: t, ttl: onlineGossipTtl}] gossipOriginHops))
1939+ (sendBatchedGossip result.peerServers [{eventType: GossipOnline {}, pubkey: pubkey, server: serverHostN, timestamp: t, ttl: onlineGossipTtl}] gossipOriginHops))
16911940 yield (respond {token: token, expiresAt: expiry})
16921941 else yield (respond {token: "", expiresAt: 0})
16931942
16941943-- POST /poll: Long-poll for new messages. Cursor-based: messages at or before cursor
16951944-- are acked (removed from mailbox). Blocks until new messages arrive or heartbeat fires.
16961945-- Acks for federated messages are sent to origin servers in the background.
1697−handlePoll = \cursor authorization respond -> do
1946+handlePoll = \cursor authorization clientIp respond -> do
16981947 t <- now
16991948 draining <- *draining
17001949 if draining then yield (respond {events: [], cursor: 0, authorized: True {}, rateLimited: False {}, draining: True {}})
@@ -1704,7 +1953,7 @@ handlePoll = \cursor authorization respond -> do
17041953 else do
17051954 result <- case extractToken authorization of
17061955 Nothing {} -> yield {authorized: False {}, events: [], cursor: 0, pendingAcks: [], evictedEntries: []}
1707− Just {value: token} -> pollOnce token cursor
1956+ Just {value: token} -> pollOnce token cursor clientIp
17081957 fork (when federationEnabled (do
17091958 sendAcksAndClean result.pendingAcks
17101959 retryEvictedAcks result.evictedEntries result.pendingAcks))
@@ -1713,7 +1962,7 @@ handlePoll = \cursor authorization respond -> do
17131962-- POST /messages: Send an encrypted message to a recipient.
17141963-- Deduplicates by messageId across messages, ackedIds, and deliveredIds.
17151964-- Rate-limited per sender key. If recipient is known on a remote server, forwards via federation.
1716−handleSendMessage = \to encryptedBlob timestamp messageId authorization respond -> do
1965+handleSendMessage = \to encryptedBlob timestamp messageId authorization clientIp respond -> do
17171966 t <- now
17181967 globalRateLimited <- checkGlobalRate t
17191968 if globalRateLimited then yield (respond {status: MsRateLimited {}, messageId: ""})
@@ -1723,7 +1972,7 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization respond
17231972 result <- atomic do
17241973 draining <- *draining
17251974 if draining then yield {send: SendFailed {status: MsDraining {}}, senderRateLimited: False {}, globalCapReached: False {}}
1726− else withSessionAuth authorization t
1975+ else withSessionAuth authorization t clientIp
17271976 (yield {send: SendFailed {status: MsUnauthorized {}}, senderRateLimited: False {}, globalCapReached: False {}})
17281977 (\senderKey -> do
17291978 -- Reject self-sends (sender == recipient wastes a mailbox slot)
@@ -1732,8 +1981,9 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization respond
17321981 messages <- *messages
17331982 ackedIds <- *ackedIds
17341983 deliveredIds <- *deliveredIds
1735− let msgCheck = checkMessageStore messages messageId to t
1736− if messageId != "" && (msgCheck.isDup || hasMessage messageId to ackedIds || hasMessage messageId to deliveredIds)
1984+ -- Dedup check (only when messageId is non-empty; skips O(n) scan otherwise)
1985+ let isDup = checkMessageDup messages messageId to || (messageId != "" && (hasMessage messageId to ackedIds || hasMessage messageId to deliveredIds))
1986+ if isDup
17371987 then yield {send: SendDuplicate {}, senderRateLimited: False {}, globalCapReached: False {}}
17381988 else do
17391989 sendRates <- *sendRates
@@ -1757,18 +2007,23 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization respond
17572007 totalMsgCount <- *totalMessageCount
17582008 if totalMsgCount >= maxMessages
17592009 then yield {send: SendFailed {status: MsRejected {}}, senderRateLimited: False {}, globalCapReached: True {}}
1760− else if msgCheck.recipientCount >= maxMessagesPerRecipient
1761− then yield {send: SendFailed {status: MsRejected {}}, senderRateLimited: False {}, globalCapReached: False {}}
17622010 else do
1763− presence <- *presence
1764− seqCounter <- *seqCounter
1765− let seqTs = nextSeqTs t seqCounter
1766− set *seqCounter = seqTs
1767− set *messages = union messages
1768− [{messageId: effectiveMsgId, toKey: to,
1769− encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: "", receivedAt: seqTs}]
1770− set *totalMessageCount = totalMsgCount + 1
1771− signalMessageForKey to seqTs
2011+ -- Recipient count from maintained counts (O(unique_recipients) instead of O(all_messages))
2012+ recipientCounts <- *recipientMsgCounts
2013+ let recipientCount = getRecipientCount to recipientCounts
2014+ if recipientCount >= maxMessagesPerRecipient
2015+ then yield {send: SendFailed {status: MsRejected {}}, senderRateLimited: False {}, globalCapReached: False {}}
2016+ else do
2017+ presence <- *presence
2018+ seqCounter <- *seqCounter
2019+ let seqTs = nextSeqTs t seqCounter
2020+ set *seqCounter = seqTs
2021+ set *messages = union messages
2022+ [{messageId: effectiveMsgId, toKey: to,
2023+ encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: "", receivedAt: seqTs}]
2024+ set *totalMessageCount = totalMsgCount + 1
2025+ set *recipientMsgCounts = adjustRecipientCount recipientCounts to 1
2026+ signalMessageForKey to seqTs
17722027 let toPresence = presenceForKey presence to t
17732028 let localOnline = hasLocalPresence toPresence
17742029 let remoteServers = collectPeerServers toPresence (\p -> p.server)
@@ -1795,7 +2050,7 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization respond
17952050 forwards <- *forwards
17962051 yield (any forwards (\f -> f.messageId == effectiveMsgId && f.toServer == server))
17972052 when stillPending (do
1798− fwdResult <- fetch (federationUrl server) (ForwardMessage {
2053+ fwdResult <- fetchFederation server (ForwardMessage {
17992054 toKey: to, encryptedBlob: encryptedBlob,
18002055 timestamp: clampedTs, originServer: serverHostN, messageId: effectiveMsgId,
18012056 fromServer: serverHostN, authorization: federationAuthForServer server})
@@ -1814,7 +2069,7 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization respond
18142069
18152070-- GET /lookup/:pubkey: Look up the best known server for a pubkey.
18162071-- Returns local server if key is online locally, otherwise the remote with latest expiry.
1817−handleLookupPresence = \pubkey authorization respond -> do
2072+handleLookupPresence = \pubkey authorization clientIp respond -> do
18182073 t <- now
18192074 globalRateLimited <- checkGlobalRate t
18202075 if globalRateLimited then yield (respond {server: "", authorized: False {}})
@@ -1822,7 +2077,7 @@ handleLookupPresence = \pubkey authorization respond -> do
18222077 result <- atomic do
18232078 draining <- *draining
18242079 if draining then yield {server: "", authorized: False {}}
1825− else withSessionAuth authorization t
2080+ else withSessionAuth authorization t clientIp
18262081 (yield {server: "", authorized: False {}})
18272082 (\senderKey -> do
18282083 lookupRates <- *lookupRates
@@ -1868,7 +2123,10 @@ handleRecvGossip = \events hopsRemaining fromServer authorization respond -> do
18682123 freshEvents <- atomic do
18692124 recentGossip <- *recentGossip
18702125 let active = do r <- recentGossip; where t - r.seenAt < gossipDedupWindow; yield r
1871− let fresh = do e <- effectiveEvents; where none active (\r -> r.pubkey == e.pubkey && r.server == e.server && r.timestamp == e.timestamp); yield e
2126+ -- Pre-filter active entries to only those matching pubkeys in this batch (reduces O(n*m) inner scan)
2127+ let batchPubkeys = fold (\acc e -> if any acc (\a -> a == e.pubkey) then acc else union acc [e.pubkey]) [] effectiveEvents
2128+ let relevantActive = do r <- active; where any batchPubkeys (\pk -> pk == r.pubkey); yield r
2129+ let fresh = do e <- effectiveEvents; where none relevantActive (\r -> r.pubkey == e.pubkey && r.server == e.server && r.timestamp == e.timestamp); yield e
18722130 let newRecent = do e <- fresh; yield {pubkey: e.pubkey, server: e.server, timestamp: e.timestamp, seenAt: t}
18732131 let updated = union active newRecent
18742132 set *recentGossip = if count updated > maxRecentGossip then do r <- updated; where t - r.seenAt < gossipDedupWindow; yield r else updated
@@ -1892,35 +2150,39 @@ handleForwardMessage = \toKey encryptedBlob timestamp originServer messageId fro
18922150 messages <- *messages
18932151 ackedIds <- *ackedIds
18942152 deliveredIds <- *deliveredIds
1895− let msgCheck = checkMessageStore messages messageId toKey t
1896− if msgCheck.isDup || hasMessage messageId toKey ackedIds || hasMessage messageId toKey deliveredIds
2153+ let isDup = checkMessageDup messages messageId toKey || hasMessage messageId toKey ackedIds || hasMessage messageId toKey deliveredIds
2154+ if isDup
18972155 then yield {forward: ForwardDuplicate {}, globalCapReached: False {}, recvRateLimited: False {}}
18982156 else do
18992157 totalMsgCount <- *totalMessageCount
19002158 if totalMsgCount >= maxMessages
19012159 then yield {forward: ForwardRejected {}, globalCapReached: True {}, recvRateLimited: False {}}
1902− else if msgCheck.recipientCount >= maxMessagesPerRecipient
1903− then yield {forward: ForwardRejected {}, globalCapReached: False {}, recvRateLimited: False {}}
19042160 else do
1905− recvRates <- *recvRates
1906− let recvResult = checkRateWindow recvRates (\r -> r.pubkey == toKey)
1907− (\ts -> {pubkey: toKey, windowStart: ts, count: 1})
1908− (\r -> {pubkey: r.pubkey, windowStart: r.windowStart, count: r.count + 1})
1909− maxRecvRatePerKey recvRateWindowMs maxRecvRateEntries t
1910− if recvResult.limited
1911− then yield {forward: ForwardRejected {}, globalCapReached: False {}, recvRateLimited: True {}}
2161+ recipientCounts <- *recipientMsgCounts
2162+ let recipientCount = getRecipientCount toKey recipientCounts
2163+ if recipientCount >= maxMessagesPerRecipient
2164+ then yield {forward: ForwardRejected {}, globalCapReached: False {}, recvRateLimited: False {}}
19122165 else do
1913− set *recvRates = recvResult.entries
1914− seqCounter <- *seqCounter
1915− let seqTs = nextSeqTs t seqCounter
1916− set *seqCounter = seqTs
1917− let updatedMessages = union messages
1918− [{messageId: messageId, toKey: toKey,
1919− encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: originServer, receivedAt: seqTs}]
1920− set *messages = updatedMessages
1921− set *totalMessageCount = totalMsgCount + 1
1922− signalMessageForKey toKey seqTs
1923− yield {forward: ForwardStored {}, globalCapReached: False {}, recvRateLimited: False {}}
2166+ recvRates <- *recvRates
2167+ let recvResult = checkRateWindow recvRates (\r -> r.pubkey == toKey)
2168+ (\ts -> {pubkey: toKey, windowStart: ts, count: 1})
2169+ (\r -> {pubkey: r.pubkey, windowStart: r.windowStart, count: r.count + 1})
2170+ maxRecvRatePerKey recvRateWindowMs maxRecvRateEntries t
2171+ if recvResult.limited
2172+ then yield {forward: ForwardRejected {}, globalCapReached: False {}, recvRateLimited: True {}}
2173+ else do
2174+ set *recvRates = recvResult.entries
2175+ seqCounter <- *seqCounter
2176+ let seqTs = nextSeqTs t seqCounter
2177+ set *seqCounter = seqTs
2178+ let updatedMessages = union messages
2179+ [{messageId: messageId, toKey: toKey,
2180+ encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: originServer, receivedAt: seqTs}]
2181+ set *messages = updatedMessages
2182+ set *totalMessageCount = totalMsgCount + 1
2183+ set *recipientMsgCounts = adjustRecipientCount recipientCounts toKey 1
2184+ signalMessageForKey toKey seqTs
2185+ yield {forward: ForwardStored {}, globalCapReached: False {}, recvRateLimited: False {}}
19242186 case result.forward of
19252187 ForwardRejected {} ->
19262188 if result.globalCapReached
@@ -1954,10 +2216,19 @@ handleFederationAck = \messageIds originServer fromServer authorization respond
19542216 yield f
19552217 -- Only remove messages stored long enough for local clients to have polled them (grace period)
19562218 let pastGrace = do f <- removable; where any messages (\m -> m.messageId == f.messageId && m.toKey == f.toKey && t - m.receivedAt >= ackGracePeriod); yield f
1957− let updatedMessages = do m <- messages; where none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey); yield m
1958− set *messages = updatedMessages
2219+ -- Single-pass fold: filter messages and count removals together
2220+ let msgResult = fold (\acc m ->
2221+ if none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey)
2222+ then {kept: union acc.kept [m], removed: acc.removed}
2223+ else {kept: acc.kept, removed: acc.removed + 1}
2224+ ) {kept: [], removed: 0} messages
2225+ set *messages = msgResult.kept
19592226 totalMsgCount <- *totalMessageCount
1960− set *totalMessageCount = totalMsgCount - count pastGrace
2227+ set *totalMessageCount = totalMsgCount - msgResult.removed
2228+ -- Adjust per-recipient counts for removed messages
2229+ recipientCounts <- *recipientMsgCounts
2230+ let adjustedCounts = fold (\acc f -> adjustRecipientCount acc f.toKey (0 - 1)) recipientCounts pastGrace
2231+ set *recipientMsgCounts = adjustedCounts
19612232 yield (respond {ok: True {}, error: ""}))
19622233
19632234-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).
@@ -2010,7 +2281,7 @@ handleHealth = \respond -> do
20102281 sa <- *startedAt
20112282 presence <- *presence
20122283 forwards <- *forwards
2013− sessions <- *sessions
2284+ sessionCount <- countAllSessions
20142285 activePolls <- *activePolls
20152286 messages <- *messages
20162287 lastFedSuccess <- *lastFederationSuccessAt
@@ -2020,7 +2291,7 @@ handleHealth = \respond -> do
20202291 let capacityOk = count messages * 10 <= maxMessages * 9
20212292 && count presence * 10 <= maxPresenceEntries * 9
20222293 && count forwards * 10 <= maxForwardsEntries * 9
2023− && count sessions * 10 <= maxSessions * 9
2294+ && sessionCount * 10 <= maxSessions * 9
20242295 && count activePolls * 10 <= maxActivePolls * 9
20252296 -- Federation: healthy if disabled, in warmup, no peer issues, or recent success
20262297 let fedOk = not federationEnabled || warmup
@@ -2036,7 +2307,7 @@ handleHealth = \respond -> do
20362307 capacityHealthy: result.capacityOk,
20372308 federationHealthy: result.federationOk})
20382309
2039−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, ok: False {}}
2310+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, ok: False {}}
20402311
20412312handleStats = \authorization respond -> do
20422313 let authorized = verifyMaintenanceAuth authorization
@@ -2045,7 +2316,7 @@ handleStats = \authorization respond -> do
20452316 t <- now
20462317 -- Read all refs atomically for a consistent snapshot, count outside to minimize lock time
20472318 snapshot <- atomic do
2048− s <- *sessions
2319+ sc <- countAllSessions
20492320 m <- *messages
20502321 p <- *presence
20512322 f <- *forwards
@@ -2070,9 +2341,11 @@ handleStats = \authorization respond -> do
20702341 dr <- *draining
20712342 rg <- *recentGossip
20722343 tmc <- *totalMessageCount
2073− yield {s: s, 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}
2074− let result = {sessions: count snapshot.s, 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}
2075− 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, ok: True {}})
2344+ rc <- *recipientMsgCounts
2345+ fi <- *federationInFlight
2346+ 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}
2347+ 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}
2348+ 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, ok: True {}})
20762349 else yield (respond zeroStats)
20772350
20782351handleStatsPeers = \authorization respond -> do
@@ -2093,12 +2366,12 @@ handleStatsKey = \pubkey authorization respond -> do
20932366 result <- atomic do
20942367 messages <- *messages
20952368 presence <- *presence
2096− sessions <- *sessions
2369+ allSessions <- readAllSessions
20972370 forwards <- *forwards
20982371 activePolls <- *activePolls
20992372 let keyMessages = count (do m <- messages; where m.toKey == pubkey; yield m)
21002373 let keyPresence = do p <- presence; where p.pubkey == pubkey && p.expiresAt > t; yield {server: p.server, expiresAt: p.expiresAt}
2101− let keySessions = count (do s <- sessions; where s.pubkey == pubkey; yield s)
2374+ let keySessions = count (do s <- allSessions; where s.pubkey == pubkey; yield s)
21022375 let keyForwards = count (do f <- forwards; where f.toKey == pubkey; yield f)
21032376 let keyPolls = count (do p <- activePolls; where p.pubkey == pubkey; yield p)
21042377 yield {messages: keyMessages, presence: keyPresence, sessions: keySessions, forwards: keyForwards, activePolls: keyPolls}
@@ -2108,9 +2381,9 @@ handleStatsKey = \pubkey authorization respond -> do
21082381serve = \req -> case req of
21092382 ReqChallenge {pubkey, clientIp, respond} -> handleChallenge pubkey clientIp respond
21102383 VerifyAuth {pubkey, challenge, signature, revokeOthers, clientIp, respond} -> handleVerifyAuth pubkey challenge signature revokeOthers clientIp respond
2111− Poll {cursor, authorization, respond} -> handlePoll cursor authorization respond
2112− SendMessage {to, encryptedBlob, timestamp, messageId, authorization, respond} -> handleSendMessage to encryptedBlob timestamp messageId authorization respond
2113− LookupPresence {pubkey, authorization, respond} -> handleLookupPresence pubkey authorization respond
2384+ Poll {cursor, authorization, clientIp, respond} -> handlePoll cursor authorization clientIp respond
2385+ SendMessage {to, encryptedBlob, timestamp, messageId, authorization, clientIp, respond} -> handleSendMessage to encryptedBlob timestamp messageId authorization clientIp respond
2386+ LookupPresence {pubkey, authorization, clientIp, respond} -> handleLookupPresence pubkey authorization clientIp respond
21142387 RecvGossip {events, hopsRemaining, fromServer, authorization, respond} -> handleRecvGossip events hopsRemaining fromServer authorization respond
21152388 ForwardMessage {toKey, encryptedBlob, timestamp, originServer, messageId, fromServer, authorization, respond} -> handleForwardMessage toKey encryptedBlob timestamp originServer messageId fromServer authorization respond
21162389 FederationAck {messageIds, originServer, fromServer, authorization, respond} -> handleFederationAck messageIds originServer fromServer authorization respond