diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index a444e09c6b3..63d6f2c95e8 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -413,12 +413,27 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, ) - elif original_exception.status_code == 403 and "spending-limit" in error_str: - raise RateLimitError( - message=f"RateLimitError: {exception_provider} - {message}", - model=model, + elif original_exception.status_code == 403: + from litellm.llms.xai.oauth import is_xai_spending_limit_error + + if is_xai_spending_limit_error( + custom_llm_provider=custom_llm_provider, + status_code=original_exception.status_code, + error_str=error_str, + ): + raise RateLimitError( + message=f"RateLimitError: {exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), + litellm_debug_info=extra_information, + ) + raise APIError( + status_code=original_exception.status_code, + message=f"APIError: {exception_provider} - {message}", llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + model=model, + request=getattr(original_exception, "request", None), litellm_debug_info=extra_information, ) elif original_exception.status_code == 404: diff --git a/litellm/llms/xai/oauth.py b/litellm/llms/xai/oauth.py index ae8d47c2042..3476228fd85 100644 --- a/litellm/llms/xai/oauth.py +++ b/litellm/llms/xai/oauth.py @@ -2,6 +2,7 @@ import base64 import hashlib import json import os +import re import secrets import sys import threading @@ -30,6 +31,7 @@ XAI_OAUTH_REDIRECT_PORT: Final = 56121 XAI_OAUTH_REDIRECT_PATH: Final = "/callback" XAI_OAUTH_EXPIRY_SKEW_SECONDS: Final = 120 XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS: Final = 180 +_XAI_OAUTH_ACCOUNT_NAME_RE: Final = re.compile(r"^[A-Za-z0-9_-]+$") _XAI_OAUTH_REFRESH_LOCK: Final = threading.Lock() @@ -75,6 +77,30 @@ class XAIOAuthLoginRequiredError(XAIOAuthError): pass +def _default_xai_oauth_token_dir() -> str: + return get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser("~/.config/litellm/xai_oauth") + + +def resolve_xai_oauth_auth_file(auth_file: str | None, token_dir: str) -> str: + requested: Final = auth_file or get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json" + candidate: Final = requested if os.path.isabs(requested) else os.path.join(token_dir, requested) + resolved: Final = os.path.realpath(candidate) + allowed: Final = os.path.realpath(token_dir) + if resolved == allowed or resolved.startswith(allowed + os.sep): + return resolved + raise XAIOAuthError("xAI OAuth auth file must stay inside the token directory") + + +def oauth_auth_file_for_account(account: str, token_dir: str) -> str: + if not _XAI_OAUTH_ACCOUNT_NAME_RE.fullmatch(account): + raise ValueError("xAI OAuth account must match ^[A-Za-z0-9_-]+$") + return resolve_xai_oauth_auth_file(f"auth-{account}.json", token_dir) + + +def is_xai_spending_limit_error(*, custom_llm_provider: str, status_code: int, error_str: str) -> bool: + return custom_llm_provider == "xai" and status_code == 403 and "spending-limit" in error_str + + class _CallbackHandler(BaseHTTPRequestHandler): server: "_CallbackServer" # pyright: ignore[reportIncompatibleVariableOverride] # stdlib stubs type server as BaseServer; _CallbackServer is the only server this handler is registered on @@ -127,16 +153,8 @@ class XAIOAuthAuthenticator: auth_file: str | None = None, token_dir: str | None = None, ) -> None: - self.token_dir = ( - token_dir or get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser("~/.config/litellm/xai_oauth") - ) - if auth_file: - self.auth_file = auth_file if os.path.isabs(auth_file) else os.path.join(self.token_dir, auth_file) - else: - self.auth_file = os.path.join( - self.token_dir, - get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json", - ) + self.token_dir = token_dir or _default_xai_oauth_token_dir() + self.auth_file = resolve_xai_oauth_auth_file(auth_file, self.token_dir) self.http_client = http_client def get_api_base(self) -> str: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index f78c4221f5a..3e9925136ea 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -324,6 +324,11 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # caller-supplied value is the same exfil shape as # ``aws_web_identity_token`` on the Bedrock path. "azure_ad_token", + # xAI SuperGrok OAuth token file. The xAI transformer reads + # ``xai_oauth_token_file`` and authenticates as that local account, + # so a caller-supplied path is the same credential-selector shape as + # ``aws_profile_name`` on the Bedrock path. + "xai_oauth_token_file", # Endpoint-targeting fields that retarget the outbound request or # an observability callback. An attacker-controlled value either # exfiltrates the request payload (incl. messages + admin-set diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index f5ba98e1dde..592807fbb22 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -998,7 +998,10 @@ def run_server( ): if cli_args: if len(cli_args) >= 2 and cli_args[0] == "xai-oauth" and cli_args[1] == "login": - from litellm.llms.xai.oauth import XAIOAuthAuthenticator + from litellm.llms.xai.oauth import ( + XAIOAuthAuthenticator, + oauth_auth_file_for_account, + ) from litellm.secret_managers.main import get_secret_str account: Final = cli_args[2] if len(cli_args) >= 3 else None @@ -1006,7 +1009,7 @@ def run_server( "~/.config/litellm/xai_oauth" ) authenticator: Final = ( - XAIOAuthAuthenticator(auth_file=os.path.join(token_dir, f"auth-{account}.json")) + XAIOAuthAuthenticator(auth_file=oauth_auth_file_for_account(account, token_dir), token_dir=token_dir) if account else XAIOAuthAuthenticator() ) diff --git a/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py b/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py index 5f3bcde819f..254b6cd151d 100644 --- a/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py +++ b/tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py @@ -8,7 +8,7 @@ from click.testing import CliRunner import litellm from litellm.litellm_core_utils.exception_mapping_utils import _map_openai_exception from litellm.llms.xai.chat.transformation import XAIChatConfig -from litellm.llms.xai.oauth import XAIOAuthAuthenticator +from litellm.llms.xai.oauth import XAIOAuthAuthenticator, XAIOAuthError from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig from litellm.types.router import GenericLiteLLMParams @@ -26,9 +26,9 @@ def test_authenticator_accepts_explicit_absolute_auth_file(tmp_path): ) ) - auth = XAIOAuthAuthenticator(auth_file=str(custom)) + auth = XAIOAuthAuthenticator(auth_file=str(custom), token_dir=str(custom.parent)) - assert auth.auth_file == str(custom) + assert auth.auth_file == os.path.realpath(str(custom)) assert auth.get_access_token() == "alice-token" @@ -48,10 +48,62 @@ def test_authenticator_relative_auth_file_joins_token_dir(tmp_path, monkeypatch) auth = XAIOAuthAuthenticator(auth_file="auth-bob.json") - assert auth.auth_file == str(token_dir / "auth-bob.json") + assert auth.auth_file == os.path.realpath(str(token_dir / "auth-bob.json")) assert auth.get_access_token() == "bob-token" +def test_authenticator_rejects_relative_path_escaping_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + with pytest.raises(XAIOAuthError, match="token directory"): + XAIOAuthAuthenticator(auth_file="../secret.json") + + +def test_authenticator_rejects_absolute_path_outside_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + outside = tmp_path / "outside.json" + outside.write_text("{}") + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + with pytest.raises(XAIOAuthError, match="token directory"): + XAIOAuthAuthenticator(auth_file=str(outside)) + + +def test_authenticator_rejects_dotdot_absolute_path_outside_token_dir(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + escaped = str((token_dir / ".." / "outside.json").resolve()) + + with pytest.raises(XAIOAuthError, match="token directory"): + XAIOAuthAuthenticator(auth_file=escaped) + + +def test_oauth_auth_file_for_account_builds_path_inside_token_dir(tmp_path): + from litellm.llms.xai.oauth import oauth_auth_file_for_account + + assert oauth_auth_file_for_account("alice", str(tmp_path)) == os.path.realpath( + os.path.join(str(tmp_path), "auth-alice.json") + ) + assert oauth_auth_file_for_account("Bob_1-2", str(tmp_path)) == os.path.realpath( + os.path.join(str(tmp_path), "auth-Bob_1-2.json") + ) + + +@pytest.mark.parametrize( + "account", + ["", ".", "..", "foo/bar", "../alice", "alice.json", "foo\\bar", "alice/../bob"], +) +def test_oauth_auth_file_for_account_rejects_unsafe_names(account, tmp_path): + from litellm.llms.xai.oauth import oauth_auth_file_for_account + + with pytest.raises(ValueError, match="account"): + oauth_auth_file_for_account(account, str(tmp_path)) + + def test_authenticator_explicit_auth_file_overrides_env(tmp_path, monkeypatch): env_file = tmp_path / "env.json" env_file.write_text( @@ -75,12 +127,13 @@ def test_authenticator_explicit_auth_file_overrides_env(tmp_path, monkeypatch): ) monkeypatch.setenv("XAI_OAUTH_AUTH_FILE", str(env_file)) - auth = XAIOAuthAuthenticator(auth_file=str(explicit_file)) + auth = XAIOAuthAuthenticator(auth_file=str(explicit_file), token_dir=str(tmp_path)) assert auth.get_access_token() == "explicit-token" -def test_chat_config_uses_per_deployment_token_file(tmp_path): +def test_chat_config_uses_per_deployment_token_file(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) alice_file = tmp_path / "auth-alice.json" alice_file.write_text( json.dumps( @@ -107,7 +160,8 @@ def test_chat_config_uses_per_deployment_token_file(tmp_path): assert headers["Authorization"] == "Bearer alice-chat-token" -def test_chat_config_multi_deployment_token_isolation(tmp_path): +def test_chat_config_multi_deployment_token_isolation(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) alice = tmp_path / "auth-alice.json" alice.write_text( json.dumps( @@ -156,7 +210,8 @@ def test_chat_config_multi_deployment_token_isolation(tmp_path): assert headers_bob["Authorization"] == "Bearer bob-token" -def test_responses_config_uses_per_deployment_token_file(tmp_path): +def test_responses_config_uses_per_deployment_token_file(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) bob = tmp_path / "auth-bob.json" bob.write_text( json.dumps( @@ -180,28 +235,14 @@ def test_responses_config_uses_per_deployment_token_file(tmp_path): assert headers["Authorization"] == "Bearer bob-responses-token" -def test_proxy_cli_xai_oauth_login_with_account(monkeypatch, tmp_path): +def test_proxy_cli_xai_oauth_login_rejects_unsafe_account(monkeypatch, tmp_path): from litellm.proxy.proxy_cli import run_server - captured: dict[str, str | None] = {} - - class FakeAuthenticator: - def __init__(self, http_client=None, auth_file=None, token_dir=None): - captured["auth_file"] = auth_file - self.auth_file = auth_file or "/tmp/xai-oauth-auth.json" - - def login(self): - return {"expires_at": 1234567890} - - monkeypatch.setattr("litellm.llms.xai.oauth.XAIOAuthAuthenticator", FakeAuthenticator) monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path)) + result = CliRunner().invoke(run_server, ["xai-oauth", "login", "../alice"]) - result = CliRunner().invoke(run_server, ["xai-oauth", "login", "alice"]) - - assert result.exit_code == 0 - expected = os.path.join(str(tmp_path), "auth-alice.json") - assert captured["auth_file"] == expected - assert expected in result.output + assert result.exit_code != 0 + assert result.exception is not None def test_spending_limit_403_maps_to_rate_limit_error(): @@ -222,3 +263,25 @@ def test_spending_limit_403_maps_to_rate_limit_error(): exception_provider="XaiException", extra_information="", ) + + +def test_spending_limit_403_on_non_xai_does_not_map_to_rate_limit_error(): + class FakeHTTPError(Exception): + def __init__(self): + self.status_code = 403 + self.message = "personal-team-blocked:spending-limit" + self.response = None + super().__init__(self.message) + + with pytest.raises(litellm.APIError) as exc_info: + _map_openai_exception( + model="gpt-4", + original_exception=FakeHTTPError(), + custom_llm_provider="openai", + error_str="personal-team-blocked:spending-limit", + exception_type="APIStatusError", + exception_provider="OpenAIException", + extra_information="", + ) + assert not isinstance(exc_info.value, litellm.RateLimitError) + assert exc_info.value.status_code == 403 diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index aaf630ad29b..9dc5a191993 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3554,3 +3554,53 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: ) is True ) + + +class TestIsRequestBodySafeBlocksXaiOauthTokenFile: + """A caller must not select another local xAI OAuth account. Router kwargs + override deployment params, so ``xai_oauth_token_file`` in the request body + would authenticate as any token file readable on the proxy host. + """ + + def test_xai_oauth_token_file_is_in_banned_set(self): + from litellm.proxy.auth.auth_utils import _BANNED_REQUEST_BODY_PARAMS + + assert "xai_oauth_token_file" in set(_BANNED_REQUEST_BODY_PARAMS) + + def test_xai_oauth_token_file_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="xai_oauth_token_file"): + is_request_body_safe( + request_body={ + "model": "grok-4", + "xai_oauth_token_file": "/etc/passwd", + }, + general_settings={}, + llm_router=None, + model="grok-4", + ) + + def test_xai_oauth_token_file_under_extra_body_is_rejected(self): + with pytest.raises(ValueError, match="xai_oauth_token_file"): + is_request_body_safe( + request_body={ + "model": "grok-4", + "extra_body": {"xai_oauth_token_file": "/etc/passwd"}, + }, + general_settings={}, + llm_router=None, + model="grok-4", + ) + + def test_xai_oauth_token_file_allowed_under_proxy_wide_opt_in(self): + assert ( + is_request_body_safe( + request_body={ + "model": "grok-4", + "xai_oauth_token_file": "auth-alice.json", + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="grok-4", + ) + is True + ) diff --git a/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py b/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py index e87b206a40a..6c621b03cba 100644 --- a/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py +++ b/tests/test_litellm/proxy/auth/test_banned_params_extra_body.py @@ -28,6 +28,7 @@ from litellm.proxy.auth.auth_utils import is_request_body_safe # noqa: E402 "base_url", "vertex_credentials", "azure_ad_token", + "xai_oauth_token_file", ], ) def test_banned_param_under_extra_body_is_rejected(banned_param):