test(ec2): ensure that private key is not written to telemetry

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

@@ -2,7 +2,8 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { fs } from '../../shared'
+import os from 'os'
+import { fs, globals } from '../../shared'
 import { ToolkitError } from '../../shared/errors'
 import { tryRun } from '../../shared/utilities/pathFind'
 import { Timeout } from '../../shared/utilities/timeoutUtils'
@@ -32,12 +33,24 @@ export class SshKeyPair {
         return new SshKeyPair(keyPath, lifetime)
     }
 
-    public static async generateSshKeyPair(keyPath: string): Promise<void> {
-        const keyGenerated = await this.tryKeyTypes(keyPath, ['ed25519', 'rsa'])
+    private static async assertGenerated(keyPath: string, keyGenerated: boolean): Promise<never | void> {
         if (!keyGenerated) {
-            throw new ToolkitError('ec2: Unable to generate ssh key pair')
+            throw new ToolkitError('ec2: Unable to generate ssh key pair with either ed25519 or rsa')
+        }
+
+        if (!(await fs.exists(keyPath))) {
+            throw new ToolkitError(`ec2: Failed to generate keys, resulting key not found at ${keyPath}`)
+        }
+    }
+
+    public static async generateSshKeyPair(keyPath: string): Promise<void> {
+        const keyGenerated = await SshKeyPair.tryKeyTypes(keyPath, ['ed25519', 'rsa'])
+        await SshKeyPair.assertGenerated(keyPath, keyGenerated)
+        // Should already be the case, but just in case we assert permissions.
+        // skip on Windows since it only allows write permission to be changed.
+        if (!globals.isWeb && os.platform() !== 'win32') {
+            await fs.chmod(keyPath, 0o600)
         }
-        await fs.chmod(keyPath, 0o600)
     }
     /**
      * Attempts to generate an ssh key pair. Returns true if successful, false otherwise.
@@ -72,8 +85,8 @@ export class SshKeyPair {
     }
 
     public async delete(): Promise<void> {
-        await fs.delete(this.keyPath)
-        await fs.delete(this.publicKeyPath)
+        await fs.delete(this.keyPath, { force: true })
+        await fs.delete(this.publicKeyPath, { force: true })
 
         if (!this.lifeTimeout.completed) {
             this.lifeTimeout.cancel()

packages/core/src/shared/telemetry/telemetryLogger.ts

@@ -118,4 +118,11 @@ export class TelemetryLogger {
     public queryFull(query: MetricQuery): MetricDatum[] {
         return this._metrics.filter((m) => m.MetricName === query.metricName)
     }
+
+    /**
+     * Queries telemetry for metrics with metadata key or value matching the given regex.
+     */
+    public queryRegex(re: RegExp | string): MetricDatum[] {
+        return this._metrics.filter((m) => m.Metadata?.some((md) => md.Value?.match(re) || md.Key?.match(re)))
+    }
 }

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

@@ -13,6 +13,9 @@ import { ToolkitError } from '../../../shared/errors'
 import { IAM } from 'aws-sdk'
 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'
 
 describe('Ec2ConnectClient', function () {
     let client: Ec2ConnectionManager
@@ -144,6 +147,25 @@ describe('Ec2ConnectClient', function () {
             sinon.assert.calledWith(sendCommandStub, testSelection.instanceId, 'AWS-RunShellScript')
             sinon.restore()
         })
+
+        it('avoids writing the keys to any telemetry metrics', async function () {
+            sinon.stub(SsmClient.prototype, 'sendCommandAndWait')
+
+            const testSelection = {
+                instanceId: 'test-id',
+                region: 'test-region',
+            }
+            const testWorkspaceFolder = await createTestWorkspaceFolder()
+            const keyPath = path.join(testWorkspaceFolder.uri.fsPath, 'key')
+            const keys = await SshKeyPair.getSshKeyPair(keyPath, 60000)
+            await client.sendSshKeyToInstance(testSelection, keys, 'test-user')
+            const privKey = await fs.readFileText(keyPath)
+            assertNoTelemetryMatch(privKey)
+            sinon.restore()
+
+            await keys.delete()
+            await fs.delete(testWorkspaceFolder.uri, { force: true })
+        })
     })
 
     describe('getRemoteUser', async function () {

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

@@ -2,7 +2,6 @@
  * 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 nodefs from 'fs' // eslint-disable-line no-restricted-imports
 import * as sinon from 'sinon'
@@ -14,7 +13,7 @@ import { InstalledClock } from '@sinonjs/fake-timers'
 import { ChildProcess } from '../../../shared/utilities/processUtils'
 import { fs, globals } from '../../../shared'
 
-describe('SshKeyUtility', async function () {
+describe('SshKeyPair', async function () {
     let temporaryDirectory: string
     let keyPath: string
     let keyPair: SshKeyPair
@@ -41,14 +40,14 @@ describe('SshKeyUtility', async function () {
     })
 
     it('generates key in target file', async function () {
-        const contents = await fs.readFileBytes(vscode.Uri.file(keyPath))
+        const contents = await fs.readFileBytes(keyPath)
         assert.notStrictEqual(contents.length, 0)
     })
 
     it('generates unique key each time', async function () {
-        const beforeContent = await fs.readFileBytes(vscode.Uri.file(keyPath))
+        const beforeContent = await fs.readFileBytes(keyPath)
         keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
-        const afterContent = await fs.readFileBytes(vscode.Uri.file(keyPath))
+        const afterContent = await fs.readFileBytes(keyPath)
         assert.notStrictEqual(beforeContent, afterContent)
     })
 
@@ -90,10 +89,10 @@ describe('SshKeyUtility', async function () {
 
     it('does overwrite existing keys on get call', async function () {
         const generateStub = sinon.spy(SshKeyPair, 'generateSshKeyPair')
-        const keyBefore = await fs.readFileBytes(vscode.Uri.file(keyPath))
+        const keyBefore = await fs.readFileBytes(keyPath)
         keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
 
-        const keyAfter = await fs.readFileBytes(vscode.Uri.file(keyPath))
+        const keyAfter = await fs.readFileBytes(keyPath)
         sinon.assert.calledOnce(generateStub)
 
         assert.notStrictEqual(keyBefore, keyAfter)

packages/core/src/test/testUtil.ts

@@ -295,6 +295,13 @@ export function partialDeepCompare<T>(actual: unknown, expected: T, message?: st
     const partial = selectFrom(actual, ...keys(expected as object))
     assert.deepStrictEqual(partial, expected, message)
 }
+/**
+ * Asserts that no metrics metadata (key OR value) matches the given regex.
+ * @param keyword target substring to search for
+ */
+export function assertNoTelemetryMatch(re: RegExp | string): void | never {
+    return assert.ok(globals.telemetry.logger.queryRegex(re).length === 0)
+}
 
 /**
  * Finds the emitted telemetry metrics with the given `name`, then checks if the metadata fields