Skip to content

Commit ea6a170

Browse files
yuxianrzyuxianrz
committed
fix: sync with changes to invalidateStsCredentials in runtimes
1 parent 5eb4ecd commit ea6a170

File tree

4 files changed

+80
-39
lines changed

4 files changed

+80
-39
lines changed

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

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -292,11 +292,6 @@ async function postStartLanguageServer(
292292
sso: {
293293
startUrl: AuthUtil.instance.connection?.startUrl,
294294
},
295-
// Add IAM credentials metadata
296-
iam: {
297-
region: AuthUtil.instance.connection?.region,
298-
accesskey: AuthUtil.instance.connection?.accessKey,
299-
},
300295
}
301296
})
302297

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

Lines changed: 37 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,6 @@ import {
5151
SsoSession,
5252
GetMfaCodeParams,
5353
getMfaCodeRequestType,
54-
5554
} from '@aws/language-server-runtimes/protocol'
5655
import { LanguageClient } from 'vscode-languageclient'
5756
import { getLogger } from '../shared/logger/logger'
@@ -73,9 +72,7 @@ export const notificationTypes = {
7372
getConnectionMetadata: new RequestType<undefined, ConnectionMetadata, Error>(
7473
getConnectionMetadataRequestType.method
7574
),
76-
getMfaCode: new RequestType<GetMfaCodeParams, ResponseMessage, Error>(
77-
getMfaCodeRequestType.method
78-
)
75+
getMfaCode: new RequestType<GetMfaCodeParams, ResponseMessage, Error>(getMfaCodeRequestType.method),
7976
}
8077

8178
export type AuthState = 'notConnected' | 'connected' | 'expired'
@@ -101,7 +98,10 @@ export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoToken
10198
*/
10299
export class LanguageClientAuth {
103100
readonly #ssoCacheWatcher = getCacheFileWatcher(getCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq))
104-
readonly #stsCacheWatcher = getCacheFileWatcher(getStsCacheDir(), getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq))
101+
readonly #stsCacheWatcher = getCacheFileWatcher(
102+
getStsCacheDir(),
103+
getFlareCacheFileName(VSCODE_EXTENSION_ID.amazonq)
104+
)
105105

