scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(xai): treat oauth token file as operator-only config

Browse source
This commit is contained in:
hx 2026-09-09 15:14:46 +08:00
parent dfab8e2969
commit 53b12558fe
7 changed files with 198 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -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:

38
litellm/llms/xai/oauth.py
View file

@ -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:

5
litellm/proxy/auth/auth_utils.py
View file

@ -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

7
litellm/proxy/proxy_cli.py
View file

@ -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()
)

115
tests/test_litellm/llms/xai/test_xai_oauth_multi_account.py
View file

@ -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

50
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -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
)

1
tests/test_litellm/proxy/auth/test_banned_params_extra_body.py
View file

@ -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):