Skip to content

Commit 87492de

Browse files
yuxianrzyuxianrz
committed
add sts credentials and webview, add autofill for IAM and STS
1 parent d6ab092 commit 87492de

File tree

13 files changed

+236
-80
lines changed

13 files changed

+236
-80
lines changed

‎packages/amazonq/src/extensionNode.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
100100
async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
101101
const state = AuthUtil.instance.getAuthState()
102102

103-
if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
104-
getLogger().error('Current Amazon Q connection is not SSO')
103+
if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || AuthUtil.instance.isIamSession() || isSageMaker())) {
104+
getLogger().error('Current Amazon Q connection is not SSO nor IAM')
105105
}
106106

107107
return {

‎packages/amazonq/src/lsp/client.ts‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,9 @@ export async function startLanguageServer(
164164
},
165165
credentials: {
166166
providesBearerToken: true,
167+
// Add IAM credentials support
168+
providesIamCredentials: true,
169+
supportsAssumeRole: true,
167170
},
168171
},
169172
/**
@@ -211,9 +214,10 @@ export async function startLanguageServer(
211214

212215
/** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */
213216
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
217+
const activeProfile = AuthUtil.instance.regionProfileManager.activeRegionProfile
214218
void pushConfigUpdate(client, {
215219
type: 'profile',
216-
profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
220+
profileArn: activeProfile?.arn,
217221
})
218222
})
219223

@@ -286,6 +290,11 @@ async function postStartLanguageServer(
286290
sso: {
287291
startUrl: AuthUtil.instance.connection?.startUrl,
288292
},
293+
// Add IAM credentials metadata
294+
iam: {
295+
region: AuthUtil.instance.connection?.region,
296+
accesskey: AuthUtil.instance.connection?.accessKey,
297+
},
289298
}
290299
})
291300

‎packages/core/src/auth/auth2.ts‎

