fix: issue with aws credentials not being passed in correctly #1266

seawatts · 2024-12-21T17:02:42Z

Issue

Lendflow caught a bug where they were getting the error The security token included in the request is invalid when calling the baml client functions within the playground as well as their local environment while trying to use an assumed role. They were setting the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY but did not have the ability to also set AWS_SESSION_TOKEN which is required in the aws credentials chain.

Root cause and fix

When assuming an IAM role, it’s required by the AWS credentials chain to supply the AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, as well as AWS_SESSION_TOKEN. Previously, we did not pass the AWS_SESSION_TOKEN into the credentials chain. I updated the code to do so.

Additionally, I fixed another issue with AWS_PROFILE and SSO. There is now separate logic for when AWS_PROFILE is passed in that will use the credentials chain to grab the SSO token from disk (not available in WASM).

Important

Fix AWS credentials handling by including AWS_SESSION_TOKEN and adding AWS_PROFILE logic for SSO, with updates to integration tests.

Behavior:
- Include AWS_SESSION_TOKEN in credentials chain in aws_bedrock.rs and aws_client.rs.
- Add logic for AWS_PROFILE handling for SSO in aws_client.rs.
Code Changes:
- Update UnresolvedAwsBedrock and ResolvedAwsBedrock structs to include session_token and profile.
- Modify AwsClient::client_anyhow() to set credentials provider with session_token.
Tests:
- Add integration tests for invalid AWS credentials and profiles in integ-tests/typescript/tests/integ-tests.test.ts.
Misc:
- Remove patchelf dependency from pyproject.toml.

^{This description was created by}^{for 9a3efc3. It will automatically update as commits are pushed.}

vercel · 2024-12-21T17:02:45Z

The latest updates on your projects. Learn more about Vercel for Git ↗︎

Name	Status	Preview	Comments	Updated (UTC)
baml	✅ Ready (Inspect)	Visit Preview	💬 Add feedback	Dec 31, 2024 5:40pm

ellipsis-dev

👍 Looks good to me! Reviewed everything up to bf6fc64 in 13 seconds

More details

Looked at 150 lines of code in 3 files
Skipped 0 files when reviewing.
Skipped posting 1 drafted comments based on config settings.

1. integ-tests/python/pyproject.toml:23

Draft comment:
The patchelf dependency was removed, but this change is not mentioned in the PR description. Ensure that this removal is intentional and won't affect the build or tests.
Reason this comment was not posted:
Comment did not seem useful.

Workflow ID: wflow_3s2dnsisE8SqGCOR

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

Findings: When assuming an IAM role, it’s required by the AWS credentials chain to supply the AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, as well as AWS_SESSION_TOKEN. Previously, we did not pass the AWS_SESSION_TOKEN into the credentials chain. I updated the code to do so. Additionally, fix another issue with AWS_PROFILE and SSO. There is now separate logic for when AWS_PROFILE is passed in that will use the credentials chain to grab the SSO token from disk (not available in WASM).

aws

ellipsis-dev

👍 Looks good to me! Incremental review on b7a18c2 in 47 seconds

More details

Looked at 413 lines of code in 6 files
Skipped 0 files when reviewing.
Skipped posting 2 drafted comments based on config settings.

1. engine/baml-lib/llm-client/src/clients/aws_bedrock.rs:100

Draft comment:
Consider using unwrap_or instead of unwrap_or_else for string literals to improve readability. This applies to other similar instances in this file.
Reason this comment was not posted:
Confidence changes required: 20%
The code in aws_bedrock.rs and aws_client.rs has multiple instances where unwrap_or_else is used with a closure that returns a string. This can be simplified using unwrap_or. This is a minor optimization but improves readability.

2. engine/baml-runtime/src/internal/llm_client/primitive/aws/aws_client.rs:192

Draft comment:
Remove the redundant comma after await in the loader.credentials_provider call.
Reason this comment was not posted:
Confidence changes required: 10%
In aws_client.rs, the client_anyhow function has a redundant comma in the await call. This is a minor issue but should be corrected for consistency.

Workflow ID: wflow_lRFF8wsMhgDmak34

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

ellipsis-dev

👍 Looks good to me! Incremental review on 9a3efc3 in 53 seconds

More details

Looked at 1522 lines of code in 17 files
Skipped 0 files when reviewing.
Skipped posting 1 drafted comments based on config settings.

1. integ-tests/typescript/tests/integ-tests.test.ts:45

Draft comment:
The test for handling invalid AWS profile is currently skipped. Ensure there's a valid reason for this, or consider enabling it to verify the functionality.
Reason this comment was not posted:
Confidence changes required: 50%
The test case for handling invalid AWS profile is currently skipped. This might be intentional, but it's important to ensure that all test cases are active unless there's a specific reason to skip them. The test should be reviewed to determine if it can be enabled or if there's a valid reason for it to remain skipped.

Workflow ID: wflow_PnhJNYHLiMXRAAxL

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

ellipsis-dev

Skipped PR review on 6370553 because no changed files had a supported extension. If you think this was in error, please contact us and we'll fix it right away.

@Vaibhav

