Skip to content

Commit 5eb4ecd

Browse files
yuxianrzyuxianrz
committed
fix: sync with IAM PR changes
1 parent 1fb9e4e commit 5eb4ecd

File tree

4 files changed

+275
-34
lines changed

4 files changed

+275
-34
lines changed

‎packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts‎

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -427,18 +427,18 @@ describe('AuthUtil', async function () {
427427
login: sinon.stub().resolves(mockResponse),
428428
loginType: 'iam',
429429
}
430-
430+
431431
sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
432432

433433
const response = await auth.loginIam('accessKey', 'secretKey', 'sessionToken')
434434

435435
assert.ok(mockIamLogin.login.calledOnce)
436-
assert.ok(mockIamLogin.login.calledWith({
437-
accessKey: 'accessKey',
438-
secretKey: 'secretKey',
439-
sessionToken: 'sessionToken',
440-
roleArn: undefined,
441-
}))
436+
assert.ok(
437+
mockIamLogin.login.calledWith({
438+
accessKey: 'accessKey',
439+
secretKey: 'secretKey',
440+
})
441+
)
442442
assert.strictEqual(response, mockResponse)
443443
})
444444

@@ -462,7 +462,7 @@ describe('AuthUtil', async function () {
462462

463463
sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
464464

465-
const response = await auth.login_iam('accessKey', 'secretKey', 'sessionToken', 'arn:aws:iam::123456789012:role/TestRole')
465+
const response = await auth.loginIam('accessKey', 'secretKey', 'sessionToken', 'arn:aws:iam::123456789012:role/TestRole')
466466

467467
assert.ok(mockIamLogin.login.calledOnce)
468468
assert.ok(mockIamLogin.login.calledWith({
@@ -568,23 +568,5 @@ describe('AuthUtil', async function () {
568568
assert.ok(mockLspAuth.updateIamCredential.called)
569569
assert.strictEqual(mockLspAuth.updateIamCredential.firstCall.args[0].data, 'fake-data')
570570
})
571-
572-
it('cleans up IAM credential when connection expires', async function () {
573-
const mockSession = { loginType: 'iam' }
574-
;(auth as any).session = mockSession
575-
576-
await (auth as any).stateChangeHandler({ state: 'expired' })
577-
578-
assert.ok(mockLspAuth.deleteIamCredential.called)
579-
})
580-
581-
it('deletes IAM credential when disconnected', async function () {
582-
const mockSession = { loginType: 'iam' }
583-
;(auth as any).session = mockSession
584-
585-
await (auth as any).stateChangeHandler({ state: 'notConnected' })
586-
587-
assert.ok(mockLspAuth.deleteIamCredential.called)
588-
})
589571
})
590572
})

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

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -109,11 +109,11 @@ export class AuthUtil implements IAuthProvider {
109109
}
110110

111111
isSsoSession(): boolean {
112-
return this.session?.loginType === LoginTypes.SSO || this.session instanceof SsoLogin
112+
return this.session instanceof SsoLogin
113113
}
114114

115115
isIamSession(): boolean {
116-
return this.session?.loginType === LoginTypes.IAM || this.session instanceof IamLogin
116+
return this.session instanceof IamLogin
117117
}
118118

119119
/**
@@ -176,11 +176,7 @@ export class AuthUtil implements IAuthProvider {
176176
}
177177

178178
// Log in using IAM or STS credentials
179-
async loginIam(
180-
accessKey: string,
181-
secretKey: string,
182-
sessionToken?: string
183-
): Promise<GetIamCredentialResult | undefined> {
179+
async loginIam(accessKey: string, secretKey: string, sessionToken?: string, roleArn?: string): Promise<GetIamCredentialResult | undefined> {
184180
let response: GetIamCredentialResult | undefined
185181
// Create IAM login session
186182
if (!this.isIamSession()) {

‎packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
206206
await globals.globalState.update('recentRoleArn', { roleArn: roleArn })
207207
const runAuth = async (): Promise<AuthError | undefined> => {
208208
try {
209-
await AuthUtil.instance.loginIam(accessKey, secretKey)
209+
await AuthUtil.instance.loginIam(accessKey, secretKey, sessionToken, roleArn)
210210
} catch (e) {
211211
getLogger().error('Failed submitting credentials %O', e)
212212
const message = e instanceof Error ? e.message : e as string

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

Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,79 @@ describe('LanguageClientAuth', () => {
182182
})
183183
})
184184

185+
describe('updateIamCredential', () => {
186+
it('sends request', async () => {
187+
const updateParams: UpdateCredentialsParams = {
188+
data: 'credential-data',
189+
encrypted: true,
190+
}
191+
192+
await auth.updateIamCredential(updateParams)
193+
194+
sinon.assert.calledOnce(client.sendRequest)
195+
sinon.assert.calledWith(client.sendRequest, iamCredentialsUpdateRequestType.method, updateParams)
196+
})
197+
})
198+
199+
describe('deleteIamCredential', () => {
200+
it('sends notification', async () => {
201+
auth.deleteIamCredential()
202+
203+
sinon.assert.calledOnce(client.sendNotification)
204+
sinon.assert.calledWith(client.sendNotification, iamCredentialsDeleteNotificationType.method)
205+
})
206+
})
207+
208+
describe('getIamCredential', () => {
209+
it('sends correct request parameters', async () => {
210+
await auth.getIamCredential(profileName, true)
211+
212+
sinon.assert.calledOnce(client.sendRequest)
213+
sinon.assert.calledWith(
214+
client.sendRequest,
215+
sinon.match.any,
216+
sinon.match({
217+
profileName: profileName,
218+
options: {
219+
callStsOnInvalidIamCredential: true,
220+
},
221+
})
222+
)
223+
})
224+
})
225+
226+
describe('invalidateStsCredential', () => {
227+
it('sends request', async () => {
228+
client.sendRequest.resolves({ success: true })
229+
const result = await auth.invalidateStsCredential(profileName)
230+
231+
sinon.assert.calledOnce(client.sendRequest)
232+
sinon.assert.calledWith(client.sendRequest, invalidateStsCredentialRequestType.method, { profileName: profileName })
233+
sinon.assert.match(result, { success: true })
234+
})
235+
})
236+
237+
describe('registerStsCredentialChangedHandler', () => {
238+
it('registers the handler correctly', () => {
239+
const handler = sinon.spy()
240+
241+
auth.registerStsCredentialChangedHandler(handler)
242+
243+
sinon.assert.calledOnce(client.onNotification)
244+
sinon.assert.calledWith(client.onNotification, stsCredentialChangedRequestType.method, sinon.match.func)
245+
246+
const credentialChangedParams: StsCredentialChangedParams = {
247+
kind: StsCredentialChangedKind.Refreshed,
248+
stsCredentialId: 'test-credential-id',
249+
}
250+
const registeredHandler = client.onNotification.firstCall.args[1]
251+
registeredHandler(credentialChangedParams)
252+
253+
sinon.assert.calledOnce(handler)
254+
sinon.assert.calledWith(handler, credentialChangedParams)
255+
})
256+
})
257+
185258
describe('invalidateSsoToken', () => {
186259
it('sends request', async () => {
187260
client.sendRequest.resolves({ success: true })
@@ -551,3 +624,193 @@ describe('SsoLogin', () => {
551624
})
552625
})
553626
})
627+
628+
describe('IamLogin', () => {
629+
let lspAuth: sinon.SinonStubbedInstance<LanguageClientAuth>
630+
let iamLogin: IamLogin
631+
let eventEmitter: vscode.EventEmitter<any>
632+
let fireEventSpy: sinon.SinonSpy
633+
634+
const loginOpts = {
635+
accessKey: 'test-access-key',
636+
secretKey: 'test-secret-key',
637+
sessionToken: 'test-session-token',
638+
}
639+
640+
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',
646+
},
647+
updateCredentialsParams: {
648+
data: 'credential-data',
649+
},
650+
}
651+
652+
beforeEach(() => {
653+
lspAuth = sinon.createStubInstance(LanguageClientAuth)
654+
eventEmitter = new vscode.EventEmitter()
655+
fireEventSpy = sinon.spy(eventEmitter, 'fire')
656+
iamLogin = new IamLogin(profileName, lspAuth as any, eventEmitter)
657+
;(iamLogin as any).eventEmitter = eventEmitter
658+
;(iamLogin as any).connectionState = 'notConnected'
659+
})
660+
661+
afterEach(() => {
662+
sinon.restore()
663+
eventEmitter.dispose()
664+
})
665+
666+
describe('login', () => {
667+
it('updates profile and returns IAM credential', async () => {
668+
lspAuth.updateIamProfile.resolves()
669+
lspAuth.getIamCredential.resolves(mockGetIamCredentialResponse)
670+
671+
const response = await iamLogin.login(loginOpts)
672+
673+
sinon.assert.calledOnce(lspAuth.updateIamProfile)
674+
sinon.assert.calledWith(lspAuth.updateIamProfile, profileName, loginOpts.accessKey, loginOpts.secretKey)
675+
sinon.assert.calledOnce(lspAuth.getIamCredential)
676+
sinon.assert.match(iamLogin.getConnectionState(), 'connected')
677+
sinon.assert.match(response.id, 'test-credential-id')
678+
})
679+
})
680+
681+
describe('reauthenticate', () => {
682+
it('throws when not connected', async () => {
683+
;(iamLogin as any).connectionState = 'notConnected'
684+
try {
685+
await iamLogin.reauthenticate()
686+
sinon.assert.fail('Should have thrown an error')
687+
} catch (err) {
688+
sinon.assert.match((err as Error).message, 'Cannot reauthenticate when not connected.')
689+
}
690+
})
691+
692+
it('returns new IAM credential when connected', async () => {
693+
;(iamLogin as any).connectionState = 'connected'
694+
lspAuth.getIamCredential.resolves(mockGetIamCredentialResponse)
695+
696+
const response = await iamLogin.reauthenticate()
697+
698+
sinon.assert.calledOnce(lspAuth.getIamCredential)
699+
sinon.assert.match(iamLogin.getConnectionState(), 'connected')
700+
sinon.assert.match(response.id, 'test-credential-id')
701+
})
702+
})
703+
704+
describe('logout', () => {
705+
it('invalidates credential and updates state', async () => {
706+
;(iamLogin as any).iamCredentialId = 'test-credential-id'
707+
lspAuth.invalidateStsCredential.resolves({ success: true })
708+
lspAuth.updateIamProfile.resolves()
709+
710+
await iamLogin.logout()
711+
712+
sinon.assert.calledOnce(lspAuth.invalidateStsCredential)
713+
sinon.assert.calledWith(lspAuth.invalidateStsCredential, 'test-credential-id')
714+
sinon.assert.match(iamLogin.getConnectionState(), 'notConnected')
715+
sinon.assert.match(iamLogin.data, undefined)
716+
})
717+
})
718+
719+
describe('restore', () => {
720+
it('restores connection state', async () => {
721+
lspAuth.getIamCredential.resolves(mockGetIamCredentialResponse)
722+
723+
await iamLogin.restore()
724+
725+
sinon.assert.calledOnce(lspAuth.getIamCredential)
726+
sinon.assert.calledWith(lspAuth.getIamCredential, profileName, false)
727+
sinon.assert.match(iamLogin.getConnectionState(), 'connected')
728+
})
729+
})
730+
731+
describe('_getIamCredential', () => {
732+
const testErrorHandling = async (errorCode: string, expectedState: string) => {
733+
const error = new Error('Credential error')
734+
;(error as any).data = { awsErrorCode: errorCode }
735+
lspAuth.getIamCredential.rejects(error)
736+
737+
try {
738+
await (iamLogin as any)._getIamCredential(false)
739+
sinon.assert.fail('Should have thrown an error')
740+
} catch (err) {
741+
sinon.assert.match(err, error)
742+
}
743+
744+
sinon.assert.match(iamLogin.getConnectionState(), expectedState)
745+
}
746+
747+
const notConnectedErrors = [
748+
AwsErrorCodes.E_CANCELLED,
749+
AwsErrorCodes.E_INVALID_PROFILE,
750+
AwsErrorCodes.E_PROFILE_NOT_FOUND,
751+
AwsErrorCodes.E_CANNOT_CREATE_STS_CREDENTIAL,
752+
AwsErrorCodes.E_INVALID_STS_CREDENTIAL,
753+
]
754+
755+
for (const errorCode of notConnectedErrors) {
756+
it(`handles ${errorCode} error`, async () => {
757+
await testErrorHandling(errorCode, 'notConnected')
758+
})
759+
}
760+
761+
it('returns correct response and updates state', async () => {
762+
lspAuth.getIamCredential.resolves(mockGetIamCredentialResponse)
763+
764+
const response = await (iamLogin as any)._getIamCredential(true)
765+
766+
sinon.assert.calledWith(lspAuth.getIamCredential, profileName, true)
767+
sinon.assert.match(response, mockGetIamCredentialResponse)
768+
sinon.assert.match(iamLogin.getConnectionState(), 'connected')
769+
// Note: iamCredentialId is commented out in the implementation
770+
// sinon.assert.match((iamLogin as any).iamCredentialId, 'test-credential-id')
771+
})
772+
})
773+
774+
describe('stsCredentialChangedHandler', () => {
775+
beforeEach(() => {
776+
;(iamLogin as any).iamCredentialId = 'test-credential-id'
777+
;(iamLogin as any).connectionState = 'connected'
778+
})
779+
780+
it('updates state when credential expires', () => {
781+
;(iamLogin as any).stsCredentialChangedHandler({
782+
kind: StsCredentialChangedKind.Expired,
783+
stsCredentialId: 'test-credential-id',
784+
})
785+
786+
sinon.assert.match(iamLogin.getConnectionState(), 'expired')
787+
sinon.assert.calledOnce(fireEventSpy)
788+
sinon.assert.calledWith(fireEventSpy, {
789+
id: profileName,
790+
state: 'expired',
791+
})
792+
})
793+
794+
it('emits refresh event when credential is refreshed', () => {
795+
;(iamLogin as any).stsCredentialChangedHandler({
796+
kind: StsCredentialChangedKind.Refreshed,
797+
stsCredentialId: 'test-credential-id',
798+
})
799+
800+
sinon.assert.calledOnce(fireEventSpy)
801+
sinon.assert.calledWith(fireEventSpy, {
802+
id: profileName,
803+
state: 'refreshed',
804+
})
805+
})
806+
807+
it('does not emit event for different credential ID', () => {
808+
;(iamLogin as any).stsCredentialChangedHandler({
809+
kind: StsCredentialChangedKind.Refreshed,
810+
stsCredentialId: 'different-credential-id',
811+
})
812+
813+
sinon.assert.notCalled(fireEventSpy)
814+
})
815+
})
816+
})

0 commit comments

Comments
 (0)