Skip to content

Commit a79939a

Browse files
liramon1liramon1
committed
Move IAM profile options into typed object
1 parent 26139e8 commit a79939a

File tree

7 files changed

+99
-123
lines changed

7 files changed

+99
-123
lines changed

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

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { createTestAuthUtil, TestFolder } from 'aws-core-vscode/test'
1111
import { constants, cache } from 'aws-core-vscode/auth'
1212
import { auth2 } from 'aws-core-vscode/auth'
1313
import { mementoUtils, fs } from 'aws-core-vscode/shared'
14+
import { GetIamCredentialResult } from '@aws/language-server-runtimes/protocol'
1415

1516
describe('AuthUtil', async function () {
1617
let auth: any
@@ -431,7 +432,11 @@ describe('AuthUtil', async function () {
431432

432433
sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
433434

434-
const response = await auth.loginIam('accessKey', 'secretKey', 'sessionToken')
435+
const response = await auth.loginIam({
436+
accessKey: 'accessKey',
437+
secretKey: 'secretKey',
438+
sessionToken: 'sessionToken',
439+
})
435440

436441
assert.ok(mockIamLogin.login.calledOnce)
437442
assert.ok(
@@ -446,12 +451,15 @@ describe('AuthUtil', async function () {
446451
})
447452

448453
it('creates IAM session with role ARN', async function () {
449-
const mockResponse = {
450-
id: 'test-credential-id',
451-
credentials: {
452-
accessKeyId: 'encrypted-access-key',
453-
secretAccessKey: 'encrypted-secret-key',
454-
sessionToken: 'encrypted-session-token',
454+
const mockResponse: GetIamCredentialResult = {
455+
credential: {
456+
id: 'test-credential-id',
457+
kinds: [],
458+
credentials: {
459+
accessKeyId: 'encrypted-access-key',
460+
secretAccessKey: 'encrypted-secret-key',
461+
sessionToken: 'encrypted-session-token',
462+
},
455463
},
456464
updateCredentialsParams: {
457465
data: 'credential-data',
@@ -465,22 +473,17 @@ describe('AuthUtil', async function () {
465473

466474
sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
467475

468-
const response = await auth.loginIam(
469-
'accessKey',
470-
'secretKey',
471-
'sessionToken',
472-
'arn:aws:iam::123456789012:role/TestRole'
473-
)
476+
const opts: auth2.IamProfileOptions = {
477+
accessKey: 'myAccessKey',
478+
secretKey: 'mySecretKey',
479+
sessionToken: 'mySessionToken',
480+
roleArn: 'arn:aws:iam::123456789012:role/MyTestRole',
481+
}
482+
483+
const response = await auth.loginIam(opts)
474484

475485
assert.ok(mockIamLogin.login.calledOnce)
476-
assert.ok(
477-
mockIamLogin.login.calledWith({
478-
accessKey: 'accessKey',
479-
secretKey: 'secretKey',
480-
sessionToken: 'sessionToken',
481-
roleArn: 'arn:aws:iam::123456789012:role/TestRole',
482-
})
483-
)
486+
assert.ok(mockIamLogin.login.calledWith(opts))
484487
assert.strictEqual(response, mockResponse)
485488
})
486489
})

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

Lines changed: 58 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,22 @@ export type Login = SsoLogin | IamLogin
9494

9595
export type TokenSource = IamIdentityCenterSsoTokenSource | AwsBuilderIdSsoTokenSource
9696

97+
export type IamProfileOptions = {
98+
accessKey?: string
99+
secretKey?: string
100+
sessionToken?: string
101+
roleArn?: string
102+
sourceProfile?: string
103+
}
104+
105+
const IamProfileOptionsDefaults = {
106+
accessKey: '',
107+
secretKey: '',
108+
sessionToken: '',
109+
roleArn: '',
110+
sourceProfile: '',
111+
} satisfies IamProfileOptions
112+
97113
/**
98114
* Handles auth requests to the Identity Server in the Amazon Q LSP.
99115
*/
@@ -216,56 +232,32 @@ export class LanguageClientAuth {
216232
return this.client.sendRequest(updateProfileRequestType.method, params)
217233
}
218234

219-
async updateIamProfile(
220-
profileName: string,
221-
accessKey: string,
222-
secretKey: string,
223-
sessionToken?: string,
224-
roleArn?: string,
225-
sourceProfile?: string
226-
): Promise<UpdateProfileResult> {
227-
// Add credentials and delete SSO settings from profile
228-
let profile: Profile
229-
if (roleArn && sourceProfile) {
230-
profile = {
231-
kinds: [ProfileKind.IamSourceProfileProfile],
232-
name: profileName,
233-
settings: {
234-
sso_session: '',
235-
aws_access_key_id: '',
236-
aws_secret_access_key: '',
237-
aws_session_token: '',
238-
role_arn: roleArn,
239-
source_profile: sourceProfile,
240-
},
241-
}
242-
} else if (accessKey && secretKey) {
243-
profile = {
244-
kinds: [ProfileKind.IamCredentialsProfile],
245-
name: profileName,
246-
settings: {
247-
sso_session: '',
248-
aws_access_key_id: accessKey,
249-
aws_secret_access_key: secretKey,
250-
aws_session_token: sessionToken,
251-
role_arn: '',
252-
source_profile: '',
253-
},
254-
}
235+
async updateIamProfile(profileName: string, opts: IamProfileOptions): Promise<UpdateProfileResult> {
236+
// Substitute missing fields for defaults
237+
const fields = { ...IamProfileOptionsDefaults, ...opts }
238+
// Get the profile kind matching the provided fields
239+
let kind: ProfileKind
240+
if (fields.roleArn && fields.sourceProfile) {
241+
kind = ProfileKind.IamSourceProfileProfile
242+
} else if (fields.accessKey && fields.secretKey) {
243+
kind = ProfileKind.IamCredentialsProfile
255244
} else {
256-
profile = {
257-
kinds: [ProfileKind.Unknown],
245+
kind = ProfileKind.Unknown
246+
}
247+
248+
const params = await this.encrypt({
249+
profile: {
250+
kinds: [kind],
258251
name: profileName,
259252
settings: {
260-
aws_access_key_id: '',
261-
aws_secret_access_key: '',
262-
aws_session_token: '',
263-
role_arn: '',
264-
source_profile: '',
253+
aws_access_key_id: fields.accessKey,
254+
aws_secret_access_key: fields.secretKey,
255+
aws_session_token: fields.sessionToken,
256+
role_arn: fields.roleArn,
257+
source_profile: fields.sourceProfile,
265258
},
266-
}
267-
}
268-
const params = await this.encrypt({ profile: profile })
259+
},
260+
})
269261
return this.client.sendRequest(updateProfileRequestType.method, params)
270262
}
271263

@@ -567,7 +559,7 @@ export class IamLogin extends BaseLogin {
567559
lspAuth.registerGetMfaCodeHandler((params: GetMfaCodeParams) => this.getMfaCodeHandler(params))
568560
}
569561

570-
async login(opts: { accessKey: string; secretKey: string; sessionToken?: string; roleArn?: string }) {
562+
async login(opts: IamProfileOptions) {
571563
await this.updateProfile(opts)
572564
return this._getIamCredential(true)
573565
}
@@ -583,34 +575,33 @@ export class IamLogin extends BaseLogin {
583575
if (this.iamCredentialId) {
584576
await this.lspAuth.invalidateStsCredential(this.iamCredentialId)
585577
}
586-
await this.lspAuth.updateIamProfile(this.profileName, '', '', '', '', '')
587-
await this.lspAuth.updateIamProfile(this.profileName + '-source', '', '', '', '', '')
578+
await this.lspAuth.updateIamProfile(this.profileName, {})
579+
await this.lspAuth.updateIamProfile(this.profileName + '-source', {})
588580
this.updateConnectionState('notConnected')
589581
this._data = undefined
590582
// TODO: DeleteProfile api in Identity Service (this doesn't exist yet)
591583
}
592584

593-
async updateProfile(opts: { accessKey: string; secretKey: string; sessionToken?: string; roleArn?: string }) {
585+
async updateProfile(opts: IamProfileOptions) {
594586
if (opts.roleArn) {
587+
// Create the source and target profiles
595588
const sourceProfile = this.profileName + '-source'
596-
await this.lspAuth.updateIamProfile(
597-
sourceProfile,
598-
opts.accessKey,
599-
opts.secretKey,
600-
opts.sessionToken,
601-
'',
602-
''
603-
)
604-
await this.lspAuth.updateIamProfile(this.profileName, '', '', '', opts.roleArn, sourceProfile)
589+
await this.lspAuth.updateIamProfile(sourceProfile, {
590+
accessKey: opts.accessKey,
591+
secretKey: opts.secretKey,
592+
sessionToken: opts.sessionToken,
593+
})
594+
await this.lspAuth.updateIamProfile(this.profileName, {
595+
roleArn: opts.roleArn,
596+
sourceProfile: sourceProfile,
597+
})
605598
} else {
606-
await this.lspAuth.updateIamProfile(
607-
this.profileName,
608-
opts.accessKey,
609-
opts.secretKey,
610-
opts.sessionToken,
611-
'',
612-
''
613-
)
599+
// Create the target profile
600+
await this.lspAuth.updateIamProfile(this.profileName, {
601+
accessKey: opts.accessKey,
602+
secretKey: opts.secretKey,
603+
sessionToken: opts.sessionToken,
604+
})
614605
}
615606
}
616607

‎packages/core/src/auth/credentials/utils.ts‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -105,11 +105,8 @@ const errorMessageUserCancelled = localize('AWS.error.mfa.userCancelled', 'User
105105
/**
106106
* @description Prompts user for MFA serial number
107107
*
108-
* Entered token is passed to the callback.
109-
* If user cancels out, the callback is passed an error with a fixed message string.
110-
*
108+
* @param defaultSerial Default value for the serial number input
111109
* @param profileName Name of Credentials profile we are asking an MFA Token for
112-
* @param callback tokens/errors are passed through here
113110
*/
114111
export async function getMfaSerialFromUser(defaultSerial: string, profileName: string): Promise<string> {
115112
const inputBox = createInputBox({
@@ -133,12 +130,8 @@ export async function getMfaSerialFromUser(defaultSerial: string, profileName: s
133130
/**
134131
* @description Prompts user for MFA token
135132
*
136-
* Entered token is passed to the callback.
137-
* If user cancels out, the callback is passed an error with a fixed message string.
138-
*
139133
* @param mfaSerial Serial arn of MFA device
140134
* @param profileName Name of Credentials profile we are asking an MFA Token for
141-
* @param callback tokens/errors are passed through here
142135
*/
143136
export async function getMfaTokenFromUser(mfaSerial: string, profileName: string): Promise<string> {
144137
const inputBox = createInputBox({

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ 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')
39+
const defaultStsCacheDir = () => path.join(fs.getUserHomeDir(), '.aws/flare/cache')
4040
export const getCacheDir = () => DevSettings.instance.get('ssoCacheDirectory', defaultCacheDir())
4141
export const getStsCacheDir = () => DevSettings.instance.get('stsCacheDirectory', defaultStsCacheDir())
4242

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

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import {
4040
IamLogin,
4141
AuthState,
4242
LoginTypes,
43+
IamProfileOptions,
4344
} from '../../auth/auth2'
4445
import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants'
4546
import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
@@ -195,22 +196,12 @@ export class AuthUtil implements IAuthProvider {
195196
}
196197

197198
// Log in using IAM or STS credentials
198-
async loginIam(
199-
accessKey: string,
200-
secretKey: string,
201-
sessionToken?: string,
202-
roleArn?: string
203-
): Promise<GetIamCredentialResult | undefined> {
199+
async loginIam(opts: IamProfileOptions): Promise<GetIamCredentialResult | undefined> {
204200
// Create IAM login session
205201
if (!this.isIamSession()) {
206202
this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
207203
}
208-
const response = await (this.session as IamLogin).login({
209-
accessKey: accessKey,
210-
secretKey: secretKey,
211-
sessionToken: sessionToken,
212-
roleArn: roleArn,
213-
})
204+
const response = await (this.session as IamLogin).login(opts)
214205
await showAmazonQWalkthroughOnce()
215206
return response
216207
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
215215

216216
const runAuth = async (): Promise<AuthError | undefined> => {
217217
try {
218-
await AuthUtil.instance.loginIam(accessKey, secretKey, sessionToken, roleArn)
218+
await AuthUtil.instance.loginIam({ accessKey, secretKey, sessionToken, roleArn })
219219
} catch (e) {
220220
getLogger().error('Failed submitting credentials %O', e)
221221
const message = e instanceof Error ? e.message : (e as string)

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

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,11 @@ describe('LanguageClientAuth', () => {
104104
})
105105

106106
it('sends correct IAM profile update parameters', async () => {
107-
await auth.updateIamProfile(profileName, 'accessKey', 'secretKey', 'sessionToken')
107+
await auth.updateIamProfile(profileName, {
108+
accessKey: 'myAccessKey',
109+
secretKey: 'mySecretKey',
110+
sessionToken: 'mySessionToken',
111+
})
108112

109113
sinon.assert.calledOnce(client.sendRequest)
110114
const requestParams = client.sendRequest.firstCall.args[1]
@@ -113,9 +117,11 @@ describe('LanguageClientAuth', () => {
113117
kinds: [ProfileKind.IamCredentialsProfile],
114118
})
115119
sinon.assert.match(requestParams.profile.settings, {
116-
aws_access_key_id: 'accessKey',
117-
aws_secret_access_key: 'secretKey',
118-
aws_session_token: 'sessionToken',
120+
aws_access_key_id: 'myAccessKey',
121+
aws_secret_access_key: 'mySecretKey',
122+
aws_session_token: 'mySessionToken',
123+
role_arn: '',
124+
source_profile: '',
119125
})
120126
})
121127
})
@@ -676,15 +682,7 @@ describe('IamLogin', () => {
676682
const response = await iamLogin.login(loginOpts)
677683

678684
sinon.assert.calledOnce(lspAuth.updateIamProfile)
679-
sinon.assert.calledWith(
680-
lspAuth.updateIamProfile,
681-
profileName,
682-
loginOpts.accessKey,
683-
loginOpts.secretKey,
684-
loginOpts.sessionToken,
685-
'',
686-
''
687-
)
685+
sinon.assert.calledWith(lspAuth.updateIamProfile, profileName, loginOpts)
688686
sinon.assert.calledOnce(lspAuth.getIamCredential)
689687
sinon.assert.match(iamLogin.getConnectionState(), 'connected')
690688
sinon.assert.match(response.credential.id, 'test-credential-id')

0 commit comments

Comments
 (0)