Change aws credential loading logic: - if any of `access_key_id`, `secret_access_key`, or `session_token` are set, all 3 are loaded explicitly (either from the `.baml` client definition or the dynamic client properties) - if, and only if, none of the 3 are set, all 3 are loaded from, respectively, `AWS_ACCESS_KEY_ID` `AWS_SECRET_ACCESS_KEY` `AWS_SESSION_TOKEN` This most closely matches the behavior of the AWS SDKs (Python, TS, and Rust). See [slack thread](https://gloo-global.slack.com/archives/C03KV1PJ6EM/p1736215459393779) which is copied below: > OK, so chris and i figured out what happened with bedrock/ethan: > > in #1266, chris correctly added support for aws session token so that if a user set it in aws.baml as properties { session_token env.AWS_SESSION_TOKEN }, baml would respect that (prior to #1266 baml would not) > > - however, 1266 also introduced an implicit default: if AWS_SESSION_TOKEN is set in the process' environment, but the user only set properties { access_key_id ... ; secret_access_key ... ; } then baml would construct the aws creds using access_key_id secret_access_key and session_token > - in ethan's case this is problematic, because he uses custom values for the access_key_id secret_access_key pair from his lambda secrets, but the aws lambda environment also sets AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN > - as a result, after updating past #1266, his (access_key_id, secret_access_key) did not agree with his AWS_SESSION_TOKEN and caused a runtime failure > - he also has no way to opt out of this behavior, because we do not currently provide a way to force session_token to null: it is always inferred from the environment by default > > to solve this, we're going to use the following logic for aws creds: > - if any of access_key_id, secret_access_key, or session_token are set in baml client properties, we will never magically infer a value for any of the 3 from the environment > - but if none of the 3 are set, we will read all 3 from the env > - this behavior feels most in line with how credential init in the aws sdk normally works > - in ts, an AwsCredentialIdentityProvider is any function that returns { accessKeyId: string, secretAccessKey: string, sessionToken?: string, expiration?: Date} (docs) > - in python, if you set any of the 3, a creds object is constructed using the explicitly provided values of all 3 (impl callsite, impl source code) > - in rust, this is what constructing Credentials::new does when you override the credentials_loader > > stepping back, this is a mix of the two approaches: (1) session_token: unset defaults to reading AWS_SESSION_TOKEN from the env, and user is allowed to explicitly set session_token: null or (2) session_token: unset never reads from the env, and the user must always set it > > NB: this does not explain why multiple other customers are complaining about not being able to figure out how to use aws bedrock. so @Vaibhav Gupta we still need to see what those other complaints are  ---- > [!IMPORTANT] > Updates AWS credential loading logic to use explicit credentials if set, otherwise defaults to environment variables, aligning with AWS SDK behavior. > > - **Behavior**: > - Updates AWS credential loading logic in `aws_bedrock.rs` and `aws_client.rs`. > - If any of `access_key_id`, `secret_access_key`, or `session_token` are set, all are loaded explicitly. > - If none are set, all are loaded from environment variables `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_SESSION_TOKEN`. > - **Error Handling**: > - Adds error checks for environment variable placeholders in `aws_client.rs`. > - **Misc**: > - Adjusts credential provider logic in `aws_client.rs` to use `DefaultCredentialsChain` when no credentials are provided. > > <sup>This description was created by </sup>[<img alt="Ellipsis" src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup> for 8199507. It will automatically update as commits are pushed.</sup>  --------- Co-authored-by: Chris Watts <chris.watts.t@gmail.com>

ellipsis-dev bot reviewed Dec 21, 2024

View reviewed changes

vercel bot deployed to Preview December 21, 2024 17:03 View deployment

seawatts added 2 commits December 30, 2024 15:31

fix: another issue where we were not setting the env vars correctly for

b7a18c2

aws

seawatts force-pushed the fix-aws-credentials branch from bf6fc64 to b7a18c2 Compare December 30, 2024 23:51

ellipsis-dev bot reviewed Dec 30, 2024

View reviewed changes

fix issue with added code

036dba7

vercel bot deployed to Preview December 30, 2024 23:57 View deployment

fix: make sure to handl wasm differently for env vars

9a3efc3

ellipsis-dev bot reviewed Dec 31, 2024

View reviewed changes

vercel bot deployed to Preview December 31, 2024 01:39 View deployment

fix lockfile

6370553

aaronvg enabled auto-merge December 31, 2024 17:37

ellipsis-dev bot reviewed Dec 31, 2024

View reviewed changes

aaronvg added this pull request to the merge queue Dec 31, 2024

vercel bot deployed to Preview December 31, 2024 17:40 View deployment

Merged via the queue into canary with commit 7b79ac4 Dec 31, 2024
11 checks passed

aaronvg deleted the fix-aws-credentials branch December 31, 2024 17:40

sxlijin mentioned this pull request Jan 7, 2025

feat: provide saner semantics around aws_session_token #1295

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: issue with aws credentials not being passed in correctly #1266

fix: issue with aws credentials not being passed in correctly #1266

seawatts commented Dec 21, 2024 •

edited by ellipsis-dev bot

Loading

vercel bot commented Dec 21, 2024 •

edited

Loading

ellipsis-dev bot left a comment

ellipsis-dev bot left a comment

ellipsis-dev bot left a comment

ellipsis-dev bot left a comment

fix: issue with aws credentials not being passed in correctly #1266

fix: issue with aws credentials not being passed in correctly #1266

Conversation

seawatts commented Dec 21, 2024 • edited by ellipsis-dev bot Loading

Issue

Root cause and fix

vercel bot commented Dec 21, 2024 • edited Loading

ellipsis-dev bot left a comment

Choose a reason for hiding this comment

ellipsis-dev bot left a comment

Choose a reason for hiding this comment

ellipsis-dev bot left a comment

Choose a reason for hiding this comment

ellipsis-dev bot left a comment

Choose a reason for hiding this comment

seawatts commented Dec 21, 2024 •

edited by ellipsis-dev bot

Loading

vercel bot commented Dec 21, 2024 •

edited

Loading