106106
constructor(
107107
private readonly client: LanguageClient,
@@ -279,7 +279,7 @@ export class LanguageClientAuth {
279279

280280
invalidateStsCredential(tokenId: string) {
281281
return this.client.sendRequest(invalidateStsCredentialRequestType.method, {
282-
profileName: tokenId,
282+
iamCredentialId: tokenId,
283283
} satisfies InvalidateStsCredentialParams) as Promise<InvalidateStsCredentialResult>
284284
}
285285

@@ -309,7 +309,9 @@ export abstract class BaseLogin {
309309
protected loginType: LoginType | undefined
310310
protected connectionState: AuthState = 'notConnected'
311311
protected cancellationToken: CancellationTokenSource | undefined
312-
protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string } | undefined
312+
protected _data:
313+
| { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string }
314+
| undefined
313315

314316
constructor(
315317
public readonly profileName: string,
@@ -536,7 +538,7 @@ export class IamLogin extends BaseLogin {
536538
)
537539
}
538540

539-
async login(opts: { accessKey: string; secretKey: string, sessionToken?: string, roleArn?: string }) {
541+
async login(opts: { accessKey: string; secretKey: string; sessionToken?: string; roleArn?: string }) {
540542
await this.updateProfile(opts)
541543
return this._getIamCredential(true)
542544
}
@@ -559,13 +561,27 @@ export class IamLogin extends BaseLogin {
559561
// TODO: DeleteProfile api in Identity Service (this doesn't exist yet)
560562
}
561563

562-
async updateProfile(opts: { accessKey: string; secretKey: string, sessionToken?: string, roleArn?: string }) {
564+
async updateProfile(opts: { accessKey: string; secretKey: string; sessionToken?: string; roleArn?: string }) {
563565
if (opts.roleArn) {
564566
const sourceProfile = this.profileName + '-source'
565-
await this.lspAuth.updateIamProfile(sourceProfile, opts.accessKey, opts.secretKey, opts.sessionToken, '', '')
567+
await this.lspAuth.updateIamProfile(
568+
sourceProfile,
569+
opts.accessKey,
570+
opts.secretKey,
571+
opts.sessionToken,
572+
'',
573+
''
574+
)
566575
await this.lspAuth.updateIamProfile(this.profileName, '', '', '', opts.roleArn, sourceProfile)
567576
} else {
568-
await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey, opts.sessionToken, '', '')
577+
await this.lspAuth.updateIamProfile(
578+
this.profileName,
579+
opts.accessKey,
580+
opts.secretKey,
581+
opts.sessionToken,
582+
'',
583+
''
584+
)
569585
}
570586
}
571587

@@ -587,10 +603,10 @@ export class IamLogin extends BaseLogin {
587603
async getCredential() {
588604
const response = await this._getIamCredential(false)
589605
const credentials: IamCredentials = {
590-
accessKeyId: await this.decrypt(response.credentials.accessKeyId),
591-
secretAccessKey: await this.decrypt(response.credentials.secretAccessKey),
592-
sessionToken: response.credentials.sessionToken
593-
? await this.decrypt(response.credentials.sessionToken)
606+
accessKeyId: await this.decrypt(response.credential.credentials.accessKeyId),
607+
secretAccessKey: await this.decrypt(response.credential.credentials.secretAccessKey),
608+
sessionToken: response.credential.credentials.sessionToken
609+
? await this.decrypt(response.credential.credentials.sessionToken)
594610
: undefined,
595611
}
596612
return {
@@ -629,13 +645,13 @@ export class IamLogin extends BaseLogin {
629645
}
630646

631647
// Update cached credentials and credential ID
632-
if (response.credentials.accessKeyId && response.credentials.secretAccessKey) {
648+
if (response.credential.credentials.accessKeyId && response.credential.credentials.secretAccessKey) {
633649
this._data = {
634-
accessKey: response.credentials.accessKeyId,
635-
secretKey: response.credentials.secretAccessKey,
636-
sessionToken: response.credentials.sessionToken,
650+
accessKey: response.credential.credentials.accessKeyId,
651+
secretKey: response.credential.credentials.secretAccessKey,
652+
sessionToken: response.credential.credentials.sessionToken,
637653
}
638-
this.iamCredentialId = response.id
654+
this.iamCredentialId = response.credential.id
639655
}
640656
this.updateConnectionState('connected')
641657
return response
@@ -647,7 +663,7 @@ export class IamLogin extends BaseLogin {
647663
this.updateConnectionState('expired')
648664
return
649665
} else if (params.kind === StsCredentialChangedKind.Refreshed) {
650-
this.eventEmitter.fire({ id: this.profileName, state: 'refreshed' })
666+
this.eventEmitter.fire({ id: this.iamCredentialId, state: 'refreshed' })
651667
}
652668
}
653669
}

‎packages/core/src/codewhisperer/util/authUtil.ts‎

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,16 @@ import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthr
3030
import { setContext } from '../../shared/vscode/setContext'
3131
import { openUrl } from '../../shared/utilities/vsCodeUtils'
3232
import { telemetry } from '../../shared/telemetry/telemetry'
33-
import { AuthStateEvent, cacheChangedEvent, stsCacheChangedEvent, LanguageClientAuth, Login, SsoLogin, IamLogin, LoginTypes } from '../../auth/auth2'
33+
import {
34+
AuthStateEvent,
35+
cacheChangedEvent,
36+
stsCacheChangedEvent,
37+
LanguageClientAuth,
38+
Login,
39+
SsoLogin,
40+
IamLogin,
41+
LoginTypes,
42+
} from '../../auth/auth2'
3443
import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants'
3544
import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
3645
import { RegionProfileManager } from '../region/regionProfileManager'
@@ -64,7 +73,13 @@ export interface IAuthProvider {
6473
getToken(): Promise<string>
6574
getIamCredential(): Promise<IamCredentials>
6675
readonly profileName: string
67-
readonly connection?: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string; sessionToken?: string }
76+
readonly connection?: {
77+
startUrl?: string
78+
region?: string
79+
accessKey?: string
80+
secretKey?: string
81+
sessionToken?: string
82+
}
6883
}
6984

7085
/**
@@ -176,13 +191,23 @@ export class AuthUtil implements IAuthProvider {
176191
}
177192

178193
// Log in using IAM or STS credentials
179-
async loginIam(accessKey: string, secretKey: string, sessionToken?: string, roleArn?: string): Promise<GetIamCredentialResult | undefined> {
194+
async loginIam(
195+
accessKey: string,
196+
secretKey: string,
197+
sessionToken?: string,
198+
roleArn?: string
199+
): Promise<GetIamCredentialResult | undefined> {
180200
let response: GetIamCredentialResult | undefined
181201
// Create IAM login session
182202
if (!this.isIamSession()) {
183203
this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
184204
}
185-
response = await (this.session as IamLogin).login({ accessKey: accessKey, secretKey: secretKey, sessionToken: sessionToken, roleArn: roleArn })
205+
response = await (this.session as IamLogin).login({
206+
accessKey: accessKey,
207+
secretKey: secretKey,
208+
sessionToken: sessionToken,
209+
roleArn: roleArn,
210+
})
186211
await showAmazonQWalkthroughOnce()
187212
return response
188213
}
@@ -216,7 +241,7 @@ export class AuthUtil implements IAuthProvider {
216241
if (this.session) {
217242
const credential = (await this.session.getCredential()).credential
218243
if (typeof credential !== 'object') {
219-
throw new ToolkitError('Cannot get token with SSO session')
244+
throw new ToolkitError('Cannot get credential with SSO session')
220245
}
221246
return credential
222247
} else {

‎packages/core/src/test/credentials/auth2.test.ts‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,9 @@ describe('LanguageClientAuth', () => {
229229
const result = await auth.invalidateStsCredential(profileName)
230230

231231
sinon.assert.calledOnce(client.sendRequest)
232-
sinon.assert.calledWith(client.sendRequest, invalidateStsCredentialRequestType.method, { profileName: profileName })
232+
sinon.assert.calledWith(client.sendRequest, invalidateStsCredentialRequestType.method, {
233+
profileName: profileName,
234+
})
233235
sinon.assert.match(result, { success: true })
234236
})
235237
})
@@ -638,11 +640,14 @@ describe('IamLogin', () => {
638640
}
639641

640642
const mockGetIamCredentialResponse: GetIamCredentialResult = {
641-
id: 'test-credential-id',
642-
credentials: {
643-
accessKeyId: 'encrypted-access-key',
644-
secretAccessKey: 'encrypted-secret-key',
645-
sessionToken: 'encrypted-session-token',
643+
credential: {
644+
id: 'test-credential-id',
645+
kinds: [],
646+
credentials: {
647+
accessKeyId: 'encrypted-access-key',
648+
secretAccessKey: 'encrypted-secret-key',
649+
sessionToken: 'encrypted-session-token',
650+
},
646651
},
647652
updateCredentialsParams: {
648653
data: 'credential-data',
@@ -674,7 +679,7 @@ describe('IamLogin', () => {
674679
sinon.assert.calledWith(lspAuth.updateIamProfile, profileName, loginOpts.accessKey, loginOpts.secretKey)
675680
sinon.assert.calledOnce(lspAuth.getIamCredential)
676681
sinon.assert.match(iamLogin.getConnectionState(), 'connected')
677-
sinon.assert.match(response.id, 'test-credential-id')
682+
sinon.assert.match(response.credential.id, 'test-credential-id')
678683
})
679684
})
680685

@@ -697,7 +702,7 @@ describe('IamLogin', () => {
697702

698703
sinon.assert.calledOnce(lspAuth.getIamCredential)
699704
sinon.assert.match(iamLogin.getConnectionState(), 'connected')
700-
sinon.assert.match(response.id, 'test-credential-id')
705+
sinon.assert.match(response.credential.id, 'test-credential-id')
701706
})
702707
})
703708

0 commit comments

Comments
 (0)