fix(ec2): avoid wiping authorized_keys files on each connection

docs/arch_features.md

@@ -41,9 +41,11 @@ For connecting a new VSCode _terminal_, remote connect works like this:
 
 For EC2 specifically, there are a few additional steps:
 
+1. Remote window connections are only supported for EC2 instances running a linux based OS such as Amazon Linux or Ubuntu. However, the terminal option is supported by all OS, and will open a Powershell-based terminal for Windows instances.
 1. If connecting to EC2 instance via remote window, the toolkit generates temporary SSH keys (30 second lifetime), with the public key sent to the remote instance.
     - Key type is ed25519 if supported, or RSA otherwise.
-    - This connection will overwrite the `.ssh/authorized_keys` file on the remote machine with each connection.
+    - Lines in `.ssh/authorized_keys` marked with the comment `#AWSToolkitForVSCode` will be removed by AWS Toolkit.
+    - Assumes `.sss/authorized_keys` can be found under `/home/ec2-user/` on Amazon Linux and `/home/ubuntu/` on Ubuntu.
 1. If insufficient permissions are detected on the attached IAM role, toolkit will prompt to add an inline policy with the necessary actions.
 1. If SSM sessions remain open after closing the window/terminal, the toolkit will terminate them on-shutdown, or when starting another session to the same instance.
 

packages/core/src/awsService/ec2/model.ts

@@ -44,6 +44,12 @@ export interface Ec2RemoteEnv extends VscodeRemoteConnection {
     ssmSession: SSM.StartSessionResponse
 }
 
+export type Ec2OS = 'Amazon Linux' | 'Ubuntu' | 'macOS'
+interface RemoteUser {
+    os: Ec2OS
+    name: string
+}
+
 export class Ec2Connecter implements vscode.Disposable {
     protected ssmClient: SsmClient
     protected ec2Client: Ec2Client
@@ -198,10 +204,16 @@ export class Ec2Connecter implements vscode.Disposable {
                 remoteEnv.SessionProcess,
                 remoteEnv.hostname,
                 remoteEnv.sshPath,
-                remoteUser,
+                remoteUser.name,
                 testSession
             )
-            await startVscodeRemote(remoteEnv.SessionProcess, remoteEnv.hostname, '/', remoteEnv.vscPath, remoteUser)
+            await startVscodeRemote(
+                remoteEnv.SessionProcess,
+                remoteEnv.hostname,
+                '/',
+                remoteEnv.vscPath,
+                remoteUser.name
+            )
         } catch (err) {
             const message = err instanceof SshError ? 'Testing SSH connection to instance failed' : ''
             this.throwConnectionError(message, selection, err as Error)
@@ -210,7 +222,10 @@ export class Ec2Connecter implements vscode.Disposable {
         }
     }
 
