From 4ef7d0815b73f208c9b98bbcbc79c156b313c242 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 22 Jun 2026 15:28:43 -0700 Subject: [PATCH] fix(bedrock): only expand config-sourced AWS credential references (#30867) AWS auth parameters in the Bedrock and SageMaker path could be expanded against the process environment when credentials were built. Config-sourced references are already expanded at load time, so restrict expansion to that path: a reference still present at request time is treated as caller-supplied input and is left as-is, and the web-identity helper rejects environment-variable references before resolving the token. Also rework the ambient AWS_* fallback as a single pass that pairs each value with its own env-var name, fixing a latent index misalignment that left AWS_EXTERNAL_ID unresolved. Adds regression tests covering the resolution behavior. --- litellm/llms/bedrock/base_aws_llm.py | 62 +++++---- .../llms/bedrock/test_base_aws_llm.py | 128 ++++++++++++++++++ 2 files changed, 162 insertions(+), 28 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 2c9ea187912..c31462a735b 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -10,7 +10,6 @@ from typing import ( Callable, ClassVar, Dict, - List, Literal, Optional, Tuple, @@ -210,32 +209,11 @@ class BaseAWSLLM: """ Return a boto3.Credentials object """ - ## CHECK IS 'os.environ/' passed in - params_to_check: List[Optional[str]] = [ - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - aws_region_name, - aws_session_name, - aws_profile_name, - aws_role_name, - aws_web_identity_token, - aws_sts_endpoint, - aws_external_id, - ] - - # Iterate over parameters and update if needed - for i, param in enumerate(params_to_check): - if param and param.startswith("os.environ/"): - _v = get_secret(param) - if _v is not None and isinstance(_v, str): - params_to_check[i] = _v - elif param is None: # check if uppercase value in env - key = self.aws_authentication_params[i] - if key.upper() in os.environ: - params_to_check[i] = os.getenv(key.upper()) - - # Assign updated values back to parameters + # Only config-sourced credentials are expanded against the environment. + # os.environ/ references in the model config are resolved at load time, + # so any reference still present at this point is caller-supplied input and is + # left as-is rather than expanded into a process environment variable. Each + # unset param falls back to its matching fixed AWS_* ambient env var. ( aws_access_key_id, aws_secret_access_key, @@ -247,7 +225,21 @@ class BaseAWSLLM: aws_web_identity_token, aws_sts_endpoint, aws_external_id, - ) = params_to_check + ) = tuple( + value if value is not None else os.getenv(env_var) + for value, env_var in ( + (aws_access_key_id, "AWS_ACCESS_KEY_ID"), + (aws_secret_access_key, "AWS_SECRET_ACCESS_KEY"), + (aws_session_token, "AWS_SESSION_TOKEN"), + (aws_region_name, "AWS_REGION_NAME"), + (aws_session_name, "AWS_SESSION_NAME"), + (aws_profile_name, "AWS_PROFILE_NAME"), + (aws_role_name, "AWS_ROLE_NAME"), + (aws_web_identity_token, "AWS_WEB_IDENTITY_TOKEN"), + (aws_sts_endpoint, "AWS_STS_ENDPOINT"), + (aws_external_id, "AWS_EXTERNAL_ID"), + ) + ) verbose_logger.debug( "in get credentials\n" @@ -845,6 +837,20 @@ class BaseAWSLLM: f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}" ) + # get_secret() expands environment-variable references (an os.environ/ + # prefix, or a bare name matching an environment variable). Config-sourced + # references are expanded at load time, so such a reference reaching here is + # caller-supplied input; reject it rather than expanding a process-environment + # value for use as the token. + if ( + aws_web_identity_token.startswith("os.environ/") + or aws_web_identity_token in os.environ + ): + raise AwsAuthError( + message="Invalid web identity token reference.", + status_code=400, + ) + oidc_token = get_secret(aws_web_identity_token) if oidc_token is None: diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 3f91f6ac26e..2d5242d510f 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -163,6 +163,134 @@ def test_aws_profile_path_not_cached_in_iam_cache(): assert mock_profile.call_count == 2 +def test_get_credentials_does_not_expand_request_env_reference(): + """ + A parameter of the form os.environ/ reaching get_credentials is left as-is + rather than expanded against the process environment, so the downstream auth + helper only ever receives the literal value. + """ + env = _os_environ_without_aws_keys() + env["SERVER_ONLY_VALUE"] = "config-managed-value" + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch.object( + base, + "_auth_with_aws_profile", + return_value=(Credentials("ak", "sk", None), None), + ) as mock_profile: + base.get_credentials(aws_profile_name="os.environ/SERVER_ONLY_VALUE") + + assert mock_profile.call_args.args[0] == "os.environ/SERVER_ONLY_VALUE" + assert "config-managed-value" not in str(mock_profile.call_args) + + +def test_get_credentials_falls_back_to_ambient_aws_profile_name_env(): + """ + The fixed AWS_* ambient fallback keeps working: an unset aws_profile_name + resolves from the AWS_PROFILE_NAME environment variable. + """ + env = _os_environ_without_aws_keys() + env["AWS_PROFILE_NAME"] = "ambient-profile" + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch.object( + base, + "_auth_with_aws_profile", + return_value=(Credentials("ak", "sk", None), None), + ) as mock_profile: + base.get_credentials(aws_profile_name=None) + + assert mock_profile.call_args.args[0] == "ambient-profile" + + +def test_get_credentials_ambient_fallback_resolves_aws_external_id(): + """ + Each unset param falls back to its own AWS_* env var. Regression for an index + misalignment between the value list and the env-name list, which left + AWS_EXTERNAL_ID unresolved. + """ + env = _os_environ_without_aws_keys() + env["AWS_EXTERNAL_ID"] = "ext-from-env" + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch.object( + base, + "_auth_with_aws_role", + return_value=(Credentials("ak", "sk", "tok"), None), + ) as mock_role: + base.get_credentials( + aws_role_name="arn:aws:iam::123456789012:role/x", + aws_session_name="s", + ) + + assert mock_role.call_args.kwargs["aws_external_id"] == "ext-from-env" + + +def _capturing_sts_client(captured: Dict[str, Any]) -> MagicMock: + sts = MagicMock() + + def _assume(**params): + captured["WebIdentityToken"] = params.get("WebIdentityToken") + return { + "Credentials": { + "AccessKeyId": "AKIA", + "SecretAccessKey": "sk", + "SessionToken": "tok", + }, + "PackedPolicySize": 10, + } + + sts.assume_role_with_web_identity.side_effect = _assume + return sts + + +@pytest.mark.parametrize( + "token_ref", + ["os.environ/SERVER_ONLY_VALUE", "SERVER_ONLY_VALUE"], + ids=["os_environ_prefix", "bare_env_name"], +) +def test_web_identity_token_env_reference_not_expanded(token_ref): + """ + A web-identity token that is an environment-variable reference (an os.environ/ + prefix, or a bare name matching an env var) is rejected rather than expanded, so + the process-environment value is never used as the token. + """ + env = _os_environ_without_aws_keys() + env["SERVER_ONLY_VALUE"] = "server-only-value" + captured: Dict[str, Any] = {} + base = BaseAWSLLM() + with patch.dict(os.environ, env, clear=True), patch( + "boto3.client", side_effect=lambda *a, **k: _capturing_sts_client(captured) + ), patch("boto3.Session", return_value=MagicMock()): + with pytest.raises(AwsAuthError): + base.get_credentials( + aws_web_identity_token=token_ref, + aws_role_name="arn:aws:iam::123456789012:role/x", + aws_session_name="s", + aws_sts_endpoint="https://custom-sts.example", + ) + + assert "server-only-value" not in str(captured) + + +def test_web_identity_token_oidc_reference_still_resolved(): + """ + The env-reference guard does not over-reject: an oidc/ reference still flows to + get_secret (mocked to None here), surfacing the existing 401 rather than the 400 + used for rejected env-var references. + """ + base = BaseAWSLLM() + env = _os_environ_without_aws_keys() + with patch.dict(os.environ, env, clear=True), patch( + "litellm.llms.bedrock.base_aws_llm.get_secret", return_value=None + ): + with pytest.raises(AwsAuthError) as exc: + base.get_credentials( + aws_web_identity_token="oidc/circleci/", + aws_role_name="arn:aws:iam::123456789012:role/x", + aws_session_name="s", + ) + + assert exc.value.status_code == 401 + + def test_web_identity_path_not_cached_in_iam_cache(): base = BaseAWSLLM() with patch.object(