fix: harden the debug command, link cache and rate limiter

- commands: /bot_dict dumped the whole chat state to any member of the chat
  and could exceed Telegram's 4096-char message limit (the send then failed
  and bubbled up as a handler error). It is now admin-only and capped at
  MAX_DEBUG_DUMP_CHARS; README, README.en and the /help description updated.
- send: the edit-before-forward template buttons were built from a HashMap
  walk, so their order changed between prompts. Now sorted by name.
- link_cache: an unparseable payload (older schema) was reported as a miss
  but left in place, re-failing the parse on every later hit; the row is
  dropped on read.
- handlers: a link handed to the URL workers after the channel closed
  (shutdown) was discarded silently; it is now logged.
- rate_limit: LIMITERS kept one bucket per chat that ever sent media,
  forever. The periodic sweep now drops buckets that are idle (refilled to
  capacity) and not held by an in-flight sender; acquire's refill was
  factored into a shared helper used by the idle check.

Tests: +3 (corrupted row dropped, sorted markup, idle-bucket pruning); the
cache one was verified to fail before the fix. fmt/clippy clean, 55 + 69.
This commit is contained in:
2026-09-16 21:16:14 +08:00
parent 0a577600fd
commit 475cfd18f9
8 changed files with 152 additions and 19 deletions
+56 -8
View File
@@ -48,6 +48,19 @@ impl TokenBucket {
}
}
/// Applies the elapsed refill to `state`. Shared by [`Self::acquire`] and
/// the idle check so the two cannot drift apart.
fn refill(&self, state: &mut State) {
let now = tokio::time::Instant::now();
let elapsed = now
.saturating_duration_since(state.last_refill)
.as_secs_f64();
// Refill up to the capacity; a debt (negative balance) is repaid
// before any surplus accumulates.
state.tokens = (state.tokens + elapsed * self.refill_per_sec).min(self.capacity);
state.last_refill = now;
}
/// Waits until `n` tokens are available, consuming them. The wait is
/// bounded: the deficit is committed as debt and repaid over time, so a
/// large acquire returns once its share of the refill budget has passed.
@@ -57,14 +70,7 @@ impl TokenBucket {
// would make the future !Send).
let wait = {
let mut state = self.state.lock();
let now = tokio::time::Instant::now();
let elapsed = now
.saturating_duration_since(state.last_refill)
.as_secs_f64();
// Refill up to the capacity; a debt (negative balance) is repaid
// before any surplus accumulates.
state.tokens = (state.tokens + elapsed * self.refill_per_sec).min(self.capacity);
state.last_refill = now;
self.refill(&mut state);
if state.tokens >= n {
state.tokens -= n;
return;
@@ -77,6 +83,14 @@ impl TokenBucket {
};
tokio::time::sleep(Duration::from_secs_f64(wait)).await;
}
/// True when the bucket has refilled to capacity: no debt outstanding, so
/// the chat has not sent anything recently.
fn is_idle(&self) -> bool {
let mut state = self.state.lock();
self.refill(&mut state);
state.tokens >= self.capacity
}
}
/// One limiter per chat, created on first use. Per-chat so one chat's burst
@@ -93,6 +107,18 @@ pub fn limiter_for(chat_id: i64) -> Arc<TokenBucket> {
.clone()
}
/// Drops limiters that are idle (refilled to capacity, so the chat has not
/// sent recently) and are not still held by an in-flight sender. The map
/// would otherwise keep one bucket per chat that ever sent media, forever.
/// Called from the periodic sweep; returns how many were dropped.
pub fn prune_idle() -> usize {
let mut limiters = LIMITERS.lock();
let before = limiters.len();
// Lock order map → bucket, the only order taken anywhere.
limiters.retain(|_, bucket| Arc::strong_count(bucket) > 1 || !bucket.is_idle());
before - limiters.len()
}
#[cfg(test)]
mod tests {
use super::*;
@@ -133,4 +159,26 @@ mod tests {
start.elapsed()
);
}
#[tokio::test(start_paused = true)]
async fn prune_idle_drops_full_unheld_buckets_only() {
// Held by this task: kept even at full capacity, a sender has it.
let held = limiter_for(9_001);
assert!(held.is_idle(), "a fresh bucket is full");
// Only the map holds this one and it is full → dropped.
limiter_for(9_002);
// Mid-debt (an acquire larger than the capacity): kept.
{
let bucket = Arc::new(TokenBucket::new(CAPACITY, REFILL_PER_SEC));
bucket.state.lock().tokens = -1.0;
LIMITERS.lock().insert(9_003, bucket);
}
assert!(prune_idle() >= 1);
let limiters = LIMITERS.lock();
assert!(limiters.contains_key(&9_001), "held bucket pruned");
assert!(!limiters.contains_key(&9_002), "idle unheld bucket kept");
assert!(limiters.contains_key(&9_003), "indebted bucket pruned");
}
}