fix(proxy): preserve AWS session tokens for RDS IAM auth (#46559)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/auth/rds_iam_token.py

@@ -24,6 +24,7 @@ def init_rds_client(
     aws_role_name: str | None = None,
     aws_web_identity_token: str | None = None,
     timeout: float | httpx.Timeout | None = None,
+    aws_session_token: str | None = None,
 ):
     from litellm.secret_managers.main import get_secret
 
@@ -108,6 +109,7 @@ def init_rds_client(
             "sts",
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
         )
 
         sts_response = sts_client.assume_role(RoleArn=aws_role_name, RoleSessionName=aws_session_name)
@@ -128,6 +130,7 @@ def init_rds_client(
             service_name="rds",
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
             region_name=region_name,
             config=config,
         )
@@ -235,6 +238,7 @@ def generate_iam_auth_token(
             aws_region_name=os.getenv("AWS_REGION_NAME"),
             aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
             aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
+            aws_session_token=os.getenv("AWS_SESSION_TOKEN"),
             aws_session_name=os.getenv("AWS_SESSION_NAME"),
             aws_profile_name=os.getenv("AWS_PROFILE_NAME"),
             aws_role_name=os.getenv("AWS_ROLE_NAME", os.getenv("AWS_ROLE_ARN")),

tests/unit/proxy/auth/test_rds_iam_token.py

@@ -1,10 +1,13 @@
+from io import BytesIO
 from pathlib import Path
 from typing import Final
 from urllib.parse import parse_qs, unquote, urlsplit
 
 import boto3
 import pytest
+from botocore.awsrequest import AWSPreparedRequest, AWSResponse
 from botocore.exceptions import ClientError
+from urllib3.response import HTTPResponse
 
 from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token, init_rds_client
 
@@ -168,3 +171,75 @@ def test_init_rds_client_surfaces_the_sts_refusal_when_there_is_no_distinct_irsa
 
     with pytest.raises(ClientError, match="AssumeRoleWithWebIdentity"):
         _init(monkeypatch, web_identity_token_file, roles_trusting_oidc=frozenset())
+
+
+@pytest.mark.parametrize("assume_role", [False, True])
+@pytest.mark.parametrize("session_token", [None, "temporary/session+token=="])
+def test_environment_credentials_retain_session_token_when_signing(
+    assume_role: bool, session_token: str | None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    for name in (
+        "AWS_PROFILE",
+        "AWS_DEFAULT_PROFILE",
+        "AWS_PROFILE_NAME",
+        "AWS_ROLE_NAME",
+        "AWS_ROLE_ARN",
+        "AWS_SESSION_NAME",
+        "AWS_WEB_IDENTITY_TOKEN",
+        "AWS_WEB_IDENTITY_TOKEN_FILE",
+        "AWS_SESSION_TOKEN",
+    ):
+        monkeypatch.delenv(name, raising=False)
+    monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
+    monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIDEXAMPLE")
+    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "x" * 40)
+    if session_token is not None:
+        monkeypatch.setenv("AWS_SESSION_TOKEN", session_token)
+    if assume_role:
+        monkeypatch.setenv("AWS_ROLE_NAME", "arn:aws:iam::123456789012:role/database")
+        monkeypatch.setenv("AWS_SESSION_NAME", "database-iam-auth")
+    session: Final = boto3.Session(region_name="us-east-1")
+    monkeypatch.setattr(boto3, "DEFAULT_SESSION", session)
+
+    def assume_role_response(request: AWSPreparedRequest, **kwargs: object) -> AWSResponse:
+        assert request.headers.get("X-Amz-Security-Token") == (
+            session_token.encode() if session_token is not None else None
+        )
+        assert b"Credential=AKIDEXAMPLE/" in request.headers["Authorization"]
+        payload: Final = b"""<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
+          <AssumeRoleResult><Credentials>
+            <AccessKeyId>ASIAASSUMED</AccessKeyId><SecretAccessKey>assumed-secret</SecretAccessKey>
+            <SessionToken>assumed/session+token==</SessionToken>
+          </Credentials></AssumeRoleResult>
+        </AssumeRoleResponse>"""
+        return AWSResponse(request.url, 200, {}, HTTPResponse(body=BytesIO(payload), preload_content=False))
+
+    session.events.register("before-send.sts.AssumeRole", assume_role_response)
+    token: Final = generate_iam_auth_token("db.example.com", "5432", "litellm")
+    query: Final = parse_qs(urlsplit(unquote(token)).query)
+    expected_token: Final = "assumed/session+token==" if assume_role else session_token
+
+    assert query.get("X-Amz-Security-Token") == ([expected_token] if expected_token is not None else None)
+    assert query["X-Amz-Credential"][0].split("/")[0] == ("ASIAASSUMED" if assume_role else "AKIDEXAMPLE")
+    assert query["DBUser"] == ["litellm"]
+
+
+@pytest.mark.parametrize("explicit_credentials", [False, True])
+def test_explicit_credentials_take_precedence_over_the_default_chain(
+    explicit_credentials: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    session: Final = boto3.Session(
+        aws_access_key_id="DEFAULTKEY", aws_secret_access_key="default-secret", aws_session_token="default-token"
+    )
+    monkeypatch.setattr(boto3, "DEFAULT_SESSION", session)
+    client: Final = init_rds_client(
+        aws_region_name="us-east-1",
+        aws_access_key_id="EXPLICITKEY" if explicit_credentials else None,
+        aws_secret_access_key="explicit-secret" if explicit_credentials else None,
+        aws_session_token="explicit-token" if explicit_credentials else None,
+    )
+    token: Final = generate_iam_auth_token("db.example.com", "5432", "litellm", client=client)
+    query: Final = parse_qs(urlsplit(unquote(token)).query)
+
+    assert query["X-Amz-Security-Token"] == (["explicit-token"] if explicit_credentials else ["default-token"])
+    assert query["X-Amz-Credential"][0].split("/")[0] == ("EXPLICITKEY" if explicit_credentials else "DEFAULTKEY")