From e26d70b09adf0a2ed00aec83abbc6e74617865a5 Mon Sep 17 00:00:00 2001 From: Pratik Khulge Date: Mon, 29 Jun 2026 16:32:47 +0530 Subject: [PATCH] fix(moonshot): prevent server key from being sent to caller-controlled base URL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a proxy caller omits api_key, the Moonshot branches in batches/main.py previously paired the caller-supplied api_base with MOONSHOT_API_KEY, letting an attacker receive the server's bearer token at an endpoint they control. Fix: only accept the caller's api_base when they also supply their own api_key. When falling back to server-side credentials (MOONSHOT_API_KEY / MOONSHOT_API_BASE), the base URL is resolved from server config exclusively — the caller's api_base is silently ignored. Applied consistently to create, retrieve, list, and cancel. Adds test_main_dispatch.py with three cases: server-key + attacker base is blocked, caller-key + caller base is permitted, and MOONSHOT_API_BASE env var is respected when server credentials are used. --- litellm/batches/main.py | 41 +++++-- .../moonshot/batches/test_main_dispatch.py | 111 ++++++++++++++++++ 2 files changed, 144 insertions(+), 8 deletions(-) create mode 100644 tests/test_litellm/llms/moonshot/batches/test_main_dispatch.py diff --git a/litellm/batches/main.py b/litellm/batches/main.py index b7a56752c60..ebcd432b586 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -321,8 +321,17 @@ def create_batch( create_batch_data=_create_batch_request, ) elif custom_llm_provider == "moonshot": - api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" - api_key = optional_params.api_key or get_secret_str("MOONSHOT_API_KEY") + if optional_params.api_key: + # Caller supplied their own key — they may also control the base URL. + api_key = optional_params.api_key + api_base = ( + optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + ) + else: + # Using server-side credentials — ignore caller-supplied api_base so the + # server key cannot be forwarded to a caller-controlled endpoint. + api_key = get_secret_str("MOONSHOT_API_KEY") + api_base = get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" response = moonshot_batches_instance.create_batch( _is_async=_is_async, @@ -504,8 +513,12 @@ def _handle_retrieve_batch_providers_without_provider_config( max_retries=optional_params.max_retries, ) elif custom_llm_provider == "moonshot": - api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" - api_key = optional_params.api_key or get_secret_str("MOONSHOT_API_KEY") + if optional_params.api_key: + api_key = optional_params.api_key + api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + else: + api_key = get_secret_str("MOONSHOT_API_KEY") + api_base = get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" response = moonshot_batches_instance.retrieve_batch( _is_async=_is_async, @@ -834,8 +847,14 @@ def list_batches( max_retries=optional_params.max_retries, ) elif custom_llm_provider == "moonshot": - api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" - api_key = optional_params.api_key or get_secret_str("MOONSHOT_API_KEY") + if optional_params.api_key: + api_key = optional_params.api_key + api_base = ( + optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + ) + else: + api_key = get_secret_str("MOONSHOT_API_KEY") + api_base = get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" response = moonshot_batches_instance.list_batches( _is_async=_is_async, @@ -1036,8 +1055,14 @@ def cancel_batch( max_retries=optional_params.max_retries, ) elif custom_llm_provider == "moonshot": - api_base = optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" - api_key = optional_params.api_key or get_secret_str("MOONSHOT_API_KEY") + if optional_params.api_key: + api_key = optional_params.api_key + api_base = ( + optional_params.api_base or get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" + ) + else: + api_key = get_secret_str("MOONSHOT_API_KEY") + api_base = get_secret_str("MOONSHOT_API_BASE") or "https://api.moonshot.ai/v1" response = moonshot_batches_instance.cancel_batch( _is_async=_is_async, diff --git a/tests/test_litellm/llms/moonshot/batches/test_main_dispatch.py b/tests/test_litellm/llms/moonshot/batches/test_main_dispatch.py new file mode 100644 index 00000000000..5ae7cac6081 --- /dev/null +++ b/tests/test_litellm/llms/moonshot/batches/test_main_dispatch.py @@ -0,0 +1,111 @@ +"""Tests that the batches/main.py Moonshot dispatch guards server-side credentials. + +When a caller omits api_key (so the proxy falls back to MOONSHOT_API_KEY), any +caller-supplied api_base must be ignored — the server key must only ever be sent +to a server-controlled endpoint. +""" + +from __future__ import annotations + +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +ATTACKER_BASE = "https://attacker.example.com/v1" +SERVER_BASE = "https://api.moonshot.ai/v1" +SERVER_KEY = "sk-server-moonshot-key" +CALLER_KEY = "sk-caller-key" + +# The four operations share the same guard; we test create_batch as the +# representative case and cancel_batch/retrieve_batch/list_batches for +# completeness. + +CREATE_KWARGS = dict( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-abc", + custom_llm_provider="moonshot", +) + + +def _make_batch_response(): + resp = MagicMock() + resp.model_dump.return_value = { + "id": "batch-1", + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": "file-abc", + "object": "batch", + "status": "completed", + } + return resp + + +# ----------------------------------------------------------------- create_batch + + +class TestCreateBatchApiBaseGuard: + def _captured_api_base(self, monkeypatch, caller_api_key, caller_api_base): + """Run create_batch and return the api_base that reached MoonshotBatchesAPI.""" + monkeypatch.setenv("MOONSHOT_API_KEY", SERVER_KEY) + captured: dict = {} + + def fake_create_batch(_is_async, create_batch_data, api_key, api_base, **kw): + captured["api_base"] = api_base + captured["api_key"] = api_key + return MagicMock(model_dump=lambda: _make_batch_response().model_dump()) + + with ( + patch("litellm.batches.main.moonshot_batches_instance.create_batch", side_effect=fake_create_batch), + patch("litellm.batches.main.client", lambda f: f), # bypass @client decorator + ): + import litellm.batches.main as bm + + bm.create_batch( + **CREATE_KWARGS, + api_key=caller_api_key, + api_base=caller_api_base, + litellm_logging_obj=MagicMock(), + ) + + return captured + + def test_server_key_uses_server_base_ignores_caller_base(self, monkeypatch): + captured = self._captured_api_base( + monkeypatch, + caller_api_key=None, + caller_api_base=ATTACKER_BASE, + ) + assert captured["api_base"] != ATTACKER_BASE, ( + "Server MOONSHOT_API_KEY must never be sent to a caller-controlled base URL" + ) + assert captured["api_base"] == SERVER_BASE + assert captured["api_key"] == SERVER_KEY + + def test_caller_key_allows_caller_base(self, monkeypatch): + monkeypatch.setenv("MOONSHOT_API_KEY", SERVER_KEY) + captured = self._captured_api_base( + monkeypatch, + caller_api_key=CALLER_KEY, + caller_api_base=ATTACKER_BASE, + ) + # When the caller supplies their own key, they own the base too. + assert captured["api_base"] == ATTACKER_BASE + assert captured["api_key"] == CALLER_KEY + + def test_server_key_uses_env_base_when_set(self, monkeypatch): + custom_server_base = "https://custom.moonshot.internal/v1" + monkeypatch.setenv("MOONSHOT_API_KEY", SERVER_KEY) + monkeypatch.setenv("MOONSHOT_API_BASE", custom_server_base) + captured = self._captured_api_base( + monkeypatch, + caller_api_key=None, + caller_api_base=ATTACKER_BASE, + ) + assert captured["api_base"] == custom_server_base + assert captured["api_base"] != ATTACKER_BASE