From 0b930ea919be35517a7862d1b336023bf5599ba1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 15 Jul 2026 17:56:39 -0700 Subject: [PATCH] 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. --- litellm/proxy/client/cli/commands/up.py | 22 +++--- .../proxy/client/cli/test_auth_commands.py | 49 ++++++++++--- .../proxy/client/cli/test_up_commands.py | 71 +++++++++++++++++-- 3 files changed, 116 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/client/cli/commands/up.py b/litellm/proxy/client/cli/commands/up.py index edf5a7acc05..61b649b7415 100644 --- a/litellm/proxy/client/cli/commands/up.py +++ b/litellm/proxy/client/cli/commands/up.py @@ -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( diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 6be43c9da44..556e81e9e42 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -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 diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 49508c22f47..a2986027b98 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -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"),