mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(cli): bind lite up's apiKeyHelper to the proxy it was started against
_ensure_fresh_login only checked token freshness, not which proxy the cached token belonged to, and resolve_api_key_helper built a bare `lite auth print-token` command with no --base-url. A user logged into proxy A who ran `up --base-url proxy-b` (or LITELLM_PROXY_URL=proxy-b) would silently get proxy A's real token wired into Claude Code's apiKeyHelper; since apiKeyHelper is invoked bare, print-token's existing origin check never engaged, so proxy B -- attacker-controlled or not -- received every subsequent request's Authorization header carrying proxy A's credential. _ensure_fresh_login now requires the cached token's base_url to match before treating it as usable, forcing a fresh login for the selected proxy otherwise. resolve_api_key_helper now takes that base_url and threads it through as an explicit --base-url, so print-token's existing (but previously unreachable in the apiKeyHelper flow) base_url_explicit check actually enforces the match at request time too.
This commit is contained in:
parent
eb240d864c
commit
0b930ea919
3 changed files with 116 additions and 26 deletions
|
|
@ -127,11 +127,16 @@ def restore_claude_settings(settings_path: Path | None = None, backup_path: Path
|
|||
return record
|
||||
|
||||
|
||||
def resolve_api_key_helper() -> str:
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from.
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
Code would otherwise use) makes `print-token` enforce that the cached
|
||||
token was actually issued for this proxy -- without it, a token minted
|
||||
for a different, previously-logged-into proxy would be handed to
|
||||
whichever server `up` currently points at.
|
||||
"""
|
||||
lite_path = shutil.which("lite")
|
||||
if lite_path is None:
|
||||
|
|
@ -139,24 +144,25 @@ def resolve_api_key_helper() -> str:
|
|||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs "
|
||||
"an absolute path to it, so `lite up` cannot continue."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} auth print-token"
|
||||
return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}"
|
||||
|
||||
|
||||
def _ensure_fresh_login(ctx: click.Context) -> None:
|
||||
base_url = ctx.obj["base_url"].rstrip("/")
|
||||
token_data = load_token()
|
||||
if token_data and is_cli_token_fresh(token_data):
|
||||
if token_data and token_data.get("base_url") == base_url and is_cli_token_fresh(token_data):
|
||||
return
|
||||
|
||||
if not sys.stdin.isatty():
|
||||
raise UpError(
|
||||
"No fresh LiteLLM login found. Run `lite login` first (apiKeyHelper "
|
||||
"No fresh LiteLLM login found for this proxy. Run `lite login` first (apiKeyHelper "
|
||||
"reads this token on every Claude Code request)."
|
||||
)
|
||||
|
||||
click.echo("No fresh LiteLLM login found; starting login...")
|
||||
click.echo("No fresh LiteLLM login found for this proxy; starting login...")
|
||||
ctx.invoke(login)
|
||||
token_data = load_token()
|
||||
if not token_data or not is_cli_token_fresh(token_data):
|
||||
if not token_data or token_data.get("base_url") != base_url or not is_cli_token_fresh(token_data):
|
||||
raise UpError("Login did not produce a usable token; cannot start `lite up`.")
|
||||
|
||||
|
||||
|
|
@ -195,7 +201,7 @@ def up(ctx: click.Context) -> None:
|
|||
"running (or crashed without cleanup). Run `lite down` first."
|
||||
)
|
||||
|
||||
api_key_helper = resolve_api_key_helper()
|
||||
api_key_helper = resolve_api_key_helper(base_url)
|
||||
original_existed = CLAUDE_SETTINGS_PATH.exists()
|
||||
original_settings = load_json_or_empty(CLAUDE_SETTINGS_PATH)
|
||||
write_backup(
|
||||
|
|
|
|||
|
|
@ -664,14 +664,18 @@ class TestPrintTokenCommand:
|
|||
verbatim as the bearer token, so any diagnostic text on stdout would
|
||||
corrupt authentication.
|
||||
|
||||
apiKeyHelper is configured as a bare command (managed-settings.json sets
|
||||
just `"apiKeyHelper": "lite auth print-token"`, no --base-url flag) --
|
||||
so in the common case ctx.obj has no explicit base_url at all, and the
|
||||
command must resolve the server from whatever `lite login` stored in
|
||||
token.json, not from a CLI default. `--base-url`/`LITELLM_PROXY_URL`
|
||||
only matters when a caller explicitly overrides it (tracked via
|
||||
ctx.obj["base_url_explicit"], set by the `cli` group from
|
||||
click's ParameterSource).
|
||||
`lite up` now writes `apiKeyHelper` with an explicit `--base-url` bound
|
||||
to whatever proxy it was pointed at (resolve_api_key_helper), so
|
||||
print-token enforces that the cached token was actually issued for that
|
||||
server -- a token minted for a different, previously-logged-into proxy
|
||||
must never be handed to whichever server the helper is invoked for.
|
||||
Settings patched by an older `lite up`, or a manually-configured
|
||||
apiKeyHelper, can still invoke this bare (no --base-url at all); that
|
||||
case falls back to trusting whatever `lite login` stored in token.json,
|
||||
since there is no explicit target to check it against. `--base-url`/
|
||||
`LITELLM_PROXY_URL` only enforces the match when a caller explicitly
|
||||
passes it (tracked via ctx.obj["base_url_explicit"], set by the `cli`
|
||||
group from click's ParameterSource).
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
|
|
@ -685,8 +689,9 @@ class TestPrintTokenCommand:
|
|||
assert "Not authenticated" in result.output
|
||||
|
||||
def test_bare_invocation_resolves_server_from_stored_token(self):
|
||||
"""The apiKeyHelper's real invocation shape: no --base-url given at
|
||||
all. Must use token.json's own base_url, not a hardcoded default."""
|
||||
"""The legacy/manual invocation shape: no --base-url given at all
|
||||
(e.g. settings patched before resolve_api_key_helper started binding
|
||||
one). Must use token.json's own base_url, not a hardcoded default."""
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
|
|
@ -706,7 +711,10 @@ class TestPrintTokenCommand:
|
|||
|
||||
def test_explicit_base_url_mismatch_fails_cleanly(self):
|
||||
"""When the caller *does* explicitly pass --base-url, a token issued
|
||||
for a different server must never be printed."""
|
||||
for a different server must never be printed. This is the exact
|
||||
scenario `lite up`'s own bound --base-url now guards against: a
|
||||
token minted for proxy A must not reach a helper invocation aimed
|
||||
at proxy B, even though the token itself is otherwise fresh."""
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
return_value={
|
||||
|
|
@ -723,6 +731,25 @@ class TestPrintTokenCommand:
|
|||
assert result.exit_code != 0
|
||||
assert "sk-should-not-print" not in result.output
|
||||
|
||||
def test_explicit_base_url_match_prints_token(self):
|
||||
"""`lite up`'s own bound invocation shape: --base-url matching the token's origin
|
||||
must succeed exactly like the bare/legacy invocation does."""
|
||||
with patch(
|
||||
"litellm.proxy.client.cli.commands.auth.load_token",
|
||||
return_value={
|
||||
"base_url": "http://localhost:4000",
|
||||
"key": "sk-matches",
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
):
|
||||
result = self.runner.invoke(
|
||||
print_token,
|
||||
obj={"base_url": "http://localhost:4000", "base_url_explicit": True},
|
||||
)
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert result.output.strip() == "sk-matches"
|
||||
|
||||
def test_fresh_cached_key_printed_without_network_call(self):
|
||||
"""A recently-issued key should be printed straight from cache -- no
|
||||
refresh call on every single invocation (apiKeyHelper gets called
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import stat
|
|||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import click
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
|
|
@ -12,6 +13,7 @@ from litellm.proxy.client.cli.commands.agents import AgentRunError
|
|||
from litellm.proxy.client.cli.commands.up import (
|
||||
BackupRecord,
|
||||
UpError,
|
||||
_ensure_fresh_login,
|
||||
down,
|
||||
load_json_or_empty,
|
||||
merge_claude_settings,
|
||||
|
|
@ -156,15 +158,70 @@ class TestBackupRoundTrip:
|
|||
|
||||
|
||||
class TestResolveApiKeyHelper:
|
||||
def test_returns_helper_command_when_lite_found(self, monkeypatch):
|
||||
def test_returns_helper_command_bound_to_the_selected_proxy(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite")
|
||||
helper = resolve_api_key_helper()
|
||||
assert helper == "/usr/local/bin/lite auth print-token"
|
||||
helper = resolve_api_key_helper("http://localhost:4000")
|
||||
assert helper == "/usr/local/bin/lite auth print-token --base-url http://localhost:4000"
|
||||
|
||||
def test_quotes_a_base_url_containing_shell_metacharacters(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite")
|
||||
helper = resolve_api_key_helper("http://example.com/path; rm -rf /")
|
||||
assert helper == "/usr/local/bin/lite auth print-token --base-url 'http://example.com/path; rm -rf /'"
|
||||
|
||||
def test_raises_when_lite_not_on_path(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: None)
|
||||
with pytest.raises(UpError, match="Could not find `lite`"):
|
||||
resolve_api_key_helper()
|
||||
resolve_api_key_helper("http://localhost:4000")
|
||||
|
||||
|
||||
def _make_ctx(base_url):
|
||||
return click.Context(click.Command("test"), obj={"base_url": base_url})
|
||||
|
||||
|
||||
class TestEnsureFreshLogin:
|
||||
"""A token that is fresh but was issued for a *different* proxy must not be trusted: without
|
||||
this check, a user logged into proxy A who runs `up --base-url proxy-b` would silently get an
|
||||
apiKeyHelper wired up around proxy A's real token, which print-token would then hand to proxy B."""
|
||||
|
||||
def test_reuses_a_fresh_token_issued_for_the_same_proxy(self, monkeypatch):
|
||||
monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"})
|
||||
monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True)
|
||||
login_calls = []
|
||||
monkeypatch.setattr(up_module, "login", lambda ctx: login_calls.append(ctx))
|
||||
|
||||
_ensure_fresh_login(_make_ctx("http://proxy-a:4000"))
|
||||
|
||||
assert login_calls == []
|
||||
|
||||
def test_forces_a_fresh_login_when_the_cached_token_is_for_a_different_proxy(self, monkeypatch):
|
||||
monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: True)
|
||||
tokens = iter(
|
||||
[
|
||||
{"key": "sk-a", "base_url": "http://proxy-a:4000"},
|
||||
{"key": "sk-b", "base_url": "http://proxy-b:4000"},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(up_module, "load_token", lambda: next(tokens))
|
||||
monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True)
|
||||
login_calls = []
|
||||
|
||||
@click.pass_context
|
||||
def fake_login(ctx):
|
||||
login_calls.append(ctx.obj["base_url"])
|
||||
|
||||
monkeypatch.setattr(up_module, "login", fake_login)
|
||||
|
||||
_ensure_fresh_login(_make_ctx("http://proxy-b:4000"))
|
||||
|
||||
assert login_calls == ["http://proxy-b:4000"]
|
||||
|
||||
def test_fails_cleanly_non_interactively_when_only_a_different_proxys_token_is_cached(self, monkeypatch):
|
||||
monkeypatch.setattr(up_module.sys.stdin, "isatty", lambda: False)
|
||||
monkeypatch.setattr(up_module, "load_token", lambda: {"key": "sk-a", "base_url": "http://proxy-a:4000"})
|
||||
monkeypatch.setattr(up_module, "is_cli_token_fresh", lambda token_data: True)
|
||||
|
||||
with pytest.raises(UpError, match="lite login"):
|
||||
_ensure_fresh_login(_make_ctx("http://proxy-b:4000"))
|
||||
|
||||
|
||||
class TestUpCommand:
|
||||
|
|
@ -177,7 +234,7 @@ class TestUpCommand:
|
|||
backup_path.write_text(json.dumps(existing_backup))
|
||||
|
||||
with (
|
||||
patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh"}),
|
||||
patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}),
|
||||
patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True),
|
||||
patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"),
|
||||
patch(f"{UP_MODULE}.verify_proxy_key"),
|
||||
|
|
@ -204,7 +261,7 @@ class TestUpCommand:
|
|||
_patch_paths(monkeypatch, tmp_path)
|
||||
|
||||
with (
|
||||
patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh"}),
|
||||
patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}),
|
||||
patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True),
|
||||
patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"),
|
||||
patch(
|
||||
|
|
@ -230,7 +287,7 @@ class TestUpCommand:
|
|||
return True
|
||||
|
||||
with (
|
||||
patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh"}),
|
||||
patch(f"{UP_MODULE}.load_token", return_value={"key": "sk-fresh", "base_url": "http://localhost:4000"}),
|
||||
patch(f"{UP_MODULE}.is_cli_token_fresh", return_value=True),
|
||||
patch(f"{UP_MODULE}.resolve_api_key", return_value="sk-fresh"),
|
||||
patch(f"{UP_MODULE}.verify_proxy_key"),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue