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:
Krrish Dholakia 2026-07-15 17:56:39 -07:00
parent eb240d864c
commit 0b930ea919
3 changed files with 116 additions and 26 deletions

View file

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

View file

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

View file

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