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

liramon1 · web-flow · commit 6276983b2372 · 2025-07-17T09:42:21.000-04:00
This change adds request handlers and their associated types to manage IAM and STS credentials from aws-lsp-identity.
diff --git a/runtimes/protocol/identity-management.ts b/runtimes/protocol/identity-management.ts
@@ -1,4 +1,5 @@
 import {
+    IamCredentials,
     LSPErrorCodes,
     ProgressType,
     ProtocolNotificationType,
@@ -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',
     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',
@@ -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
 
@@ -64,6 +82,18 @@ 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
+        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
     }
 }
 
@@ -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
+    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<
+    GetMfaCodeParams,
+    GetMfaCodeResult,
+    never,
+    AwsResponseError,
+    void
+>('aws/identity/getMfaCode')
+
 // invalidateSsoToken
 export interface InvalidateSsoTokenParams {
     ssoTokenId: SsoTokenId
@@ -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<
+    InvalidateStsCredentialParams,
+    InvalidateStsCredentialResult,
+    never,
+    AwsResponseError,
+    void
+>('aws/identity/invalidateStsCredential')
+
 // ssoTokenChanged
 export type Expired = 'Expired'
 export type Refreshed = 'Refreshed'
@@ -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>(
+    'aws/identity/stsCredentialChanged'
+)
diff --git 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<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:
diff --git a/runtimes/runtimes/base-runtime.ts b/runtimes/runtimes/base-runtime.ts
@@ -85,11 +85,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'
@@ -208,8 +212,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
diff --git a/runtimes/runtimes/standalone.ts b/runtimes/runtimes/standalone.ts
@@ -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,
@@ -50,6 +56,7 @@ import {
     getSsoTokenRequestType,
     IdentityManagement,
     invalidateSsoTokenRequestType,
+    invalidateStsCredentialRequestType,
     listProfilesRequestType,
     ssoTokenChangedRequestType,
     updateProfileRequestType,
@@ -302,31 +309,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()
diff --git a/runtimes/server-interface/identity-management.ts b/runtimes/server-interface/identity-management.ts
@@ -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'
 
@@ -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>
 }
diff --git 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 = {

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,7 @@ export type IamCredentials = {`
`2`	`2`	`readonly accessKeyId: string`
`3`	`3`	`readonly secretAccessKey: string`
`4`	`4`	`readonly sessionToken?: string`
	`5`	`+ readonly expiration?: Date`
`5`	`6`	`}`
`6`	`7`
`7`	`8`	`export type BearerCredentials = {`