This commit is contained in:
Vineeth Sai Varikuntla 2026-09-28 19:26:13 -04:00 • committed by GitHub
commit f955fd8514
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 111 additions and 3 deletions

View file

@ -795,6 +795,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
return credentials, aws_region_name
@staticmethod
def _request_api_key_is_a_bedrock_credential(request_data: Mapping[str, Any] | None) -> bool:
"""Whether ``request_data["api_key"]`` can be a Bedrock credential at all.
That field holds the key for the **LLM** call, not for this guardrail's own
ApplyGuardrail request. Credential-override routing
(``_apply_credential_overrides_from_model_config``) writes the caller's BYOK
provider key into it, so on a non-Bedrock deployment it is someone else's
secret: signing with it fails, and it would be sent to AWS as a bearer token.
Only a request that is itself Bedrock-routed can carry a key that is also
valid here. When the provider cannot be determined the old behaviour is kept,
so a clientside Bedrock key keeps working.
"""
if not request_data:
return False
litellm_params: Final = request_data.get("litellm_params")
nested_provider: Final = (
litellm_params.get("custom_llm_provider") if isinstance(litellm_params, Mapping) else None
)
provider: Final = request_data.get("custom_llm_provider") or nested_provider
if provider:
return provider == "bedrock"
model: Final[str] = request_data.get("model") or ""
# A bare name gives nothing to route on, so leave it to the caller's config.
return "/" not in model or model.startswith("bedrock/")
def _prepare_request(
self,
credentials,
@ -937,7 +964,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
if key not in _BEDROCK_DYNAMIC_BODY_DENYLIST
}
)
if request_data.get("api_key") is not None:
if request_data.get("api_key") is not None and self._request_api_key_is_a_bedrock_credential(request_data):
api_key = request_data["api_key"]
event_type: Final = (
@ -1914,7 +1941,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does.
return BedrockGuardrailResponse()
api_key: Final[str | None] = request_data.get("api_key") if request_data else None
api_key: Final[str | None] = (
request_data.get("api_key")
if request_data and self._request_api_key_is_a_bedrock_credential(request_data)
else None
)
credentials, aws_region_name = await run_aws_signing(
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
)

View file

@ -12,7 +12,6 @@ import httpx
import pytest
from fastapi import HTTPException
import litellm
from litellm.caching.caching import DualCache
from litellm.exceptions import ModifyResponseException
@ -6068,6 +6067,83 @@ async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
assert response["action"] == "NONE"
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
class TestRequestApiKeyIsNotAlwaysABedrockCredential:
"""request_data["api_key"] is the LLM call's key, not this guardrail's.
Credential-override routing writes the caller's BYOK provider key into that
field, so on a non-Bedrock deployment using it as an ApplyGuardrail bearer
token both fails signing with a 403 and sends the caller's third-party secret
to AWS (issue #37872).
"""
@staticmethod
def _guardrail():
return BedrockGuardrail(guardrailIdentifier="gid", guardrailVersion="DRAFT")
@pytest.mark.parametrize(
"request_data, expected",
[
# BYOK credential override on a non-Bedrock deployment: never a Bedrock key.
({"model": "nvidia_nim/meta/llama-3.1-8b", "api_key": "nvapi-xxx"}, False),
({"model": "openrouter/anthropic/claude-3", "api_key": "sk-or-xxx"}, False),
({"custom_llm_provider": "nvidia_nim", "model": "meta/llama", "api_key": "nvapi-xxx"}, False),
# Bedrock-routed requests keep the existing clientside-key behaviour.
({"model": "bedrock/anthropic.claude-3", "api_key": "bedrock-key"}, True),
({"custom_llm_provider": "bedrock", "model": "anthropic.claude-3", "api_key": "k"}, True),
({"litellm_params": {"custom_llm_provider": "bedrock"}, "model": "x/y", "api_key": "k"}, True),
# Nothing to route on: unchanged, so a clientside Bedrock key still works.
({"model": "claude-sonnet", "api_key": "k"}, True),
({"api_key": "k"}, True),
# No request data at all.
(None, False),
({}, False),
],
)
def test_only_a_bedrock_routed_request_yields_a_bedrock_credential(self, request_data, expected):
assert self._guardrail()._request_api_key_is_a_bedrock_credential(request_data) is expected
def test_byok_key_is_not_sent_to_aws_as_a_bearer_token(self):
"""End to end through _prepare_request: the NIM key must not become Authorization."""
guardrail = self._guardrail()
credentials = MagicMock()
credentials.access_key, credentials.secret_key, credentials.token = "ak", "sk", None
request_data = {"model": "nvidia_nim/meta/llama-3.1-8b", "api_key": "nvapi-SECRET"}
api_key = request_data["api_key"] if guardrail._request_api_key_is_a_bedrock_credential(request_data) else None
with patch.dict(os.environ, {}, clear=False):
os.environ.pop("AWS_BEARER_TOKEN_BEDROCK", None)
prepped = guardrail._prepare_request(
credentials=credentials,
data={"source": "INPUT", "content": [{"text": {"text": "hi"}}]},
optional_params={},
aws_region_name="us-east-1",
api_key=api_key,
)
assert "nvapi-SECRET" not in str(dict(prepped.headers))
assert prepped.headers.get("Authorization", "") != "Bearer nvapi-SECRET"
def test_an_explicit_bedrock_key_still_becomes_the_bearer_token(self):
"""Control: _prepare_request is untouched, so it passes with or without the fix.
Written against _prepare_request directly rather than through the new helper,
so it proves the capability the api_key field exists for is not regressed.
"""
guardrail = self._guardrail()
credentials = MagicMock()
credentials.access_key, credentials.secret_key, credentials.token = "ak", "sk", None
api_key = "bedrock-KEY"
prepped = guardrail._prepare_request(
credentials=credentials,
data={"source": "INPUT", "content": [{"text": {"text": "hi"}}]},
optional_params={},
aws_region_name="us-east-1",
api_key=api_key,
)
assert prepped.headers.get("Authorization") == "Bearer bedrock-KEY"
@pytest.mark.asyncio
async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch):
@ -6100,3 +6176,4 @@ async def test_apply_guardrail_signs_off_the_event_loop(monkeypatch):
assert response["action"] == "NONE"
assert probe.served_during_refresh is True
import os