From 22f77b025969df10bbc75b4fb59a2e554387c3d1 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 7 Jul 2025 12:03:34 -0400 Subject: [PATCH 01/19] feat(runtimes): add support for IAM and STS credentials management --- runtimes/protocol/identity-management.ts | 87 +++++++++++++++++-- runtimes/runtimes/base-runtime.ts | 7 ++ runtimes/runtimes/standalone.ts | 42 +++++++++ .../server-interface/identity-management.ts | 19 ++++ types/auth.ts | 1 + 5 files changed, 150 insertions(+), 6 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 1c19d89b..73890dcb 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -1,4 +1,5 @@ import { + IamCredentials, LSPErrorCodes, ProgressType, ProtocolNotificationType, @@ -15,20 +16,26 @@ 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', 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_TIMEOUT: 'E_TIMEOUT', E_UNKNOWN: 'E_UNKNOWN', @@ -47,14 +54,16 @@ export class AwsResponseError extends ResponseError { } // listProfiles -export type ProfileKind = 'Unknown' | 'SsoTokenProfile' +export type ProfileKind = 'Unknown' | 'SsoTokenProfile' | 'IamCredentialProfile' | 'EmptyProfile' export const ProfileKind = { SsoTokenProfile: 'SsoTokenProfile', + IamCredentialProfile: 'IamCredentialProfile', + EmptyProfile: 'EmptyProfile', Unknown: 'Unknown', } as const -// Profile and SsoSession use 'settings' property as namescope for their settings to avoid future +// Profile and Session use 'settings' property as namescope for their settings to avoid future // name conflicts with 'kinds', 'name', and future properties as well as making some setting // iteration operations easier. @@ -64,6 +73,10 @@ export interface Profile { settings?: { region?: string sso_session?: string + aws_access_key_id?: string + aws_secret_access_key?: string + aws_session_token?: string + role_arn?: string } } @@ -131,7 +144,7 @@ export const updateProfileRequestType = new ProtocolRequestType< >('aws/identity/updateProfile') // getSsoToken -export type SsoTokenId = string // Opaque identifier +export type CredentialId = string // Opaque identifier export type IamIdentityCenterSsoTokenSourceKind = 'IamIdentityCenter' export type AwsBuilderIdSsoTokenSourceKind = 'AwsBuilderId' @@ -199,7 +212,7 @@ export interface GetSsoTokenParams { } export interface SsoToken { - id: SsoTokenId + id: CredentialId accessToken: string // This field is encrypted with JWT like 'update' // Additional fields captured in token cache file may be added here in the future } @@ -218,9 +231,37 @@ export const getSsoTokenRequestType = new ProtocolRequestType< void >('aws/identity/getSsoToken') +// getIamCredential +export interface GetIamCredentialOptions { + generateOnInvalidStsCredential?: boolean +} + +export const getIamCredentialOptionsDefaults = { + generateOnInvalidStsCredential: true, +} satisfies GetIamCredentialOptions + +export interface GetIamCredentialParams { + profileName: string + options?: GetIamCredentialOptions +} + +export interface GetIamCredentialResult { + id: CredentialId + credentials: IamCredentials + updateCredentialsParams: UpdateCredentialsParams +} + +export const getIamCredentialRequestType = new ProtocolRequestType< + GetIamCredentialParams, + GetIamCredentialResult, + never, + AwsResponseError, + void +>('aws/identity/getIamCredential') + // invalidateSsoToken export interface InvalidateSsoTokenParams { - ssoTokenId: SsoTokenId + ssoTokenId: CredentialId } export interface InvalidateSsoTokenResult { @@ -236,6 +277,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< + InvalidateStsCredentialParams, + InvalidateStsCredentialResult, + never, + AwsResponseError, + void +>('aws/identity/invalidateStsCredential') + // ssoTokenChanged export type Expired = 'Expired' export type Refreshed = 'Refreshed' @@ -249,9 +307,26 @@ export const SsoTokenChangedKind = { export interface SsoTokenChangedParams { kind: SsoTokenChangedKind - ssoTokenId: SsoTokenId + ssoTokenId: CredentialId } export const ssoTokenChangedRequestType = new ProtocolNotificationType( 'aws/identity/ssoTokenChanged' ) + +// stsCredentialChanged +export type StsCredentialChangedKind = Refreshed | Expired + +export const StsCredentialChangedKind = { + Expired: 'Expired', + Refreshed: 'Refreshed', +} as const + +export interface StsCredentialChangedParams { + kind: StsCredentialChangedKind + stsCredentialId: CredentialId +} + +export const stsCredentialChangedRequestType = new ProtocolNotificationType( + 'aws/identity/stsCredentialChanged' +) diff --git a/runtimes/runtimes/base-runtime.ts b/runtimes/runtimes/base-runtime.ts index debc5508..ea21a265 100644 --- a/runtimes/runtimes/base-runtime.ts +++ b/runtimes/runtimes/base-runtime.ts @@ -83,11 +83,15 @@ import { observe } from './lsp' import { LspRouter } from './lsp/router/lspRouter' import { LspServer } from './lsp/router/lspServer' import { + AwsResponseError, + getIamCredentialRequestType, getSsoTokenRequestType, + invalidateStsCredentialRequestType, invalidateSsoTokenRequestType, listProfilesRequestType, ssoTokenChangedRequestType, updateProfileRequestType, + stsCredentialChangedRequestType, } from '../protocol/identity-management' import { IdentityManagement } from '../server-interface/identity-management' import { WebBase64Encoding } from './encoding' @@ -202,8 +206,11 @@ 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), } // Set up auth without encryption diff --git a/runtimes/runtimes/standalone.ts b/runtimes/runtimes/standalone.ts index 53a10921..3a36f25b 100644 --- a/runtimes/runtimes/standalone.ts +++ b/runtimes/runtimes/standalone.ts @@ -28,8 +28,12 @@ import { didWriteFileNotificationType, didAppendFileNotificationType, didCreateDirectoryNotificationType, + getIamCredentialRequestType, + GetIamCredentialParams, + IamCredentials, ShowOpenDialogParams, ShowOpenDialogRequestType, + stsCredentialChangedRequestType, } from '../protocol' import { ProposedFeatures, createConnection } from 'vscode-languageserver/node' import { @@ -50,6 +54,7 @@ import { getSsoTokenRequestType, IdentityManagement, invalidateSsoTokenRequestType, + invalidateStsCredentialRequestType, listProfilesRequestType, ssoTokenChangedRequestType, updateProfileRequestType, @@ -317,11 +322,48 @@ export const standalone = (props: RuntimeProps) => { } } + return result + } + ), + onGetIamCredential: handler => + lspConnection.onRequest( + getIamCredentialRequestType, + async (params: GetIamCredentialParams, token: CancellationToken) => { + const result = await handler(params, token) + + // Encrypt the IAM credential before sending to client + if (result && !(result instanceof Error) && encryptionKey) { + result.credentials = { + accessKeyId: await encryptObjectWithKey(result.credentials.accessKeyId, encryptionKey), + secretAccessKey: await encryptObjectWithKey( + result.credentials.secretAccessKey, + encryptionKey + ), + ...(result.credentials.sessionToken + ? { + sessionToken: await encryptObjectWithKey( + result.credentials.sessionToken, + encryptionKey + ), + } + : {}), + } + if (!result.updateCredentialsParams.encrypted) { + result.updateCredentialsParams.data = await encryptObjectWithKey( + { data: result.updateCredentialsParams.data }, + encryptionKey + ) + result.updateCredentialsParams.encrypted = true + } + } + 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), } const credentialsProvider: CredentialsProvider = auth.getCredentialsProvider() diff --git a/runtimes/server-interface/identity-management.ts b/runtimes/server-interface/identity-management.ts index 5dd1a72a..28e79ebb 100644 --- a/runtimes/server-interface/identity-management.ts +++ b/runtimes/server-interface/identity-management.ts @@ -1,12 +1,17 @@ import { AwsResponseError, + GetIamCredentialParams, + GetIamCredentialResult, GetSsoTokenParams, GetSsoTokenResult, InvalidateSsoTokenParams, InvalidateSsoTokenResult, + InvalidateStsCredentialParams, + InvalidateStsCredentialResult, ListProfilesParams, ListProfilesResult, SsoTokenChangedParams, + StsCredentialChangedParams, UpdateProfileParams, UpdateProfileResult, } from '../protocol/identity-management' @@ -27,9 +32,23 @@ export type IdentityManagement = { handler: RequestHandler ) => void + onGetIamCredential: ( + handler: RequestHandler + ) => void + onInvalidateSsoToken: ( handler: RequestHandler ) => void + onInvalidateStsCredential: ( + handler: RequestHandler< + InvalidateStsCredentialParams, + InvalidateStsCredentialResult | undefined | null, + AwsResponseError + > + ) => void + sendSsoTokenChanged: (params: SsoTokenChangedParams) => void + + sendStsCredentialChanged: (params: StsCredentialChangedParams) => void } diff --git a/types/auth.ts b/types/auth.ts index 86a26f76..ead8e1d4 100644 --- a/types/auth.ts +++ b/types/auth.ts @@ -2,6 +2,7 @@ export type IamCredentials = { readonly accessKeyId: string readonly secretAccessKey: string readonly sessionToken?: string + readonly expiration?: Date } export type BearerCredentials = { From d9c4753fe84f6b8f64e17338d386e5d82e8bb401 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 7 Jul 2025 12:21:40 -0400 Subject: [PATCH 02/19] chore: undo unnecessary changes --- runtimes/protocol/identity-management.ts | 2 +- runtimes/runtimes/base-runtime.ts | 1 - types/auth.ts | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 73890dcb..61c83008 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -63,7 +63,7 @@ export const ProfileKind = { Unknown: 'Unknown', } as const -// Profile and Session use 'settings' property as namescope for their settings to avoid future +// Profile and SsoSession use 'settings' property as namescope for their settings to avoid future // name conflicts with 'kinds', 'name', and future properties as well as making some setting // iteration operations easier. diff --git a/runtimes/runtimes/base-runtime.ts b/runtimes/runtimes/base-runtime.ts index ea21a265..3deb3cf9 100644 --- a/runtimes/runtimes/base-runtime.ts +++ b/runtimes/runtimes/base-runtime.ts @@ -83,7 +83,6 @@ import { observe } from './lsp' import { LspRouter } from './lsp/router/lspRouter' import { LspServer } from './lsp/router/lspServer' import { - AwsResponseError, getIamCredentialRequestType, getSsoTokenRequestType, invalidateStsCredentialRequestType, diff --git a/types/auth.ts b/types/auth.ts index ead8e1d4..86a26f76 100644 --- a/types/auth.ts +++ b/types/auth.ts @@ -2,7 +2,6 @@ export type IamCredentials = { readonly accessKeyId: string readonly secretAccessKey: string readonly sessionToken?: string - readonly expiration?: Date } export type BearerCredentials = { From db626ab8fa1b0cd57b7003581f52e9d7cd8a5758 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 7 Jul 2025 18:09:11 -0400 Subject: [PATCH 03/19] fix: re-add IamCredentials change --- types/auth.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/types/auth.ts b/types/auth.ts index 86a26f76..ead8e1d4 100644 --- a/types/auth.ts +++ b/types/auth.ts @@ -2,6 +2,7 @@ export type IamCredentials = { readonly accessKeyId: string readonly secretAccessKey: string readonly sessionToken?: string + readonly expiration?: Date } export type BearerCredentials = { From 306c2b0290eb17545f926a89ef99db8479dcb237 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 9 Jul 2025 09:56:45 -0400 Subject: [PATCH 04/19] chore: revert EmptyProfile change --- runtimes/protocol/identity-management.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 61c83008..dfb63fe4 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -54,16 +54,15 @@ export class AwsResponseError extends ResponseError { } // listProfiles -export type ProfileKind = 'Unknown' | 'SsoTokenProfile' | 'IamCredentialProfile' | 'EmptyProfile' +export type ProfileKind = 'Unknown' | 'SsoTokenProfile' | 'IamCredentialProfile' export const ProfileKind = { SsoTokenProfile: 'SsoTokenProfile', IamCredentialProfile: 'IamCredentialProfile', - EmptyProfile: 'EmptyProfile', Unknown: 'Unknown', } as const -// Profile and SsoSession use 'settings' property as namescope for their settings to avoid future +// Profile and Session use 'settings' property as namescope for their settings to avoid future // name conflicts with 'kinds', 'name', and future properties as well as making some setting // iteration operations easier. @@ -77,6 +76,7 @@ export interface Profile { aws_secret_access_key?: string aws_session_token?: string role_arn?: string + credential_process?: string } } From fbfb2d622eaa755e3f46eacab28a81175db3f5ab Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 10 Jul 2025 14:08:07 -0400 Subject: [PATCH 05/19] refactor: split IamCredentialProfile into multiple profiles --- runtimes/protocol/identity-management.ts | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index dfb63fe4..fb0e1d63 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -54,11 +54,20 @@ export class AwsResponseError extends ResponseError { } // listProfiles -export type ProfileKind = 'Unknown' | 'SsoTokenProfile' | 'IamCredentialProfile' +export type ProfileKind = + | 'Unknown' + | 'SsoTokenProfile' + | 'IamUserProfile' + | 'RoleSourceProfile' + | 'RoleInstanceProfile' + | 'ProcessProfile' export const ProfileKind = { SsoTokenProfile: 'SsoTokenProfile', - IamCredentialProfile: 'IamCredentialProfile', + IamUserProfile: 'IamUserProfile', + RoleSourceProfile: 'RoleSourceProfile', + RoleInstanceProfile: 'RoleInstanceProfile', + ProcessProfile: 'ProcessProfile', Unknown: 'Unknown', } as const @@ -76,7 +85,11 @@ export interface Profile { 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 } } From 56fc78a44c5887460201eee38c00c7b7e0bc65e7 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 10 Jul 2025 14:12:49 -0400 Subject: [PATCH 06/19] chore: revert comment --- runtimes/protocol/identity-management.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index fb0e1d63..87127bad 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -71,7 +71,7 @@ export const ProfileKind = { Unknown: 'Unknown', } as const -// Profile and Session use 'settings' property as namescope for their settings to avoid future +// Profile and SsoSession use 'settings' property as namescope for their settings to avoid future // name conflicts with 'kinds', 'name', and future properties as well as making some setting // iteration operations easier. From 939241ba2366ca79656b5a038dcebd2d3b6b0a5d Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 10 Jul 2025 17:36:54 -0400 Subject: [PATCH 07/19] feat: add mfaCode to getIamCredentialParams --- runtimes/protocol/identity-management.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 87127bad..f5bb8f03 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -255,6 +255,7 @@ export const getIamCredentialOptionsDefaults = { export interface GetIamCredentialParams { profileName: string + mfaCode?: string options?: GetIamCredentialOptions } From 52b509b51d6c928f496f21ef920c0d8ed0edf0cb Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 10 Jul 2025 18:09:56 -0400 Subject: [PATCH 08/19] refactor: prefix IAM-related profiles with 'Iam' --- runtimes/protocol/identity-management.ts | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index f5bb8f03..39df6ea5 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -58,16 +58,16 @@ export type ProfileKind = | 'Unknown' | 'SsoTokenProfile' | 'IamUserProfile' - | 'RoleSourceProfile' - | 'RoleInstanceProfile' - | 'ProcessProfile' + | 'IamRoleSourceProfile' + | 'IamRoleInstanceProfile' + | 'IamProcessProfile' export const ProfileKind = { SsoTokenProfile: 'SsoTokenProfile', IamUserProfile: 'IamUserProfile', - RoleSourceProfile: 'RoleSourceProfile', - RoleInstanceProfile: 'RoleInstanceProfile', - ProcessProfile: 'ProcessProfile', + IamRoleSourceProfile: 'IamRoleSourceProfile', + IamRoleInstanceProfile: 'IamRoleInstanceProfile', + IamProcessProfile: 'IamProcessProfile', Unknown: 'Unknown', } as const From d38dbb72f173db1a6da967609c2d7c93906b0458 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 11 Jul 2025 10:09:33 -0400 Subject: [PATCH 09/19] refactor: move encryption into separate function and split credential id --- runtimes/protocol/identity-management.ts | 14 ++--- .../runtimes/auth/standalone/encryption.ts | 43 ++++++++++++++++ runtimes/runtimes/standalone.ts | 51 +++---------------- 3 files changed, 57 insertions(+), 51 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 39df6ea5..973c9bb3 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -157,7 +157,7 @@ export const updateProfileRequestType = new ProtocolRequestType< >('aws/identity/updateProfile') // getSsoToken -export type CredentialId = string // Opaque identifier +export type SsoTokenId = string // Opaque identifier export type IamIdentityCenterSsoTokenSourceKind = 'IamIdentityCenter' export type AwsBuilderIdSsoTokenSourceKind = 'AwsBuilderId' @@ -225,7 +225,7 @@ export interface GetSsoTokenParams { } export interface SsoToken { - id: CredentialId + id: SsoTokenId accessToken: string // This field is encrypted with JWT like 'update' // Additional fields captured in token cache file may be added here in the future } @@ -245,6 +245,8 @@ export const getSsoTokenRequestType = new ProtocolRequestType< >('aws/identity/getSsoToken') // getIamCredential +export type IamCredentialId = string // Opaque identifier + export interface GetIamCredentialOptions { generateOnInvalidStsCredential?: boolean } @@ -260,7 +262,7 @@ export interface GetIamCredentialParams { } export interface GetIamCredentialResult { - id: CredentialId + id: IamCredentialId credentials: IamCredentials updateCredentialsParams: UpdateCredentialsParams } @@ -275,7 +277,7 @@ export const getIamCredentialRequestType = new ProtocolRequestType< // invalidateSsoToken export interface InvalidateSsoTokenParams { - ssoTokenId: CredentialId + ssoTokenId: SsoTokenId } export interface InvalidateSsoTokenResult { @@ -321,7 +323,7 @@ export const SsoTokenChangedKind = { export interface SsoTokenChangedParams { kind: SsoTokenChangedKind - ssoTokenId: CredentialId + ssoTokenId: SsoTokenId } export const ssoTokenChangedRequestType = new ProtocolNotificationType( @@ -338,7 +340,7 @@ export const StsCredentialChangedKind = { export interface StsCredentialChangedParams { kind: StsCredentialChangedKind - stsCredentialId: CredentialId + stsCredentialId: IamCredentialId } export const stsCredentialChangedRequestType = new ProtocolNotificationType( diff --git a/runtimes/runtimes/auth/standalone/encryption.ts b/runtimes/runtimes/auth/standalone/encryption.ts index ebe576dd..4f9dbb5b 100644 --- a/runtimes/runtimes/auth/standalone/encryption.ts +++ b/runtimes/runtimes/auth/standalone/encryption.ts @@ -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') @@ -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 { + 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 { + 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: diff --git a/runtimes/runtimes/standalone.ts b/runtimes/runtimes/standalone.ts index 3a36f25b..74a33071 100644 --- a/runtimes/runtimes/standalone.ts +++ b/runtimes/runtimes/standalone.ts @@ -30,15 +30,16 @@ import { didCreateDirectoryNotificationType, getIamCredentialRequestType, GetIamCredentialParams, - IamCredentials, ShowOpenDialogParams, ShowOpenDialogRequestType, stsCredentialChangedRequestType, } from '../protocol' import { ProposedFeatures, createConnection } from 'vscode-languageserver/node' import { + encryptIamResultWithKey, EncryptionInitialization, encryptObjectWithKey, + encryptSsoResultWithKey, readEncryptionDetails, shouldWaitForEncryptionKey, validateEncryptionDetails, @@ -302,26 +303,10 @@ 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 } ), @@ -329,34 +314,10 @@ export const standalone = (props: RuntimeProps) => { lspConnection.onRequest( getIamCredentialRequestType, async (params: GetIamCredentialParams, token: CancellationToken) => { - const result = await handler(params, token) - - // Encrypt the IAM credential before sending to client + let result = await handler(params, token) if (result && !(result instanceof Error) && encryptionKey) { - result.credentials = { - accessKeyId: await encryptObjectWithKey(result.credentials.accessKeyId, encryptionKey), - secretAccessKey: await encryptObjectWithKey( - result.credentials.secretAccessKey, - encryptionKey - ), - ...(result.credentials.sessionToken - ? { - sessionToken: await encryptObjectWithKey( - result.credentials.sessionToken, - encryptionKey - ), - } - : {}), - } - if (!result.updateCredentialsParams.encrypted) { - result.updateCredentialsParams.data = await encryptObjectWithKey( - { data: result.updateCredentialsParams.data }, - encryptionKey - ) - result.updateCredentialsParams.encrypted = true - } + result = await encryptIamResultWithKey(result, encryptionKey) } - return result } ), From 8adfdeed029cfea719c724ba4618d5e48d02c3e5 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 11 Jul 2025 11:17:53 -0400 Subject: [PATCH 10/19] feat: add external_id field to profile --- runtimes/protocol/identity-management.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 973c9bb3..c2643674 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -90,6 +90,7 @@ export interface Profile { credential_source?: string source_profile?: string mfa_serial?: string + external_id?: string } } From bb6d6a25440238e291777d54323b399962221084 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 14 Jul 2025 10:53:37 -0400 Subject: [PATCH 11/19] feat: optionalize permission validation in GetIamCredentialOptions --- runtimes/protocol/identity-management.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index c2643674..481ad02d 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -250,10 +250,12 @@ export type IamCredentialId = string // Opaque identifier export interface GetIamCredentialOptions { generateOnInvalidStsCredential?: boolean + validatePermissions?: boolean } export const getIamCredentialOptionsDefaults = { generateOnInvalidStsCredential: true, + validatePermissions: true, } satisfies GetIamCredentialOptions export interface GetIamCredentialParams { From 7dc451c6b151120c0a3a9e5934814922ea3d2b61 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 14 Jul 2025 16:23:15 -0400 Subject: [PATCH 12/19] refactor: move MFA code retrieval into separate request --- runtimes/protocol/identity-management.ts | 19 ++++++++++++++++++- runtimes/runtimes/base-runtime.ts | 2 ++ runtimes/runtimes/standalone.ts | 2 ++ .../server-interface/identity-management.ts | 4 ++++ 4 files changed, 26 insertions(+), 1 deletion(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 481ad02d..a8b3b3d7 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -37,6 +37,7 @@ export const AwsErrorCodes = { 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', @@ -260,7 +261,6 @@ export const getIamCredentialOptionsDefaults = { export interface GetIamCredentialParams { profileName: string - mfaCode?: string options?: GetIamCredentialOptions } @@ -278,6 +278,23 @@ export const getIamCredentialRequestType = new ProtocolRequestType< void >('aws/identity/getIamCredential') +// getMfaCode +export interface GetMfaCodeParams { + // Intentionally left blank +} + +export interface GetMfaCodeResult { + code: string +} + +export const getMfaCodeRequestType = new ProtocolRequestType< + GetMfaCodeParams, + GetMfaCodeResult, + never, + AwsResponseError, + void +>('aws/identity/getMfaCode') + // invalidateSsoToken export interface InvalidateSsoTokenParams { ssoTokenId: SsoTokenId diff --git a/runtimes/runtimes/base-runtime.ts b/runtimes/runtimes/base-runtime.ts index 3deb3cf9..116808d4 100644 --- a/runtimes/runtimes/base-runtime.ts +++ b/runtimes/runtimes/base-runtime.ts @@ -91,6 +91,7 @@ import { ssoTokenChangedRequestType, updateProfileRequestType, stsCredentialChangedRequestType, + getMfaCodeRequestType, } from '../protocol/identity-management' import { IdentityManagement } from '../server-interface/identity-management' import { WebBase64Encoding } from './encoding' @@ -210,6 +211,7 @@ export const baseRuntime = (connections: { reader: MessageReader; writer: Messag 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 diff --git a/runtimes/runtimes/standalone.ts b/runtimes/runtimes/standalone.ts index 74a33071..60c79b4c 100644 --- a/runtimes/runtimes/standalone.ts +++ b/runtimes/runtimes/standalone.ts @@ -33,6 +33,7 @@ import { ShowOpenDialogParams, ShowOpenDialogRequestType, stsCredentialChangedRequestType, + getMfaCodeRequestType, } from '../protocol' import { ProposedFeatures, createConnection } from 'vscode-languageserver/node' import { @@ -325,6 +326,7 @@ export const standalone = (props: RuntimeProps) => { 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() diff --git a/runtimes/server-interface/identity-management.ts b/runtimes/server-interface/identity-management.ts index 28e79ebb..6e0de534 100644 --- a/runtimes/server-interface/identity-management.ts +++ b/runtimes/server-interface/identity-management.ts @@ -10,10 +10,12 @@ import { InvalidateStsCredentialResult, ListProfilesParams, ListProfilesResult, + GetMfaCodeParams, SsoTokenChangedParams, StsCredentialChangedParams, UpdateProfileParams, UpdateProfileResult, + GetMfaCodeResult, } from '../protocol/identity-management' import { RequestHandler } from '../protocol' @@ -51,4 +53,6 @@ export type IdentityManagement = { sendSsoTokenChanged: (params: SsoTokenChangedParams) => void sendStsCredentialChanged: (params: StsCredentialChangedParams) => void + + sendGetMfaCode: (params: GetMfaCodeParams) => Promise } From 38961c943c07c82017d70413f705afc9213fe921 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Mon, 14 Jul 2025 17:55:38 -0400 Subject: [PATCH 13/19] fix: add parameters to mfa request --- runtimes/protocol/identity-management.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index a8b3b3d7..89d6a20c 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -280,7 +280,8 @@ export const getIamCredentialRequestType = new ProtocolRequestType< // getMfaCode export interface GetMfaCodeParams { - // Intentionally left blank + mfaSerial: string + profileName: string } export interface GetMfaCodeResult { From 7948b97d5bb0475139fa9a15f535052ec9a12f37 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 15 Jul 2025 14:53:13 -0400 Subject: [PATCH 14/19] fix: naming changes --- runtimes/protocol/identity-management.ts | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 89d6a20c..a6c6613e 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -58,17 +58,17 @@ export class AwsResponseError extends ResponseError { export type ProfileKind = | 'Unknown' | 'SsoTokenProfile' - | 'IamUserProfile' - | 'IamRoleSourceProfile' - | 'IamRoleInstanceProfile' - | 'IamProcessProfile' + | 'IamCredentialsProfile' + | 'IamSourceProfileProfile' + | 'IamCredentialSourceProfile' + | 'IamCredentialProcessProfile' export const ProfileKind = { SsoTokenProfile: 'SsoTokenProfile', - IamUserProfile: 'IamUserProfile', - IamRoleSourceProfile: 'IamRoleSourceProfile', - IamRoleInstanceProfile: 'IamRoleInstanceProfile', - IamProcessProfile: 'IamProcessProfile', + IamCredentialsProfile: 'IamCredentialsProfile', + IamSourceProfileProfile: 'IamSourceProfileProfile', + IamCredentialSourceProfile: 'IamCredentialSourceProfile', + IamCredentialProcessProfile: 'IamCredentialProcessProfile', Unknown: 'Unknown', } as const @@ -92,6 +92,8 @@ export interface Profile { source_profile?: string mfa_serial?: string external_id?: string + credential_cache?: string + credential_cache_location?: string } } @@ -250,12 +252,12 @@ export const getSsoTokenRequestType = new ProtocolRequestType< export type IamCredentialId = string // Opaque identifier export interface GetIamCredentialOptions { - generateOnInvalidStsCredential?: boolean + callStsOnInvalidIamCredential?: boolean validatePermissions?: boolean } export const getIamCredentialOptionsDefaults = { - generateOnInvalidStsCredential: true, + callStsOnInvalidIamCredential: true, validatePermissions: true, } satisfies GetIamCredentialOptions From 93e856df1a07c00a6538d12772c0590733555631 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 17 Jul 2025 12:29:57 -0400 Subject: [PATCH 15/19] fix: incorporate PR feedback --- runtimes/protocol/identity-management.ts | 34 ++++++++++++++++--- .../runtimes/auth/standalone/encryption.ts | 10 +++--- 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index a6c6613e..99616127 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -38,11 +38,33 @@ export const AwsErrorCodes = { 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_PERMISSION_DENIED: 'E_PERMISSION_DENIED', E_TIMEOUT: 'E_TIMEOUT', E_UNKNOWN: 'E_UNKNOWN', E_CANCELLED: 'E_CANCELLED', } as const +// Permissions +export const PermissionSets = { + Q: [ + 'q:StartConversation', + 'q:SendMessage', + 'q:GetConversation', + 'q:ListConversations', + 'q:UpdateConversation', + 'q:DeleteConversation', + 'q:PassRequest', + 'q:StartTroubleshootingAnalysis', + 'q:StartTroubleshootingResolutionExplanation', + 'q:GetTroubleshootingResults', + 'q:UpdateTroubleshootingCommandResult', + 'q:GetIdentityMetaData', + 'q:GenerateCodeFromCommands', + 'q:UsePlugin', + 'codewhisperer:GenerateRecommendations', + ], +} + export interface AwsResponseErrorData { awsErrorCode: string } @@ -253,12 +275,12 @@ export type IamCredentialId = string // Opaque identifier export interface GetIamCredentialOptions { callStsOnInvalidIamCredential?: boolean - validatePermissions?: boolean + permissionSet?: string[] } export const getIamCredentialOptionsDefaults = { callStsOnInvalidIamCredential: true, - validatePermissions: true, + permissionSet: PermissionSets.Q, } satisfies GetIamCredentialOptions export interface GetIamCredentialParams { @@ -266,9 +288,13 @@ export interface GetIamCredentialParams { options?: GetIamCredentialOptions } -export interface GetIamCredentialResult { - id: IamCredentialId +export interface IamCredential { + id: string credentials: IamCredentials +} + +export interface GetIamCredentialResult { + credential: IamCredential updateCredentialsParams: UpdateCredentialsParams } diff --git a/runtimes/runtimes/auth/standalone/encryption.ts b/runtimes/runtimes/auth/standalone/encryption.ts index 4f9dbb5b..f5f3475b 100644 --- a/runtimes/runtimes/auth/standalone/encryption.ts +++ b/runtimes/runtimes/auth/standalone/encryption.ts @@ -124,11 +124,11 @@ export async function encryptIamResultWithKey( request: GetIamCredentialResult, key: string ): Promise { - 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) } + request.credential.credentials = { + accessKeyId: await encryptObjectWithKey(request.credential.credentials.accessKeyId, key), + secretAccessKey: await encryptObjectWithKey(request.credential.credentials.secretAccessKey, key), + ...(request.credential.credentials.sessionToken + ? { sessionToken: await encryptObjectWithKey(request.credential.credentials.sessionToken, key) } : {}), } if (!request.updateCredentialsParams.encrypted) { From 9d1ad6b61210746faa9d65c60dc1c4b92cbb97d2 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 18 Jul 2025 10:56:44 -0400 Subject: [PATCH 16/19] fix: change IamCredential fields --- runtimes/protocol/identity-management.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 99616127..0152fd7b 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -289,7 +289,8 @@ export interface GetIamCredentialParams { } export interface IamCredential { - id: string + id: IamCredentialId + kinds: ProfileKind[] credentials: IamCredentials } @@ -344,7 +345,7 @@ export const invalidateSsoTokenRequestType = new ProtocolRequestType< // invalidateStsCredential export interface InvalidateStsCredentialParams { - profileName: string + iamCredentialId: IamCredentialId } export interface InvalidateStsCredentialResult { From ec9c7c4328071810b7b4cb19679e2f2f6f144f41 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Tue, 22 Jul 2025 14:00:49 -0400 Subject: [PATCH 17/19] fix: add error message for failed caller identity --- runtimes/protocol/identity-management.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 0152fd7b..39430c07 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -37,6 +37,7 @@ export const AwsErrorCodes = { 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_CALLER_IDENTITY_NOT_FOUND: 'E_CALLER_IDENTITY_NOT_FOUND', E_MFA_REQUIRED: 'E_MFA_REQUIRED', E_PERMISSION_DENIED: 'E_PERMISSION_DENIED', E_TIMEOUT: 'E_TIMEOUT', From 3444e0ad83c37f29d843a31b0f92fc13ec509ec7 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 25 Jul 2025 09:59:26 -0400 Subject: [PATCH 18/19] fix: add mfaSerial to MfaCode output --- runtimes/protocol/identity-management.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 39430c07..1e0c7ed7 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -310,12 +310,13 @@ export const getIamCredentialRequestType = new ProtocolRequestType< // getMfaCode export interface GetMfaCodeParams { - mfaSerial: string profileName: string + mfaSerial?: string } export interface GetMfaCodeResult { code: string + mfaSerial: string } export const getMfaCodeRequestType = new ProtocolRequestType< From 05618366fe81def939b82b8df5823d80ac658d28 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 25 Jul 2025 17:20:54 -0400 Subject: [PATCH 19/19] feat: add credential override --- runtimes/protocol/identity-management.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts index 1e0c7ed7..69f93bc5 100644 --- a/runtimes/protocol/identity-management.ts +++ b/runtimes/protocol/identity-management.ts @@ -277,11 +277,13 @@ export type IamCredentialId = string // Opaque identifier export interface GetIamCredentialOptions { callStsOnInvalidIamCredential?: boolean permissionSet?: string[] + credentialOverride?: IamCredentials } export const getIamCredentialOptionsDefaults = { callStsOnInvalidIamCredential: true, permissionSet: PermissionSets.Q, + credentialOverride: undefined, } satisfies GetIamCredentialOptions export interface GetIamCredentialParams {