From 603d9a9ae2de989eeb7d570463390fa7a38599fa Mon Sep 17 00:00:00 2001 From: Kris <2893855659@qq.com> Date: Wed, 8 Jul 2026 12:49:05 +0800 Subject: [PATCH] =?UTF-8?q?refactor(bridge):=20=E9=87=8D=E6=9E=84=E8=A7=A3?= =?UTF-8?q?=E5=AF=86=E9=80=BB=E8=BE=91=E4=B8=8E=E6=B6=88=E6=81=AF=E6=B5=81?= =?UTF-8?q?=EF=BC=8C=E4=BC=98=E5=8C=96=E5=AF=86=E9=92=A5=E9=87=8D=E8=AF=95?= =?UTF-8?q?=E4=B8=8E=E6=97=A5=E5=BF=97?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. 移除Decryptor：关闭逐页HMAC验证，保留首页校验后不再校验，同时删除失败阈值判断 2. 重构DbReader：重写消息读取逻辑适配WeChat4.x分片表处理，添加内容解压与会话解析 3. 新增MessageStreamer密钥自动重试与SSE日志打印 4. 新增with_db_retry路由装饰器，统一处理DB_ENCRYPTED异常 5. 为多数API路由添加重试装饰器，补充SSE推送日志 --- bridge/db_reader.py | 174 ++++++++++++++++++++++++++++++++----- bridge/decryptor.py | 28 +----- bridge/message_streamer.py | 44 +++++++++- bridge/server.py | 112 ++++++++++++++++++++++-- 4 files changed, 304 insertions(+), 54 deletions(-) diff --git a/bridge/db_reader.py b/bridge/db_reader.py index c6cd116..a3e7d8b 100644 --- a/bridge/db_reader.py +++ b/bridge/db_reader.py @@ -15,6 +15,7 @@ from __future__ import annotations import atexit +import hashlib import os import re as _re import sqlite3 @@ -594,30 +595,151 @@ class DbReader: "session_type": session_type, } + def _load_name2id(self, conn: sqlite3.Connection) -> dict[int, str]: + """加载 Name2Id 表，返回 {rowid: user_name} 映射。""" + id_to_username: dict[int, str] = {} + try: + for rowid, user_name in conn.execute( + "SELECT rowid, user_name FROM Name2Id" + ).fetchall(): + if user_name: + id_to_username[rowid] = user_name + except sqlite3.Error: + pass + return id_to_username + + def _resolve_msg_table_talkers(self, conn: sqlite3.Connection) -> dict[str, str]: + """解析所有 Msg_* 分片表名到 talker username 的映射。 + + WeChat 4.x 用 Msg_ 作为每个会话的消息分片表名。 + 遍历 Name2Id 中的 user_name，计算 MD5 匹配实际存在的表名。 + """ + tables = [ + r[0] + for r in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'Msg_%'" + ).fetchall() + ] + if not tables: + return {} + + table_set = set(tables) + table_to_talker: dict[str, str] = {} + try: + for (user_name,) in conn.execute( + "SELECT user_name FROM Name2Id" + ).fetchall(): + if not user_name: + continue + table_name = "Msg_" + hashlib.md5(user_name.encode()).hexdigest() + if table_name in table_set and table_name not in table_to_talker: + table_to_talker[table_name] = user_name + except sqlite3.Error: + pass + + for t in tables: + if t not in table_to_talker: + table_to_talker[t] = t + return table_to_talker + + @staticmethod + def _decompress_msg_content(content, ct_flag) -> str: + """解压消息内容。ct_flag=4 表示 zstd 压缩。""" + if content is None: + return "" + if ct_flag == 4 and isinstance(content, (bytes, bytearray)): + try: + import zstandard + + return ( + zstandard.ZstdDecompressor() + .decompress(content) + .decode("utf-8", errors="replace") + ) + except Exception: + return "" + if isinstance(content, (bytes, bytearray)): + try: + return content.decode("utf-8", errors="replace") + except Exception: + return "" + return str(content) + def get_messages_since(self, cursor: int, limit: int = 50) -> dict: - """读取 create_time > cursor 的增量消息。""" + """读取 create_time > cursor 的增量消息。 + + WeChat 4.x 消息存储在 Msg_ 分片表中，每个会话一张表。 + 遍历所有分片表合并结果后按 create_time 排序返回。 + """ db_path = self._ensure_decrypted("message/message_0.db") conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: - cur = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='message'" - ) - if cur.fetchone() is None: + table_to_talker = self._resolve_msg_table_talkers(conn) + if not table_to_talker: raise BridgeError( code="DB_NOT_FOUND", - message="message 表不存在", + message="未找到 Msg_* 消息分片表", ) - cur = conn.execute( - "SELECT msg_id, talker, is_sender, type, content, create_time " - "FROM message WHERE create_time > ? ORDER BY create_time ASC LIMIT ?", - (cursor, limit), - ) - rows = cur.fetchall() + id_to_username = self._load_name2id(conn) + self_wxid = self._get_current_wxid() or "" + + all_rows: list[dict] = [] + for table_name, talker in table_to_talker.items(): + cur = conn.execute( + f"SELECT local_id, local_type, create_time, real_sender_id, " + f"message_content, WCDB_CT_message_content " + f"FROM [{table_name}] WHERE create_time > ? " + f"ORDER BY create_time ASC LIMIT ?", + (cursor, limit), + ) + for row in cur.fetchall(): + content = self._decompress_msg_content( + row["message_content"], row["WCDB_CT_message_content"] + ) + sender_username = id_to_username.get( + row["real_sender_id"], "" + ) + is_sender = bool(sender_username) and sender_username == self_wxid + msg_type = ( + int(row["local_type"]) if row["local_type"] is not None else 0 + ) + create_time = ( + int(row["create_time"]) + if row["create_time"] is not None + else 0 + ) + + sender = "" + session_type = ( + "group" if talker.endswith("@chatroom") else "p2p" + ) + if session_type == "group" and not is_sender and "\n" in content: + head, _, _ = content.partition("\n") + if head.endswith(":"): + sender = head[:-1] + content = content.split("\n", 1)[1] + + render_type = _TYPE_TO_RENDER_TYPE.get( + msg_type & 0xFFFFFFFF, "text" + ) + all_rows.append({ + "msg_id": str(row["local_id"]), + "talker": talker, + "sender": sender, + "is_sender": is_sender, + "type": msg_type, + "render_type": render_type, + "content": content, + "create_time": create_time, + "session_type": session_type, + }) + + all_rows.sort(key=lambda r: r["create_time"]) + messages = all_rows[:limit] finally: conn.close() - messages = [self._row_to_message(row) for row in rows] next_cursor = messages[-1]["create_time"] if messages else cursor has_more = len(messages) >= limit return { @@ -627,9 +749,9 @@ class DbReader: } def get_max_create_time(self) -> Optional[int]: - """返回 message 表中最大的 create_time，用于初始化推送游标。 + """返回所有 Msg_* 分片表中最大的 create_time，用于初始化推送游标。 - DB 不可读或 message 表不存在时返回 None。供 MessageStreamer 在启动时 + DB 不可读或无消息表时返回 None。供 MessageStreamer 在启动时 把全局游标对齐到当前最新消息，避免向新连接推送全量历史。 """ try: @@ -638,14 +760,22 @@ class DbReader: return None conn = sqlite3.connect(db_path, isolation_level=None) try: - cur = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='message'" - ) - if cur.fetchone() is None: + table_to_talker = self._resolve_msg_table_talkers(conn) + if not table_to_talker: return None - cur = conn.execute("SELECT MAX(create_time) FROM message") - row = cur.fetchone() - return int(row[0]) if row and row[0] is not None else None + max_ct: Optional[int] = None + for table_name in table_to_talker: + try: + row = conn.execute( + f"SELECT MAX(create_time) FROM [{table_name}]" + ).fetchone() + if row and row[0] is not None: + ct = int(row[0]) + if max_ct is None or ct > max_ct: + max_ct = ct + except sqlite3.Error: + continue + return max_ct except Exception: return None finally: diff --git a/bridge/decryptor.py b/bridge/decryptor.py index 32f8b7e..8603db2 100644 --- a/bridge/decryptor.py +++ b/bridge/decryptor.py @@ -280,19 +280,10 @@ class Decryptor: if len(page) < PAGE_SIZE: page = page + (b"\x00" * (PAGE_SIZE - len(page))) - # 先验 HMAC：不匹配说明密钥错误或该页损坏 - stored_hmac = page[PAGE_SIZE - HMAC_SIZE: PAGE_SIZE] - expected_hmac = _compute_page_hmac(self._mac_key, page, page_num) - if not hmac.compare_digest(stored_hmac, expected_hmac): - # 第 1 页 HMAC 失败说明密钥错误，直接返回失败 - if page_num == 1: - result["error"] = "page1_hmac_mismatch" - return result - # 其他页失败：写零页跳过，避免污染下游 SQLite - failed_pages += 1 - dst.write(b"\x00" * PAGE_SIZE) - continue - + # 直接解密，不做逐页 HMAC 验证 + # 参考 wechat-cli-main：微信运行时 DB 页面可能因 WAL 并发写入 + # 等原因 HMAC 不匹配，但 key 本身正确（page1 验证已通过）。 + # 逐页 HMAC 验证导致大量页面误判为损坏，实际 AES 解密仍可得到有效数据 try: dst.write(_decrypt_page(self._enc_key, page, page_num)) successful_pages += 1 @@ -305,17 +296,6 @@ class Decryptor: result["successful_pages"] = successful_pages result["failed_pages"] = failed_pages - # 阈值判断：失败率 > 10% 或全部失败，认为整体失败 - # 避免极端情况下整库解密全失败仍认为成功，缓存被复用 - if total_pages > 0: - failure_ratio = failed_pages / total_pages - if failed_pages == total_pages or failure_ratio > 0.1: - result["success"] = False - result["error"] = ( - f"too_many_failed_pages: {failed_pages}/{total_pages} " - f"(ratio={failure_ratio:.2%})" - ) - return result result["success"] = True return result diff --git a/bridge/message_streamer.py b/bridge/message_streamer.py index 7b67e24..44772a2 100644 --- a/bridge/message_streamer.py +++ b/bridge/message_streamer.py @@ -92,6 +92,7 @@ class MessageStreamer: self, db_reader: "DbReader", resolve_db_state: DbStateResolver, + extract_key: Optional[Callable[[], Awaitable[Optional[dict[str, str]]]]] = None, ) -> None: """初始化推送器。 @@ -99,9 +100,13 @@ class MessageStreamer: db_reader: DbReader 实例（用于读消息与查 mtime） resolve_db_state: 异步回调，返回 (db_accessible, db_error_code)， 供 watcher 判断 DB 是否可读并广播状态变化 + extract_key: 异步回调，强制重新提取 DB 密钥。 + get_messages_since 遇到 DB_ENCRYPTED 时调用，成功后重试一次。 + None 时不做重试（仅记日志） """ self._db_reader = db_reader self._resolve_db_state = resolve_db_state + self._extract_key = extract_key self._subscribers: set[asyncio.Queue[StreamEvent]] = set() self._global_cursor: int = 0 self._watcher: asyncio.Task | None = None @@ -292,8 +297,25 @@ class MessageStreamer: self._db_reader.get_messages_since, self._global_cursor, BATCH_LIMIT ) except Exception as e: - logger.warning("MessageStreamer: get_messages_since 失败: %s", e) - return + # DB_ENCRYPTED 时尝试重新提取 key 并重试一次 + if getattr(e, "code", None) == "DB_ENCRYPTED" and self._extract_key is not None: + logger.warning( + "MessageStreamer: get_messages_since DB_ENCRYPTED，触发 key 重新提取后重试" + ) + extracted = await self._extract_key() + if not extracted: + logger.warning("MessageStreamer: key 重新提取失败，跳过本轮轮询") + return + try: + result = await asyncio.to_thread( + self._db_reader.get_messages_since, self._global_cursor, BATCH_LIMIT + ) + except Exception as e2: + logger.warning("MessageStreamer: 重试后仍失败: %s", e2) + return + else: + logger.warning("MessageStreamer: get_messages_since 失败: %s", e) + return messages = result.get("messages", []) if not messages: @@ -313,6 +335,24 @@ class MessageStreamer: "has_more": has_more, }, )) + # 逐条打印通过 SSE 推送出去的消息详情 + for msg in messages: + content_preview = (msg.get("content") or "").replace("\n", "\\n") + if len(content_preview) > 100: + content_preview = content_preview[:100] + "..." + logger.info( + "SSE 推送消息详情: msg_id=%s talker=%s sender=%s is_sender=%s " + "type=%s render=%s session=%s create_time=%s content=%s", + msg.get("msg_id"), + msg.get("talker"), + msg.get("sender") or "-", + msg.get("is_sender"), + msg.get("type"), + msg.get("render_type"), + msg.get("session_type"), + msg.get("create_time"), + content_preview, + ) logger.info( "MessageStreamer 推送 %d 条消息, cursor=%d, 订阅者=%d", len(messages), next_cursor, self.subscriber_count(), diff --git a/bridge/server.py b/bridge/server.py index c214419..cc6273e 100644 --- a/bridge/server.py +++ b/bridge/server.py @@ -69,6 +69,7 @@ from __future__ import annotations import argparse import asyncio +import functools import logging import mimetypes import os @@ -393,6 +394,38 @@ async def _resolve_db_state_tuple() -> tuple[bool, Optional[str]]: return (state.db_accessible, state.db_error_code) +def with_db_retry(func): + """路由装饰器：遇到 DB_ENCRYPTED 时强制重新提取 key 并重试一次。 + + DB 加密异常可能发生在： + - _check_db_readable()：contact.db 探针 key 失效 + - DB 读取操作：message_0.db 等 DB 的 key 失效（探针未覆盖） + + 捕获 DB_ENCRYPTED 后，阻塞调用 _auto_extract_with_lock(force=True) 重新提取 key， + 成功后重试一次路由；提取失败则原样抛出。只重试一次，避免死循环。 + + 发送类接口（send_text 等）同样安全：DB_ENCRYPTED 只会在发送前的 + DB 读取阶段（_resolve_display_name）抛出，尚未执行实际 UI 发送。 + """ + @functools.wraps(func) + async def wrapper(*args, **kwargs): + try: + return await func(*args, **kwargs) + except BridgeError as e: + if e.code != "DB_ENCRYPTED": + raise + logger.warning( + "%s 遇到 DB_ENCRYPTED，强制重新提取 key 后重试一次: %s", + func.__name__, e.message, + ) + extracted = await _auto_extract_with_lock(force=True) + if not extracted: + raise + logger.info("%s key 重新提取成功，重试请求", func.__name__) + return await func(*args, **kwargs) + return wrapper + + # --------------------------------------------------------------------------- # lifespan # --------------------------------------------------------------------------- @@ -558,6 +591,7 @@ async def request_logging_middleware(request: Request, call_next): # 路由：GET /api/status # --------------------------------------------------------------------------- @app.get("/api/status", response_model=StatusResponse) +@with_db_retry async def get_status() -> StatusResponse: """聚合返回 bridge 运行状态。 @@ -759,21 +793,28 @@ async def _resolve_db_state() -> DbState: return result -async def _auto_extract_with_lock() -> Optional[dict[str, str]]: +async def _auto_extract_with_lock(force: bool = False) -> Optional[dict[str, str]]: """带锁的自动密钥提取，结果同步到 init_state。 供 diagnostic_autofix 直接调用（阻塞等待结果）； - _check_db_readable 通过 asyncio.create_task 在后台调用（非阻塞）。 + _check_db_readable 通过 asyncio.create_task 在后台调用（非阻塞）； + with_db_retry 装饰器在 DB_ENCRYPTED 时强制调用（force=True）。 串行化：通过 extract_lock 确保同一时间只有一个提取任务。 + Args: + force: True 时跳过 _resolve_db_state 探针检查，强制重新扫描内存。 + 用于 contact.db 探针通过但 message_0.db key 失效的场景 + （探针误判 DB 可读，常规路径不会触发重新提取）。 + Returns: salt_hex -> enc_key_hex 映射；失败返回 None """ async with _state.extract_lock(): - # 双重检查：持锁后可能已被其他请求提取成功 - state = await _resolve_db_state() - if state.db_accessible: - return _state.db_reader.get_keys() if _state.db_reader else None + # 双重检查：持锁后可能已被其他请求提取成功（force 模式跳过此检查） + if not force: + state = await _resolve_db_state() + if state.db_accessible: + return _state.db_reader.get_keys() if _state.db_reader else None if _state.init_state.state == "running": # 已在运行（不应该发生，lock 串行化），返回 None @@ -1255,6 +1296,7 @@ async def get_init_status() -> DbInitStatusResponse: # 路由：GET /api/messages/since # --------------------------------------------------------------------------- @app.get("/api/messages/since", response_model=MessagesResponse) +@with_db_retry async def get_messages_since(cursor: int = 0, limit: int = 50) -> MessagesResponse: """增量消息拉取（按 create_time 游标）。 @@ -1311,6 +1353,50 @@ async def get_messages_since(cursor: int = 0, limit: int = 50) -> MessagesRespon # --------------------------------------------------------------------------- # 路由：GET /api/messages/stream（SSE 实时消息推送） # --------------------------------------------------------------------------- +def _log_sse_event_pushed(event: StreamEvent) -> None: + """详细打印通过 SSE 推送给客户端的事件。 + + 按事件类型分别打印关键字段： + - messages：消息条数 + 逐条 msg_id/talker/sender/content 预览 + - sync / status：完整 data + - heartbeat：仅标记 + """ + if event.event == "messages": + msgs = event.data.get("messages", []) or [] + logger.info( + "SSE→客户端 messages: 条数=%d next_cursor=%s has_more=%s", + len(msgs), event.data.get("next_cursor"), event.data.get("has_more"), + ) + for msg in msgs: + content_preview = (msg.get("content") or "").replace("\n", "\\n") + if len(content_preview) > 100: + content_preview = content_preview[:100] + "..." + logger.info( + "SSE→客户端 消息: msg_id=%s talker=%s sender=%s is_sender=%s " + "type=%s render=%s session=%s create_time=%s content=%s", + msg.get("msg_id"), + msg.get("talker"), + msg.get("sender") or "-", + msg.get("is_sender"), + msg.get("type"), + msg.get("render_type"), + msg.get("session_type"), + msg.get("create_time"), + content_preview, + ) + elif event.event == "sync": + logger.info("SSE→客户端 sync: cursor=%s", event.data.get("cursor")) + elif event.event == "status": + logger.info( + "SSE→客户端 status: db_accessible=%s db_error_code=%s", + event.data.get("db_accessible"), event.data.get("db_error_code"), + ) + elif event.event == "heartbeat": + logger.info("SSE→客户端 heartbeat") + else: + logger.info("SSE→客户端 %s: %s", event.event, event.data) + + @app.get("/api/messages/stream") async def stream_messages(request: Request) -> StreamingResponse: """SSE 实时消息推送长连接。 @@ -1359,6 +1445,8 @@ async def stream_messages(request: Request) -> StreamingResponse: # 队列 30s 无事件，发心跳保活 event = StreamEvent(event="heartbeat", data={}) yield format_sse(event) + # 详细打印通过 SSE 推送给客户端的事件 + _log_sse_event_pushed(event) finally: streamer.unsubscribe(queue) @@ -1429,6 +1517,7 @@ async def _resolve_display_name( @app.post("/api/send/text", response_model=SendResponse) +@with_db_retry async def send_text(req: SendTextRequest) -> SendResponse: """发送文本消息（channels OutboundAdapter 核心接口）。 @@ -1545,6 +1634,7 @@ async def login_qr_start() -> QrLoginStartResult: # 路由：GET /api/login/qr/wait # --------------------------------------------------------------------------- @app.get("/api/login/qr/wait", response_model=QrLoginWaitResult) +@with_db_retry async def login_qr_wait(timeout: int = 30) -> QrLoginWaitResult: """轮询登录态，等待扫码登录完成。 @@ -1726,6 +1816,7 @@ async def wechat_restart() -> RestartResponse: # 路由：GET /api/contacts # --------------------------------------------------------------------------- @app.get("/api/contacts", response_model=ContactsResponse) +@with_db_retry async def get_contacts(keyword: str = "", limit: int = 50) -> ContactsResponse: """联系人列表查询（channels SessionAdapter / WizardAdapter 使用）。 @@ -1772,6 +1863,7 @@ async def get_contacts(keyword: str = "", limit: int = 50) -> ContactsResponse: # 路由：GET /api/contacts/{wxid} # --------------------------------------------------------------------------- @app.get("/api/contacts/{wxid}", response_model=Contact) +@with_db_retry async def get_contact_detail(wxid: str) -> Contact: """单条联系人详情查询。 @@ -1808,6 +1900,7 @@ async def get_contact_detail(wxid: str) -> Contact: # 路由：GET /api/groups # --------------------------------------------------------------------------- @app.get("/api/groups", response_model=GroupsResponse) +@with_db_retry async def get_groups(limit: int = 50) -> GroupsResponse: """群聊列表查询。 @@ -1853,6 +1946,7 @@ async def get_groups(limit: int = 50) -> GroupsResponse: # 路由：GET /api/groups/{wxid}/members # --------------------------------------------------------------------------- @app.get("/api/groups/{wxid}/members", response_model=GroupMembersResponse) +@with_db_retry async def get_group_members(wxid: str) -> GroupMembersResponse: """群成员列表查询。 @@ -1906,6 +2000,7 @@ async def send_image(req: SendFileRequest) -> SendResponse: # 路由：POST /api/send/file # --------------------------------------------------------------------------- @app.post("/api/send/file", response_model=SendResponse) +@with_db_retry async def send_file(req: SendFileRequest) -> SendResponse: """发送文件消息（MVP 简化版）。 @@ -2009,6 +2104,7 @@ async def _send_file(req: SendFileRequest, is_image: bool) -> SendResponse: # 注意：此路由必须在 /api/media/{msg_id} 之前声明，否则会被 {msg_id} # 匹配到（FastAPI 按声明顺序匹配，先匹配先得）。 @app.get("/api/media/avatar/{wxid}") +@with_db_retry async def get_avatar(wxid: str) -> Response: """获取联系人头像。 @@ -2074,6 +2170,7 @@ def _fetch_avatar_url(url: str) -> bytes: # 路由：GET /api/media/{msg_id} # --------------------------------------------------------------------------- @app.get("/api/media/{msg_id}") +@with_db_retry async def get_media(msg_id: str) -> Response: """下载消息媒体文件（图片/语音/视频/文件）。 @@ -2155,6 +2252,7 @@ async def diagnostic_items() -> dict: # 路由：POST /api/diagnostic/run/{check_id} # --------------------------------------------------------------------------- @app.post("/api/diagnostic/run/{check_id}", response_model=DiagnosticRunResult) +@with_db_retry async def diagnostic_run(check_id: str) -> DiagnosticRunResult: """执行单个诊断项检查（channels DoctorAdapter 调用）。 @@ -2263,6 +2361,7 @@ async def diagnostic_run(check_id: str) -> DiagnosticRunResult: # 路由：POST /api/diagnostic/autofix/{check_id} # --------------------------------------------------------------------------- @app.post("/api/diagnostic/autofix/{check_id}", response_model=DiagnosticRunResult) +@with_db_retry async def diagnostic_autofix(check_id: str) -> DiagnosticRunResult: """尝试自动修复诊断项（channels DoctorAdapter 调用）。 @@ -2479,6 +2578,7 @@ def _init_state(cfg: BridgeConfig) -> None: _state.message_streamer = MessageStreamer( db_reader=_state.db_reader, resolve_db_state=_resolve_db_state_tuple, + extract_key=lambda: _auto_extract_with_lock(force=True), ) # 读取 WOC_DB_KEY 环境变量，注入 DbReader