From 475cfd18f9ca47b9e514d2d1a51e8eba6e60856e Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Wed, 16 Sep 2026 21:16:14 +0800 Subject: [PATCH] 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. --- README.en.md | 2 +- README.md | 2 +- crates/xmedia-bot/src/handlers/commands.rs | 29 +++++++++- crates/xmedia-bot/src/handlers/mod.rs | 7 ++- crates/xmedia-bot/src/link_cache.rs | 38 ++++++++++++- crates/xmedia-bot/src/main.rs | 4 ++ crates/xmedia-bot/src/rate_limit.rs | 64 +++++++++++++++++++--- crates/xmedia-bot/src/send.rs | 25 ++++++++- 8 files changed, 152 insertions(+), 19 deletions(-) diff --git a/README.en.md b/README.en.md index 53f4c9d..3d55c50 100644 --- a/README.en.md +++ b/README.en.md @@ -113,7 +113,7 @@ Telegram only accepts ports 443/80/88/8443. | `/set_template ` | Reply to a message containing `[]` to save it as a named template; `[]` is replaced by the original post link when forwarding (used with "edit before forward") | | `/set_format ` | Customize the caption format for one site. Sites: `twitter` / `bsky` / `pixiv` / `misskey`. Placeholders: `{url}` `{author}` `{author_url}` `{title}` `{tags}` | | `/clear_cache [link]` | Clear the link cache (admin only); with a link only that entry, otherwise everything | -| `/bot_dict` | Show the current chat state (debugging) | +| `/bot_dict` | Show the current chat state (debugging; admin only) | | `/test ` | Debug: parse a link and report the parse result only (site, title, author, tags, media list) — no media is sent | Link processing works only in private chats; commands work in any chat. diff --git a/README.md b/README.md index 5ac3c6f..f7ac499 100644 --- a/README.md +++ b/README.md @@ -113,7 +113,7 @@ Telegram 只接受 443/80/88/8443 端口。 | `/set_template <名称>` | 回复一条含 `[]` 的消息，将其保存为命名模板；转发时 `[]` 会被替换为原帖链接（配合「转发前编辑」使用） | | `/set_format <站点> <格式>` | 自定义某站点的 caption 格式。站点：`twitter` / `bsky` / `pixiv` / `misskey`。占位符：`{url}` `{author}` `{author_url}` `{title}` `{tags}` | | `/clear_cache [链接]` | 清空链接缓存（仅管理员）；带链接只清该条，否则清空全部 | -| `/bot_dict` | 查看当前聊天状态（调试用） | +| `/bot_dict` | 查看当前聊天状态（调试用；仅管理员） | | `/test <链接>` | 调试：只解析链接并返回解析结果（站点、标题、作者、标签、媒体列表），不发送任何媒体 | 链接处理仅限私聊；命令在任意聊天可用。 diff --git a/crates/xmedia-bot/src/handlers/commands.rs b/crates/xmedia-bot/src/handlers/commands.rs index 4ecfc17..19224e9 100644 --- a/crates/xmedia-bot/src/handlers/commands.rs +++ b/crates/xmedia-bot/src/handlers/commands.rs @@ -31,7 +31,7 @@ pub(crate) enum Command { parse_with = "split" )] SetTemplate(String), - #[command(description = "Show chat state (debug)")] + #[command(description = "Show chat state (debug; admin only)")] BotDict, #[command(description = "Set site caption format", parse_with = "split")] SetFormat(String), @@ -226,9 +226,28 @@ pub(crate) async fn execute_command( reply(bot, message.chat.id.0, message.id, text).await?; } Command::BotDict => { + // Debug dump of the chat's persisted state: admin only (it echoes + // forward-channel ids and templates to whoever asks). + let sender_id = message + .from + .as_ref() + .map(|user| user.id.0 as i64) + .unwrap_or(-1); + if !CONFIG.admin_ids.contains(&sender_id) { + reply(bot, message.chat.id.0, message.id, "Admin only.").await?; + return Ok(()); + } let chat_data = CHAT_STORE.get(message.chat.id.0).await; - let debug = format!("{chat_data:?}"); - let text = html_escape::encode_text(&debug).into_owned(); + let debug = html_escape::encode_text(&format!("{chat_data:?}")).into_owned(); + // A chat with many templates/edit records exceeds Telegram's 4096 + // char message limit; the dump is plain text (no parse mode), so a + // plain byte-boundary cut is safe. + let end = debug.floor_char_boundary(MAX_DEBUG_DUMP_CHARS.min(debug.len())); + let text = if end < debug.len() { + format!("{}…", &debug[..end]) + } else { + debug + }; reply(bot, message.chat.id.0, message.id, text).await?; } Command::SetFormat(arg) => { @@ -389,6 +408,10 @@ pub async fn register_commands(bot: &Bot) -> Result<(), RequestError> { /// it even for very large threads (many media lines + a long caption). const MAX_TEST_REPORT_CHARS: usize = 4000; +/// Cap for the `/bot_dict` debug dump: the state is echoed as one plain-text +/// message, so it must stay under Telegram's 4096-char limit. +const MAX_DEBUG_DUMP_CHARS: usize = 3500; + /// Builds the HTML report for the `/test` command: what the parser produced /// for a link (site, canonical URL, title/author/tags, caption and the media /// list) — no media is sent and nothing is cached or forwarded. Sent with diff --git a/crates/xmedia-bot/src/handlers/mod.rs b/crates/xmedia-bot/src/handlers/mod.rs index 9eb9b33..a3e6310 100644 --- a/crates/xmedia-bot/src/handlers/mod.rs +++ b/crates/xmedia-bot/src/handlers/mod.rs @@ -157,7 +157,12 @@ pub async fn message_handler(bot: Bot, message: Message) -> Result<(), RequestEr log::warn!("url workers not started; dropping link"); break; }; - let _ = tx.send((message.clone(), url)).await; + // A closed channel means the workers are stopping (shutdown): + // report the dropped link instead of losing it silently. + if tx.send((message.clone(), url)).await.is_err() { + log::warn!("url workers stopped; dropping link"); + break; + } } } respond(()) diff --git a/crates/xmedia-bot/src/link_cache.rs b/crates/xmedia-bot/src/link_cache.rs index 0e55479..fc86510 100644 --- a/crates/xmedia-bot/src/link_cache.rs +++ b/crates/xmedia-bot/src/link_cache.rs @@ -78,9 +78,15 @@ impl LinkCache { conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?; return Ok(None); } - Ok(Some(serde_json::from_str::(&payload).map_err( - |e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)), - )?)) + match serde_json::from_str::(&payload) { + Ok(post) => Ok(Some(post)), + Err(e) => { + // Unreadable payload (e.g. an older schema): drop it + // instead of re-failing the parse on every later hit. + conn.execute("DELETE FROM link_cache WHERE url = ?1", params![key])?; + Err(rusqlite::Error::ToSqlConversionFailure(Box::new(e))) + } + } }) .await; match result { @@ -228,6 +234,32 @@ mod tests { ); } + #[tokio::test] + async fn unreadable_entry_is_dropped_on_read() { + // A payload from an older schema must not be re-parsed on every hit: + // the row is removed and the read reports a miss. + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("c.db"); + let cache = LinkCache::new(crate::db::open_store(db_path.to_str().unwrap()).unwrap()); + { + let conn = rusqlite::Connection::open(&db_path).unwrap(); + conn.execute( + "INSERT INTO link_cache (url, payload, created_at) VALUES (?1, ?2, ?3)", + params!["twitter:1", "{not json", now_f64()], + ) + .unwrap(); + } + + assert!( + cache + .get("twitter:1", Duration::from_secs(3600)) + .await + .is_none() + ); + // Dropped, not left behind for the next hit. + assert_eq!(cache.clear(None).await, 0, "corrupted row still present"); + } + #[tokio::test] async fn remove_and_prune() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/xmedia-bot/src/main.rs b/crates/xmedia-bot/src/main.rs index 65994e6..d4c4556 100644 --- a/crates/xmedia-bot/src/main.rs +++ b/crates/xmedia-bot/src/main.rs @@ -108,6 +108,10 @@ async fn main() { if pruned > 0 { log::info!("link cache: pruned {pruned} expired entr(ies)"); } + let idle_limiters = crate::rate_limit::prune_idle(); + if idle_limiters > 0 { + log::debug!("rate limiter: dropped {idle_limiters} idle bucket(s)"); + } for (chat_id, prompt_message_id) in removed { // If the prompt was already deleted, this fails with a // 400 "message to edit not found" — log and ignore. diff --git a/crates/xmedia-bot/src/rate_limit.rs b/crates/xmedia-bot/src/rate_limit.rs index 034b012..30adea0 100644 --- a/crates/xmedia-bot/src/rate_limit.rs +++ b/crates/xmedia-bot/src/rate_limit.rs @@ -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 { .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"); + } } diff --git a/crates/xmedia-bot/src/send.rs b/crates/xmedia-bot/src/send.rs index 8358982..375b8c9 100644 --- a/crates/xmedia-bot/src/send.rs +++ b/crates/xmedia-bot/src/send.rs @@ -1146,9 +1146,13 @@ pub async fn forward_messages(sender: &dyn MediaSender, task: &Task) -> Result<( } /// One button per template name (column layout), then the confirm button. +/// Sorted by name: the templates live in a `HashMap`, so an unsorted walk +/// would reshuffle the buttons between prompts. pub fn build_edit_markup(templates: &HashMap) -> InlineKeyboardMarkup { - let mut rows = Vec::new(); - for name in templates.keys() { + let mut names: Vec<&String> = templates.keys().collect(); + names.sort(); + let mut rows = Vec::with_capacity(names.len() + 1); + for name in names { rows.push(vec![InlineKeyboardButton::callback( name.clone(), format!("template|{name}"), @@ -1488,6 +1492,23 @@ mod tests { } } + #[test] + fn edit_markup_lists_templates_sorted_then_confirm() { + // Six names: a HashMap walk would land on this order by chance only + // 1 time in 720. + let templates: HashMap = ["z", "a", "m", "q", "b", "y"] + .into_iter() + .map(|name| (name.to_string(), "[]".to_string())) + .collect(); + let labels: Vec = build_edit_markup(&templates) + .inline_keyboard + .iter() + .flatten() + .map(|button| button.text.clone()) + .collect(); + assert_eq!(labels, ["a", "b", "m", "q", "y", "z", "↩️ Confirm"]); + } + #[test] fn is_media_fetch_failure_matches_markers() { for description in [