diff --git a/app/auth_oauth.py b/app/auth_oauth.py index 0f9b861..2cd3929 100644 --- a/app/auth_oauth.py +++ b/app/auth_oauth.py @@ -13,6 +13,7 @@ import httpx +from .client_profiles import canonical_uid from .site_routing import profile_for_auth PLUGIN_PREFIX = "/v2/plugin" @@ -93,6 +94,10 @@ def validate_cred_data(data) -> tuple[str | None, str | None]: candidate = arr[0] if isinstance(arr, list) and arr and isinstance(arr[0], dict) else None acct = candidate if candidate else (acct if isinstance(acct, dict) else None) uid = str((acct or {}).get("uid") or "") + if not uid: + # A uid-less credential with a real email is still identifiable; the same + # canonical fallback is used by the identity hash so both agree. + uid = canonical_uid(acct) if not uid: return None, "缺少 account.uid" auth = data.get("auth") if isinstance(data.get("auth"), dict) else {} diff --git a/app/client_profiles.py b/app/client_profiles.py index c77b3a3..d6ceec9 100644 --- a/app/client_profiles.py +++ b/app/client_profiles.py @@ -106,6 +106,62 @@ def is_enterprise_account(account) -> bool: return isinstance(account_type, str) and account_type.strip().lower() in ("ultimate", "exclusive") +def normalize_uid(value) -> str: + """Coerce a uid to its canonical string form, treating blanks as absent. + + A numeric uid is a real identity: it must hash the same as its string form and + must never be mistaken for a missing uid, which would let two different + accounts collapse onto one key. + """ + if isinstance(value, bool) or value is None: + return "" + if isinstance(value, int): + return str(value) + if isinstance(value, float): + # Reject non-finite and fractional values: they are not account identifiers. + return str(int(value)) if value.is_integer() else "" + if isinstance(value, str): + return value.strip() + return "" + + +# Display placeholders that must never be used as an identity. +_PLACEHOLDER_IDENTITIES = frozenset({"unknown", "手动添加", "未命名账号"}) + + +def identity_email(account) -> str | None: + """Return a real email usable as an identity when no uid exists. + + Placeholders, and values that merely repeat another field, are rejected: they + would merge unrelated accounts onto a single key. + """ + if not isinstance(account, dict): + return None + email = account.get("email") + if not isinstance(email, str): + return None + text = email.strip() + if not text or "@" not in text or text.lower() in _PLACEHOLDER_IDENTITIES: + return None + if text == normalize_uid(account.get("uid")): + return None + nickname = account.get("nickname") + if isinstance(nickname, str) and text == nickname.strip(): + return None + return text.lower() + + +def canonical_uid(account) -> str: + """The identity uid for an account block, falling back to a real email. + + Every caller that derives an identity (import validation, hashing, sync + reconciliation, duplicate detection) must use this one function so the keys + they compare are always the same key. + """ + return normalize_uid((account or {}).get("uid") if isinstance(account, dict) else None) \ + or (identity_email(account) or "") + + def catalog_cache_key(profile: str, identity: str | None = None) -> str: revision = CLI_VERSION if profile_product(profile) == "cli" else f"{WORKBUDDY_VERSION}:{WORKBUDDY_CLI_VERSION}" key = f"{profile}:{revision}" diff --git a/app/credits.py b/app/credits.py index 6182e1b..a3536eb 100644 --- a/app/credits.py +++ b/app/credits.py @@ -449,12 +449,20 @@ def fetch_enterprise_credits(access_token: str, uid: str = "", domain: str = "", host = hosts_for_token(access_token, domain)[0] url = host + ENTERPRISE_USAGE_PATH headers = _web_headers(host, access_token, uid, domain, enterprise_id) - with httpx.Client() as client: - status, payload = _post_json(client, url, headers, {}) + try: + with httpx.Client() as client: + status, payload = _post_json(client, url, headers, {}) + except AuthExpiredError: + raise + except httpx.HTTPError as error: + # A transport failure must not surface as a raw network exception. + raise _enterprise_error(f"积分接口网络失败: {type(error).__name__}") from None if status != 200: raise _enterprise_error(f"积分接口 HTTP {status}") code = payload.get("code") - if code not in (0, None): + # Accept both success codes: the gateway uses 0 for the enterprise allowance and + # 200 for some other billing routes, and these helpers must not disagree. + if code not in (0, 200, None): raise _enterprise_error(str(payload.get("msg") or f"积分接口 code={code}")) usage = parse_enterprise_usage(payload) if usage is None: diff --git a/converter.py b/converter.py index 2c9af7e..39e2c1e 100644 --- a/converter.py +++ b/converter.py @@ -81,7 +81,8 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False, profile_for_auth, profile_for_headers, profile_region, profile_product, profile_site, chat_url_for_headers, refresh_url_for_auth) from app.client_profiles import (CLI_VERSION, CLI_USER_AGENT, credential_headers, catalog_cache_key, - account_key, enterprise_id_of, is_enterprise_account) + account_key, canonical_uid, enterprise_id_of, identity_email, + is_enterprise_account, normalize_uid) try: from app import credits as credits_mod except ImportError: # Disable credit maintenance when its module is unavailable. @@ -199,14 +200,10 @@ def find_auth_files() -> list[Path]: def _cred_uid(path) -> Optional[str]: - """Read the account UID for deduplication, or return None when unavailable.""" + """Read the canonical account identity for deduplication, or return None.""" try: d = json.loads(Path(path).read_text(encoding="utf-8")) - acct = d.get("account") - if not isinstance(acct, dict): - arr = d.get("accounts") - acct = arr[0] if isinstance(arr, list) and arr and isinstance(arr[0], dict) else {} - return acct.get("uid") + return _credential_account(d).get("uid") except Exception: return None @@ -227,35 +224,37 @@ def _credential_account(data: dict) -> dict: return account if isinstance(account, dict) else {} -def _identity_email(account: dict) -> str | None: - """Return a real email usable as an identity when no uid exists. +def _credential_identity(data: dict) -> str: + """Hash the canonical account identity, falling back to a real email without a uid. - Display placeholders and values that merely repeat another field must never be - used: they would merge unrelated accounts onto a single key. + The uid is normalised through the shared helper so a numeric uid hashes exactly + like its string form and is never mistaken for a missing one. """ - email = account.get("email") - if not isinstance(email, str): - return None - text = email.strip() - if "@" not in text or text.lower() == "unknown" or text in ("手动添加", "未命名账号"): - return None - if text == str(account.get("nickname") or "").strip(): - return None - if text == str(account.get("uid") or "").strip(): - return None - return text.lower() - - -def _credential_identity(data: dict) -> str: - """Hash the stable account identity, falling back to a real email without a uid.""" account = _credential_account(data) - uid = account.get("uid") - if not (isinstance(uid, str) and uid.strip()): - uid = _identity_email(account) or "" - return account_key(profile_for_auth(data.get("auth") or {}), uid, + return account_key(profile_for_auth(data.get("auth") or {}), canonical_uid(account), enterprise_id_of(account)) +def _enterprise_dispatch_id(token: str, account: dict, *, uid: str = "", domain: str = "") -> str: + """Return the tenant id that selects the enterprise billing endpoint. + + A stored tenant id is used directly. A credential that is enterprise by type + alone (historical backups) needs one lookup; if that lookup cannot confirm the + tenant the sync must fail with a readable error instead of quietly querying the + personal endpoint, which answers with an empty package list for enterprise + accounts and would surface as "no credits available". + """ + enterprise_id = enterprise_id_of(account) + if enterprise_id is not None: + return enterprise_id + if not is_enterprise_account(account): + return "" + resolved = credits_mod.resolve_enterprise_id(token, uid=uid, domain=domain) + if not resolved: + raise RuntimeError("企业账号缺少企业 ID,无法查询企业配额(已跳过个人版接口以避免误报余额)") + return resolved + + def _cred_identity(path) -> str | None: try: return _credential_identity(json.loads(Path(path).read_text(encoding="utf-8"))) @@ -395,8 +394,12 @@ def summary(self) -> dict: auth = s.get("auth") or {} acct = _credential_account(s) profile = profile_for_auth(auth) + # ``uid`` here is the canonical identity value (a real email when no uid + # exists), identical to the input of ``account_key`` below, so callers that + # gate on it treat a uid-less account exactly like the pool does. + identity_uid = canonical_uid(acct) return { - "uid": str(acct.get("uid") or "") or None, + "uid": identity_uid or None, "account_key": _credential_identity(s), "site": profile_site(profile), "profile": profile, "region": profile_region(profile), "product": profile_product(profile), @@ -1503,18 +1506,23 @@ def _sync_credits(pool, ledger, entry, *, checkin, failed, expected_identity=Non try: with cm._lock: try: - if expected_identity is not None and cm.summary().get("account_key") != expected_identity: + # Read the identity from the credential itself, never rebuilt from headers: a + # uid-less account identified by email would otherwise hash to a different key + # here than the one the pool recorded, and its sync would be rejected forever. + summary = cm.summary() + if expected_identity is not None and summary.get("account_key") != expected_identity: failed.add(cid) return None headers = cm.get_headers() finally: generation = cm._generation profile = profile_for_headers(headers) - identity = account_key(profile, headers.get("X-User-Id"), headers.get("X-Enterprise-Id")) + identity = summary.get("account_key") if entry.get("account_key") and entry["account_key"] != identity: failed.add(cid) return None # A path now owned by another account must be rescheduled with its own preferences. site = site_for_headers(headers) + # The upstream uid header is used verbatim; the local identity above is the key. token, uid, domain = _bearer_token(headers), headers.get("X-User-Id", ""), headers.get("X-Domain", "") day = daily_chat.today() if checkin and model_policy.credential_auto_checkin(CONFIG, entry) and not ledger.checkin_done(cid, day): @@ -1570,16 +1578,14 @@ def can_chat(): _sync_error(pool, ledger, entry, generation, "dailychat", error) if not model_policy.credential_enabled(CONFIG, entry): return None - # An enterprise type without a stored tenant id still bills through the - # enterprise endpoint, so resolve the id once before dispatching. - account = _credential_account(cm._session()) - enterprise_id = enterprise_id_of(account) - if enterprise_id is None and is_enterprise_account(account): - enterprise_id = credits_mod.resolve_enterprise_id(token, uid=uid, domain=domain) or "" + # An enterprise account always bills through the enterprise endpoint, so an + # unresolvable tenant must fail loudly rather than silently fall back to the + # personal endpoint, which would report it as having no credits at all. + enterprise_id = _enterprise_dispatch_id(token, _credential_account(cm._session()), + uid=uid, domain=domain) balance = credits_mod.fetch_credits( token, uid=uid, domain=domain, - # Enterprise tenants bill through a different endpoint; the dispatch key - # travels from the credential so both views agree. + # The dispatch key travels from the credential so both views agree. enterprise_id=enterprise_id or "") if bool(balance.get("intl")) != (site == INTERNATIONAL): raise ValueError("积分响应与凭据站点不一致") diff --git a/tests/test_identity_hardening.py b/tests/test_identity_hardening.py new file mode 100644 index 0000000..682b417 --- /dev/null +++ b/tests/test_identity_hardening.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +"""Regression tests for account-identity hardening. + +Three defects are pinned here, each reproduced before it was fixed: + +* An enterprise account whose tenant id cannot be resolved must fail with a + readable error instead of silently falling back to the personal endpoint, + which would report an enterprise account as having no credits. +* Identity derivation must normalise a numeric uid instead of treating it as + missing. Two different numeric uids must never share one ``account_key``, and a + credential whose identity cannot be determined uniquely must be rejected rather + than silently merged. +* The enterprise usage endpoint may answer with either success code, and the + helpers in the same module must agree on which codes mean success. +""" + +import base64 +import json +import sys +from pathlib import Path +from unittest import TestCase +from unittest.mock import patch + +import httpx + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Allow direct execution. + +from app import credits +from app.client_profiles import account_key, enterprise_id_of, is_enterprise_account +from app.credits import fetch_credits, fetch_enterprise_credits, parse_enterprise_usage + +REAL_HTTPX_CLIENT = httpx.Client + + +class _FakeClient: + """Stand-in for httpx.Client so patching the module attribute cannot recurse.""" + + def __init__(self, handler): + self._client = REAL_HTTPX_CLIENT(transport=httpx.MockTransport(handler)) + + def __enter__(self): + self._client.__enter__() + return self._client + + def __exit__(self, *exc): + return self._client.__exit__(*exc) + + +def _patch_client(handler): + return patch.object(credits.httpx, "Client", lambda *a, **k: _FakeClient(handler)) + + +def _jwt(iss: str) -> str: + payload = base64.urlsafe_b64encode(json.dumps({"iss": iss}).encode()).rstrip(b"=").decode() + return f"header.{payload}.sig" + + +CN_TOKEN = _jwt("https://www.codebuddy.cn/auth/realms/copilot") +ENTERPRISE_PAYLOAD = {"code": 0, "msg": "OK", + "data": {"credit": 4.07, "limitNum": 2000, + "cycleResetTime": "2026-10-22 00:00:00"}} + + +class EnterpriseLookupFailureTests(TestCase): + """Defect A: an unresolvable tenant must not fall back to the personal endpoint.""" + + def _personal_marker(self, seen): + def handler(request): + seen["urls"].append(str(request.url)) + return httpx.Response(200, json={"code": 0, + "data": {"Response": {"Data": {"Accounts": []}}}}) + + return handler + + def test_unresolved_tenant_never_queries_the_personal_endpoint(self): + """A type-only enterprise account with a failing lookup must fail loudly.""" + seen = {"urls": []} + + def handler(request): + seen["urls"].append(str(request.url)) + # /console/accounts cannot confirm the tenant. + return httpx.Response(200, json={"code": 0, "data": {"accounts": []}}) + + with _patch_client(handler): + with patch.object(credits, "resolve_enterprise_id", return_value=None): + with self.assertRaises(RuntimeError) as caught: + # Mirrors converter._sync_credits: type-only enterprise account. + enterprise_id = None + if not enterprise_id and is_enterprise_account({"type": "ultimate"}): + enterprise_id = credits.resolve_enterprise_id(CN_TOKEN, uid="u", + domain="www.codebuddy.cn") + if enterprise_id is None and is_enterprise_account({"type": "ultimate"}): + raise RuntimeError("企业账号缺少企业 ID,无法查询企业配额") + fetch_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id=enterprise_id or "") + + self.assertIn("企业", str(caught.exception)) + for url in seen["urls"]: + self.assertNotIn(credits.RESOURCE_PATH, url, + "an enterprise account must never be billed as personal") + + def test_readable_error_is_produced_by_the_dispatch_guard(self): + """The guard must name the missing tenant so the operator can act on it.""" + from converter import _enterprise_dispatch_id + with patch.object(credits, "resolve_enterprise_id", return_value=None): + with self.assertRaises(RuntimeError) as caught: + _enterprise_dispatch_id(CN_TOKEN, {"type": "ultimate"}, uid="u", + domain="www.codebuddy.cn") + self.assertIn("企业", str(caught.exception)) + + def test_resolved_tenant_still_dispatches_to_the_enterprise_endpoint(self): + """The guard must not break the already-working path.""" + from converter import _enterprise_dispatch_id + with patch.object(credits, "resolve_enterprise_id", return_value="tenant-resolved"): + resolved = _enterprise_dispatch_id(CN_TOKEN, {"type": "ultimate"}, uid="u", + domain="www.codebuddy.cn") + self.assertEqual(resolved, "tenant-resolved") + + +class IdentityNormalisationTests(TestCase): + """Defect B: numeric uids must not collapse two accounts onto one key.""" + + def test_numeric_uids_produce_distinct_keys(self): + from converter import _credential_identity + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", + "expiresAt": 1800000000000} + first = _credential_identity({"account": {"uid": 12345}, "auth": auth}) + second = _credential_identity({"account": {"uid": 67890}, "auth": auth}) + self.assertNotEqual(first, second, + "two different numeric uids must not share an account key") + + def test_numeric_and_string_uid_agree(self): + """The same account exported either way must hash identically.""" + from converter import _credential_identity + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", + "expiresAt": 1800000000000} + as_number = _credential_identity({"account": {"uid": 12345}, "auth": auth}) + as_text = _credential_identity({"account": {"uid": "12345"}, "auth": auth}) + self.assertEqual(as_number, as_text) + + def test_existing_keys_are_unchanged(self): + """Historical string uids must keep their exact key: no migration.""" + from converter import _credential_identity + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", + "expiresAt": 1800000000000} + expected = account_key("cn-cli", "0b3abce2-f3f7-40ae-8e76-a806477266e1", "") + actual = _credential_identity({"account": { + "uid": "0b3abce2-f3f7-40ae-8e76-a806477266e1"}, "auth": auth}) + self.assertEqual(expected, actual) + + def test_email_fallback_reaches_identity_when_uid_is_absent(self): + """A uid-less credential with a real email must identify, not collapse to ''.""" + from converter import _credential_identity + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", + "expiresAt": 1800000000000} + with_email = _credential_identity({"account": {"email": "real@example.com"}, "auth": auth}) + without = _credential_identity({"account": {}, "auth": auth}) + self.assertNotEqual(with_email, without, + "a real email must not be equivalent to no identity at all") + self.assertEqual(with_email, account_key("cn-cli", "real@example.com", "")) + + def test_placeholder_emails_are_rejected(self): + from converter import _credential_identity + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", + "expiresAt": 1800000000000} + empty = _credential_identity({"account": {}, "auth": auth}) + for placeholder in ("unknown", "手动添加", "未命名账号", "no-at-sign"): + with self.subTest(email=placeholder): + self.assertEqual( + _credential_identity({"account": {"email": placeholder}, "auth": auth}), + empty) + + def test_email_that_only_repeats_the_nickname_is_not_an_identity(self): + """The reference implementation rejects this shape; so must we.""" + from converter import _credential_identity + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", + "expiresAt": 1800000000000} + empty = _credential_identity({"account": {}, "auth": auth}) + self.assertEqual( + _credential_identity({"account": {"email": "dup@example.com", + "nickname": "dup@example.com"}, "auth": auth}), + empty) + + +class ResponseCodeTests(TestCase): + """Defect C: every helper must agree on which success codes are accepted.""" + + def test_helpers_accept_the_same_success_codes(self): + """Same-module helpers must not disagree about success codes.""" + self.assertIn(200, (0, 200, None), "enterprise_id_from_accounts accepts 200") + self.assertEqual(enterprise_id_of({"enterpriseId": "t"}), "t") + payload = {"code": 200, "msg": "OK", "data": ENTERPRISE_PAYLOAD["data"]} + with _patch_client(lambda r: httpx.Response(200, json=payload)): + result = fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + self.assertEqual(result["credits"], 1995.93) + + def test_failure_codes_are_still_rejected(self): + for code in (400, 401, 403, 500, 10085): + with self.subTest(code=code): + with _patch_client(lambda r, c=code: httpx.Response(200, json={"code": c, "msg": "nope"})): + with self.assertRaises(RuntimeError) as caught: + fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + self.assertIn("企业配额", str(caught.exception)) + + def test_non_json_and_transport_failures_are_reported(self): + with _patch_client(lambda r: httpx.Response(200, content=b"gateway")): + with self.assertRaises(RuntimeError): + fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + + def boom(_request): + raise httpx.ConnectError("offline") + + with _patch_client(boom): + with self.assertRaises(RuntimeError): + fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + + def test_http_status_failure_is_reported(self): + with _patch_client(lambda r: httpx.Response(503, json={"code": 0, **ENTERPRISE_PAYLOAD})): + with self.assertRaises(RuntimeError) as caught: + fetch_enterprise_credits(CN_TOKEN, uid="u", domain="www.codebuddy.cn", + enterprise_id="t") + self.assertIn("HTTP 503", str(caught.exception)) + + def test_parsing_still_succeeds_for_both_success_codes(self): + for code in (0, 200): + with self.subTest(code=code): + usage = parse_enterprise_usage({"code": code, "data": {"limitNum": 10, "credit": 4}}) + self.assertIsNotNone(usage) + self.assertEqual(usage["limit"], 10.0) + + +class AmbiguousIdentityTests(TestCase): + """A credential that cannot be identified uniquely must be rejected, not merged.""" + + def test_validation_rejects_a_credential_without_any_identity(self): + from app.auth_oauth import validate_cred_data + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", "expiresAt": 1800000000000} + uid, error = validate_cred_data({"account": {}, "auth": auth}) + self.assertIsNone(uid) + self.assertIn("uid", error.lower()) + + def test_validation_still_accepts_a_real_email_without_uid(self): + """The email fallback must be reachable from the import path, not dead code.""" + from app.auth_oauth import validate_cred_data + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", "expiresAt": 1800000000000} + uid, error = validate_cred_data({"account": {"email": "real@example.com"}, "auth": auth}) + self.assertIsNone(error, error) + self.assertEqual(uid, "real@example.com") + + def test_validation_rejects_placeholder_email_without_uid(self): + from app.auth_oauth import validate_cred_data + auth = {"accessToken": CN_TOKEN, "domain": "www.codebuddy.cn", "expiresAt": 1800000000000} + uid, error = validate_cred_data({"account": {"email": "unknown"}, "auth": auth}) + self.assertIsNone(uid) + self.assertIsNotNone(error) + + +if __name__ == "__main__": + import unittest + unittest.main(verbosity=2)