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

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

@@ -33,6 +33,7 @@ export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMC
 
 interface Ec2RemoteEnv extends VscodeRemoteConnection {
     selection: Ec2Selection
+    keyPair: SshKeyPair
 }
 
 export class Ec2ConnectionManager {
@@ -195,9 +196,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()) {
@@ -224,6 +225,7 @@ export class Ec2ConnectionManager {
             vscPath: vsc,
             SessionProcess,
             selection,
+            keyPair,
         }
     }
 
@@ -233,11 +235,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

@@ -5,19 +5,32 @@
 import * as fs from 'fs-extra'
 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.pathExists(keyPath)) {
+            await fs.remove(keyPath)
         }
-        return new SshKeyPair(keyPath)
+        await SshKeyPair.generateSshKeyPair(keyPath)
+        return new SshKeyPair(keyPath, lifetime)
     }
 
     public static async generateSshKeyPair(keyPath: string): Promise<void> {
@@ -26,14 +39,38 @@ export class SshKeyPair {
         if (result.exitCode !== 0) {
             throw new ToolkitError('ec2: Failed to generate ssh key', { details: { stdout: result.stdout } })
         }
+        await fs.chmod(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')
         return contents
     }
+
+    public async delete(): Promise<void> {
+        await fs.remove(this.publicKeyPath)
+        await fs.remove(this.keyPath)
+
+        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('', 0)
             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,33 +2,51 @@
  * 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'
 
 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()
+    })
+
+    it('generates key in target file', async function () {
+        const contents = await vscode.workspace.fs.readFile(vscode.Uri.file(keyPath))
+        assert.notStrictEqual(contents.length, 0)
     })
 
-    describe('generateSshKeys', async function () {
-        it('generates key in target file', async function () {
-            const contents = await fs.readFile(keyPath, 'utf-8')
-            assert.notStrictEqual(contents.length, 0)
-        })
+    it('sets key permission to read/write by owner', async function () {
+        const fileMode = (await fs.stat(keyPath)).mode
+        // Mode is in decimal, so convert to decimal with bitmask
+        // source: https://github.com/nodejs/node-v0.x-archive/issues/3045
+        assert.strictEqual(fileMode & 0o777, 0o600)
     })
 
     it('properly names the public key', function () {
@@ -40,10 +58,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.pathExists(keyPair.getPublicKeyPath())
+        const privateKeyExistsBefore = await fs.pathExists(keyPair.getPrivateKeyPath())
+
+        await keyPair.delete()
+
+        const pubKeyExistsAfter = await fs.pathExists(keyPair.getPublicKeyPath())
+        const privateKeyExistsAfter = await fs.pathExists(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()
     })
 })