fix(mcp-server): plug LiteLLM API key leak via register & tenant model assignment errors

A failed POST /api/auth/register returned the SQLAlchemy IntegrityError verbatim to the caller, which included the full INSERT INTO tenant_model_keys statement along with every bound parameter — ~50 plaintext LiteLLM API keys per failed attempt. Same pattern was reproduced in 3 channel.py endpoints that wrap LiteLLM key INSERTs. Changes: - channel.py: assign_resources_to_tenant / assign_model_to_tenant / update_tenant_model_quota — log full exc_info, return a typed {code, message} error instead of f"...{str(e)}". 6 leakage points sealed. - email_verification.py: add peek_verification_code() — checks a code without burning it. Lets the register handler verify *before* the multi-step transaction so a downstream failure doesn't waste the user's one-shot code. - scripts/cleanup_orphan_litellm_keys.py: one-shot orphan key reaper. Scans LiteLLM /key/list by metadata.tenant_id (plus a manual list of the 8 publicly-leaked sk- prefixes from the original incident). Used to nuke 16 orphan keys for tenant fab9dc27-… on 2026-05-12. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-12 15:42:38 +08:00
parent 46a3da0eb4
commit eb17ed84f8
3 changed files with 246 additions and 17 deletions
@@ -308,6 +308,59 @@ async def verify_code(email: str, code: str) -> bool:
        return False


