diff --git a/backend/app/routers/upstreams.py b/backend/app/routers/upstreams.py index 71c9696..eaabbcd 100644 --- a/backend/app/routers/upstreams.py +++ b/backend/app/routers/upstreams.py @@ -586,6 +586,7 @@ def _ensure_group_key( db.refresh(row) return _key_response(row, include_value=True) except Exception as exc: + db.rollback() logger.exception("ensure group key failed for upstream=%s group=%s", upstream.id, gid) return GeneratedUpstreamKeyResponse( upstream_id=upstream.id, diff --git a/backend/test_upstream_key_sync.py b/backend/test_upstream_key_sync.py index 1292b3b..36812b4 100644 --- a/backend/test_upstream_key_sync.py +++ b/backend/test_upstream_key_sync.py @@ -1735,3 +1735,96 @@ def test_generate_keys_by_groups_warm_cache_failure_fails_fast(db_session, monke assert excinfo.value.status_code == 502 assert "上游连接或缓存预热失败" in excinfo.value.detail + +def test_ensure_group_key_rollback_prevents_leak(db_session): + """验证 _ensure_group_key 在失败时执行 rollback,防止脏状态(如 status='replaced')泄漏。""" + from app.routers.upstreams import _ensure_group_key + from app.schemas.upstream import GenerateKeysByGroupsRequest + from app.models.upstream_key import UpstreamGeneratedKey + from app.services.upstream_client import UpstreamError + + upstream = Upstream( + name="NoxAPI", + base_url="http://nox.local", + api_prefix="", + auth_type="nox_token", + auth_config_json=json.dumps({"token": "abc", "new_api_user": "7"}), + groups_endpoint="/api/user/self/groups", + rate_endpoint="/api/user/self/groups", + ) + db_session.add(upstream) + db_session.commit() + db_session.refresh(upstream) + + # 1. 本地有两条记录:vip 和 normal。 + # 状态都为 exists。 + old_key_name_vip = f"SmartUp-{upstream.id}-VIP-vip" + row_vip = UpstreamGeneratedKey( + upstream_id=upstream.id, + group_id="vip", + group_name="VIP", + key_name=old_key_name_vip, + key_value="sk-vip-local", + managed_prefix="SmartUp", + key_id="vip-123", + status="exists", + ) + db_session.add(row_vip) + + old_key_name_normal = f"SmartUp-{upstream.id}-Normal-normal" + row_normal = UpstreamGeneratedKey( + upstream_id=upstream.id, + group_id="normal", + group_name="Normal", + key_name=old_key_name_normal, + key_value="sk-normal-local", + managed_prefix="SmartUp", + key_id="normal-123", + status="exists", + ) + db_session.add(row_normal) + db_session.commit() + + # 2. 我们执行确保 vip 分组: + # Step 1 发现远端不存在此 key (返回 None) -> 本地 row_vip 设为 replaced + # Step 2 抛出 429 异常 -> _ensure_group_key 应该捕获并 db.rollback(), + # 否则 row_vip.status 将为 replaced 并遗留在 Session 中。 + calls = {"vip": 0} + + class MockClientVip: + def find_smartup_group_key(self, group_id, name, prefix): + calls["vip"] += 1 + if calls["vip"] == 1: + return None # Step 1: 远端不存在旧 token + # Step 2: 遇到 429 + raise UpstreamError("429 Too Many Requests") + + def create_api_key(self, *args, **kwargs): + raise AssertionError("should not call create") + + body = GenerateKeysByGroupsRequest(group_ids=["vip", "normal"], name_prefix="SmartUp", quota=0) + group_vip = {"id": "vip", "name": "VIP"} + res_vip = _ensure_group_key(db_session, MockClientVip(), upstream, group_vip, "SmartUp", body) + assert res_vip.status == "failed" + + # 3. 接下来执行正常的分组: + # 模拟能够正常同步,内部会执行 db.commit() + class MockClientNormal: + def find_smartup_group_key(self, group_id, name, prefix): + return {"id": "normal-123", "name": old_key_name_normal, "key": "sk-normal-remote"} + + group_normal = {"id": "normal", "name": "Normal"} + res_normal = _ensure_group_key(db_session, MockClientNormal(), upstream, group_normal, "SmartUp", body) + assert res_normal.status == "exists" + + # 4. 最终验证: + # 虽然在同一个 Session 之后提交了 normal,但 vip 的状态绝对不能变成 replaced! + # 如果泄露了,这里 vip 的状态就会变成 replaced。 + db_session.expire_all() + vip_db = db_session.query(UpstreamGeneratedKey).filter_by( + upstream_id=upstream.id, group_id="vip" + ).one() + assert vip_db.status == "exists", "VIP group status was incorrectly modified to replaced!" + assert vip_db.key_value == "sk-vip-local" + +