Skip to content

feat(runtimes): add handlers for IAM and STS credentials management #599

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
liramon1 merged 14 commits into from
Jul 17, 2025
Merged

feat(runtimes): add handlers for IAM and STS credentials management #599

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
22f77b0
feat(runtimes): add support for IAM and STS credentials management
liramon1 Jul 7, 2025
d9c4753
chore: undo unnecessary changes
liramon1 Jul 7, 2025
db626ab
fix: re-add IamCredentials change
liramon1 Jul 7, 2025
306c2b0
chore: revert EmptyProfile change
liramon1 Jul 9, 2025
fbfb2d6
refactor: split IamCredentialProfile into multiple profiles
liramon1 Jul 10, 2025
56fc78a
chore: revert comment
liramon1 Jul 10, 2025
939241b
feat: add mfaCode to getIamCredentialParams
liramon1 Jul 10, 2025
52b509b
refactor: prefix IAM-related profiles with 'Iam'
liramon1 Jul 10, 2025
d38dbb7
refactor: move encryption into separate function and split credential id
liramon1 Jul 11, 2025
8adfdee
feat: add external_id field to profile
liramon1 Jul 11, 2025
bb6d6a2
feat: optionalize permission validation in GetIamCredentialOptions
liramon1 Jul 14, 2025
7dc451c
refactor: move MFA code retrieval into separate request
liramon1 Jul 14, 2025
38961c9
fix: add parameters to mfa request
liramon1 Jul 14, 2025
7948b97
fix: naming changes
liramon1 Jul 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 115 additions & 1 deletion runtimes/protocol/identity-management.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import {
IamCredentials,
LSPErrorCodes,
ProgressType,
ProtocolNotificationType,
Expand All @@ -15,21 +16,28 @@ export const AwsErrorCodes = {
E_CANNOT_OVERWRITE_SSO_SESSION: 'E_CANNOT_OVERWRITE_SSO_SESSION',
E_CANNOT_READ_SHARED_CONFIG: 'E_CANNOT_READ_SHARED_CONFIG',
E_CANNOT_READ_SSO_CACHE: 'E_CANNOT_READ_SSO_CACHE',
E_CANNOT_READ_STS_CACHE: 'E_CANNOT_READ_STS_CACHE',
Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You should be reading and writing credentials into the shared config file ~/.aws/credentials. You should use E_CANNOT_READ_SHARED_CONFIG (and the WRITE variation) for these errors. STS is just a service to retrieve and only one way of getting IAM credentials. The naming convention here should be "IAM" instead of "STS".

Copy link
Contributor Author

@liramon1 liramon1 Jul 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

According to the AWS CLI docs, credentials obtained from AssumeRole (which Flare will call if the profile contains a role_arn and some credential source) are cached in ~/.aws/cli/cache, while all other IAM credentials live in ~/.aws/credentials. Should we still only consider ~/.aws/credentials?

Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, I'll send you an SEP offline that goes into more detail and should be considered canon.

Copy link
Contributor Author

@liramon1 liramon1 Jul 16, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case, I think the naming convention can stay as "STS". Technically, the SEP uses the term "temporary credentials", but "STS" is more concise and distinct (and can't be confused with SSO tokens which are also temporary). The cache only stores credentials obtained from STS AssumeRole and not any other IAM credential, so an "STS" prefix fits.

E_CANNOT_REFRESH_SSO_TOKEN: 'E_CANNOT_REFRESH_SSO_TOKEN',
E_CANNOT_REFRESH_STS_CREDENTIAL: 'E_CANNOT_REFRESH_STS_CREDENTIAL',
E_CANNOT_REGISTER_CLIENT: 'E_CANNOT_REGISTER_CLIENT',
E_CANNOT_CREATE_SSO_TOKEN: 'E_CANNOT_CREATE_SSO_TOKEN',
E_CANNOT_CREATE_STS_CREDENTIAL: 'E_CANNOT_CREATE_STS_CREDENTIAL',
E_CANNOT_WRITE_SHARED_CONFIG: 'E_CANNOT_WRITE_SHARED_CONFIG',
E_CANNOT_WRITE_SSO_CACHE: 'E_CANNOT_WRITE_SSO_CACHE',
E_CANNOT_WRITE_STS_CACHE: 'E_CANNOT_WRITE_STS_CACHE',
E_ENCRYPTION_REQUIRED: 'E_ENCRYPTION_REQUIRED',
E_INVALID_PROFILE: 'E_INVALID_PROFILE',
E_INVALID_SSO_CLIENT: 'E_INVALID_SSO_CLIENT',
E_INVALID_SSO_SESSION: 'E_INVALID_SSO_SESSION',
E_INVALID_SSO_TOKEN: 'E_INVALID_SSO_TOKEN',
E_INVALID_STS_CREDENTIAL: 'E_INVALID_STS_CREDENTIAL',
E_PROFILE_NOT_FOUND: 'E_PROFILE_NOT_FOUND',
E_RUNTIME_NOT_SUPPORTED: 'E_RUNTIME_NOT_SUPPORTED',
E_SSO_SESSION_NOT_FOUND: 'E_SSO_SESSION_NOT_FOUND',
E_SSO_TOKEN_EXPIRED: 'E_SSO_TOKEN_EXPIRED',
E_STS_CREDENTIAL_EXPIRED: 'E_STS_CREDENTIAL_EXPIRED',
E_SSO_TOKEN_SOURCE_NOT_SUPPORTED: 'E_SSO_TOKEN_SOURCE_NOT_SUPPORTED',
E_MFA_REQUIRED: 'E_MFA_REQUIRED',
E_TIMEOUT: 'E_TIMEOUT',
E_UNKNOWN: 'E_UNKNOWN',
E_CANCELLED: 'E_CANCELLED',
Expand All @@ -47,10 +55,20 @@ export class AwsResponseError extends ResponseError<AwsResponseErrorData> {
}

// listProfiles
export type ProfileKind = 'Unknown' | 'SsoTokenProfile'
export type ProfileKind =
| 'Unknown'
| 'SsoTokenProfile'
| 'IamCredentialsProfile'
| 'IamSourceProfileProfile'
| 'IamCredentialSourceProfile'
| 'IamCredentialProcessProfile'

export const ProfileKind = {
SsoTokenProfile: 'SsoTokenProfile',
IamCredentialsProfile: 'IamCredentialsProfile',
IamSourceProfileProfile: 'IamSourceProfileProfile',
IamCredentialSourceProfile: 'IamCredentialSourceProfile',
IamCredentialProcessProfile: 'IamCredentialProcessProfile',
Unknown: 'Unknown',
} as const

Expand All @@ -64,6 +82,18 @@ export interface Profile {
settings?: {
region?: string
sso_session?: string
aws_access_key_id?: string
Copy link

@nkomonen-amazon nkomonen-amazon Jul 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it make sense to separate Profile in to a "superset" type which consists of separate types SsoProfile and CredentialProfile?

Copy link
Contributor Author

@liramon1 liramon1 Jul 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably not if profiles can be any combination of profile types. Making these fields optional would allow the combined profile type to be decided at runtime and checked using the profile service's current duck typer.

Copy link
Contributor

@ctlai95 ctlai95 Jul 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might still make sense to separate the settings types while keeping them optional for better readability. Right now it's hard to determine which settings are for which type of profile

Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd recommend reaching out to connorstw@ here as he is working on a related project the generate client code from an API definition for Flare. As that would be the primary consumer of this information, there may be better ways to capture it than information-only interfaces.

aws_secret_access_key?: string
aws_session_token?: string
role_arn?: string
role_session_name?: string
credential_process?: string
credential_source?: string
source_profile?: string
mfa_serial?: string
external_id?: string
credential_cache?: string
credential_cache_location?: string
}
}

Expand Down Expand Up @@ -218,6 +248,56 @@ export const getSsoTokenRequestType = new ProtocolRequestType<
void
>('aws/identity/getSsoToken')

// getIamCredential
export type IamCredentialId = string // Opaque identifier

export interface GetIamCredentialOptions {
callStsOnInvalidIamCredential?: boolean
validatePermissions?: boolean
}

export const getIamCredentialOptionsDefaults = {
callStsOnInvalidIamCredential: true,
validatePermissions: true,
} satisfies GetIamCredentialOptions

export interface GetIamCredentialParams {
profileName: string
options?: GetIamCredentialOptions
}

export interface GetIamCredentialResult {
id: IamCredentialId
credentials: IamCredentials
Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was encrypting these in-flight considered or discussed with AppSec similar to SSO tokens?

Copy link
Contributor Author

@liramon1 liramon1 Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The credentials are encrypted in-flight using this code, but I have not discussed this with AppSec.

Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, trying to recall the reasoning for the additional encryption for SSO tokens, but this is fine for now.

updateCredentialsParams: UpdateCredentialsParams
}

export const getIamCredentialRequestType = new ProtocolRequestType<
GetIamCredentialParams,
GetIamCredentialResult,
never,
AwsResponseError,
void
>('aws/identity/getIamCredential')

// getMfaCode
export interface GetMfaCodeParams {
mfaSerial: string
profileName: string
}

export interface GetMfaCodeResult {
code: string
}

export const getMfaCodeRequestType = new ProtocolRequestType<
Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How does this work?

Copy link
Contributor Author

@liramon1 liramon1 Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this PR, the identityService will request an MFA code from the language client if getIamCredential needs to call AssumeRole and the parent credentials' permission to assume roles is locked behind MFA. This assumes the language clients will implement a handler for getMfaCode.

Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I'll review in more detail in that PR when it is posted.

GetMfaCodeParams,
GetMfaCodeResult,
never,
AwsResponseError,
void
>('aws/identity/getMfaCode')

// invalidateSsoToken
export interface InvalidateSsoTokenParams {
ssoTokenId: SsoTokenId
Expand All @@ -236,6 +316,23 @@ export const invalidateSsoTokenRequestType = new ProtocolRequestType<
void
>('aws/identity/invalidateSsoToken')

// invalidateStsCredential
export interface InvalidateStsCredentialParams {
profileName: string
}

export interface InvalidateStsCredentialResult {
// Intentionally left blank
}

export const invalidateStsCredentialRequestType = new ProtocolRequestType<
Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the plan for implementing this (i.e. what API(s) will you call)?

Copy link
Contributor Author

@liramon1 liramon1 Jul 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this PR, invalidateStsCredential will delete the credential from the cache folder using fs unlink.

Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I'll review in more detail in that PR when it is posted.

InvalidateStsCredentialParams,
InvalidateStsCredentialResult,
never,
AwsResponseError,
void
>('aws/identity/invalidateStsCredential')

// ssoTokenChanged
export type Expired = 'Expired'
export type Refreshed = 'Refreshed'
Expand All @@ -255,3 +352,20 @@ export interface SsoTokenChangedParams {
export const ssoTokenChangedRequestType = new ProtocolNotificationType<SsoTokenChangedParams, void>(
'aws/identity/ssoTokenChanged'
)

// stsCredentialChanged
export type StsCredentialChangedKind = Refreshed | Expired

export const StsCredentialChangedKind = {
Expired: 'Expired',
Refreshed: 'Refreshed',
} as const

export interface StsCredentialChangedParams {
kind: StsCredentialChangedKind
stsCredentialId: IamCredentialId
}

export const stsCredentialChangedRequestType = new ProtocolNotificationType<StsCredentialChangedParams, void>(
Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What is the implementation plan for this?

Copy link
Contributor Author

@liramon1 liramon1 Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this PR, stsCredentialChanged will be used similarly to ssoTokenChanged. It will fire a notification whenever an IAM credential obtained from AssumedRole is refreshed.

Copy link
Contributor

@floralph floralph Jul 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I'll review in more detail in that PR when it is posted.

'aws/identity/stsCredentialChanged'
)
43 changes: 43 additions & 0 deletions runtimes/runtimes/auth/standalone/encryption.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { Readable } from 'stream'
import { CompactEncrypt } from 'jose'
import { GetIamCredentialResult, GetSsoTokenResult } from '../../../protocol'

export function shouldWaitForEncryptionKey(): boolean {
return process.argv.some(arg => arg === '--set-credentials-encryption-key')
Expand Down Expand Up @@ -98,6 +99,48 @@ export function encryptObjectWithKey(request: Object, key: string, alg?: string,
.encrypt(keyBuffer)
}

/**
* Encrypts the SSO access tokens inside the result object with the provided key
*/
export async function encryptSsoResultWithKey(request: GetSsoTokenResult, key: string): Promise<GetSsoTokenResult> {
if (request.ssoToken.accessToken) {
request.ssoToken.accessToken = await encryptObjectWithKey(request.ssoToken.accessToken, key)
}
if (request.updateCredentialsParams.data && !request.updateCredentialsParams.encrypted) {
request.updateCredentialsParams.data = await encryptObjectWithKey(
// decodeCredentialsRequestToken expects nested 'data' fields
{ data: request.updateCredentialsParams.data },
key
)
request.updateCredentialsParams.encrypted = true
}
return request
}

/**
* Encrypts the IAM credentials inside the result object with the provided key
*/
export async function encryptIamResultWithKey(
request: GetIamCredentialResult,
key: string
): Promise<GetIamCredentialResult> {
request.credentials = {
accessKeyId: await encryptObjectWithKey(request.credentials.accessKeyId, key),
secretAccessKey: await encryptObjectWithKey(request.credentials.secretAccessKey, key),
...(request.credentials.sessionToken
? { sessionToken: await encryptObjectWithKey(request.credentials.sessionToken, key) }
: {}),
}
if (!request.updateCredentialsParams.encrypted) {
request.updateCredentialsParams.data = await encryptObjectWithKey(
{ data: request.updateCredentialsParams.data },
key
)
request.updateCredentialsParams.encrypted = true
}
return request
}

/**
* Check if a message is an encrypted JWE message with the provided key management algorithm and encoding
* As per RFC-7516:
Expand Down
8 changes: 8 additions & 0 deletions runtimes/runtimes/base-runtime.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,15 @@ import { observe } from './lsp'
import { LspRouter } from './lsp/router/lspRouter'
import { LspServer } from './lsp/router/lspServer'
import {
getIamCredentialRequestType,
getSsoTokenRequestType,
invalidateStsCredentialRequestType,
invalidateSsoTokenRequestType,
listProfilesRequestType,
ssoTokenChangedRequestType,
updateProfileRequestType,
stsCredentialChangedRequestType,
getMfaCodeRequestType,
} from '../protocol/identity-management'
import { IdentityManagement } from '../server-interface/identity-management'
import { WebBase64Encoding } from './encoding'
Expand Down Expand Up @@ -202,8 +206,12 @@ export const baseRuntime = (connections: { reader: MessageReader; writer: Messag
onListProfiles: handler => lspConnection.onRequest(listProfilesRequestType, handler),
onUpdateProfile: handler => lspConnection.onRequest(updateProfileRequestType, handler),
onGetSsoToken: handler => lspConnection.onRequest(getSsoTokenRequestType, handler),
onGetIamCredential: handler => lspConnection.onRequest(getIamCredentialRequestType, handler),
onInvalidateSsoToken: handler => lspConnection.onRequest(invalidateSsoTokenRequestType, handler),
onInvalidateStsCredential: handler => lspConnection.onRequest(invalidateStsCredentialRequestType, handler),
sendSsoTokenChanged: params => lspConnection.sendNotification(ssoTokenChangedRequestType, params),
sendStsCredentialChanged: params => lspConnection.sendNotification(stsCredentialChangedRequestType, params),
sendGetMfaCode: params => lspConnection.sendRequest(getMfaCodeRequestType, params),
}

// Set up auth without encryption
Expand Down
41 changes: 23 additions & 18 deletions runtimes/runtimes/standalone.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,19 @@ import {
didWriteFileNotificationType,
didAppendFileNotificationType,
didCreateDirectoryNotificationType,
getIamCredentialRequestType,
GetIamCredentialParams,
ShowOpenDialogParams,
ShowOpenDialogRequestType,
stsCredentialChangedRequestType,
getMfaCodeRequestType,
} from '../protocol'
import { ProposedFeatures, createConnection } from 'vscode-languageserver/node'
import {
encryptIamResultWithKey,
EncryptionInitialization,
encryptObjectWithKey,
encryptSsoResultWithKey,
readEncryptionDetails,
shouldWaitForEncryptionKey,
validateEncryptionDetails,
Expand All @@ -50,6 +56,7 @@ import {
getSsoTokenRequestType,
IdentityManagement,
invalidateSsoTokenRequestType,
invalidateStsCredentialRequestType,
listProfilesRequestType,
ssoTokenChangedRequestType,
updateProfileRequestType,
Expand Down Expand Up @@ -297,31 +304,29 @@ export const standalone = (props: RuntimeProps) => {
lspConnection.onRequest(
getSsoTokenRequestType,
async (params: GetSsoTokenParams, token: CancellationToken) => {
const result = await handler(params, token)

// Encrypt SsoToken.accessToken before sending to client
let result = await handler(params, token)
if (result && !(result instanceof Error) && encryptionKey) {
if (result.ssoToken.accessToken) {
result.ssoToken.accessToken = await encryptObjectWithKey(
result.ssoToken.accessToken,
encryptionKey
)
}
if (result.updateCredentialsParams.data && !result.updateCredentialsParams.encrypted) {
result.updateCredentialsParams.data = await encryptObjectWithKey(
// decodeCredentialsRequestToken expects nested 'data' fields
{ data: result.updateCredentialsParams.data },
encryptionKey
)
result.updateCredentialsParams.encrypted = true
}
result = await encryptSsoResultWithKey(result, encryptionKey)
}
return result
}
),
onGetIamCredential: handler =>
lspConnection.onRequest(
getIamCredentialRequestType,
async (params: GetIamCredentialParams, token: CancellationToken) => {
let result = await handler(params, token)
if (result && !(result instanceof Error) && encryptionKey) {
result = await encryptIamResultWithKey(result, encryptionKey)
}

return result
}
),
onInvalidateSsoToken: handler => lspConnection.onRequest(invalidateSsoTokenRequestType, handler),
onInvalidateStsCredential: handler => lspConnection.onRequest(invalidateStsCredentialRequestType, handler),
sendSsoTokenChanged: params => lspConnection.sendNotification(ssoTokenChangedRequestType, params),
sendStsCredentialChanged: params => lspConnection.sendNotification(stsCredentialChangedRequestType, params),
sendGetMfaCode: params => lspConnection.sendRequest(getMfaCodeRequestType, params),
}

const credentialsProvider: CredentialsProvider = auth.getCredentialsProvider()
Expand Down
23 changes: 23 additions & 0 deletions runtimes/server-interface/identity-management.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,21 @@
import {
AwsResponseError,
GetIamCredentialParams,
GetIamCredentialResult,
GetSsoTokenParams,
GetSsoTokenResult,
InvalidateSsoTokenParams,
InvalidateSsoTokenResult,
InvalidateStsCredentialParams,
InvalidateStsCredentialResult,
ListProfilesParams,
ListProfilesResult,
GetMfaCodeParams,
SsoTokenChangedParams,
StsCredentialChangedParams,
UpdateProfileParams,
UpdateProfileResult,
GetMfaCodeResult,
} from '../protocol/identity-management'
import { RequestHandler } from '../protocol'

Expand All @@ -27,9 +34,25 @@ export type IdentityManagement = {
handler: RequestHandler<GetSsoTokenParams, GetSsoTokenResult | undefined | null, AwsResponseError>
) => void

onGetIamCredential: (
handler: RequestHandler<GetIamCredentialParams, GetIamCredentialResult | undefined | null, AwsResponseError>
) => void

onInvalidateSsoToken: (
handler: RequestHandler<InvalidateSsoTokenParams, InvalidateSsoTokenResult | undefined | null, AwsResponseError>
) => void

onInvalidateStsCredential: (
handler: RequestHandler<
InvalidateStsCredentialParams,
InvalidateStsCredentialResult | undefined | null,
AwsResponseError
>
) => void

sendSsoTokenChanged: (params: SsoTokenChangedParams) => void

sendStsCredentialChanged: (params: StsCredentialChangedParams) => void

sendGetMfaCode: (params: GetMfaCodeParams) => Promise<GetMfaCodeResult>
}
1 change: 1 addition & 0 deletions types/auth.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ export type IamCredentials = {
readonly accessKeyId: string
readonly secretAccessKey: string
readonly sessionToken?: string
readonly expiration?: Date
}

export type BearerCredentials = {
Expand Down