@@ -615,6 +615,12 @@ gateAuth = \checkDrain auth ip t respond mkErr onAuth -> do
615615 if g then yield (respond (mkErr "rate_limited"))
616616 else withSessionAuth auth t ip (yield (respond (mkErr "unauthorized"))) onAuth
617617
618+-- gateAuth + per-key rate gate. `rateResp` is the response value when the per-key gate trips.
619+gateRated = \checkDrain auth ip rateRef maxRate windowMs t respond mkErr rateResp onAuth ->
620+ gateAuth checkDrain auth ip t respond mkErr (\pk -> do
621+ lim <- checkPubkeyRateGate rateRef pk maxRate windowMs t
622+ if lim then yield (respond rateResp) else onAuth pk)
623+
618624authenticateSession : [Session] -> Text -> Timestamp -> AuthResult
619625authenticateSession = \sessions token t ->
620626 case findFirst sessions (\s -> s.expiresAt > t && s.token == token) of
@@ -865,6 +871,16 @@ recordAcksAndDeliveries = \toAck t -> do
865871 appendCapped senderDeliveries localDel 10000 (\d -> d.deliveredAt)
866872 yield ackedFed
867873
874+-- Common cancel block: drop matching messages + forwards, record cancel notification.
875+-- Returns the federated forward servers that had the message (for caller follow-up).
876+applyCancel = \msgs matches to mid t -> do
877+ set *messages = removeWhere msgs matches
878+ fwds <- *forwards
879+ let fwdServers = collectPeers (filter fwds matches) (\f -> f.toServer)
880+ pruneRef forwards (\f -> not (matches f))
881+ appendCapped cancelNotifications [{toKey: to, messageId: mid, cancelledAt: t}] 10000 (\c -> c.cancelledAt)
882+ yield fwdServers
883+
868884sendBatchedGossip = \peers events hops -> forEach peers id (\server -> do
869885 t <- now
870886 failures <- *forwardFailures
@@ -1253,20 +1269,16 @@ forkForwardsTo = \effMsgId to encryptedBlob clampedTs servers ->
12531269handleLookupPresence = \pubkey authorization clientIp respond -> do
12541270 t <- now
12551271 let mkErr = \e -> {server: "", authorized: False {}, error: if e == "rate_limited" then "rate_limited" else ""}
1256− gateAuth (False {}) authorization clientIp t respond mkErr (\senderKey -> do
1257− lim <- checkPubkeyRateGate lookupRates senderKey maxLookupRatePerKey lookupRateWindowMs t
1258− if lim then yield (respond {server: "", authorized: True {}, error: "rate_limited"})
1259− else do
1272+ gateRated (False {}) authorization clientIp lookupRates maxLookupRatePerKey lookupRateWindowMs t respond mkErr
1273+ {server: "", authorized: True {}, error: "rate_limited"} (\senderKey -> do
12601274 pres <- *presence
12611275 yield (respond {server: maskedServer pres pubkey t, authorized: True {}, error: ""}))
12621276
12631277handleBatchLookupPresence = \pubkeys authorization clientIp respond -> do
12641278 t <- now
12651279 let mkErr = \e -> {results: [], authorized: False {}, error: if e == "rate_limited" then "rate_limited" else ""}
1266− gateAuth (False {}) authorization clientIp t respond mkErr (\senderKey -> do
1267− lim <- checkPubkeyRateGate lookupRates senderKey maxLookupRatePerKey lookupRateWindowMs t
1268− if lim then yield (respond {results: [], authorized: True {}, error: "rate_limited"})
1269− else do
1280+ gateRated (False {}) authorization clientIp lookupRates maxLookupRatePerKey lookupRateWindowMs t respond mkErr
1281+ {results: [], authorized: True {}, error: "rate_limited"} (\senderKey -> do
12701282 pres <- *presence
12711283 let results = do
12721284 k <- takeRelation 50 pubkeys
@@ -1295,10 +1307,8 @@ handleMessageStatus = \to messageId authorization clientIp respond -> do
12951307handleAckMessages = \messageIds authorization clientIp respond -> do
12961308 t <- now
12971309 let mkErr = \e -> {ok: False {}, authorized: False {}, ackedCount: 0, error: e}
1298− gateAuth (True {}) authorization clientIp t respond mkErr (\pubkey -> do
1299− lim <- checkPubkeyRateGate clientWriteRates pubkey maxClientWriteRatePerKey clientWriteRateWindowMs t
1300− if lim then yield (respond {ok: False {}, authorized: True {}, ackedCount: 0, error: "rate_limited"})
1301− else do
1310+ gateRated (True {}) authorization clientIp clientWriteRates maxClientWriteRatePerKey clientWriteRateWindowMs t respond mkErr
1311+ {ok: False {}, authorized: True {}, ackedCount: 0, error: "rate_limited"} (\pubkey -> do
13021312 let dedupedIds = dedupMessageIds messageIds
13031313 let matches = \x -> x.toKey == pubkey && elem dedupedIds x.messageId
13041314 result <- atomic do
@@ -1314,10 +1324,8 @@ handleAckMessages = \messageIds authorization clientIp respond -> do
13141324handleListMessages = \offset authorization clientIp respond -> do
13151325 t <- now
13161326 let mkErr = \e -> {messages: [], authorized: False {}, count: 0, error: if e == "rate_limited" then "rate_limited" else ""}
1317− gateAuth (True {}) authorization clientIp t respond mkErr (\pubkey -> do
1318− lim <- checkPubkeyRateGate lookupRates pubkey maxLookupRatePerKey lookupRateWindowMs t
1319− if lim then yield (respond {messages: [], authorized: True {}, count: 0, error: "rate_limited"})
1320− else do
1327+ gateRated (True {}) authorization clientIp lookupRates maxLookupRatePerKey lookupRateWindowMs t respond mkErr
1328+ {messages: [], authorized: True {}, count: 0, error: "rate_limited"} (\pubkey -> do
13211329 msgs <- *messages
13221330 let myMsgs = sortBy (\m -> m.receivedAt) (filter msgs (\m -> m.toKey == pubkey && m.expiresAt > t))
13231331 let total = count myMsgs
@@ -1329,10 +1337,8 @@ handleListMessages = \offset authorization clientIp respond -> do
13291337handleDeleteMessages = \messageIds authorization clientIp respond -> do
13301338 t <- now
13311339 let mkErr = \e -> {ok: False {}, authorized: False {}, deletedCount: 0, error: if e == "draining" then "" else e}
1332− gateAuth (True {}) authorization clientIp t respond mkErr (\pubkey -> do
1333− lim <- checkPubkeyRateGate clientWriteRates pubkey maxClientWriteRatePerKey clientWriteRateWindowMs t
1334− if lim then yield (respond {ok: False {}, authorized: True {}, deletedCount: 0, error: "rate_limited"})
1335− else do
1340+ gateRated (True {}) authorization clientIp clientWriteRates maxClientWriteRatePerKey clientWriteRateWindowMs t respond mkErr
1341+ {ok: False {}, authorized: True {}, deletedCount: 0, error: "rate_limited"} (\pubkey -> do
13361342 let dedupedIds = dedupMessageIds messageIds
13371343 let matches = \x -> x.toKey == pubkey && elem dedupedIds x.messageId
13381344 result <- atomic do
@@ -1349,10 +1355,8 @@ handleDeleteMessages = \messageIds authorization clientIp respond -> do
13491355handleCancelMessage = \messageId to authorization clientIp respond -> do
13501356 t <- now
13511357 let mkErr = \e -> {ok: False {}, authorized: False {}, cancelled: False {}, error: if e == "draining" then "" else e}
1352− gateAuth (True {}) authorization clientIp t respond mkErr (\senderKey -> do
1353− lim <- checkPubkeyRateGate clientWriteRates senderKey maxClientWriteRatePerKey clientWriteRateWindowMs t
1354− if lim then yield (respond {ok: False {}, authorized: True {}, cancelled: False {}, error: "rate_limited"})
1355− else do
1358+ gateRated (True {}) authorization clientIp clientWriteRates maxClientWriteRatePerKey clientWriteRateWindowMs t respond mkErr
1359+ {ok: False {}, authorized: True {}, cancelled: False {}, error: "rate_limited"} (\senderKey -> do
13561360 let matches = \x -> x.messageId == messageId && x.toKey == to
13571361 result <- atomic do
13581362 msgs <- *messages
@@ -1362,13 +1366,9 @@ handleCancelMessage = \messageId to authorization clientIp respond -> do
13621366 if msg.originServer != "" || msg.senderKey != senderKey || msg.senderKey == ""
13631367 then yield {cancelled: False {}, cancelServers: []}
13641368 else do
1365− set *messages = removeWhere msgs matches
1366− fwds <- *forwards
1367− let fwdServers = collectPeers (filter fwds matches) (\f -> f.toServer)
1368− pruneRef forwards (\f -> not (matches f))
1369+ fwdServers <- applyCancel msgs matches to messageId t
13691370 pres <- *presence
13701371 let presServers = collectPeers (presenceForKey pres to t) (\p -> p.server)
1371− appendCapped cancelNotifications [{toKey: to, messageId: messageId, cancelledAt: t}] 10000 (\c -> c.cancelledAt)
13721372 yield {cancelled: True {}, cancelServers: unique (union presServers fwdServers)}
13731373 when (result.cancelled && count result.cancelServers > 0 && federationEnabled)
13741374 (fork (sendCancelToEach messageId to result.cancelServers))
@@ -1389,10 +1389,8 @@ sendCancelToEach = \messageId to servers ->
13891389handleCancelBatch = \cancels authorization clientIp respond -> do
13901390 t <- now
13911391 let mkErr = \e -> {results: [], authorized: False {}, error: if e == "draining" then "" else e}
1392− gateAuth (True {}) authorization clientIp t respond mkErr (\senderKey -> do
1393− lim <- checkPubkeyRateGate clientWriteRates senderKey maxClientWriteRatePerKey clientWriteRateWindowMs t
1394− if lim then yield (respond {results: [], authorized: True {}, error: "rate_limited"})
1395− else do
1392+ gateRated (True {}) authorization clientIp clientWriteRates maxClientWriteRatePerKey clientWriteRateWindowMs t respond mkErr
1393+ {results: [], authorized: True {}, error: "rate_limited"} (\senderKey -> do
13961394 results <- mapEach (takeRelation maxBatchSize cancels) (\e -> e.messageId) (\e -> do
13971395 let matches = \x -> x.messageId == e.messageId && x.toKey == e.to
13981396 r <- atomic do
@@ -1403,9 +1401,7 @@ handleCancelBatch = \cancels authorization clientIp respond -> do
14031401 if msg.originServer != "" || msg.senderKey != senderKey || msg.senderKey == ""
14041402 then yield False
14051403 else do
1406− set *messages = removeWhere msgs matches
1407− pruneRef forwards (\f -> not (matches f))
1408− appendCapped cancelNotifications [{toKey: e.to, messageId: e.messageId, cancelledAt: t}] 10000 (\c -> c.cancelledAt)
1404+ _ <- applyCancel msgs matches e.to e.messageId t
14091405 yield True
14101406 yield {messageId: e.messageId, cancelled: r})
14111407 yield (respond {results: results, authorized: True {}, error: ""}))
@@ -1551,9 +1547,7 @@ handleFederationCancel = \messageId toKey fromServer requestTimestamp correlatio
15511547 case findFirst msgs (\m -> matches m && m.originServer == fromServer) of
15521548 Nothing {} -> yield False
15531549 Just {} -> do
1554− set *messages = removeWhere msgs matches
1555− pruneRef forwards (\f -> not (matches f))
1556− appendCapped cancelNotifications [{toKey: toKey, messageId: messageId, cancelledAt: t}] 10000 (\c -> c.cancelledAt)
1550+ _ <- applyCancel msgs matches toKey messageId t
15571551 yield True
15581552 yield (respond {ok: True {}, error: "", cancelled: cancelled, retryAfterMs: 0}))
15591553
@@ -1593,7 +1587,11 @@ handleFederationDeliveryReceipt = \messageIds toKeys fromServer requestTimestamp
15931587 if count messageIds != count toKeys
15941588 then yield (respond {ok: False {}, error: "mismatched_lengths", receivedCount: 0, retryAfterMs: 0})
15951589 else do
1596− let pairs = zipPairs messageIds toKeys
1590+ let pairs = do
1591+ mp <- withIndex messageIds
1592+ kp <- withIndex toKeys
1593+ where mp.i == kp.i
1594+ yield {messageId: mp.v, toKey: kp.v}
15971595 atomic do
15981596 msgs <- *messages
15991597 let receipts = do
@@ -1607,10 +1605,6 @@ handleFederationDeliveryReceipt = \messageIds toKeys fromServer requestTimestamp
16071605-- Index a list (1-based by virtue of `count` of the partial fold).
16081606withIndex = \xs -> fold (\acc x -> union acc [{i: count acc, v: x}]) [] xs
16091607
1610−-- Pair parallel lists by index (O(n*m) cross-filter; trivial code).
1611−zipPairs = \mids keys ->
1612− do mp <- withIndex mids; kp <- withIndex keys; where mp.i == kp.i; yield {messageId: mp.v, toKey: kp.v}
1613−
16141608-- ============================================
16151609-- Maintenance handlers
16161610-- ============================================
@@ -1682,10 +1676,12 @@ handleUnblockPeer = \server authorization respond -> withMaintAuth authorization
16821676 (yield (respond {ok: False {}, unblocked: False {}}))
16831677 (do removed <- removeServerFromList blockedPeers server; yield (respond {ok: True {}, unblocked: removed}))
16841678
1685−handleListBlocked = \authorization respond -> do
1679+handleListPeers = \listRef authorization respond -> do
16861680 ok <- verifyMaintenanceAuth authorization
1687− blocked <- *blockedPeers
1688− yield (respond {peers: if ok then blocked else [], ok: ok})
1681+ list <- *listRef
1682+ yield (respond {peers: if ok then list else [], ok: ok})
1683+
1684+handleListBlocked = \authorization respond -> handleListPeers blockedPeers authorization respond
16891685
16901686handleAllowPeer = \server authorization respond -> withMaintAuth authorization
16911687 (yield (respond {ok: False {}, allowed: 0}))
@@ -1695,10 +1691,7 @@ handleUnallowPeer = \server authorization respond -> withMaintAuth authorization
16951691 (yield (respond {ok: False {}, removed: False {}}))
16961692 (do removed <- removeServerFromList allowedPeers server; yield (respond {ok: True {}, removed: removed}))
16971693
1698−handleListAllowed = \authorization respond -> do
1699− ok <- verifyMaintenanceAuth authorization
1700− allowed <- *allowedPeers
1701− yield (respond {peers: if ok then allowed else [], ok: ok})
1694+handleListAllowed = \authorization respond -> handleListPeers allowedPeers authorization respond
17021695
17031696handlePingPeer = \server authorization respond -> withMaintAuth authorization
17041697 (yield (respond {ok: False {}, latencyMs: 0, error: "unauthorized"}))