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

This commit is contained in:
liumangmang
2026-06-30 17:33:43 +08:00
parent edc00444dc
commit 6065e824e3
2 changed files with 94 additions and 0 deletions
+1
View File
@@ -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,
+93
View File
@@ -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"