Skip to content

Commit 8b4d088

Browse files
kzr-at-amazonkzr-at-amazonashishrp-aws
authored
fix(smus): Look for profile in both config and credentials files when adding region to profile (#8522)
## Problem - If a user picks a profile that does not have region entry from the config file, smus tries to update the profile with selected region. But smus is only looking for the profile in credentials file. ## Solution - Look for profile in both config and credentials files when adding region to profile - Use shared parsing method to handle profiles with `profile` prefix ## Test Updating profile with `profile` prefix in config file before ``` [profile configWithProfilePrefix] AWS_ACCESS_KEY_ID=xyz AWS_SECRET_ACCESS_KEY=xyz AWS_SESSION_TOKEN=xyz ``` after ``` [profile configWithProfilePrefix] AWS_ACCESS_KEY_ID=xyz AWS_SECRET_ACCESS_KEY=xyz AWS_SESSION_TOKEN=xyz region = ap-east-1 ``` --- - 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. Co-authored-by: kzr-at-amazon <build@amazon.com> Co-authored-by: invictus <149003065+ashishrp-aws@users.noreply.github.com>
1 parent e9cc11a commit 8b4d088

File tree

2 files changed

+208
-48
lines changed

2 files changed

+208
-48
lines changed

‎packages/core/src/sagemakerunifiedstudio/auth/ui/iamProfileSelection.ts‎

Lines changed: 72 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,10 @@
55

66
import * as vscode from 'vscode'
77
import * as path from 'path'
8+
import * as os from 'os'
89
import { getLogger } from '../../../shared/logger/logger'
910
import { ToolkitError } from '../../../shared/errors'
10-
import { loadSharedCredentialsProfiles } from '../../../auth/credentials/sharedCredentials'
11+
import { loadSharedCredentialsProfiles, parseIni } from '../../../auth/credentials/sharedCredentials'
1112
import { getCredentialsFilename, getConfigFilename } from '../../../auth/credentials/sharedCredentialsFile'
1213
import { SmusErrorCodes, DataZoneServiceId } from '../../shared/smusUtils'
1314
import globals from '../../../shared/extensionGlobals'
@@ -1189,7 +1190,7 @@ export class SmusIamProfileSelector {
11891190
}
11901191

11911192
// Parse the file line by line to handle profile replacement properly
1192-
const lines = content.split('\n')
1193+
const lines = content.split(os.EOL)
11931194
const newLines: string[] = []
11941195
let inTargetProfile = false
11951196
let profileFound = false
@@ -1230,7 +1231,7 @@ export class SmusIamProfileSelector {
12301231
}
12311232

12321233
// Update content with the new lines
1233-
content = newLines.join('\n')
1234+
content = newLines.join(os.EOL)
12341235

12351236
// Write back to file
12361237
await fs.writeFile(credentialsPath, content)
@@ -1245,62 +1246,85 @@ export class SmusIamProfileSelector {
12451246
try {
12461247
logger.debug(`Updating profile ${profileName} with region ${region}`)
12471248

1248-
const credentialsPath = getCredentialsFilename()
1249+
// Check both config and credential files
1250+
const filepathsToCheck = [getCredentialsFilename(), getConfigFilename()]
12491251

1250-
if (!(await fs.existsFile(credentialsPath))) {
1251-
throw new ToolkitError('Credentials file not found', { code: 'CredentialsFileNotFound' })
1252-
}
1252+
let profileUpdated = false
12531253

1254-
// Read the current credentials file
1255-
const content = await fs.readFileText(credentialsPath)
1254+
for (const filePath of filepathsToCheck) {
1255+
// File does not exist, try next file
1256+
if (!(await fs.existsFile(filePath))) {
1257+
continue
1258+
}
12561259

1257-
// Find the profile section
1258-
const profileSectionRegex = new RegExp(`^\\[${profileName}\\]$`, 'm')
1259-
const profileMatch = content.match(profileSectionRegex)
1260+
const content = await fs.readFileText(filePath)
1261+
const sections = parseIni(content, vscode.Uri.file(filePath))
12601262

1261-
if (!profileMatch) {
1262-
throw new ToolkitError(`Profile ${profileName} not found in credentials file`, {
1263-
code: 'ProfileNotFound',
1264-
})
1265-
}
1263+
// Find the profile section in this file
1264+
const profileSection = sections.find(
1265+
(section) => section.type === 'profile' && section.name === profileName
1266+
)
12661267

1267-
// Find the next profile section or end of file
1268-
const profileStartIndex = profileMatch.index!
1269-
const nextProfileMatch = content.slice(profileStartIndex + 1).match(/^\[.*\]$/m)
1270-
const profileEndIndex = nextProfileMatch ? profileStartIndex + 1 + nextProfileMatch.index! : content.length
1271-
1272-
// Extract the profile section
1273-
const profileSection = content.slice(profileStartIndex, profileEndIndex)
1274-
1275-
// Check if region already exists in the profile
1276-
let updatedProfileSection: string
1277-
1278-
if (this.regionLinePattern.test(profileSection)) {
1279-
// Replace existing region
1280-
updatedProfileSection = profileSection.replace(this.regionLinePattern, `region = ${region}`)
1281-
} else {
1282-
// Add region to the profile (before any empty lines at the end)
1283-
const lines = profileSection.split('\n')
1284-
// Find the last non-empty line index (compatible with older JS versions)
1285-
let lastNonEmptyIndex = -1
1286-
for (let i = lines.length - 1; i >= 0; i--) {
1287-
if (lines[i].trim() !== '') {
1288-
lastNonEmptyIndex = i
1268+
// Profile not in this file, try next file
1269+
if (!profileSection) {
1270+
continue
1271+
}
1272+
1273+
// Find the profile section boundaries using the startLines from parsed section
1274+
const profileStartLine = profileSection.startLines[0]
1275+
const lines = content.split(os.EOL)
1276+
1277+
// Find the next profile section or end of file
1278+
let profileEndLine = lines.length
1279+
for (let i = profileStartLine + 1; i < lines.length; i++) {
1280+
if (lines[i].match(/^\s*\[([^\[\]]+)]\s*$/)) {
1281+
profileEndLine = i
12891282
break
12901283
}
12911284
}
1292-
lines.splice(lastNonEmptyIndex + 1, 0, `region = ${region}`)
1293-
updatedProfileSection = lines.join('\n')
1294-
}
12951285

1296-
// Replace the profile section in the content
1297-
const updatedContent =
1298-
content.slice(0, profileStartIndex) + updatedProfileSection + content.slice(profileEndIndex)
1286+
// Extract the profile section lines
1287+
const profileLines = lines.slice(profileStartLine, profileEndLine)
1288+
1289+
// Check if region already exists in the profile
1290+
const regionLineIndex = profileLines.findIndex((line) => this.regionLinePattern.test(line))
12991291

1300-
// Write back to file
1301-
await fs.writeFile(credentialsPath, updatedContent)
1292+
if (regionLineIndex !== -1) {
1293+
// Replace existing region
1294+
profileLines[regionLineIndex] = `region = ${region}`
1295+
} else {
1296+
// Add region to the profile (after the last non-empty line)
1297+
let lastNonEmptyIndex = -1
1298+
for (let i = profileLines.length - 1; i >= 0; i--) {
1299+
if (profileLines[i].trim() !== '') {
1300+
lastNonEmptyIndex = i
1301+
break
1302+
}
1303+
}
1304+
profileLines.splice(lastNonEmptyIndex + 1, 0, `region = ${region}`)
1305+
}
1306+
1307+
// Reconstruct the file content
1308+
const updatedLines = [
1309+
...lines.slice(0, profileStartLine),
1310+
...profileLines,
1311+
...lines.slice(profileEndLine),
1312+
]
1313+
const updatedContent = updatedLines.join(os.EOL)
13021314

1303-
logger.debug(`Successfully updated profile ${profileName} with region ${region}`)
1315+
// Write back to file
1316+
await fs.writeFile(filePath, updatedContent)
1317+
1318+
logger.debug(`Successfully updated profile ${profileName} with region ${region} in ${filePath}`)
1319+
profileUpdated = true
1320+
break
1321+
}
1322+
1323+
if (!profileUpdated) {
1324+
throw new ToolkitError(`Profile ${profileName} not found in credentials or config file`, {
1325+
code: 'ProfileNotFound',
1326+
})
1327+
}
13041328
} catch (error) {
13051329
logger.error('Failed to update profile region: %s', error)
13061330
throw new ToolkitError(`Failed to update profile region: ${(error as Error).message}`, {

‎packages/core/src/test/sagemakerunifiedstudio/auth/ui/iamProfileSelection.test.ts‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,13 @@
44
*/
55

66
import assert from 'assert'
7+
import * as path from 'path'
8+
import * as os from 'os'
9+
import * as sinon from 'sinon'
710
import { SmusIamProfileSelector } from '../../../../sagemakerunifiedstudio/auth/ui/iamProfileSelection'
11+
import { makeTemporaryToolkitFolder } from '../../../../shared/filesystemUtilities'
12+
import { fs } from '../../../../shared'
13+
import { EnvironmentVariables } from '../../../../shared/environmentVariables'
814

915
describe('SmusIamProfileSelector', function () {
1016
describe('showRegionSelection', function () {
@@ -18,4 +24,134 @@ describe('SmusIamProfileSelector', function () {
1824
assert.strictEqual(typeof SmusIamProfileSelector.showIamProfileSelection, 'function')
1925
})
2026
})
27+
28+
describe('updateProfileRegion', function () {
29+
let tempFolder: string
30+
let credentialsPath: string
31+
let configPath: string
32+
33+
beforeEach(async function () {
34+
tempFolder = await makeTemporaryToolkitFolder()
35+
credentialsPath = path.join(tempFolder, 'credentials')
36+
configPath = path.join(tempFolder, 'config')
37+
38+
// Stub environment variables to use temp files
39+
sinon.stub(process, 'env').value({
40+
AWS_SHARED_CREDENTIALS_FILE: credentialsPath,
41+
AWS_CONFIG_FILE: configPath,
42+
} as EnvironmentVariables)
43+
})
44+
45+
afterEach(async function () {
46+
await fs.delete(tempFolder, { recursive: true })
47+
sinon.restore()
48+
})
49+
50+
it('should update region in credentials file when profile exists there', async function () {
51+
// Create credentials file with a profile without region
52+
const credentialsContent = [
53+
'[test-profile]',
54+
'aws_access_key_id = XYZ',
55+
'aws_secret_access_key = XYZ',
56+
'',
57+
].join(os.EOL)
58+
await fs.writeFile(credentialsPath, credentialsContent)
59+
60+
// Call the private method using bracket notation
61+
await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'us-west-2')
62+
63+
// Verify the region was added
64+
const updatedContent = await fs.readFileText(credentialsPath)
65+
assert.ok(updatedContent.includes('region = us-west-2'))
66+
assert.ok(updatedContent.includes('[test-profile]'))
67+
assert.ok(updatedContent.includes('aws_access_key_id = XYZ'))
68+
})
69+
70+
it('should update region in config file when profile exists there', async function () {
71+
// Create config file with a profile without region
72+
const configContent = ['[profile test-profile]', 'output = json', ''].join(os.EOL)
73+
await fs.writeFile(configPath, configContent)
74+
75+
// Call the private method
76+
await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'eu-west-1')
77+
78+
// Verify the region was added
79+
const updatedContent = await fs.readFileText(configPath)
80+
assert.ok(updatedContent.includes('region = eu-west-1'))
81+
assert.ok(updatedContent.includes('[profile test-profile]'))
82+
assert.ok(updatedContent.includes('output = json'))
83+
})
84+
85+
it('should handle multiple profiles in credentials file', async function () {
86+
// Create credentials file with multiple profiles
87+
const credentialsContent = [
88+
'[default]',
89+
'aws_access_key_id = XYZ',
90+
'aws_secret_access_key = XYZ',
91+
'',
92+
'[test-profile]',
93+
'aws_access_key_id = XYZ',
94+
'aws_secret_access_key = XYZ',
95+
'',
96+
'[another-profile]',
97+
'aws_access_key_id = XYZ',
98+
'aws_secret_access_key = XYZ',
99+
'',
100+
].join(os.EOL)
101+
await fs.writeFile(credentialsPath, credentialsContent)
102+
103+
// Update the region for test-profile
104+
await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'us-west-2')
105+
106+
// Verify the region was added only to test-profile
107+
const updatedContent = await fs.readFileText(credentialsPath)
108+
const lines = updatedContent.split(os.EOL)
109+
110+
// Find test-profile section
111+
const testProfileIndex = lines.findIndex((line) => line.includes('[test-profile]'))
112+
const anotherProfileIndex = lines.findIndex((line) => line.includes('[another-profile]'))
113+
114+
// Check that region is between test-profile and another-profile
115+
const testProfileSection = lines.slice(testProfileIndex, anotherProfileIndex).join(os.EOL)
116+
assert.ok(testProfileSection.includes('region = us-west-2'))
117+
118+
// Check that other profiles are unchanged
119+
assert.ok(updatedContent.includes('[default]'))
120+
assert.ok(updatedContent.includes('[another-profile]'))
121+
})
122+
123+
it('should throw error when profile does not exist in either file', async function () {
124+
// Create both files without the target profile
125+
const credentialsContent = ['[default]', 'aws_access_key_id = XYZ', 'aws_secret_access_key = XYZ'].join(
126+
os.EOL
127+
)
128+
const configContent = ['[profile default]', 'region = us-east-1'].join(os.EOL)
129+
await fs.writeFile(credentialsPath, credentialsContent)
130+
await fs.writeFile(configPath, configContent)
131+
132+
// Attempt to update non-existent profile
133+
await assert.rejects(
134+
async () => {
135+
await (SmusIamProfileSelector as any).updateProfileRegion('non-existent-profile', 'us-west-2')
136+
},
137+
(error: Error) => {
138+
assert.ok(error.message.includes('not found'))
139+
return true
140+
}
141+
)
142+
})
143+
144+
it('should throw error when neither file exists', async function () {
145+
// Attempt to update profile
146+
await assert.rejects(
147+
async () => {
148+
await (SmusIamProfileSelector as any).updateProfileRegion('test-profile', 'us-west-2')
149+
},
150+
(error: Error) => {
151+
assert.ok(error.message.includes('not found'))
152+
return true
153+
}
154+
)
155+
})
156+
})
21157
})

0 commit comments

Comments
 (0)