Skip to content

Commit e6311d9

Browse files
yueny2020yueny2020
authored
fix(amazonq): not set IAM auth as default for SMAI remote ssh (aws#7795)
## Problem In SMAI remote ssh space, Q chat request failed as IAM creds not found ## Solution Discussed with SMAI team, for SMAI remote ssh space, not set the default auth mode as IAM. ## Test before tested with the latest version v1.87.0 https://github.com/user-attachments/assets/7a4f8b23-fb64-453c-830a-bfc58fc8394c after tested with a local build version https://github.com/user-attachments/assets/765c772c-13b2-444f-bada-eb9eb5068d55 --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent d40d74a commit e6311d9

File tree

2 files changed

+220
-3
lines changed

2 files changed

+220
-3
lines changed

‎packages/core/src/shared/lsp/utils/platform.ts‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { ToolkitError } from '../../errors'
88
import { Logger } from '../../logger/logger'
99
import { ChildProcess } from '../../utilities/processUtils'
1010
import { waitUntil } from '../../utilities/timeoutUtils'
11-
import { isDebugInstance } from '../../vscode/env'
11+
import { isDebugInstance, isRemoteWorkspace } from '../../vscode/env'
1212
import { isSageMaker } from '../../extensionUtilities'
1313
import { getLogger } from '../../logger/logger'
1414

@@ -124,8 +124,16 @@ export function createServerOptions({
124124
getLogger().info(`[SageMaker Debug] Using SSO auth mode, not setting USE_IAM_AUTH`)
125125
}
126126
} catch (err) {
127-
getLogger().warn(`[SageMaker Debug] Failed to parse SageMaker cookies, defaulting to IAM auth: ${err}`)
128-
env.USE_IAM_AUTH = 'true'
127+
if (isRemoteWorkspace() && env.SERVICE_NAME !== 'SageMakerUnifiedStudio') {
128+
getLogger().warn(
129+
`[SageMaker Debug] Failed to parse SageMaker cookies in remote space, not SMUS env, not defaulting to IAM auth: ${err}`
130+
)
131+
} else {
132+
getLogger().warn(
133+
`[SageMaker Debug] Failed to parse SageMaker cookies, defaulting to IAM auth: ${err}`
134+
)
135+
env.USE_IAM_AUTH = 'true'
136+
}
129137
}
130138

131139
// Log important environment variables for debugging

‎packages/core/src/test/shared/lsp/utils/platform.test.ts‎

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import assert from 'assert'
7+
import * as sinon from 'sinon'
8+
import * as vscode from 'vscode'
9+
import { createServerOptions } from '../../../../shared/lsp/utils/platform'
10+
import * as extensionUtilities from '../../../../shared/extensionUtilities'
11+
import * as env from '../../../../shared/vscode/env'
12+
import { ChildProcess } from '../../../../shared/utilities/processUtils'
13+
14+
describe('createServerOptions - SageMaker Authentication', function () {
15+
let sandbox: sinon.SinonSandbox
16+
let isSageMakerStub: sinon.SinonStub
17+
let isRemoteWorkspaceStub: sinon.SinonStub
18+
let executeCommandStub: sinon.SinonStub
19+
20+
beforeEach(function () {
21+
sandbox = sinon.createSandbox()
22+
23+
isSageMakerStub = sandbox.stub(extensionUtilities, 'isSageMaker')
24+
isRemoteWorkspaceStub = sandbox.stub(env, 'isRemoteWorkspace')
25+
sandbox.stub(env, 'isDebugInstance').returns(false)
26+
executeCommandStub = sandbox.stub(vscode.commands, 'executeCommand')
27+
28+
sandbox.stub(ChildProcess.prototype, 'run').resolves()
29+
sandbox.stub(ChildProcess.prototype, 'send').resolves()
30+
sandbox.stub(ChildProcess.prototype, 'proc').returns({} as any)
31+
})
32+
33+
afterEach(function () {
34+
sandbox.restore()
35+
})
36+
37+
// jscpd:ignore-start
38+
it('sets USE_IAM_AUTH=true when authMode is Iam', async function () {
39+
isSageMakerStub.returns(true)
40+
executeCommandStub.withArgs('sagemaker.parseCookies').resolves({ authMode: 'Iam' })
41+
42+
// Capture constructor arguments using sinon stub
43+
let capturedOptions: any = undefined
44+
const childProcessConstructorSpy = sandbox.stub().callsFake((command: string, args: string[], options: any) => {
45+
capturedOptions = options
46+
// Create a fake instance with the methods we need
47+
const fakeInstance = {
48+
run: sandbox.stub().resolves(),
49+
send: sandbox.stub().resolves(),
50+
proc: sandbox.stub().returns({}),
51+
pid: sandbox.stub().returns(12345),
52+
stop: sandbox.stub(),
53+
stopped: false,
54+
}
55+
return fakeInstance
56+
})
57+
58+
// Replace ChildProcess constructor
59+
sandbox.replace(
60+
require('../../../../shared/utilities/processUtils'),
61+
'ChildProcess',
62+
childProcessConstructorSpy
63+
)
64+
65+
const serverOptions = createServerOptions({
66+
encryptionKey: Buffer.from('test-key'),
67+
executable: ['node'],
68+
serverModule: 'test-module.js',
69+
execArgv: ['--stdio'],
70+
})
71+
72+
await serverOptions()
73+
74+
assert(capturedOptions, 'ChildProcess constructor should have been called')
75+
assert(capturedOptions.spawnOptions, 'spawnOptions should be defined')
76+
assert(capturedOptions.spawnOptions.env, 'spawnOptions.env should be defined')
77+
assert.equal(capturedOptions.spawnOptions.env.USE_IAM_AUTH, 'true')
78+
})
79+
80+
it('does not set USE_IAM_AUTH when authMode is Sso', async function () {
81+
isSageMakerStub.returns(true)
82+
executeCommandStub.withArgs('sagemaker.parseCookies').resolves({ authMode: 'Sso' })
83+
84+
// Capture constructor arguments using sinon stub
85+
let capturedOptions: any = undefined
86+
const childProcessConstructorSpy = sandbox.stub().callsFake((command: string, args: string[], options: any) => {
87+
capturedOptions = options
88+
// Create a fake instance with the methods we need
89+
const fakeInstance = {
90+
run: sandbox.stub().resolves(),
91+
send: sandbox.stub().resolves(),
92+
proc: sandbox.stub().returns({}),
93+
pid: sandbox.stub().returns(12345),
94+
stop: sandbox.stub(),
95+
stopped: false,
96+
}
97+
return fakeInstance
98+
})
99+
100+
// Replace ChildProcess constructor
101+
sandbox.replace(
102+
require('../../../../shared/utilities/processUtils'),
103+
'ChildProcess',
104+
childProcessConstructorSpy
105+
)
106+
107+
const serverOptions = createServerOptions({
108+
encryptionKey: Buffer.from('test-key'),
109+
executable: ['node'],
110+
serverModule: 'test-module.js',
111+
execArgv: ['--stdio'],
112+
})
113+
114+
await serverOptions()
115+
116+
assert(capturedOptions, 'ChildProcess constructor should have been called')
117+
assert(capturedOptions.spawnOptions, 'spawnOptions should be defined')
118+
assert(capturedOptions.spawnOptions.env, 'spawnOptions.env should be defined')
119+
assert.equal(capturedOptions.spawnOptions.env.USE_IAM_AUTH, undefined)
120+
})
121+
122+
it('defaults to IAM auth when parseCookies fails', async function () {
123+
isSageMakerStub.returns(true)
124+
isRemoteWorkspaceStub.returns(false)
125+
executeCommandStub.withArgs('sagemaker.parseCookies').rejects(new Error('Command failed'))
126+
127+
// Capture constructor arguments using sinon stub
128+
let capturedOptions: any = undefined
129+
const childProcessConstructorSpy = sandbox.stub().callsFake((command: string, args: string[], options: any) => {
130+
capturedOptions = options
131+
// Create a fake instance with the methods we need
132+
const fakeInstance = {
133+
run: sandbox.stub().resolves(),
134+
send: sandbox.stub().resolves(),
135+
proc: sandbox.stub().returns({}),
136+
pid: sandbox.stub().returns(12345),
137+
stop: sandbox.stub(),
138+
stopped: false,
139+
}
140+
return fakeInstance
141+
})
142+
143+
// Replace ChildProcess constructor
144+
sandbox.replace(
145+
require('../../../../shared/utilities/processUtils'),
146+
'ChildProcess',
147+
childProcessConstructorSpy
148+
)
149+
150+
const serverOptions = createServerOptions({
151+
encryptionKey: Buffer.from('test-key'),
152+
executable: ['node'],
153+
serverModule: 'test-module.js',
154+
execArgv: ['--stdio'],
155+
})
156+
157+
await serverOptions()
158+
159+
assert(capturedOptions, 'ChildProcess constructor should have been called')
160+
assert(capturedOptions.spawnOptions, 'spawnOptions should be defined')
161+
assert(capturedOptions.spawnOptions.env, 'spawnOptions.env should be defined')
162+
assert.equal(capturedOptions.spawnOptions.env.USE_IAM_AUTH, 'true')
163+
})
164+
165+
it('does not default to IAM in remote workspace without SMUS', async function () {
166+
isSageMakerStub.returns(true)
167+
isRemoteWorkspaceStub.returns(true)
168+
process.env.SERVICE_NAME = 'OtherService'
169+
executeCommandStub.withArgs('sagemaker.parseCookies').rejects(new Error('Command failed'))
170+
171+
// Capture constructor arguments using sinon stub
172+
let capturedOptions: any = undefined
173+
const childProcessConstructorSpy = sandbox.stub().callsFake((command: string, args: string[], options: any) => {
174+
capturedOptions = options
175+
// Create a fake instance with the methods we need
176+
const fakeInstance = {
177+
run: sandbox.stub().resolves(),
178+
send: sandbox.stub().resolves(),
179+
proc: sandbox.stub().returns({}),
180+
pid: sandbox.stub().returns(12345),
181+
stop: sandbox.stub(),
182+
stopped: false,
183+
}
184+
return fakeInstance
185+
})
186+
187+
// Replace ChildProcess constructor
188+
sandbox.replace(
189+
require('../../../../shared/utilities/processUtils'),
190+
'ChildProcess',
191+
childProcessConstructorSpy
192+
)
193+
194+
const serverOptions = createServerOptions({
195+
encryptionKey: Buffer.from('test-key'),
196+
executable: ['node'],
197+
serverModule: 'test-module.js',
198+
execArgv: ['--stdio'],
199+
})
200+
201+
await serverOptions()
202+
203+
assert(capturedOptions, 'ChildProcess constructor should have been called')
204+
assert(capturedOptions.spawnOptions, 'spawnOptions should be defined')
205+
assert(capturedOptions.spawnOptions.env, 'spawnOptions.env should be defined')
206+
assert.equal(capturedOptions.spawnOptions.env.USE_IAM_AUTH, undefined)
207+
})
208+
// jscpd:ignore-end
209+
})

0 commit comments

Comments
 (0)