Skip to content

Commit 754f87e

Browse files
yuxianrzliramon1
yuxianrz
and
liramon1
authored
fix(auth): fix UI and error message bug, disable inline chat for IAM, autofill access key id (#7797)
## Problem UI and error message no longer compatible after adding IAM credentials authflow IAM Access Key needs manual input every time a client log in ## Solution This is part of #7507 and is built on top of #7659. --- - License: I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: Ramon Li <[email protected]>
1 parent 3ff9966 commit 754f87e

File tree

13 files changed

+133
-62
lines changed

13 files changed

+133
-62
lines changed

‎packages/amazonq/src/extensionNode.ts

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,11 @@ 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 (
104+
AuthUtil.instance.isConnected() &&
105+
!(AuthUtil.instance.isSsoSession() || AuthUtil.instance.isIamSession() || isSageMaker())
106+
) {
107+
getLogger().error('Current Amazon Q connection is not SSO nor IAM')
105108
}
106109

107110
return {

‎packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ export class InlineChatProvider {
143143
private async generateResponse(
144144
triggerPayload: TriggerPayload & { projectContextQueryLatencyMs?: number },
145145
triggerID: string
146-
) {
146+
): Promise<GenerateAssistantResponseCommandOutput | undefined> {
147147
const triggerEvent = this.triggerEventsStorage.getTriggerEvent(triggerID)
148148
if (triggerEvent === undefined) {
149149
return
@@ -182,7 +182,12 @@ export class InlineChatProvider {
182182
let response: GenerateAssistantResponseCommandOutput | undefined = undefined
183183
session.createNewTokenSource()
184184
try {
185-
response = await session.chatSso(request)
185+
if (AuthUtil.instance.isSsoSession()) {
186+
response = await session.chatSso(request)
187+
} else {
188+
// Call sendMessage because Q Developer Streaming Client does not have generateAssistantResponse
189+
throw new ToolkitError('Inline chat is only available with SSO authentication')
190+
}
186191
getLogger().info(
187192
`response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${response.$metadata.requestId} metadata: %O`,
188193
response.$metadata

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -483,7 +483,7 @@ describe('AuthUtil', async function () {
483483
await auth.getIamCredential()
484484
assert.fail('Should have thrown an error')
485485
} catch (err) {
486-
assert.strictEqual((err as Error).message, 'Cannot get token with SSO session')
486+
assert.strictEqual((err as Error).message, 'Cannot get credential without logging in with IAM.')
487487
}
488488
})
489489

@@ -494,7 +494,7 @@ describe('AuthUtil', async function () {
494494
await auth.getIamCredential()
495495
assert.fail('Should have thrown an error')
496496
} catch (err) {
497-
assert.strictEqual((err as Error).message, 'Cannot get credential without logging in.')
497+
assert.strictEqual((err as Error).message, 'Cannot get credential without logging in with IAM.')
498498
}
499499
})
500500
})

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -344,8 +344,13 @@ export abstract class BaseLogin {
344344
* Decrypts an encrypted string, removes its quotes, and returns the resulting string
345345
*/
346346
protected async decrypt(encrypted: string): Promise<string> {
347-
const decrypted = await jose.compactDecrypt(encrypted, this.lspAuth.encryptionKey)
348-
return decrypted.plaintext.toString().replaceAll('"', '')
347+
try {
348+
const decrypted = await jose.compactDecrypt(encrypted, this.lspAuth.encryptionKey)
349+
return decrypted.plaintext.toString().replaceAll('"', '')
350+
} catch (e) {
351+
getLogger().error(`Failed to decrypt: ${encrypted}`)
352+
return encrypted
353+
}
349354
}
350355
}
351356

‎packages/core/src/codewhisperer/ui/codeWhispererNodes.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,11 @@ export function createManageSubscription(): DataQuickPickItem<'manageSubscriptio
192192
export function createSignout(): DataQuickPickItem<'signout'> {
193193
const label = localize('AWS.codewhisperer.signoutNode.label', 'Sign Out')
194194
const icon = getIcon('vscode-export')
195-
const connection = AuthUtil.instance.isBuilderIdConnection() ? 'AWS Builder ID' : 'IAM Identity Center'
195+
const connection = AuthUtil.instance.isIamSession()
196+
? 'IAM Credentials'
197+
: AuthUtil.instance.isBuilderIdConnection()
198+
? 'AWS Builder ID'
199+
: 'IAM Identity Center'
196200

197201
return {
198202
data: 'signout',

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

Lines changed: 32 additions & 26 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, LanguageClientAuth, Login, SsoLogin, IamLogin } from '../../auth/auth2'
33+
import {
34+
AuthStateEvent,
35+
cacheChangedEvent,
36+
LanguageClientAuth,
37+
Login,
38+
SsoLogin,
39+
IamLogin,
40+
AuthState,
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'
@@ -108,11 +117,11 @@ export class AuthUtil implements IAuthProvider {
108117
}
109118

110119
isSsoSession(): boolean {
111-
return this.session instanceof SsoLogin
120+
return this.session?.loginType === LoginTypes.SSO
112121
}
113122

114123
isIamSession(): boolean {
115-
return this.session instanceof IamLogin
124+
return this.session?.loginType === LoginTypes.IAM
116125
}
117126

118127
/**
@@ -204,36 +213,31 @@ export class AuthUtil implements IAuthProvider {
204213
}
205214

206215
async getToken() {
207-
if (this.session) {
208-
const token = (await this.session.getCredential()).credential
209-
if (typeof token !== 'string') {
210-
throw new ToolkitError('Cannot get token with IAM session')
211-
}
212-
return token
216+
if (this.isSsoSession()) {
217+
const response = await this.session!.getCredential()
218+
return response.credential as string
213219
} else {
214-
throw new ToolkitError('Cannot get credential without logging in.')
220+
throw new ToolkitError('Cannot get credential without logging in with SSO.')
215221
}
216222
}
217223

218224
async getIamCredential() {
219-
if (this.session) {
220-
const credential = (await this.session.getCredential()).credential
221-
if (typeof credential !== 'object') {
222-
throw new ToolkitError('Cannot get token with SSO session')
223-
}
224-
return credential
225+
if (this.isIamSession()) {
226+
const response = await this.session!.getCredential()
227+
return response.credential as IamCredentials
225228
} else {
226-
throw new ToolkitError('Cannot get credential without logging in.')
229+
throw new ToolkitError('Cannot get credential without logging in with IAM.')
227230
}
228231
}
229232

230233
get connection() {
231234
return this.session?.data
232235
}
233236

234-
getAuthState() {
235-
if (this.session) {
236-
return this.session.getConnectionState()
237+
getAuthState(): AuthState {
238+
// Check if getConnectionState exists in case of type casts
239+
if (typeof this.session?.getConnectionState === 'function') {
240+
return this.session!.getConnectionState()
237241
} else {
238242
return 'notConnected'
239243
}
@@ -356,11 +360,12 @@ export class AuthUtil implements IAuthProvider {
356360

357361
private async stateChangeHandler(e: AuthStateEvent) {
358362
if (e.state === 'refreshed') {
359-
const params = this.session ? (await this.session.getCredential()).updateCredentialsParams : undefined
360363
if (this.isSsoSession()) {
361-
await this.lspAuth.updateBearerToken(params)
364+
const params = await this.session!.getCredential()
365+
await this.lspAuth.updateBearerToken(params.updateCredentialsParams)
362366
} else if (this.isIamSession()) {
363-
await this.lspAuth.updateIamCredential(params)
367+
const params = await this.session!.getCredential()
368+
await this.lspAuth.updateIamCredential(params.updateCredentialsParams)
364369
}
365370
} else {
366371
this.logger.info(`codewhisperer: connection changed to ${e.state}`)
@@ -383,11 +388,12 @@ export class AuthUtil implements IAuthProvider {
383388
this.session = undefined
384389
}
385390
if (state === 'connected') {
386-
const params = this.session ? (await this.session.getCredential()).updateCredentialsParams : undefined
387391
if (this.isSsoSession()) {
388-
await this.lspAuth.updateBearerToken(params)
392+
const params = await this.session!.getCredential()
393+
await this.lspAuth.updateBearerToken(params.updateCredentialsParams)
389394
} else if (this.isIamSession()) {
390-
await this.lspAuth.updateIamCredential(params)
395+
const params = await this.session!.getCredential()
396+
await this.lspAuth.updateIamCredential(params.updateCredentialsParams)
391397
}
392398

393399
if (this.isIdcConnection()) {

‎packages/core/src/login/webview/vue/backend.ts

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import { getLogger } from '../../../shared/logger/logger'
3333
import { isValidUrl } from '../../../shared/utilities/uriUtils'
3434
import { RegionProfile } from '../../../codewhisperer/models/model'
3535
import { ProfileSwitchIntent } from '../../../codewhisperer/region/regionProfileManager'
36+
import { showMessage } from '../../../shared/utilities/messages'
3637

3738
export abstract class CommonAuthWebview extends VueWebview {
3839
private readonly className = 'CommonAuthWebview'
@@ -183,7 +184,7 @@ export abstract class CommonAuthWebview extends VueWebview {
183184
abstract fetchConnections(): Promise<AwsConnection[] | undefined>
184185

185186
async errorNotification(e: AuthError) {
186-
void vscode.window.showInformationMessage(`${e.text}`)
187+
await showMessage('error', e.text)
187188
}
188189

189190
abstract quitLoginScreen(): Promise<void>
@@ -296,6 +297,15 @@ export abstract class CommonAuthWebview extends VueWebview {
296297
return globals.globalState.tryGet('recentSso', Object, { startUrl: '', region: 'us-east-1' })
297298
}
298299

300+
getDefaultIamKeys(): { accessKey: string } {
301+
const devSettings = DevSettings.instance.get('autofillAccessKey', '')
302+
if (devSettings) {
303+
return { accessKey: devSettings }
304+
}
305+
306+
return globals.globalState.tryGet('recentIamKeys', Object, { accessKey: '' })
307+
}
308+
299309
cancelAuthFlow() {
300310
AuthSSOServer.lastInstance?.cancelCurrentFlow()
301311
}

‎packages/core/src/login/webview/vue/login.vue

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,7 @@
230230

231231
<template v-if="stage === 'AUTHENTICATING'">
232232
<div class="auth-container-section">
233-
<div v-if="app === 'TOOLKIT' && profileName.length > 0" class="header bottomMargin">
233+
<div v-if="selectedLoginOption === LoginOption.IAM_CREDENTIAL" class="header bottomMargin">
234234
Connecting to IAM...
235235
</div>
236236
<div v-else class="header bottomMargin">Authenticating in browser...</div>
@@ -273,7 +273,7 @@
273273
<div class="title">Secret Access Key</div>
274274
<input
275275
class="iamInput bottomMargin"
276-
type="text"
276+
type="password"
277277
id="secretKey"
278278
name="secretKey"
279279
v-model="secretKey"
@@ -330,6 +330,10 @@ interface ImportedLogin {
330330
type: number
331331
startUrl: string
332332
region: string
333+
// Add IAM credential fields
334+
profileName?: string
335+
accessKey?: string
336+
secretKey?: string // Note: storing secrets has security implications
333337
}
334338
335339
export default defineComponent({
@@ -349,6 +353,7 @@ export default defineComponent({
349353
data() {
350354
return {
351355
existingStartUrls: [] as string[],
356+
existingIamAccessKeys: [] as string[],
352357
importedLogins: [] as ImportedLogin[],
353358
selectedLoginOption: LoginOption.NONE,
354359
stage: 'START' as Stage,
@@ -368,6 +373,8 @@ export default defineComponent({
368373
const defaultSso = await this.getDefaultSso()
369374
this.startUrl = defaultSso.startUrl
370375
this.selectedRegion = defaultSso.region
376+
const defaultIamAccessKey = await this.getDefaultIamAccessKey()
377+
this.accessKey = defaultIamAccessKey.accessKey
371378
await this.emitUpdate('created')
372379
},
373380
@@ -397,6 +404,10 @@ export default defineComponent({
397404
}
398405
},
399406
handleDocumentClick(event: any) {
407+
// Only reset selection when in START stage to avoid clearing during authentication
408+
if (this.stage !== 'START') {
409+
return
410+
}
400411
const isClickInsideSelectableItems = event.target.closest('.selectable-item')
401412
if (!isClickInsideSelectableItems) {
402413
this.selectedLoginOption = 0
@@ -437,17 +448,32 @@ export default defineComponent({
437448
const selectedConnection =
438449
this.importedLogins[this.selectedLoginOption - LoginOption.IMPORTED_LOGINS]
439450
440-
// Imported connections cannot be Builder IDs, they are filtered out in the client.
441-
const error = await client.startEnterpriseSetup(
442-
selectedConnection.startUrl,
443-
selectedConnection.region,
444-
this.app
445-
)
446-
if (error) {
447-
this.stage = 'START'
448-
void client.errorNotification(error)
449-
} else {
450-
this.stage = 'CONNECTED'
451+
// Handle both SSO and IAM imported connections
452+
if (selectedConnection.type === LoginOption.ENTERPRISE_SSO) {
453+
const error = await client.startEnterpriseSetup(
454+
selectedConnection.startUrl,
455+
selectedConnection.region,
456+
this.app
457+
)
458+
if (error) {
459+
this.stage = 'START'
460+
void client.errorNotification(error)
461+
} else {
462+
this.stage = 'CONNECTED'
463+
}
464+
} else if (selectedConnection.type === LoginOption.IAM_CREDENTIAL) {
465+
// Use stored IAM credentials
466+
const error = await client.startIamCredentialSetup(
467+
selectedConnection.profileName || '',
468+
selectedConnection.accessKey || '',
469+
selectedConnection.secretKey || ''
470+
)
471+
if (error) {
472+
this.stage = 'START'
473+
void client.errorNotification(error)
474+
} else {
475+
this.stage = 'CONNECTED'
476+
}
451477
}
452478
} else if (this.selectedLoginOption === LoginOption.IAM_CREDENTIAL) {
453479
this.stage = 'AWS_PROFILE'
@@ -581,6 +607,9 @@ export default defineComponent({
581607
async getDefaultSso() {
582608
return await client.getDefaultSsoProfile()
583609
},
610+
async getDefaultIamAccessKey() {
611+
return await client.getDefaultIamKeys()
612+
},
584613
handleHelpLinkClick() {
585614
void client.emitUiClick('auth_helpLink')
586615
},

‎packages/core/src/shared/featureConfig.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ export class FeatureConfigProvider {
118118
}
119119

120120
async fetchFeatureConfigs(): Promise<void> {
121-
if (AuthUtil.instance.isConnectionExpired()) {
121+
if (AuthUtil.instance.isConnectionExpired() || AuthUtil.instance.isIamSession()) {
122122
return
123123
}
124124

‎packages/core/src/shared/globalState.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ export type globalKey =
7171
| 'lastOsStartTime'
7272
| 'recentCredentials'
7373
| 'recentSso'
74+
| 'recentIamKeys'
7475
// List of regions enabled in AWS Explorer.
7576
| 'region'
7677
// TODO: implement this via `PromptSettings` instead of globalState.

0 commit comments

Comments
 (0)