merge in remote auth fix

Author: Hweinstock

.changes/next-release/Bug Fix-b1b5475c-fc00-40d2-a05e-cc247ab69a01.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "auth: remote workspaces use the global state, showing incorrect profile info"
+}

src/auth/auth.ts

@@ -743,11 +743,21 @@ export class Auth implements AuthService, ConnectionManager {
 
     static #instance: Auth | undefined
     public static get instance() {
-        const devEnvId = getCodeCatalystDevEnvId()
-        const memento =
-            devEnvId !== undefined ? partition(globals.context.globalState, devEnvId) : globals.context.globalState
+        return (this.#instance ??= new Auth(new ProfileStore(getMemento())))
 
-        return (this.#instance ??= new Auth(new ProfileStore(memento)))
+        function getMemento() {
+            if (!vscode.env.remoteName) {
+                return globals.context.globalState
+            }
+
+            const devEnvId = getCodeCatalystDevEnvId()
+
+            if (devEnvId !== undefined) {
+                return partition(globals.context.globalState, devEnvId)
+            }
+
+            return globals.context.workspaceState
+        }
     }
 
     private getSsoProfileLabel(profile: SsoProfile) {

src/codecatalyst/model.ts

@@ -30,10 +30,9 @@ import { CodeCatalystAuthenticationProvider } from './auth'
 import { ToolkitError } from '../shared/errors'
 import { Result } from '../shared/utilities/result'
 import { VscodeRemoteConnection, ensureDependencies } from '../shared/remoteSession'
-import { VscodeRemoteSshConfig, sshLogFileLocation } from '../shared/vscodeRemoteSshConfig'
+import { SshConfig, sshLogFileLocation } from '../shared/sshConfig'
 
 export type DevEnvironmentId = Pick<DevEnvironment, 'id' | 'org' | 'project'>
-export const hostNamePrefix = 'aws-devenv-'
 export const connectScriptPrefix = 'codecatalyst_connect'
 
 export const docs = {
@@ -142,10 +141,6 @@ export function bearerTokenCacheLocation(devenvId: string): string {
     return path.join(globals.context.globalStorageUri.fsPath, `codecatalyst.${devenvId}.token`)
 }
 
-export function getHostNameFromEnv(env: DevEnvironmentId): string {
-    return `${hostNamePrefix}${env.id}`
-}
-
 export interface ConnectedDevEnv {
     readonly summary: DevEnvironment
     readonly devenvClient: DevEnvClient
@@ -209,7 +204,8 @@ export async function prepareDevEnvConnection(
     { topic, timeout }: { topic?: string; timeout?: Timeout } = {}
 ): Promise<DevEnvConnection> {
     const { ssm, vsc, ssh } = (await ensureDependencies()).unwrap()
-    const sshConfig = new VscodeRemoteSshConfig(ssh, hostNamePrefix, connectScriptPrefix)
+    const hostNamePrefix = 'aws-devenv-'
+    const sshConfig = new SshConfig(ssh, hostNamePrefix, connectScriptPrefix)
     const config = await sshConfig.ensureValid()
 
     if (config.isErr()) {
@@ -225,7 +221,7 @@ export async function prepareDevEnvConnection(
         projectName: project.name,
     })
 
-    const hostname = getHostNameFromEnv({ id, org, project })
+    const hostname = `${hostNamePrefix}${id}`
     const logPrefix = topic ? `codecatalyst ${topic} (${id})` : `codecatalyst (${id})`
     const logger = (data: string) => getLogger().verbose(`${logPrefix}: ${data}`)
     const envProvider = createCodeCatalystEnvProvider(client, ssm, runningDevEnv)
@@ -263,7 +259,7 @@ export async function openDevEnv(
         const repo = env.devenv.repositories.length == 1 ? env.devenv.repositories[0].repositoryName : undefined
         targetPath = repo ? `/projects/${repo}` : '/projects'
     }
-    await startVscodeRemote(env.SessionProcess, getHostNameFromEnv(env.devenv), targetPath, env.vscPath)
+    await startVscodeRemote(env.SessionProcess, env.hostname, targetPath, env.vscPath)
 }
 
 // The "codecatalyst_connect" metric should really be splt into two parts:

src/ec2/commands.ts

@@ -27,7 +27,7 @@ export async function openTerminal(node?: Ec2Node) {
 export async function openRemoteConnection(node?: Ec2Node) {
     const selection = await getSelection(node)
     const connectionManager = new Ec2ConnectionManager(selection.region)
-    await connectionManager.attemptToOpenRemoteConnection(selection)
+    await connectionManager.tryOpenRemoteConnection(selection)
 }
 
 export async function startInstance(node?: Ec2Node) {

src/ec2/model.ts

@@ -26,7 +26,7 @@ import { createBoundProcess } from '../codecatalyst/model'
 import { getLogger } from '../shared/logger/logger'
 import { Timeout } from '../shared/utilities/timeoutUtils'
 import { showMessageWithCancel } from '../shared/utilities/messages'
-import { VscodeRemoteSshConfig, sshLogFileLocation } from '../shared/vscodeRemoteSshConfig'
+import { SshConfig, sshLogFileLocation } from '../shared/sshConfig'
 import { SshKeyPair } from './sshKeyPair'
 import globals from '../shared/extensionGlobals'
 
@@ -36,9 +36,6 @@ interface Ec2RemoteEnv extends VscodeRemoteConnection {
     selection: Ec2Selection
 }
 
-const ec2ConnectScriptPrefix = 'ec2_connect'
-const hostNamePrefix = 'ec2-'
-
 export class Ec2ConnectionManager {
     protected ssmClient: SsmClient
     protected ec2Client: Ec2Client
@@ -177,15 +174,14 @@ export class Ec2ConnectionManager {
         }
     }
 
-    public async attemptToOpenRemoteConnection(selection: Ec2Selection): Promise<void> {
+    public async tryOpenRemoteConnection(selection: Ec2Selection): Promise<void> {
         await this.checkForStartSessionError(selection)
 
         const remoteUser = await this.getRemoteUser(selection.instanceId)
         const remoteEnv = await this.prepareEc2RemoteEnvWithProgress(selection, remoteUser)
 
-        const fullHostName = `${hostNamePrefix}${selection.instanceId}`
         try {
-            await startVscodeRemote(remoteEnv.SessionProcess, fullHostName, '/', remoteEnv.vscPath, remoteUser)
+            await startVscodeRemote(remoteEnv.SessionProcess, remoteEnv.hostname, '/', remoteEnv.vscPath, remoteUser)
         } catch (err) {
             this.throwGeneralConnectionError(selection, err as Error)
         }
@@ -202,7 +198,8 @@ export class Ec2ConnectionManager {
         const logger = this.configureRemoteConnectionLogger(selection.instanceId)
         const { ssm, vsc, ssh } = (await ensureDependencies()).unwrap()
         const keyPath = await this.configureSshKeys(selection, remoteUser)
-        const sshConfig = new VscodeRemoteSshConfig(ssh, hostNamePrefix, ec2ConnectScriptPrefix, keyPath)
+        const hostNamePrefix = 'aws-ec2-'
+        const sshConfig = new SshConfig(ssh, hostNamePrefix, 'ec2_connect', keyPath)
 
         const config = await sshConfig.ensureValid()
         if (config.isErr()) {
@@ -223,7 +220,7 @@ export class Ec2ConnectionManager {
         })
 
         return {
-            hostname: selection.instanceId,
+            hostname: `${hostNamePrefix}${selection.instanceId}`,
             envProvider,
             sshPath: ssh,
             vscPath: vsc,
@@ -250,10 +247,10 @@ export class Ec2ConnectionManager {
         sshKeyPair: SshKeyPair,
         remoteUser: string
     ): Promise<void> {
-        const sshKey = await sshKeyPair.getPublicKey()
+        const sshPubKey = await sshKeyPair.getPublicKey()
 
         const remoteAuthorizedKeysPaths = `/home/${remoteUser}/.ssh/authorized_keys`
-        const command = `echo "${sshKey}" > ${remoteAuthorizedKeysPaths}`
+        const command = `echo "${sshPubKey}" > ${remoteAuthorizedKeysPaths}`
         const documentName = 'AWS-RunShellScript'
 
         await this.ssmClient.sendCommandAndWait(selection.instanceId, documentName, {

src/shared/extensions/ssh.ts

@@ -126,7 +126,8 @@ export async function startVscodeRemote(
     vscPath: string,
     user?: string
 ): Promise<void> {
-    const workspaceUri = `vscode-remote://ssh-remote+${`${user}@` ?? ''}${hostname}${targetDirectory}`
+    const userAt = user ? `${user}@` : ''
+    const workspaceUri = `vscode-remote://ssh-remote+${userAt}${hostname}${targetDirectory}`
 
     const settings = new RemoteSshSettings()
     settings.ensureDefaultExtension(VSCODE_EXTENSION_ID.awstoolkit)

src/shared/vscodeRemoteSshConfig.ts renamed to src/shared/sshConfig.ts

@@ -20,7 +20,7 @@ import { fileExists, readFileAsString } from './filesystemUtilities'
 
 const localize = nls.loadMessageBundle()
 
-export class VscodeRemoteSshConfig {
+export class SshConfig {
     protected readonly configHostName: string
     protected readonly proxyCommandRegExp: RegExp
 

src/test/shared/vscodeRemoteSshConfig.test.ts renamed to src/test/shared/sshConfig.test.ts

@@ -10,7 +10,7 @@ import * as http from 'http'
 import { ToolkitError } from '../../shared/errors'
 import { Result } from '../../shared/utilities/result'
 import { ChildProcess, ChildProcessResult } from '../../shared/utilities/childProcess'
-import { VscodeRemoteSshConfig, ensureConnectScript, sshLogFileLocation } from '../../shared/vscodeRemoteSshConfig'
+import { SshConfig, ensureConnectScript, sshLogFileLocation } from '../../shared/sshConfig'
 import { FakeExtensionContext } from '../fakeExtensionContext'
 import { fileExists, makeTemporaryToolkitFolder } from '../../shared/filesystemUtilities'
 import {
@@ -23,7 +23,7 @@ import { StartDevEnvironmentSessionRequest } from 'aws-sdk/clients/codecatalyst'
 import { mkdir, readFile, writeFile } from 'fs-extra'
 import { SystemUtilities } from '../../shared/systemUtilities'
 
-class MockSshConfig extends VscodeRemoteSshConfig {
+class MockSshConfig extends SshConfig {
     // State variables to track logic flow.
     public testIsWin: boolean = false
     public configSection: string = ''
@@ -109,10 +109,7 @@ describe('VscodeRemoteSshConfig', async function () {
             Promise<void>
         >
         before(function () {
-            promptUserToConfigureSshConfigStub = sinon.stub(
-                VscodeRemoteSshConfig.prototype,
-                'promptUserToConfigureSshConfig'
-            )
+            promptUserToConfigureSshConfigStub = sinon.stub(SshConfig.prototype, 'promptUserToConfigureSshConfig')
         })
 
         beforeEach(function () {