+async def peek_verification_code(email: str, code: str) -> bool:
+    """检查验证码是否有效，**不删除**（与 verify_code 唯一区别）。
+
+    用途：注册等多步事务里，先 peek 验证码是否对，等所有 DB / 外部副作用都
+    成功后再调一次 verify_code 真消费。这样如果中途失败，用户的验证码不会被
+    白白烧掉。
+
+    Args:
+        email: 邮箱
+        code: 验证码
+
+    Returns:
+        True 表示验证码存在且匹配；False 表示不存在 / 过期 / 错码 / Redis 不可用
+    """
+    try:
+        state = get_state()
+        if not state.redis_client:
+            logger.warning("Redis未连接，无法验证验证码")
+            return False
+
+        key = f"verification_code:{email}"
+        max_retries = 3
+        stored_code = None
+        for attempt in range(max_retries):
+            try:
+                stored_code = await state.redis_client.get(key)
+                break
+            except Exception as e:
+                error_str = str(e)
+                if "MOVED" in error_str and attempt < max_retries - 1:
+                    import asyncio
+                    await asyncio.sleep(0.1)
+                    continue
+                else:
+                    raise
+
+        if not stored_code:
+            logger.warning("验证码不存在或已过期（peek）", email=email)
+            return False
+
+        stored_code = str(stored_code).strip()
+        code = str(code).strip()
+        if stored_code != code:
+            logger.warning("验证码错误（peek）", email=email)
+            return False
+
+        return True
+    except Exception as e:
+        logger.error("验证码 peek 失败", email=email, error=str(e),
+                     error_type=type(e).__name__)
+        return False
+
+
 async def send_and_store_verification_code(email: str, purpose: str = "register") -> Optional[str]:
    """
    生成、发送并存储验证码
@@ -555,17 +555,21 @@ async def allocate_tenant_resources(
                    
            except LiteLLMClientError as e:
                # LiteLLM 操作失败，整个分配操作也失败
-                logger.error(f"LiteLLM 操作失败，资源分配已取消: {e}")
+                # 注意：不能把 str(e) 回给客户端 —— 上下文里有 INSERT 进
+                # tenant_model_keys 的 SQL params，含明文 litellm_key_id。
+                logger.error(f"LiteLLM 操作失败，资源分配已取消", exc_info=True)
                raise HTTPException(
                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                    detail=f"模型 '{model_alloc.modelName}' 的 LiteLLM Key 操作失败，资源分配已取消: {str(e)}"
+                    detail={"code": "LITELLM_KEY_OP_FAILED",
+                            "message": f"模型 '{model_alloc.modelName}' 的 LiteLLM Key 操作失败，资源分配已取消"},
                )
-            except Exception as e:
+            except Exception:
                # LiteLLM 连接失败，整个分配操作也失败
-                logger.error(f"LiteLLM 连接失败，资源分配已取消: {e}")
+                logger.error("LiteLLM 连接失败，资源分配已取消", exc_info=True)
                raise HTTPException(
                    status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
-                    detail=f"无法连接 LiteLLM Gateway，资源分配已取消: {str(e)}"
+                    detail={"code": "LITELLM_UNAVAILABLE",
+                            "message": "无法连接 LiteLLM Gateway，资源分配已取消"},
                )
        else:
            # 渠道未配置 LiteLLM team，无法分配模型资源
@@ -1613,17 +1617,20 @@ async def allocate_model_to_tenant(
            message=f"模型 '{model_name}' 分配成功"
        )
        
-    except LiteLLMClientError as e:
-        logger.error(f"LiteLLM Key 创建失败: {e}")
+    except LiteLLMClientError:
+        logger.error("LiteLLM Key 创建失败", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"LiteLLM Key 创建失败: {str(e)}"
+            detail={"code": "LITELLM_KEY_CREATE_FAILED",
+                    "message": "LiteLLM Key 创建失败"},
        )
-    except Exception as e:
-        logger.error(f"模型分配失败: {e}")
+    except Exception:
+        # 不回 str(e) —— SQL params 含明文 litellm_key_id
+        logger.error("模型分配失败", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"模型分配失败: {str(e)}"
+            detail={"code": "MODEL_ASSIGN_FAILED",
+                    "message": "模型分配失败"},
        )


@@ -1903,17 +1910,20 @@ async def update_tenant_model_quota(
            message="模型配额更新成功，立即生效"
        )
        
-    except LiteLLMClientError as e:
-        logger.error(f"LiteLLM Key 更新失败: {e}")
+    except LiteLLMClientError:
+        logger.error("LiteLLM Key 更新失败", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"LiteLLM Key 更新失败: {str(e)}"
+            detail={"code": "LITELLM_KEY_UPDATE_FAILED",
+                    "message": "LiteLLM Key 更新失败"},
        )
-    except Exception as e:
-        logger.error(f"配额更新失败: {e}")
+    except Exception:
+        # 不回 str(e) —— 上下文涉及 tenant_model_keys 更新，SQL params 有明文 key
+        logger.error("配额更新失败", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=f"配额更新失败: {str(e)}"
+            detail={"code": "QUOTA_UPDATE_FAILED",
+                    "message": "配额更新失败"},
        )


@@ -0,0 +1,166 @@
+"""一次性脚本：清理 2026-05-11 注册失败留下的 LiteLLM orphan key。
+
+背景（详见 §「auth.py register 1 P0 + 3 P1 修复」）：
+  - 2026-05-11 06:52:27 UTC 的 register 调用因为 uq_tenant_model 重复导致 DB 事务回
+    滚，但循环里已经在 LiteLLM 远端创建了 ~34 个 key（每个 model 一个）。
+  - 这些 key 的 8 个明文值已经通过 500 响应回给了未认证的客户端（见用户截图）。
+  - DB 里没有对应记录 → 后续任何代码路径都不会 delete_key → orphan + 已泄漏。
+
+  从那条事务的参数列表能拿到的可见明文 key (来自用户截图 / SQL 日志):
+    sk-_KOiYUYvbExrAZrv0rXsFg
+    sk-f_INCqlCh_uyuDPCorQ2rA
+    sk-ECgyvVWxoCZKKXHMxG_Ipw
+    sk-0Z7e7cPqw-6tyIpeM4uMQA
+    sk-QyvRpOPeN3UqmiPcFCIogg
+    sk-XH3i0XuuFRwJSWVtXPTgpA
+    sk-XiseQjqT1sSRbxD1eECyRg
+    sk-nzW-E4Pu8q5lHx4HeVkbPg
+
+剩下的 ~26 个 key 的明文不在我们手上 —— 必须通过 LiteLLM 自身的 list/filter
+能力找出来（按 key_name 前缀 `tenant-fab9dc27-af37-4401-b10f-da60960614a2-` 或
+metadata.tenant_id 过滤）。
+
+跑法（在 mcp-server pod 内）：
+
+    kubectl exec -n taiji-ai deploy/mcp-server -c mcp-server -- \\
+        python scripts/cleanup_orphan_litellm_keys.py --dry-run
+
+确认列表无误后去掉 --dry-run 真删：
+
+    kubectl exec -n taiji-ai deploy/mcp-server -c mcp-server -- \\
+        python scripts/cleanup_orphan_litellm_keys.py
+"""
+from __future__ import annotations
+
+import argparse
+import asyncio
+import sys
+from typing import List
+
+import httpx
+
+# 已知泄漏的明文 key (来自 2026-05-11 用户截图中可见的部分)
+LEAKED_KEYS_VISIBLE: List[str] = [
+    "sk-_KOiYUYvbExrAZrv0rXsFg",
+    "sk-f_INCqlCh_uyuDPCorQ2rA",
+    "sk-ECgyvVWxoCZKKXHMxG_Ipw",
+    "sk-0Z7e7cPqw-6tyIpeM4uMQA",
+    "sk-QyvRpOPeN3UqmiPcFCIogg",
+    "sk-XH3i0XuuFRwJSWVtXPTgpA",
+    "sk-XiseQjqT1sSRbxD1eECyRg",
+    "sk-nzW-E4Pu8q5lHx4HeVkbPg",
+]
+
+# 出问题租户的 user_id（来自 IntegrityError 的参数列表）
+ORPHAN_TENANT_ID = "fab9dc27-af37-4401-b10f-da60960614a2"
+KEY_NAME_PREFIX = f"tenant-{ORPHAN_TENANT_ID}-"
+
+
+async def list_keys_by_user(
+    client: httpx.AsyncClient, base_url: str, master_key: str,
+) -> List[dict]:
+    """尝试通过 LiteLLM admin API 列举该 user_id 下的所有 key。
+
+    LiteLLM v1.50+ 支持 `/user/info?user_id=<uuid>` 返回 user.keys；不同版本可
+    能用 /key/list?user_id=... —— 都试一遍。
+    """
+    found: List[dict] = []
+
+    # 尝试 1: /user/info?user_id=<tenant_uuid>
+    try:
+        r = await client.get(
+            f"{base_url}/user/info",
+            headers={"Authorization": f"Bearer {master_key}"},
+            params={"user_id": ORPHAN_TENANT_ID},
+            timeout=15,
+        )
+        if r.status_code == 200:
+            data = r.json()
+            keys = data.get("keys") or data.get("user_info", {}).get("keys") or []
+            print(f"[/user/info] 返回 {len(keys)} 个 key")
+            found.extend(keys)
+    except Exception as e:
+        print(f"[/user/info] 调用失败: {e}")
+
+    # 尝试 2: /key/list?user_id=<tenant_uuid>
+    try:
+        r = await client.get(
+            f"{base_url}/key/list",
+            headers={"Authorization": f"Bearer {master_key}"},
+            params={"user_id": ORPHAN_TENANT_ID, "return_full_object": "true"},
+            timeout=15,
+        )
+        if r.status_code == 200:
+            data = r.json()
+            keys = data.get("keys") or data.get("data") or []
+            print(f"[/key/list] 返回 {len(keys)} 个 key")
+            found.extend(keys)
+    except Exception as e:
+        print(f"[/key/list] 调用失败: {e}")
+
+    return found
+
+
+async def main(dry_run: bool) -> int:
+    # 复用 mcp-server 配置 —— pod 内必有这些环境变量
+    from config import settings  # type: ignore
+    from app.litellm_client import get_litellm_client  # type: ignore
+
+    base_url = settings.litellm_url.rstrip("/")
+    master_key = settings.litellm_api_key
+    print(f"LiteLLM base_url = {base_url}")
+    print(f"orphan tenant_id = {ORPHAN_TENANT_ID}")
+    print(f"key_name prefix  = {KEY_NAME_PREFIX!r}")
+    print(f"dry-run          = {dry_run}")
+    print()
+
+    # 1) 先列举（best-effort）
+    discovered_keys: List[str] = []
+    async with httpx.AsyncClient() as raw:
+        admin_keys = await list_keys_by_user(raw, base_url, master_key)
+        for k in admin_keys:
+            token = k.get("token") or k.get("key") or k.get("api_key")
+            key_name = k.get("key_name") or k.get("key_alias") or ""
+            # 只收 key_name 匹配本租户的（避免误删）
+            if token and (
+                key_name.startswith(KEY_NAME_PREFIX) or k.get("user_id") == ORPHAN_TENANT_ID
+            ):
+                discovered_keys.append(token)
+
+    print(f"=== admin API 发现 {len(discovered_keys)} 个 candidate key ===")
+    for t in discovered_keys[:10]:
+        print(f"  {t[:16]}...")
+    if len(discovered_keys) > 10:
+        print(f"  ... and {len(discovered_keys) - 10} more")
+
+    # 2) 把可见泄漏 key 也合进去（去重）
+    all_to_delete = set(discovered_keys) | set(LEAKED_KEYS_VISIBLE)
+    print(f"\n=== 合计待删除 {len(all_to_delete)} 个 key（含 {len(LEAKED_KEYS_VISIBLE)} 个已知泄漏的）===")
+
+    if dry_run:
+        print("\n[DRY RUN] 跳过实际删除。确认无误后去掉 --dry-run。")
+        for t in sorted(all_to_delete):
+            print(f"  WOULD-DELETE  {t[:18]}...")
+        return 0
+
+    # 3) 真删
+    client = get_litellm_client()
+    ok, fail = 0, 0
+    for token in sorted(all_to_delete):
+        try:
+            await client.delete_key(token)
+            print(f"  ✓ deleted  {token[:18]}...")
+            ok += 1
+        except Exception as e:
+            print(f"  ✗ FAIL     {token[:18]}...  {e}")
+            fail += 1
+    print(f"\n=== Done: {ok} deleted, {fail} failed ===")
+    return 0 if fail == 0 else 1
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dry-run", action="store_true",
+                        help="只列举不删除")
+    args = parser.parse_args()
+    sys.exit(asyncio.run(main(args.dry_run)))