fix(upstreams): add db.rollback() on exception to prevent dirty replaced status leak

2026-06-30 17:33:43 +08:00
parent edc00444dc
commit 6065e824e3
2 changed files with 94 additions and 0 deletions
@@ -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"
+
+