-    public async prepareEc2RemoteEnvWithProgress(selection: Ec2Selection, remoteUser: string): Promise<Ec2RemoteEnv> {
+    public async prepareEc2RemoteEnvWithProgress(
+        selection: Ec2Selection,
+        remoteUser: RemoteUser
+    ): Promise<Ec2RemoteEnv> {
         const timeout = new Timeout(60000)
         await showMessageWithCancel('AWS: Opening remote connection...', timeout)
         const remoteEnv = await this.prepareEc2RemoteEnv(selection, remoteUser).finally(() => timeout.cancel())
@@ -223,7 +238,7 @@ export class Ec2Connecter implements vscode.Disposable {
         return ssmSession
     }
 
-    public async prepareEc2RemoteEnv(selection: Ec2Selection, remoteUser: string): Promise<Ec2RemoteEnv> {
+    public async prepareEc2RemoteEnv(selection: Ec2Selection, remoteUser: RemoteUser): Promise<Ec2RemoteEnv> {
         const logger = this.configureRemoteConnectionLogger(selection.instanceId)
         const { ssm, vsc, ssh } = (await ensureDependencies()).unwrap()
         const keyPair = await this.configureSshKeys(selection, remoteUser)
@@ -271,38 +286,78 @@ export class Ec2Connecter implements vscode.Disposable {
         return logger
     }
 
-    public async configureSshKeys(selection: Ec2Selection, remoteUser: string): Promise<SshKeyPair> {
+    public async configureSshKeys(selection: Ec2Selection, remoteUser: RemoteUser): Promise<SshKeyPair> {
         const keyPair = await SshKeyPair.getSshKeyPair(`aws-ec2-key`, 30000)
         await this.sendSshKeyToInstance(selection, keyPair, remoteUser)
         return keyPair
     }
 
+    /** Removes old key(s) that we added to the remote ~/.ssh/authorized_keys file. */
+    public async tryCleanKeys(
+        instanceId: string,
+        hintComment: string,
+        hostOS: Ec2OS,
+        remoteAuthorizedKeysPath: string
+    ) {
+        try {
+            const deleteExistingKeyCommand = getRemoveLinesCommand(hintComment, hostOS, remoteAuthorizedKeysPath)
+            await this.sendCommandAndWait(instanceId, deleteExistingKeyCommand)
+        } catch (e) {
+            getLogger().warn(`ec2: failed to clean keys: %O`, e)
+        }
+    }
+
+    private async sendCommandAndWait(instanceId: string, command: string) {
+        return await this.ssmClient.sendCommandAndWait(instanceId, 'AWS-RunShellScript', {
+            commands: [command],
+        })
+    }
+
     public async sendSshKeyToInstance(
         selection: Ec2Selection,
         sshKeyPair: SshKeyPair,
-        remoteUser: string
+        remoteUser: RemoteUser
     ): Promise<void> {
         const sshPubKey = await sshKeyPair.getPublicKey()
+        const hintComment = '#AWSToolkitForVSCode'
 
-        const remoteAuthorizedKeysPaths = `/home/${remoteUser}/.ssh/authorized_keys`
-        const command = `echo "${sshPubKey}" > ${remoteAuthorizedKeysPaths}`
-        const documentName = 'AWS-RunShellScript'
+        const remoteAuthorizedKeysPath = `/home/${remoteUser.name}/.ssh/authorized_keys`
 
-        await this.ssmClient.sendCommandAndWait(selection.instanceId, documentName, {
-            commands: [command],
-        })
+        const appendStr = (s: string) => `echo "${s}" >> ${remoteAuthorizedKeysPath}`
+        const writeKeyCommand = appendStr([sshPubKey.replace('\n', ''), hintComment].join(' '))
+
+        await this.tryCleanKeys(selection.instanceId, hintComment, remoteUser.os, remoteAuthorizedKeysPath)
+        await this.sendCommandAndWait(selection.instanceId, writeKeyCommand)
     }
 
-    public async getRemoteUser(instanceId: string) {
-        const osName = await this.ssmClient.getTargetPlatformName(instanceId)
-        if (osName === 'Amazon Linux') {
-            return 'ec2-user'
+    public async getRemoteUser(instanceId: string): Promise<RemoteUser> {
+        const os = await this.ssmClient.getTargetPlatformName(instanceId)
+        if (os === 'Amazon Linux') {
+            return { name: 'ec2-user', os }
         }
 
-        if (osName === 'Ubuntu') {
-            return 'ubuntu'
+        if (os === 'Ubuntu') {
+            return { name: 'ubuntu', os }
         }
 
-        throw new ToolkitError(`Unrecognized OS name ${osName} on instance ${instanceId}`, { code: 'UnknownEc2OS' })
+        throw new ToolkitError(`Unrecognized OS name ${os} on instance ${instanceId}`, { code: 'UnknownEc2OS' })
     }
 }
+
+/**
+ * Generate bash command (as string) to remove lines containing `pattern`.
+ * @param pattern pattern for deleted lines.
+ * @param filepath filepath (as string) to target with the command.
+ * @returns bash command to remove lines from file.
+ */
+export function getRemoveLinesCommand(pattern: string, hostOS: Ec2OS, filepath: string): string {
+    if (pattern.includes('/')) {
+        throw new ToolkitError(`ec2: cannot match pattern containing '/', given: ${pattern}`)
+    }
+    // Linux allows not passing extension to -i, whereas macOS requires zero length extension.
+    return `sed -i${isLinux(hostOS) ? '' : " ''"} /${pattern}/d ${filepath}`
+}
+
+function isLinux(os: Ec2OS): boolean {
+    return os === 'Amazon Linux' || os === 'Ubuntu'
+}

packages/core/src/shared/vscode/env.ts

@@ -149,6 +149,9 @@ export async function isCloudDesktop() {
     return (await new ChildProcess('/apollo/bin/getmyfabric').run().then((r) => r.exitCode)) === 0
 }
 
+export function isMac(): boolean {
+    return process.platform === 'darwin'
+}
 /** Returns true if OS is Windows. */
 export function isWin(): boolean {
     // if (isWeb()) {

packages/core/src/test/awsService/ec2/model.test.ts

@@ -5,7 +5,7 @@
 
 import assert from 'assert'
 import * as sinon from 'sinon'
-import { Ec2Connecter } from '../../../awsService/ec2/model'
+import { Ec2Connecter, getRemoveLinesCommand } from '../../../awsService/ec2/model'
 import { SsmClient } from '../../../shared/clients/ssmClient'
 import { Ec2Client } from '../../../shared/clients/ec2Client'
 import { Ec2Selection } from '../../../awsService/ec2/prompter'
@@ -15,6 +15,11 @@ import { SshKeyPair } from '../../../awsService/ec2/sshKeyPair'
 import { DefaultIamClient } from '../../../shared/clients/iamClient'
 import { assertNoTelemetryMatch, createTestWorkspaceFolder } from '../../testUtil'
 import { fs } from '../../../shared'
+import path from 'path'
+import { ChildProcess } from '../../../shared/utilities/processUtils'
+import { isMac, isWin } from '../../../shared/vscode/env'
+import { inspect } from '../../../shared/utilities/collectionUtils'
+import { assertLogsContain } from '../../globalSetup.test'
 
 describe('Ec2ConnectClient', function () {
     let client: Ec2Connecter
@@ -140,7 +145,7 @@ describe('Ec2ConnectClient', function () {
             }
 
             const keys = await SshKeyPair.getSshKeyPair('key', 30000)
-            await client.sendSshKeyToInstance(testSelection, keys, 'test-user')
+            await client.sendSshKeyToInstance(testSelection, keys, { name: 'test-user', os: 'Amazon Linux' })
             sinon.assert.calledWith(sendCommandStub, testSelection.instanceId, 'AWS-RunShellScript')
             sinon.restore()
         })
@@ -154,7 +159,7 @@ describe('Ec2ConnectClient', function () {
             }
             const testWorkspaceFolder = await createTestWorkspaceFolder()
             const keys = await SshKeyPair.getSshKeyPair('key', 60000)
-            await client.sendSshKeyToInstance(testSelection, keys, 'test-user')
+            await client.sendSshKeyToInstance(testSelection, keys, { name: 'test-user', os: 'Amazon Linux' })
             const privKey = await fs.readFileText(keys.getPrivateKeyPath())
             assertNoTelemetryMatch(privKey)
             sinon.restore()
@@ -178,13 +183,13 @@ describe('Ec2ConnectClient', function () {
         it('identifies the user for ubuntu as ubuntu', async function () {
             getTargetPlatformNameStub.resolves('Ubuntu')
             const remoteUser = await client.getRemoteUser('testInstance')
-            assert.strictEqual(remoteUser, 'ubuntu')
+            assert.strictEqual(remoteUser.name, 'ubuntu')
         })
 
         it('identifies the user for amazon linux as ec2-user', async function () {
             getTargetPlatformNameStub.resolves('Amazon Linux')
             const remoteUser = await client.getRemoteUser('testInstance')
-            assert.strictEqual(remoteUser, 'ec2-user')
+            assert.strictEqual(remoteUser.name, 'ec2-user')
         })
 
         it('throws error when not given known OS', async function () {
@@ -197,4 +202,94 @@ describe('Ec2ConnectClient', function () {
             }
         })
     })
+
+    describe('attemptToCleanKeys', async function () {
+        it('calls the sdk with the proper parameters', async function () {
+            const sendCommandStub = sinon.stub(SsmClient.prototype, 'sendCommandAndWait')
+
+            const testSelection = {
+                instanceId: 'test-id',
+                region: 'test-region',
+            }
+
+            await client.tryCleanKeys(testSelection.instanceId, 'hint', 'macOS', 'path/to/keys')
+            sendCommandStub.calledWith(testSelection.instanceId, 'AWS-RunShellScript', {
+                commands: [getRemoveLinesCommand('hint', 'macOS', 'path/to/keys')],
+            })
+            sinon.assert.calledWith(sendCommandStub, testSelection.instanceId, 'AWS-RunShellScript')
+            sinon.restore()
+        })
+
+        it('logs warning when sdk call fails', async function () {
+            const sendCommandStub = sinon
+                .stub(SsmClient.prototype, 'sendCommandAndWait')
+                .throws(new ToolkitError('error'))
+
+            const testSelection = {
+                instanceId: 'test-id',
+                region: 'test-region',
+            }
+
+            await client.tryCleanKeys(testSelection.instanceId, 'hint', 'macOS', 'path/to/keys')
+            sinon.assert.calledWith(sendCommandStub, testSelection.instanceId, 'AWS-RunShellScript', {
+                commands: [getRemoveLinesCommand('hint', 'macOS', 'path/to/keys')],
+            })
+            sinon.restore()
+            assertLogsContain('failed to clean keys', false, 'warn')
+        })
+    })
+})
+
+describe('getRemoveLinesCommand', async function () {
+    let tempPath: { uri: { fsPath: string } }
+
+    before(async function () {
+        tempPath = await createTestWorkspaceFolder()
+    })
+
+    after(async function () {
+        await fs.delete(tempPath.uri.fsPath, { recursive: true, force: true })
+    })
+
+    it('removes lines containing pattern', async function () {
+        if (isWin()) {
+            this.skip()
+        }
+        // For the test, we only need to distinguish mac and linux
+        const hostOS = isMac() ? 'macOS' : 'Amazon Linux'
+        const lines = ['line1', 'line2 pattern', 'line3', 'line4 pattern', 'line5', 'line6 pattern', 'line7']
+        const expected = ['line1', 'line3', 'line5', 'line7']
+
+        const lineToStr = (ls: string[]) => ls.join('\n') + '\n'
+
+        const textFile = path.join(tempPath.uri.fsPath, 'test.txt')
+        const originalContent = lineToStr(lines)
+        await fs.writeFile(textFile, originalContent)
+        const [command, ...args] = getRemoveLinesCommand('pattern', hostOS, textFile).split(' ')
+        const process = new ChildProcess(command, args, { collect: true })
+        const result = await process.run()
+        assert.strictEqual(
+            result.exitCode,
+            0,
+            `Ran command '${command} ${args.join(' ')}' and failed with result ${inspect(result)}`
+        )
+
+        const newContent = await fs.readFileText(textFile)
+        assert.notStrictEqual(newContent, originalContent)
+        assert.strictEqual(newContent, lineToStr(expected))
+    })
+
+    it('includes empty extension on macOS only', async function () {
+        const macCommand = getRemoveLinesCommand('pattern', 'macOS', 'test.txt')
+        const alCommand = getRemoveLinesCommand('pattern', 'Amazon Linux', 'test.txt')
+        const ubuntuCommand = getRemoveLinesCommand('pattern', 'Ubuntu', 'test.txt')
+
+        assert.ok(macCommand.includes("''"))
+        assert.ok(!alCommand.includes("''"))
+        assert.strictEqual(ubuntuCommand, alCommand)
+    })
+
+    it('throws when given invalid pattern', function () {
+        assert.throws(() => getRemoveLinesCommand('pat/tern', 'macOS', 'test.txt'))
+    })
 })

packages/toolkit/.changes/next-release/Bug Fix-641a1ae0-3ad4-48c8-9628-ba739a06c8e5.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "EC2: avoid overwriting authorized_keys file on remote"
+}