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

## Problem
We want to enforce that the private key generated by the toolkit is not
accidentally slipped into a telemetry metric.

## Solution
Similar to `assertTelemetry`, we implement a `assertNoTelemetryMatch`
test utility that scans all metrics to see if a specified keyword is
included in its keys or values.

With this utility, it is relatively straightforward to assert that this
private key doesn't appear in the telemetry when we generate it and
perform other operations with it.

---

<!--- 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: Justin M. Keyes <[email protected]>

Author: Hweinstock

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