Skip to content

Commit 4d2b8b0

Browse files
liramon1liramon1
committed
Add session token input to login flow
1 parent ce07bd2 commit 4d2b8b0

File tree

11 files changed

+48
-32
lines changed

11 files changed

+48
-32
lines changed

‎packages/amazonq/test/e2e/amazonq/utils/setup.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,5 @@ export async function loginToIdC() {
2222
)
2323
}
2424

25-
await AuthUtil.instance.login(startUrl, region, 'sso')
25+
await AuthUtil.instance.login_sso(startUrl, region)
2626
}

‎packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,11 @@ describe('RegionProfileManager', async function () {
2626

2727
async function setupConnection(type: 'builderId' | 'idc') {
2828
if (type === 'builderId') {
29-
await AuthUtil.instance.login(constants.builderIdStartUrl, region, 'sso')
29+
await AuthUtil.instance.login_sso(constants.builderIdStartUrl, region)
3030
assert.ok(AuthUtil.instance.isSsoSession())
3131
assert.ok(AuthUtil.instance.isBuilderIdConnection())
3232
} else if (type === 'idc') {
33-
await AuthUtil.instance.login(enterpriseSsoStartUrl, region, 'sso')
33+
await AuthUtil.instance.login_sso(enterpriseSsoStartUrl, region)
3434
assert.ok(AuthUtil.instance.isSsoSession())
3535
assert.ok(AuthUtil.instance.isIdcConnection())
3636
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ describe('showConnectionPrompt', function () {
2828
})
2929

3030
it('can select connect to AwsBuilderId', async function () {
31-
sinon.stub(AuthUtil.instance, 'login').resolves()
31+
sinon.stub(AuthUtil.instance, 'login_sso').resolves()
3232

3333
getTestWindow().onDidShowQuickPick(async (picker) => {
3434
await picker.untilReady()
@@ -44,7 +44,7 @@ describe('showConnectionPrompt', function () {
4444

4545
it('connectToAwsBuilderId calls AuthUtil login with builderIdStartUrl', async function () {
4646
sinon.stub(vscode.commands, 'executeCommand')
47-
const loginStub = sinon.stub(AuthUtil.instance, 'login').resolves()
47+
const loginStub = sinon.stub(AuthUtil.instance, 'login_sso').resolves()
4848

4949
await awsIdSignIn()
5050

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,7 @@ export class LanguageClientAuth {
174174
} satisfies UpdateProfileParams)
175175
}
176176

177-
updateIamProfile(profileName: string, accessKey: string, secretKey: string): Promise<UpdateProfileResult> {
177+
updateIamProfile(profileName: string, accessKey: string, secretKey: string, sessionToken?: string): Promise<UpdateProfileResult> {
178178
// Add credentials and delete SSO settings from profile
179179
return this.client.sendRequest(updateProfileRequestType.method, {
180180
profile: {
@@ -185,6 +185,7 @@ export class LanguageClientAuth {
185185
sso_session: '',
186186
aws_access_key_id: accessKey,
187187
aws_secret_access_key: secretKey,
188+
aws_session_token: sessionToken,
188189
},
189190
},
190191
ssoSession: {
@@ -497,7 +498,7 @@ export class IamLogin extends BaseLogin {
497498
// )
498499
}
499500

500-
async login(opts: { accessKey: string; secretKey: string }) {
501+
async login(opts: { accessKey: string; secretKey: string, sessionToken?: string }) {
501502
await this.updateProfile(opts)
502503
return this._getIamCredential(true)
503504
}
@@ -519,8 +520,8 @@ export class IamLogin extends BaseLogin {
519520
// TODO: DeleteProfile api in Identity Service (this doesn't exist yet)
520521
}
521522

522-
async updateProfile(opts: { accessKey: string; secretKey: string }) {
523-
await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey)
523+
async updateProfile(opts: { accessKey: string; secretKey: string, sessionToken?: string }) {
524+
await this.lspAuth.updateIamProfile(this.profileName, opts.accessKey, opts.secretKey, opts.sessionToken)
524525
this._data = {
525526
accessKey: opts.accessKey,
526527
secretKey: opts.secretKey,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ export function createSignIn(): DataQuickPickItem<'signIn'> {
271271
if (isWeb()) {
272272
// TODO: nkomonen, call a Command instead
273273
onClick = () => {
274-
void AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso')
274+
void AuthUtil.instance.login_sso(builderIdStartUrl, builderIdRegion)
275275
}
276276
}
277277

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

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -160,25 +160,26 @@ export class AuthUtil implements IAuthProvider {
160160
}
161161
}
162162

163-
// Log into the desired session type using the authentication parameters
164-
async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise<GetIamCredentialResult | undefined>
165-
async login(startUrl: string, region: string, loginType: 'sso'): Promise<GetSsoTokenResult | undefined>
166-
async login(
167-
first: string,
168-
second: string,
169-
loginType: 'iam' | 'sso'
170-
): Promise<GetSsoTokenResult | GetIamCredentialResult | undefined> {
171-
let response: GetSsoTokenResult | GetIamCredentialResult | undefined
172-
173-
// Start session if the current session type does not match the desired type
174-
if (loginType === 'sso' && !this.isSsoSession()) {
163+
// Log in using SSO
164+
async login_sso(startUrl: string, region: string): Promise<GetSsoTokenResult | undefined> {
165+
let response: GetSsoTokenResult | undefined
166+
// Create SSO login session
167+
if (!this.isSsoSession()) {
175168
this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter)
176-
response = await this.session.login({ startUrl: first, region: second, scopes: amazonQScopes })
177-
} else if (loginType === 'iam' && !this.isIamSession()) {
178-
this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
179-
response = await this.session.login({ accessKey: first, secretKey: second })
180169
}
170+
response = await (this.session as SsoLogin).login({ startUrl: startUrl, region: region, scopes: amazonQScopes })
171+
await showAmazonQWalkthroughOnce()
172+
return response
173+
}
181174

175+
// Log in using IAM or STS credentials
176+
async login_iam(accessKey: string, secretKey: string, sessionToken?: string): Promise<GetIamCredentialResult | undefined> {
177+
let response: GetIamCredentialResult | undefined
178+
// Create IAM login session
179+
if (!this.isIamSession()) {
180+
this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
181+
}
182+
response = await (this.session as IamLogin).login({ accessKey: accessKey, secretKey: secretKey, sessionToken: sessionToken })
182183
await showAmazonQWalkthroughOnce()
183184
return response
184185
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ export const getStartUrl = async () => {
2929

3030
export async function connectToEnterpriseSso(startUrl: string, region: Region['id']) {
3131
try {
32-
await AuthUtil.instance.login(startUrl, region, 'sso')
32+
await AuthUtil.instance.login_sso(startUrl, region)
3333
} catch (e) {
3434
throw ToolkitError.chain(e, CodeWhispererConstants.failedToConnectIamIdentityCenter, {
3535
code: 'FailedToConnect',

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ export const showCodeWhispererConnectionPrompt = async () => {
4747
export async function awsIdSignIn() {
4848
getLogger().info('selected AWS ID sign in')
4949
try {
50-
await AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso')
50+
await AuthUtil.instance.login_sso(builderIdStartUrl, builderIdRegion)
5151
} catch (e) {
5252
throw ToolkitError.chain(e, failedToConnectAwsBuilderId, { code: 'FailedToConnect' })
5353
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,8 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
196196
async startIamCredentialSetup(
197197
profileName: string,
198198
accessKey: string,
199-
secretKey: string
199+
secretKey: string,
200+
sessionToken?: string
200201
): Promise<AuthError | undefined> {
201202
getLogger().debug(`called startIamCredentialSetup()`)
202203
// Defining separate auth function to emit telemetry before returning from this method
@@ -205,7 +206,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
205206
})
206207
const runAuth = async (): Promise<AuthError | undefined> => {
207208
try {
208-
await AuthUtil.instance.login(accessKey, secretKey, 'iam')
209+
await AuthUtil.instance.login_iam(accessKey, secretKey, sessionToken)
209210
} catch (e) {
210211
getLogger().error('Failed submitting credentials %O', e)
211212
return { id: this.id, text: e as string }

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -174,7 +174,8 @@ export abstract class CommonAuthWebview extends VueWebview {
174174
abstract startIamCredentialSetup(
175175
profileName: string,
176176
accessKey: string,
177-
secretKey: string
177+
secretKey: string,
178+
sessionToken?: string,
178179
): Promise<AuthError | undefined>
179180

180181
async showResourceExplorer(): Promise<void> {

0 commit comments

Comments
 (0)