Lines changed: 67 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,20 @@ import {
1212
GetIamCredentialParams,
1313
getIamCredentialRequestType,
1414
GetIamCredentialResult,
15+
InvalidateStsCredentialResult,
1516
IamIdentityCenterSsoTokenSource,
1617
InvalidateSsoTokenParams,
18+
InvalidateStsCredentialParams,
1719
invalidateSsoTokenRequestType,
20+
invalidateStsCredentialRequestType,
1821
ProfileKind,
1922
UpdateProfileParams,
2023
updateProfileRequestType,
2124
SsoTokenChangedParams,
22-
// StsCredentialChangedParams,
25+
StsCredentialChangedParams,
26+
StsCredentialChangedKind,
2327
ssoTokenChangedRequestType,
24-
// stsCredentialChangedRequestType,
28+
stsCredentialChangedRequestType,
2529
AwsBuilderIdSsoTokenSource,
2630
UpdateCredentialsParams,
2731
AwsErrorCodes,
@@ -36,6 +40,7 @@ import {
3640
iamCredentialsDeleteNotificationType,
3741
bearerCredentialsDeleteNotificationType,
3842
bearerCredentialsUpdateRequestType,
43+
SsoTokenChangedKind,
3944
RequestType,
4045
ResponseMessage,
4146
NotificationType,
@@ -44,16 +49,12 @@ import {
4449
iamCredentialsUpdateRequestType,
4550
Profile,
4651
SsoSession,
47-
SsoTokenChangedKind,
48-
// invalidateStsCredentialRequestType,
49-
// InvalidateStsCredentialParams,
50-
// InvalidateStsCredentialResult,
5152
} from '@aws/language-server-runtimes/protocol'
5253
import { LanguageClient } from 'vscode-languageclient'
5354
import { getLogger } from '../shared/logger/logger'
5455
import { ToolkitError } from '../shared/errors'
5556
import { useDeviceFlow } from './sso/ssoAccessTokenProvider'
56-
import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName } from './sso/cache'
57+
import { getCacheDir, getCacheFileWatcher, getFlareCacheFileName, getStsCacheDir } from './sso/cache'
5758
import { VSCODE_EXTENSION_ID } from '../shared/extensions'
5859
import { IamCredentials } from '@aws/language-server-runtimes-types'
5960

@@ -83,6 +84,8 @@ export type LoginType = (typeof LoginTypes)[keyof typeof LoginTypes]
8384

8485
export type cacheChangedEvent = 'delete' | 'create'
8586

87+
export type stsCacheChangedEvent = 'delete' | 'create'
88+
8689
export type Login = SsoLogin | IamLogin
8790

8891
export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource
@@ -92,6 +95,7 @@ export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoToken
9295
*/
9396
export class LanguageClientAuth {
9497
readonly #ssoCacheWatcher = getCacheFileWatcher(getCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq))
98+
readonly #stsCacheWatcher = getCacheFileWatcher(getStsCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq))
9599

96100
constructor(
97101
private readonly client: LanguageClient,
@@ -103,6 +107,10 @@ export class LanguageClientAuth {
103107
return this.#ssoCacheWatcher
104108
}
105109

110+
public get stsCacheWatcher() {
111+
return this.#stsCacheWatcher
112+
}
113+
106114
getSsoToken(
107115
tokenSource: TokenSource,
108116
login: boolean = false,
@@ -155,6 +163,7 @@ export class LanguageClientAuth {
155163
sso_session: profileName,
156164
aws_access_key_id: '',
157165
aws_secret_access_key: '',
166+
role_arn: '',
158167
},
159168
},
160169
ssoSession: {
@@ -259,24 +268,29 @@ export class LanguageClientAuth {
259268
} satisfies InvalidateSsoTokenParams) as Promise<InvalidateSsoTokenResult>
260269
}
261270

262-
// invalidateStsCredential(tokenId: string) {
263-
// return this.client.sendRequest(invalidateStsCredentialRequestType.method, {
264-
// stsCredentialId: tokenId,
265-
// } satisfies InvalidateStsCredentialParams) as Promise<InvalidateStsCredentialResult>
266-
// }
271+
invalidateStsCredential(tokenId: string) {
272+
return this.client.sendRequest(invalidateStsCredentialRequestType.method, {
273+
profileName: tokenId,
274+
} satisfies InvalidateStsCredentialParams) as Promise<InvalidateStsCredentialResult>
275+
}
267276

268277
registerSsoTokenChangedHandler(ssoTokenChangedHandler: (params: SsoTokenChangedParams) => any) {
269278
this.client.onNotification(ssoTokenChangedRequestType.method, ssoTokenChangedHandler)
270279
}
271280

272-
// registerStsCredentialChangedHandler(stsCredentialChangedHandler: (params: StsCredentialChangedParams) => any) {
273-
// this.client.onNotification(stsCredentialChangedRequestType.method, stsCredentialChangedHandler)
274-
// }
281+
registerStsCredentialChangedHandler(stsCredentialChangedHandler: (params: StsCredentialChangedParams) => any) {
282+
this.client.onNotification(stsCredentialChangedRequestType.method, stsCredentialChangedHandler)
283+
}
275284

276285
registerCacheWatcher(cacheChangedHandler: (event: cacheChangedEvent) => any) {
277286
this.cacheWatcher.onDidCreate(() => cacheChangedHandler('create'))
278287
this.cacheWatcher.onDidDelete(() => cacheChangedHandler('delete'))
279288
}
289+
290+
registerStsCacheWatcher(stsCacheChangedHandler: (event: stsCacheChangedEvent) => any) {
291+
this.stsCacheWatcher.onDidCreate(() => stsCacheChangedHandler('create'))
292+
this.stsCacheWatcher.onDidDelete(() => stsCacheChangedHandler('delete'))
293+
}
280294
}
281295

282296
/**
@@ -286,7 +300,7 @@ export abstract class BaseLogin {
286300
protected loginType: LoginType | undefined
287301
protected connectionState: AuthState = 'notConnected'
288302
protected cancellationToken: CancellationTokenSource | undefined
289-
protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined
303+
protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string } | undefined
290304

291305
constructor(
292306
public readonly profileName: string,
@@ -504,16 +518,16 @@ export class SsoLogin extends BaseLogin {
504518
export class IamLogin extends BaseLogin {
505519
// Cached information from the identity server for easy reference
506520
override readonly loginType = LoginTypes.IAM
507-
// private iamCredentialId: string | undefined
521+
private iamCredentialId: string | undefined
508522

509523
constructor(profileName: string, lspAuth: LanguageClientAuth, eventEmitter: vscode.EventEmitter<AuthStateEvent>) {
510524
super(profileName, lspAuth, eventEmitter)
511-
// lspAuth.registerStsCredentialChangedHandler((params: StsCredentialChangedParams) =>
512-
// this.stsCredentialChangedHandler(params)
513-
// )
525+
lspAuth.registerStsCredentialChangedHandler((params: StsCredentialChangedParams) =>
526+
this.stsCredentialChangedHandler(params)
527+
)
514528
}
515529

516-
async login(opts: { accessKey: string; secretKey: string }) {
530+
async login(opts: { accessKey: string; secretKey: string, sessionToken?: string, roleArn?: string }) {
517531
await this.updateProfile(opts)
518532
return this._getIamCredential(true)
519533
}
@@ -526,36 +540,30 @@ export class IamLogin extends BaseLogin {
526540
}
527541

528542
async logout() {
529-
// if (this.iamCredentialId) {
530-
// await this.lspAuth.invalidateIamCredential(this.iamCredentialId)
531-
// }
543+
if (this.iamCredentialId) {
544+
await this.lspAuth.invalidateStsCredential(this.iamCredentialId)
545+
}
532546
await this.lspAuth.updateIamProfile(this.profileName, '', '', '', '', '')
533547
await this.lspAuth.updateIamProfile(this.profileName + '-source', '', '', '', '', '')
534548
this.updateConnectionState('notConnected')
535549
this._data = undefined
536550
// TODO: DeleteProfile api in Identity Service (this doesn't exist yet)
537551
}
538552

539-
async updateProfile(opts: { accessKey: string; secretKey: string }) {
540-
await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey)
541-
this._data = {
542-
accessKey: opts.accessKey,
543-
secretKey: opts.secretKey,
553+
async updateProfile(opts: { accessKey: string; secretKey: string, sessionToken?: string, roleArn?: string }) {
554+
if (opts.roleArn) {
555+
const sourceProfile = this.profileName + '-source'
556+
await this.lspAuth.updateIamProfile(sourceProfile, opts.accessKey, opts.secretKey, opts.sessionToken, '', '')
557+
await this.lspAuth.updateIamProfile(this.profileName, '', '', '', opts.roleArn, sourceProfile)
558+
} else {
559+
await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey, opts.sessionToken, '', '')
544560
}
545561
}
546562

547563
/**
548564
* Restore the connection state and connection details to memory, if they exist.
549565
*/
550566
async restore() {
551-
const sessionData = await this.getProfile()
552-
const credentials = sessionData?.profile?.settings
553-
if (credentials?.aws_access_key_id && credentials?.aws_secret_access_key) {
554-
this._data = {
555-
accessKey: credentials.aws_access_key_id,
556-
secretKey: credentials.aws_secret_access_key,
557-
}
558-
}
559567
try {
560568
await this._getIamCredential(false)
561569
} catch (err) {
@@ -595,8 +603,10 @@ export class IamLogin extends BaseLogin {
595603
} catch (err: any) {
596604
switch (err.data?.awsErrorCode) {
597605
case AwsErrorCodes.E_CANCELLED:
598-
case AwsErrorCodes.E_SSO_SESSION_NOT_FOUND:
606+
case AwsErrorCodes.E_INVALID_PROFILE:
599607
case AwsErrorCodes.E_PROFILE_NOT_FOUND:
608+
case AwsErrorCodes.E_CANNOT_CREATE_STS_CREDENTIAL:
609+
case AwsErrorCodes.E_INVALID_STS_CREDENTIAL:
600610
this.updateConnectionState('notConnected')
601611
break
602612
default:
@@ -609,19 +619,27 @@ export class IamLogin extends BaseLogin {
609619
this.cancellationToken = undefined
610620
}
611621

612-
// this.iamCredentialId = response.id
622+
// Update cached credentials and credential ID
623+
if (response.credentials.accessKeyId && response.credentials.secretAccessKey) {
624+
this._data = {
625+
accessKey: response.credentials.accessKeyId,
626+
secretKey: response.credentials.secretAccessKey,
627+
sessionToken: response.credentials.sessionToken,
628+
}
629+
this.iamCredentialId = response.id
630+
}
613631
this.updateConnectionState('connected')
614632
return response
615633
}
616634

617-
// private stsCredentialChangedHandler(params: StsCredentialChangedParams) {
618-
// if (params.stsCredentialId === this.iamCredentialId) {
619-
// if (params.kind === CredentialChangedKind.Expired) {
620-
// this.updateConnectionState('expired')
621-
// return
622-
// } else if (params.kind === CredentialChangedKind.Refreshed) {
623-
// this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' })
624-
// }
625-
// }
626-
// }
635+
private stsCredentialChangedHandler(params: StsCredentialChangedParams) {
636+
if (params.stsCredentialId === this.iamCredentialId) {
637+
if (params.kind === StsCredentialChangedKind.Expired) {
638+
this.updateConnectionState('expired')
639+
return
640+
} else if (params.kind === StsCredentialChangedKind.Refreshed) {
641+
this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' })
642+
}
643+
}
644+
}
627645
}

‎packages/core/src/auth/sso/cache.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,9 @@ export interface SsoCache {
3636
}
3737

3838
const defaultCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/sso/cache')
39+
const defaultStsCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/cli/cache')
3940
export const getCacheDir = () => DevSettings.instance.get('ssoCacheDirectory', defaultCacheDir())
41+
export const getStsCacheDir = () => DevSettings.instance.get('stsCacheDirectory', defaultStsCacheDir())
4042

4143
export function getCache(directory = getCacheDir()): SsoCache {
4244
return {

0 commit comments

Comments
 (0)