fix(toolkit): attach lifetime to generated ssh keys. (#5578)

## Problem

For connecting VSCode to EC2 instance, we generate an ssh key pair on
disk. This results in writing the private ssh key to VSCode global
storage, allowing the key to be potentially reused by other users on the
same machine.

## Solution  

Attach a lifetime to any key pair generated such that they wipe from
disk after X seconds. Value is currently set is 30 seconds to allow
connection to reliably establish. Also, change file permissions to
read/write owner only and change behavior to overwrite existing keys.

Unable to test file permissions, due to VSCode file system unable to
provide us with enough detail. Only provides whether it is readonly or
not (somewhat unreliably:
microsoft/vscode-discussions#673). VSCode
file system is what is used in `fs.ts` implementation.

This is in-line with how ec2 instance connect works:
https://github.com/aws/aws-ec2-instance-connect-cli/blob/master/ec2instanceconnectcli/EC2InstanceConnectKey.py


---

<!--- REMINDER: Ensure that your PR meets the guidelines in
CONTRIBUTING.md -->

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: JadenSimon <31319484+JadenSimon@users.noreply.github.com>
Co-authored-by: Justin M. Keyes <jmkeyes@amazon.com>
Co-authored-by: Weinstock <hkobew@80a9971f0a95.ant.amazon.com>

Author: 4 people

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

@@ -34,6 +34,7 @@ export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMC
 
 interface Ec2RemoteEnv extends VscodeRemoteConnection {
     selection: Ec2Selection
+    keyPair: SshKeyPair
 }
 
 export class Ec2ConnectionManager {
@@ -194,9 +195,9 @@ export class Ec2ConnectionManager {
     public async prepareEc2RemoteEnv(selection: Ec2Selection, remoteUser: string): Promise<Ec2RemoteEnv> {
         const logger = this.configureRemoteConnectionLogger(selection.instanceId)
         const { ssm, vsc, ssh } = (await ensureDependencies()).unwrap()
-        const keyPath = await this.configureSshKeys(selection, remoteUser)
+        const keyPair = await this.configureSshKeys(selection, remoteUser)
         const hostNamePrefix = 'aws-ec2-'
-        const sshConfig = new SshConfig(ssh, hostNamePrefix, 'ec2_connect', keyPath)
+        const sshConfig = new SshConfig(ssh, hostNamePrefix, 'ec2_connect', keyPair.getPrivateKeyPath())
 
         const config = await sshConfig.ensureValid()
         if (config.isErr()) {
@@ -223,6 +224,7 @@ export class Ec2ConnectionManager {
             vscPath: vsc,
             SessionProcess,
             selection,
+            keyPair,
         }
     }
 
@@ -232,11 +234,11 @@ export class Ec2ConnectionManager {
         return logger
     }
 
-    public async configureSshKeys(selection: Ec2Selection, remoteUser: string): Promise<string> {
+    public async configureSshKeys(selection: Ec2Selection, remoteUser: string): Promise<SshKeyPair> {
         const keyPath = path.join(globals.context.globalStorageUri.fsPath, `aws-ec2-key`)
-        const keyPair = await SshKeyPair.getSshKeyPair(keyPath)
+        const keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
         await this.sendSshKeyToInstance(selection, keyPair, remoteUser)
-        return keyPath
+        return keyPair
     }
 
     public async sendSshKeyToInstance(

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

@@ -2,22 +2,36 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import * as fs from 'fs-extra'
+import { fs } from '../../shared'
+import { chmodSync } from 'fs'
 import { ToolkitError } from '../../shared/errors'
 import { ChildProcess } from '../../shared/utilities/childProcess'
+import { Timeout } from '../../shared/utilities/timeoutUtils'
 
 export class SshKeyPair {
     private publicKeyPath: string
-    private constructor(keyPath: string) {
+    private lifeTimeout: Timeout
+    private deleted: boolean = false
+
+    private constructor(
+        private keyPath: string,
+        lifetime: number
+    ) {
         this.publicKeyPath = `${keyPath}.pub`
+        this.lifeTimeout = new Timeout(lifetime)
+
+        this.lifeTimeout.onCompletion(async () => {
+            await this.delete()
+        })
     }
 
-    public static async getSshKeyPair(keyPath: string) {
-        const keyExists = await fs.pathExists(keyPath)
-        if (!keyExists) {
-            await SshKeyPair.generateSshKeyPair(keyPath)
+    public static async getSshKeyPair(keyPath: string, lifetime: number) {
+        // Overwrite key if already exists
+        if (await fs.existsFile(keyPath)) {
+            await fs.delete(keyPath)
         }
-        return new SshKeyPair(keyPath)
+        await SshKeyPair.generateSshKeyPair(keyPath)
+        return new SshKeyPair(keyPath, lifetime)
     }
 
     public static async generateSshKeyPair(keyPath: string): Promise<void> {
@@ -26,14 +40,38 @@ export class SshKeyPair {
         if (result.exitCode !== 0) {
             throw new ToolkitError('ec2: Failed to generate ssh key', { details: { stdout: result.stdout } })
         }
+        chmodSync(keyPath, 0o600)
     }
 
     public getPublicKeyPath(): string {
         return this.publicKeyPath
     }
 
+    public getPrivateKeyPath(): string {
+        return this.keyPath
+    }
+
     public async getPublicKey(): Promise<string> {
-        const contents = await fs.readFile(this.publicKeyPath, 'utf-8')
+        const contents = await fs.readFileAsString(this.publicKeyPath)
         return contents
     }
+
+    public async delete(): Promise<void> {
+        await fs.delete(this.keyPath)
+        await fs.delete(this.publicKeyPath)
+
+        if (!this.lifeTimeout.completed) {
+            this.lifeTimeout.cancel()
+        }
+
+        this.deleted = true
+    }
+
+    public isDeleted(): boolean {
+        return this.deleted
+    }
+
+    public timeAlive(): number {
+        return this.lifeTimeout.elapsedTime
+    }
 }

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

@@ -139,7 +139,7 @@ describe('Ec2ConnectClient', function () {
                 instanceId: 'test-id',
                 region: 'test-region',
             }
-            const mockKeys = await SshKeyPair.getSshKeyPair('')
+            const mockKeys = await SshKeyPair.getSshKeyPair('', 30000)
             await client.sendSshKeyToInstance(testSelection, mockKeys, '')
             sinon.assert.calledWith(sendCommandStub, testSelection.instanceId, 'AWS-RunShellScript')
             sinon.restore()

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

@@ -2,35 +2,55 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-
+import * as vscode from 'vscode'
 import assert from 'assert'
-import * as fs from 'fs-extra'
 import * as sinon from 'sinon'
-import { makeTemporaryToolkitFolder, tryRemoveFolder } from '../../../shared/filesystemUtilities'
+import * as path from 'path'
 import { SshKeyPair } from '../../../awsService/ec2/sshKeyPair'
+import { createTestWorkspaceFolder, installFakeClock } from '../../testUtil'
+import { InstalledClock } from '@sinonjs/fake-timers'
 import { ChildProcess } from '../../../shared/utilities/childProcess'
+import { fs } from '../../../shared'
 
 describe('SshKeyUtility', async function () {
     let temporaryDirectory: string
     let keyPath: string
     let keyPair: SshKeyPair
+    let clock: InstalledClock
 
     before(async function () {
-        temporaryDirectory = await makeTemporaryToolkitFolder()
-        keyPath = `${temporaryDirectory}/test-key`
-        keyPair = await SshKeyPair.getSshKeyPair(keyPath)
+        temporaryDirectory = (await createTestWorkspaceFolder()).uri.fsPath
+        keyPath = path.join(temporaryDirectory, 'testKeyPair')
+        clock = installFakeClock()
+    })
+
+    beforeEach(async function () {
+        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
+    })
+
+    afterEach(async function () {
+        await keyPair.delete()
     })
 
     after(async function () {
-        await tryRemoveFolder(temporaryDirectory)
+        await keyPair.delete()
+        clock.uninstall()
+        sinon.restore()
     })
 
     describe('generateSshKeys', async function () {
         it('generates key in target file', async function () {
-            const contents = await fs.readFile(keyPath, 'utf-8')
+            const contents = await vscode.workspace.fs.readFile(vscode.Uri.file(keyPath))
             assert.notStrictEqual(contents.length, 0)
         })
 
+        it('generates unique key each time', async function () {
+            const beforeContent = await vscode.workspace.fs.readFile(vscode.Uri.file(keyPath))
+            keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
+            const afterContent = await vscode.workspace.fs.readFile(vscode.Uri.file(keyPath))
+            assert.notStrictEqual(beforeContent, afterContent)
+        })
+
         it('uses ed25519 algorithm to generate the keys', async function () {
             const process = new ChildProcess(`ssh-keygen`, ['-vvv', '-l', '-f', keyPath])
             const result = await process.run()
@@ -48,10 +68,42 @@ describe('SshKeyUtility', async function () {
         assert.notStrictEqual(key.length, 0)
     })
 
-    it('does not overwrite existing keys', async function () {
-        const generateStub = sinon.stub(SshKeyPair, 'generateSshKeyPair')
-        await SshKeyPair.getSshKeyPair(keyPath)
-        sinon.assert.notCalled(generateStub)
+    it('does overwrite existing keys on get call', async function () {
+        const generateStub = sinon.spy(SshKeyPair, 'generateSshKeyPair')
+        const keyBefore = await vscode.workspace.fs.readFile(vscode.Uri.file(keyPath))
+        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
+
+        const keyAfter = await vscode.workspace.fs.readFile(vscode.Uri.file(keyPath))
+        sinon.assert.calledOnce(generateStub)
+
+        assert.notStrictEqual(keyBefore, keyAfter)
+        sinon.restore()
+    })
+
+    it('deletes key on delete', async function () {
+        const pubKeyExistsBefore = await fs.existsFile(keyPair.getPublicKeyPath())
+        const privateKeyExistsBefore = await fs.existsFile(keyPair.getPrivateKeyPath())
+
+        await keyPair.delete()
+
+        const pubKeyExistsAfter = await fs.existsFile(keyPair.getPublicKeyPath())
+        const privateKeyExistsAfter = await fs.existsFile(keyPair.getPrivateKeyPath())
+
+        assert.strictEqual(pubKeyExistsBefore && privateKeyExistsBefore, true)
+        assert.strictEqual(pubKeyExistsAfter && privateKeyExistsAfter, false)
+        assert(keyPair.isDeleted())
+    })
+
+    it('deletes keys after timeout', async function () {
+        // Stub methods interacting with file system to avoid flaky test.
+        sinon.stub(SshKeyPair, 'generateSshKeyPair')
+        const deleteStub = sinon.stub(SshKeyPair.prototype, 'delete')
+
+        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 50)
+        await clock.tickAsync(10)
+        sinon.assert.notCalled(deleteStub)
+        await clock.tickAsync(100)
+        sinon.assert.calledOnce(deleteStub)
         sinon.restore()
     })
 })