Merge branch 'aws:master' into master

Author: laileni-aws

packages/amazonq/src/lsp/client.ts

@@ -39,6 +39,7 @@ import {
     getClientId,
     extensionVersion,
     isSageMaker,
+    DevSettings,
 } from 'aws-core-vscode/shared'
 import { processUtils } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
@@ -129,6 +130,15 @@ export async function startLanguageServer(
 
     await validateNodeExe(executable, resourcePaths.lsp, argv, logger)
 
+    const endpointOverride = DevSettings.instance.get('codewhispererService', {}).endpoint ?? undefined
+    const textDocSection = {
+        inlineEditSupport: Experiments.instance.get('amazonqLSPNEP', true),
+    } as any
+
+    if (endpointOverride) {
+        textDocSection.endpointOverride = endpointOverride
+    }
+
     // Options to control the language client
     const clientOptions: LanguageClientOptions = {
         // Register the server for json documents
@@ -177,9 +187,7 @@ export async function startLanguageServer(
                         showLogs: true,
                     },
                     textDocument: {
-                        inlineCompletionWithReferences: {
-                            inlineEditSupport: Experiments.instance.get('amazonqLSPNEP', true),
-                        },
+                        inlineCompletionWithReferences: textDocSection,
                     },
                 },
                 contextConfiguration: {

packages/core/src/auth/providers/sharedCredentialsProvider.ts

@@ -406,12 +406,28 @@ export class SharedCredentialsProvider implements CredentialsProvider {
                     `auth: Profile ${this.profileName} is missing source_profile for role assumption`
                 )
             }
-            // Use source profile to assume IAM role based on role ARN provided.
+
+            // Check if we already have resolved credentials from patchSourceCredentials
             const sourceProfile = iniData[profile.source_profile!]
-            const stsClient = new DefaultStsClient(this.getDefaultRegion() ?? 'us-east-1', {
-                accessKeyId: sourceProfile.aws_access_key_id!,
-                secretAccessKey: sourceProfile.aws_secret_access_key!,
-            })
+            let sourceCredentials: AWS.Credentials
+
+            if (sourceProfile.aws_access_key_id && sourceProfile.aws_secret_access_key) {
+                // Source credentials have already been resolved
+                sourceCredentials = {
+                    accessKeyId: sourceProfile.aws_access_key_id,
+                    secretAccessKey: sourceProfile.aws_secret_access_key,
+                    sessionToken: sourceProfile.aws_session_token,
+                }
+            } else {
+                // Source profile needs credential resolution - this should have been handled by patchSourceCredentials
+                // but if not, we need to resolve it here
+                const sourceProvider = new SharedCredentialsProvider(profile.source_profile!, this.sections)
+                sourceCredentials = await sourceProvider.getCredentials()
+            }
+
+            // Use source credentials to assume IAM role based on role ARN provided.
+            const stsClient = new DefaultStsClient(this.getDefaultRegion() ?? 'us-east-1', sourceCredentials)
+
             // Prompt for MFA Token if needed.
             const assumeRoleReq = {
                 RoleArn: profile.role_arn,

packages/core/src/test/auth/providers/sharedCredentialsProvider.test.ts

@@ -0,0 +1,79 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import { SharedCredentialsProvider } from '../../../auth/providers/sharedCredentialsProvider'
+import { createTestSections } from '../../credentials/testUtil'
+import { DefaultStsClient } from '../../../shared/clients/stsClient'
+import { oneDay } from '../../../shared/datetime'
+import sinon from 'sinon'
+import { SsoAccessTokenProvider } from '../../../auth/sso/ssoAccessTokenProvider'
+import { SsoClient } from '../../../auth/sso/clients'
+
+describe('SharedCredentialsProvider - Role Chaining with SSO', function () {
+    let sandbox: sinon.SinonSandbox
+
+    beforeEach(function () {
+        sandbox = sinon.createSandbox()
+    })
+
+    afterEach(function () {
+        sandbox.restore()
+    })
+
+    it('should handle role chaining from SSO profile', async function () {
+        // Mock the SSO authentication
+        sandbox.stub(SsoAccessTokenProvider.prototype, 'getToken').resolves({
+            accessToken: 'test-token',
+            expiresAt: new Date(Date.now() + oneDay),
+        })
+
+        // Mock SSO getRoleCredentials
+        sandbox.stub(SsoClient.prototype, 'getRoleCredentials').resolves({
+            accessKeyId: 'sso-access-key',
+            secretAccessKey: 'sso-secret-key',
+            sessionToken: 'sso-session-token',
+            expiration: new Date(Date.now() + oneDay),
+        })
+
+        // Mock STS assumeRole
+        sandbox.stub(DefaultStsClient.prototype, 'assumeRole').callsFake(async (request) => {
+            assert.strictEqual(request.RoleArn, 'arn:aws:iam::123456789012:role/dev')
+            return {
+                Credentials: {
+                    AccessKeyId: 'assumed-access-key',
+                    SecretAccessKey: 'assumed-secret-key',
+                    SessionToken: 'assumed-session-token',
+                    Expiration: new Date(Date.now() + oneDay),
+                },
+            }
+        })
+
+        const sections = await createTestSections(`
+            [sso-session aws1_session]
+            sso_start_url = https://example.awsapps.com/start
+            sso_region = us-east-1
+            sso_registration_scopes = sso:account:access
+
+            [profile Landing]
+            sso_session = aws1_session
+            sso_account_id = 111111111111
+            sso_role_name = Landing
+            region = us-east-1
+
+            [profile dev]
+            region = us-east-1
+            role_arn = arn:aws:iam::123456789012:role/dev
+            source_profile = Landing
+        `)
+
+        const provider = new SharedCredentialsProvider('dev', sections)
+        const credentials = await provider.getCredentials()
+
+        assert.strictEqual(credentials.accessKeyId, 'assumed-access-key')
+        assert.strictEqual(credentials.secretAccessKey, 'assumed-secret-key')
+        assert.strictEqual(credentials.sessionToken, 'assumed-session-token')
+    })
+})