From d670ef22f19151732dec24bd66d85aa3422c17a1 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Fri, 18 Oct 2024 13:27:27 -0400
Subject: [PATCH 01/14] start of test cases

---
 packages/core/src/awsService/ec2/sshKeyPair.ts        | 11 ++++++++++-
 .../core/src/test/awsService/ec2/sshKeyPair.test.ts   |  8 ++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/packages/core/src/awsService/ec2/sshKeyPair.ts b/packages/core/src/awsService/ec2/sshKeyPair.ts
index 0e7ab3ea3c3..9b44fead621 100644
--- a/packages/core/src/awsService/ec2/sshKeyPair.ts
+++ b/packages/core/src/awsService/ec2/sshKeyPair.ts
@@ -2,12 +2,13 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { fs } from '../../shared'
+import { fs, globals } from '../../shared'
 import { ToolkitError } from '../../shared/errors'
 import { tryRun } from '../../shared/utilities/pathFind'
 import { Timeout } from '../../shared/utilities/timeoutUtils'
 import { findAsync } from '../../shared/utilities/collectionUtils'
 import { RunParameterContext } from '../../shared/utilities/processUtils'
+import path from 'path'
 
 type sshKeyType = 'rsa' | 'ed25519'
 
@@ -28,10 +29,18 @@ export class SshKeyPair {
     }
 
     public static async getSshKeyPair(keyPath: string, lifetime: number) {
+        if (!SshKeyPair.isValidKeyPath) {
+            throw new ToolkitError(`ec2: unable to generate key outside of global storage in path ${keyPath}`)
+        }
         await SshKeyPair.generateSshKeyPair(keyPath)
         return new SshKeyPair(keyPath, lifetime)
     }
 
+    private static isValidKeyPath(keyPath: string): boolean {
+        const relative = path.relative(globals.context.globalStorageUri.fsPath, keyPath)
+        return relative !== undefined && !relative.startsWith('..') && !path.isAbsolute(relative)
+    }
+
     public static async generateSshKeyPair(keyPath: string): Promise<void> {
         const keyGenerated = await this.tryKeyTypes(keyPath, ['ed25519', 'rsa'])
         if (!keyGenerated) {
diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index 92b1f123189..43600b5174a 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -134,6 +134,14 @@ describe('SshKeyUtility', async function () {
         assert(keyPair.isDeleted())
     })
 
+    it('does not allow writing keys to non-global storage', async function () {
+        assert.throws(async () => SshKeyPair.getSshKeyPair('~/.ssh/someKey', 2000))
+
+        assert.throws(async () => SshKeyPair.getSshKeyPair('/a/path/that/isnt/real/key', 2000))
+    })
+
+    it('checks path is still valid before deleting', async function () {})
+
     describe('isDeleted', async function () {
         it('returns false if key files exist', async function () {
             assert.strictEqual(await keyPair.isDeleted(), false)

From 05e6d52d526b1229e5de98947aac109b57000d79 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Fri, 18 Oct 2024 14:03:01 -0400
Subject: [PATCH 02/14] add protections

---
 packages/core/src/awsService/ec2/sshKeyPair.ts       |  3 ++-
 .../core/src/test/awsService/ec2/sshKeyPair.test.ts  | 12 +++++++++---
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/packages/core/src/awsService/ec2/sshKeyPair.ts b/packages/core/src/awsService/ec2/sshKeyPair.ts
index 9b44fead621..b31544eb2b0 100644
--- a/packages/core/src/awsService/ec2/sshKeyPair.ts
+++ b/packages/core/src/awsService/ec2/sshKeyPair.ts
@@ -29,7 +29,8 @@ export class SshKeyPair {
     }
 
     public static async getSshKeyPair(keyPath: string, lifetime: number) {
-        if (!SshKeyPair.isValidKeyPath) {
+        const validKey = SshKeyPair.isValidKeyPath(keyPath)
+        if (!validKey) {
             throw new ToolkitError(`ec2: unable to generate key outside of global storage in path ${keyPath}`)
         }
         await SshKeyPair.generateSshKeyPair(keyPath)
diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index 43600b5174a..6e375dbfda1 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -135,12 +135,18 @@ describe('SshKeyUtility', async function () {
     })
 
     it('does not allow writing keys to non-global storage', async function () {
-        assert.throws(async () => SshKeyPair.getSshKeyPair('~/.ssh/someKey', 2000))
+        await assert.rejects(async () => await SshKeyPair.getSshKeyPair('~/.ssh/someKey', 2000))
 
-        assert.throws(async () => SshKeyPair.getSshKeyPair('/a/path/that/isnt/real/key', 2000))
+        await assert.rejects(async () => await SshKeyPair.getSshKeyPair('/a/path/that/isnt/real/key', 2000))
     })
 
-    it('checks path is still valid before deleting', async function () {})
+    it('checks the key path again before deleting', async function () {
+        let keyPath = path.join(globals.context.globalStorageUri.fsPath)
+        const keys = await SshKeyPair.getSshKeyPair(keyPath, 5000)
+        keyPath = '~/.ssh/aws-toolkit-test-key'
+        assert.ok(keys.getPrivateKeyPath(), keyPath)
+        await assert.rejects(async () => await keys.delete())
+    })
 
     describe('isDeleted', async function () {
         it('returns false if key files exist', async function () {

From 12fc398755f6714b280837a83b2a25aecd834dec Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Fri, 18 Oct 2024 14:08:47 -0400
Subject: [PATCH 03/14] refactor into its own function

---
 packages/core/src/awsService/ec2/sshKeyPair.ts | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/packages/core/src/awsService/ec2/sshKeyPair.ts b/packages/core/src/awsService/ec2/sshKeyPair.ts
index b31544eb2b0..5acaa5994da 100644
--- a/packages/core/src/awsService/ec2/sshKeyPair.ts
+++ b/packages/core/src/awsService/ec2/sshKeyPair.ts
@@ -29,10 +29,10 @@ export class SshKeyPair {
     }
 
     public static async getSshKeyPair(keyPath: string, lifetime: number) {
-        const validKey = SshKeyPair.isValidKeyPath(keyPath)
-        if (!validKey) {
-            throw new ToolkitError(`ec2: unable to generate key outside of global storage in path ${keyPath}`)
-        }
+        SshKeyPair.assertValidKeypath(
+            keyPath,
+            `ec2: unable to generate key outside of global storage in path ${keyPath}`
+        )
         await SshKeyPair.generateSshKeyPair(keyPath)
         return new SshKeyPair(keyPath, lifetime)
     }
@@ -42,6 +42,12 @@ export class SshKeyPair {
         return relative !== undefined && !relative.startsWith('..') && !path.isAbsolute(relative)
     }
 
+    private static assertValidKeypath(keyPath: string, message: string): void | never {
+        if (!SshKeyPair.isValidKeyPath(keyPath)) {
+            throw new ToolkitError(message)
+        }
+    }
+
     public static async generateSshKeyPair(keyPath: string): Promise<void> {
         const keyGenerated = await this.tryKeyTypes(keyPath, ['ed25519', 'rsa'])
         if (!keyGenerated) {
@@ -82,6 +88,10 @@ export class SshKeyPair {
     }
 
     public async delete(): Promise<void> {
+        SshKeyPair.assertValidKeypath(
+            this.keyPath,
+            `ec2: keyPath became invalid after creation, not deleting key at ${this.keyPath}`
+        )
         await fs.delete(this.keyPath)
         await fs.delete(this.publicKeyPath)
 

From e9d011a2eef2b9c3926c94f635e3c1f7653b2bf8 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Fri, 18 Oct 2024 15:27:05 -0400
Subject: [PATCH 04/14] manually create temporary directory inside global
 storage

---
 .../src/test/awsService/ec2/sshKeyPair.test.ts    | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index 6e375dbfda1..06e45176221 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -9,7 +9,7 @@ import * as sinon from 'sinon'
 import * as path from 'path'
 import * as os from 'os'
 import { SshKeyPair } from '../../../awsService/ec2/sshKeyPair'
-import { createTestWorkspaceFolder, installFakeClock } from '../../testUtil'
+import { installFakeClock } from '../../testUtil'
 import { InstalledClock } from '@sinonjs/fake-timers'
 import { ChildProcess } from '../../../shared/utilities/processUtils'
 import { fs, globals } from '../../../shared'
@@ -21,7 +21,10 @@ describe('SshKeyUtility', async function () {
     let clock: InstalledClock
 
     before(async function () {
-        temporaryDirectory = (await createTestWorkspaceFolder()).uri.fsPath
+        // Setup a temporary directory inside of globalStorage since keys need to be in dir controlled by toolkit.
+        temporaryDirectory = path.join(globals.context.globalStorageUri.fsPath, 'SshKeyUtilityTests')
+        await fs.mkdir(temporaryDirectory)
+
         keyPath = path.join(temporaryDirectory, 'testKeyPair')
         clock = installFakeClock()
     })
@@ -140,14 +143,6 @@ describe('SshKeyUtility', async function () {
         await assert.rejects(async () => await SshKeyPair.getSshKeyPair('/a/path/that/isnt/real/key', 2000))
     })
 
-    it('checks the key path again before deleting', async function () {
-        let keyPath = path.join(globals.context.globalStorageUri.fsPath)
-        const keys = await SshKeyPair.getSshKeyPair(keyPath, 5000)
-        keyPath = '~/.ssh/aws-toolkit-test-key'
-        assert.ok(keys.getPrivateKeyPath(), keyPath)
-        await assert.rejects(async () => await keys.delete())
-    })
-
     describe('isDeleted', async function () {
         it('returns false if key files exist', async function () {
             assert.strictEqual(await keyPair.isDeleted(), false)

From 0e966cb37b18913cfb06c143c54f929d19d42ce9 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Fri, 18 Oct 2024 15:30:36 -0400
Subject: [PATCH 05/14] adjust other test to avoid mocks

---
 .../core/src/test/awsService/ec2/model.test.ts     | 14 +++++++++-----
 .../src/test/awsService/ec2/sshKeyPair.test.ts     |  2 +-
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/packages/core/src/test/awsService/ec2/model.test.ts b/packages/core/src/test/awsService/ec2/model.test.ts
index 36a7895a35f..40c3c174a34 100644
--- a/packages/core/src/test/awsService/ec2/model.test.ts
+++ b/packages/core/src/test/awsService/ec2/model.test.ts
@@ -13,6 +13,8 @@ import { ToolkitError } from '../../../shared/errors'
 import { IAM } from 'aws-sdk'
 import { SshKeyPair } from '../../../awsService/ec2/sshKeyPair'
 import { DefaultIamClient } from '../../../shared/clients/iamClient'
+import path from 'path'
+import { fs, globals } from '../../../shared'
 
 describe('Ec2ConnectClient', function () {
     let client: Ec2ConnectionManager
@@ -132,17 +134,19 @@ describe('Ec2ConnectClient', function () {
         it('calls the sdk with the proper parameters', async function () {
             const sendCommandStub = sinon.stub(SsmClient.prototype, 'sendCommandAndWait')
 
-            sinon.stub(SshKeyPair, 'generateSshKeyPair')
-            sinon.stub(SshKeyPair.prototype, 'getPublicKey').resolves('test-key')
-
             const testSelection = {
                 instanceId: 'test-id',
                 region: 'test-region',
             }
-            const mockKeys = await SshKeyPair.getSshKeyPair('fakeDir', 30000)
-            await client.sendSshKeyToInstance(testSelection, mockKeys, 'test-user')
+            const temporaryDirectory = path.join(globals.context.globalStorageUri.fsPath, 'ModelTests')
+            await fs.mkdir(temporaryDirectory)
+
+            const keys = await SshKeyPair.getSshKeyPair(path.join(temporaryDirectory, 'key'), 30000)
+            await client.sendSshKeyToInstance(testSelection, keys, 'test-user')
             sinon.assert.calledWith(sendCommandStub, testSelection.instanceId, 'AWS-RunShellScript')
             sinon.restore()
+
+            await fs.delete(temporaryDirectory, { force: true, recursive: true })
         })
     })
 
diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index 06e45176221..d2bd1a143c2 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -21,7 +21,7 @@ describe('SshKeyUtility', async function () {
     let clock: InstalledClock
 
     before(async function () {
-        // Setup a temporary directory inside of globalStorage since keys need to be in dir controlled by toolkit.
+        // Setup a temporary directory inside of globalStorage since keys need to be in globalStorage subdir
         temporaryDirectory = path.join(globals.context.globalStorageUri.fsPath, 'SshKeyUtilityTests')
         await fs.mkdir(temporaryDirectory)
 

From dc2b75673e609a8ed616c35aab9ac729ba6d6fc6 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Fri, 18 Oct 2024 16:27:50 -0400
Subject: [PATCH 06/14] adjust comment

---
 packages/core/src/test/awsService/ec2/sshKeyPair.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index d2bd1a143c2..f53a4875ffe 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -21,7 +21,7 @@ describe('SshKeyUtility', async function () {
     let clock: InstalledClock
 
     before(async function () {
-        // Setup a temporary directory inside of globalStorage since keys need to be in globalStorage subdir
+        // Setup a temporary directory inside of globalStorage since keys need to be inside globalStorage
         temporaryDirectory = path.join(globals.context.globalStorageUri.fsPath, 'SshKeyUtilityTests')
         await fs.mkdir(temporaryDirectory)
 

From b1cda5fe5821adcb55221405ea774947cce1d3a4 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 21 Oct 2024 11:44:21 -0400
Subject: [PATCH 07/14] refactor key class

---
 packages/core/src/awsService/ec2/model.ts     |  5 +--
 .../core/src/awsService/ec2/sshKeyPair.ts     | 15 ++++---
 .../src/test/awsService/ec2/model.test.ts     |  8 +---
 .../test/awsService/ec2/sshKeyPair.test.ts    | 44 +++++++------------
 4 files changed, 26 insertions(+), 46 deletions(-)

diff --git a/packages/core/src/awsService/ec2/model.ts b/packages/core/src/awsService/ec2/model.ts
index 021ad623dd2..f47386ef47d 100644
--- a/packages/core/src/awsService/ec2/model.ts
+++ b/packages/core/src/awsService/ec2/model.ts
@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import * as vscode from 'vscode'
-import * as path from 'path'
 import { Session } from 'aws-sdk/clients/ssm'
 import { IAM, SSM } from 'aws-sdk'
 import { Ec2Selection } from './prompter'
@@ -28,7 +27,6 @@ import { CancellationError, Timeout } from '../../shared/utilities/timeoutUtils'
 import { showMessageWithCancel } from '../../shared/utilities/messages'
 import { SshConfig, sshLogFileLocation } from '../../shared/sshConfig'
 import { SshKeyPair } from './sshKeyPair'
-import globals from '../../shared/extensionGlobals'
 
 export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect' | 'EC2SSMAgentStatus'
 
@@ -235,8 +233,7 @@ export class Ec2ConnectionManager {
     }
 
     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, 30000)
+        const keyPair = await SshKeyPair.getSshKeyPair(`aws-ec2-key`, 30000)
         await this.sendSshKeyToInstance(selection, keyPair, remoteUser)
         return keyPair
     }
diff --git a/packages/core/src/awsService/ec2/sshKeyPair.ts b/packages/core/src/awsService/ec2/sshKeyPair.ts
index 5acaa5994da..9fd2267f9c1 100644
--- a/packages/core/src/awsService/ec2/sshKeyPair.ts
+++ b/packages/core/src/awsService/ec2/sshKeyPair.ts
@@ -20,7 +20,7 @@ export class SshKeyPair {
         private readonly keyPath: string,
         lifetime: number
     ) {
-        this.publicKeyPath = `${keyPath}.pub`
+        this.publicKeyPath = `${this.keyPath}.pub`
         this.lifeTimeout = new Timeout(lifetime)
 
         this.lifeTimeout.onCompletion(async () => {
@@ -28,18 +28,19 @@ export class SshKeyPair {
         })
     }
 
-    public static async getSshKeyPair(keyPath: string, lifetime: number) {
-        SshKeyPair.assertValidKeypath(
-            keyPath,
-            `ec2: unable to generate key outside of global storage in path ${keyPath}`
-        )
+    private static getKeypath(keyName: string): string {
+        return path.join(globals.context.globalStorageUri.fsPath, keyName)
+    }
+
+    public static async getSshKeyPair(keyName: string, lifetime: number) {
+        const keyPath = SshKeyPair.getKeypath(keyName)
         await SshKeyPair.generateSshKeyPair(keyPath)
         return new SshKeyPair(keyPath, lifetime)
     }
 
     private static isValidKeyPath(keyPath: string): boolean {
         const relative = path.relative(globals.context.globalStorageUri.fsPath, keyPath)
-        return relative !== undefined && !relative.startsWith('..') && !path.isAbsolute(relative)
+        return relative !== undefined && !relative.startsWith('..') && !path.isAbsolute(relative) && keyPath.length > 4
     }
 
     private static assertValidKeypath(keyPath: string, message: string): void | never {
diff --git a/packages/core/src/test/awsService/ec2/model.test.ts b/packages/core/src/test/awsService/ec2/model.test.ts
index 40c3c174a34..045c961003a 100644
--- a/packages/core/src/test/awsService/ec2/model.test.ts
+++ b/packages/core/src/test/awsService/ec2/model.test.ts
@@ -13,8 +13,6 @@ import { ToolkitError } from '../../../shared/errors'
 import { IAM } from 'aws-sdk'
 import { SshKeyPair } from '../../../awsService/ec2/sshKeyPair'
 import { DefaultIamClient } from '../../../shared/clients/iamClient'
-import path from 'path'
-import { fs, globals } from '../../../shared'
 
 describe('Ec2ConnectClient', function () {
     let client: Ec2ConnectionManager
@@ -138,15 +136,11 @@ describe('Ec2ConnectClient', function () {
                 instanceId: 'test-id',
                 region: 'test-region',
             }
-            const temporaryDirectory = path.join(globals.context.globalStorageUri.fsPath, 'ModelTests')
-            await fs.mkdir(temporaryDirectory)
 
-            const keys = await SshKeyPair.getSshKeyPair(path.join(temporaryDirectory, 'key'), 30000)
+            const keys = await SshKeyPair.getSshKeyPair('key', 30000)
             await client.sendSshKeyToInstance(testSelection, keys, 'test-user')
             sinon.assert.calledWith(sendCommandStub, testSelection.instanceId, 'AWS-RunShellScript')
             sinon.restore()
-
-            await fs.delete(temporaryDirectory, { force: true, recursive: true })
         })
     })
 
diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index f53a4875ffe..dfa6d531e77 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -2,11 +2,9 @@
  * 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'
-import * as path from 'path'
 import * as os from 'os'
 import { SshKeyPair } from '../../../awsService/ec2/sshKeyPair'
 import { installFakeClock } from '../../testUtil'
@@ -15,22 +13,17 @@ import { ChildProcess } from '../../../shared/utilities/processUtils'
 import { fs, globals } from '../../../shared'
 
 describe('SshKeyUtility', async function () {
-    let temporaryDirectory: string
-    let keyPath: string
-    let keyPair: SshKeyPair
     let clock: InstalledClock
+    let keyPair: SshKeyPair
+    let keyName: string
 
     before(async function () {
-        // Setup a temporary directory inside of globalStorage since keys need to be inside globalStorage
-        temporaryDirectory = path.join(globals.context.globalStorageUri.fsPath, 'SshKeyUtilityTests')
-        await fs.mkdir(temporaryDirectory)
-
-        keyPath = path.join(temporaryDirectory, 'testKeyPair')
         clock = installFakeClock()
     })
 
     beforeEach(async function () {
-        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
+        keyName = 'testKeyPair'
+        keyPair = await SshKeyPair.getSshKeyPair(keyName, 30000)
     })
 
     afterEach(async function () {
@@ -38,21 +31,20 @@ describe('SshKeyUtility', async function () {
     })
 
     after(async function () {
-        await fs.delete(temporaryDirectory, { recursive: true })
         clock.uninstall()
         sinon.restore()
     })
 
     it('generates key in target file', async function () {
-        const contents = await fs.readFileBytes(vscode.Uri.file(keyPath))
+        const contents = await fs.readFileBytes(keyPair.getPrivateKeyPath())
         assert.notStrictEqual(contents.length, 0)
     })
 
     it('generates unique key each time', async function () {
-        const beforeContent = await fs.readFileBytes(vscode.Uri.file(keyPath))
-        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
-        const afterContent = await fs.readFileBytes(vscode.Uri.file(keyPath))
-        assert.notStrictEqual(beforeContent, afterContent)
+        const keyPair2 = await SshKeyPair.getSshKeyPair(`${keyName}2`, 30000)
+        const content1 = await fs.readFileBytes(keyPair2.getPrivateKeyPath())
+        const content2 = await fs.readFileBytes(keyPair.getPrivateKeyPath())
+        assert.notStrictEqual(content1, content2)
     })
 
     it('sets permission of the file to read/write owner', async function () {
@@ -63,7 +55,7 @@ describe('SshKeyUtility', async function () {
     })
 
     it('defaults to ed25519 key type', async function () {
-        const process = new ChildProcess(`ssh-keygen`, ['-vvv', '-l', '-f', keyPath])
+        const process = new ChildProcess(`ssh-keygen`, ['-vvv', '-l', '-f', keyPair.getPrivateKeyPath()])
         const result = await process.run()
         // Check private key header for algorithm name
         assert.strictEqual(result.stdout.includes('[ED25519 256]'), true)
@@ -74,18 +66,14 @@ describe('SshKeyUtility', async function () {
         const stub = sinon.stub(SshKeyPair, 'tryKeyGen')
         stub.onFirstCall().resolves(false)
         stub.callThrough()
-        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
-        const process = new ChildProcess(`ssh-keygen`, ['-vvv', '-l', '-f', keyPath])
+        const rsaKey = await SshKeyPair.getSshKeyPair('rsa', 30000)
+        const process = new ChildProcess(`ssh-keygen`, ['-vvv', '-l', '-f', rsaKey.getPrivateKeyPath()])
         const result = await process.run()
         // Check private key header for algorithm name
         assert.strictEqual(result.stdout.includes('[RSA'), true)
         stub.restore()
     })
 
-    it('properly names the public key', function () {
-        assert.strictEqual(keyPair.getPublicKeyPath(), `${keyPath}.pub`)
-    })
-
     it('reads in public ssh key that is non-empty', async function () {
         const key = await keyPair.getPublicKey()
         assert.notStrictEqual(key.length, 0)
@@ -93,10 +81,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))
-        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 30000)
+        const keyBefore = await fs.readFileBytes(keyPair.getPrivateKeyPath())
+        keyPair = await SshKeyPair.getSshKeyPair(keyName, 30000)
 
-        const keyAfter = await fs.readFileBytes(vscode.Uri.file(keyPath))
+        const keyAfter = await fs.readFileBytes(keyPair.getPrivateKeyPath())
         sinon.assert.calledOnce(generateStub)
 
         assert.notStrictEqual(keyBefore, keyAfter)
@@ -122,7 +110,7 @@ describe('SshKeyUtility', async function () {
         sinon.stub(SshKeyPair, 'generateSshKeyPair')
         const deleteStub = sinon.stub(SshKeyPair.prototype, 'delete')
 
-        keyPair = await SshKeyPair.getSshKeyPair(keyPath, 50)
+        keyPair = await SshKeyPair.getSshKeyPair(keyName, 50)
         await clock.tickAsync(10)
         sinon.assert.notCalled(deleteStub)
         await clock.tickAsync(100)

From 593538a9816386ad975b581fc372b7b8565b4ac8 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 21 Oct 2024 15:46:35 -0400
Subject: [PATCH 08/14] clear dir inbetween each test

---
 packages/core/src/test/globalSetup.test.ts | 1 +
 1 file changed, 1 insertion(+)

diff --git a/packages/core/src/test/globalSetup.test.ts b/packages/core/src/test/globalSetup.test.ts
index c7bc947c482..a748c039a62 100644
--- a/packages/core/src/test/globalSetup.test.ts
+++ b/packages/core/src/test/globalSetup.test.ts
@@ -103,6 +103,7 @@ export const mochaHooks = {
         TelemetryDebounceInfo.instance.clear()
         // mochaGlobalSetup() set this to a fake, so it's safe to clear it here.
         await globals.globalState.clear()
+        await fs.delete(globals.context.globalStorageUri.fsPath, { recursive: true, force: true })
 
         await testUtil.closeAllEditors()
     },

From 44aa977c93c4a326cd50e0ca7eaaa3cc0701de63 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 21 Oct 2024 15:55:49 -0400
Subject: [PATCH 09/14] remove delete inbetween tests

---
 packages/core/src/test/globalSetup.test.ts | 1 -
 1 file changed, 1 deletion(-)

diff --git a/packages/core/src/test/globalSetup.test.ts b/packages/core/src/test/globalSetup.test.ts
index a748c039a62..c7bc947c482 100644
--- a/packages/core/src/test/globalSetup.test.ts
+++ b/packages/core/src/test/globalSetup.test.ts
@@ -103,7 +103,6 @@ export const mochaHooks = {
         TelemetryDebounceInfo.instance.clear()
         // mochaGlobalSetup() set this to a fake, so it's safe to clear it here.
         await globals.globalState.clear()
-        await fs.delete(globals.context.globalStorageUri.fsPath, { recursive: true, force: true })
 
         await testUtil.closeAllEditors()
     },

From 8f23d8aeb569088f88419d2203b3f449fd8f9959 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 21 Oct 2024 16:37:50 -0400
Subject: [PATCH 10/14] remove temp dir each test

---
 packages/core/src/test/globalSetup.test.ts | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/packages/core/src/test/globalSetup.test.ts b/packages/core/src/test/globalSetup.test.ts
index c7bc947c482..5c71d043ecc 100644
--- a/packages/core/src/test/globalSetup.test.ts
+++ b/packages/core/src/test/globalSetup.test.ts
@@ -101,6 +101,8 @@ export const mochaHooks = {
         globals.telemetry.clearRecords()
         globals.telemetry.logger.clear()
         TelemetryDebounceInfo.instance.clear()
+        await fs.delete(globals.context.globalStorageUri.fsPath, { recursive: true, force: true })
+
         // mochaGlobalSetup() set this to a fake, so it's safe to clear it here.
         await globals.globalState.clear()
 

From 9762798029bfd13d19b4203f58b0e031cd155b4c Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 21 Oct 2024 17:02:40 -0400
Subject: [PATCH 11/14] adjust codecatalyst to use global context in tests

---
 .../core/src/test/shared/sshConfig.test.ts    | 20 +++++++------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/packages/core/src/test/shared/sshConfig.test.ts b/packages/core/src/test/shared/sshConfig.test.ts
index 2400e7f0e44..054833293b6 100644
--- a/packages/core/src/test/shared/sshConfig.test.ts
+++ b/packages/core/src/test/shared/sshConfig.test.ts
@@ -22,6 +22,7 @@ import {
 import { StartDevEnvironmentSessionRequest } from 'aws-sdk/clients/codecatalyst'
 import { mkdir, readFile, writeFile } from 'fs/promises'
 import fs from '../../shared/fs/fs'
+import { globals } from '../../shared'
 
 class MockSshConfig extends SshConfig {
     // State variables to track logic flow.
@@ -181,22 +182,15 @@ describe('VscodeRemoteSshConfig', async function () {
 })
 
 describe('CodeCatalyst Connect Script', function () {
-    let context: FakeExtensionContext
-
     function isWithin(path1: string, path2: string): boolean {
         const rel = path.relative(path1, path2)
         return !path.isAbsolute(rel) && !rel.startsWith('..') && !!rel
     }
 
-    beforeEach(async function () {
-        context = await FakeExtensionContext.create()
-        context.globalStorageUri = vscode.Uri.file(await makeTemporaryToolkitFolder())
-    })
-
     it('can get a connect script path, adding a copy to global storage', async function () {
-        const script = (await ensureConnectScript(connectScriptPrefix, context)).unwrap().fsPath
+        const script = (await ensureConnectScript(connectScriptPrefix, globals.context)).unwrap().fsPath
         assert.ok(await fileExists(script))
-        assert.ok(isWithin(context.globalStorageUri.fsPath, script))
+        assert.ok(isWithin(globals.context.globalStorageUri.fsPath, script))
     })
 
     function createFakeServer(testDevEnv: DevEnvironmentId) {
@@ -247,8 +241,8 @@ describe('CodeCatalyst Connect Script', function () {
             server.listen({ host: 'localhost', port: 28142 }, () => resolve(`http://localhost:28142`))
         })
 
-        await writeFile(bearerTokenCacheLocation(testDevEnv.id), 'token')
-        const script = (await ensureConnectScript(connectScriptPrefix, context)).unwrap().fsPath
+        await fs.writeFile(bearerTokenCacheLocation(testDevEnv.id), 'token')
+        const script = (await ensureConnectScript(connectScriptPrefix, globals.context)).unwrap().fsPath
         const env = getCodeCatalystSsmEnv('us-weast-1', 'echo', testDevEnv)
         env.CODECATALYST_ENDPOINT = address
 
@@ -280,12 +274,12 @@ describe('CodeCatalyst Connect Script', function () {
         })
 
         it('works if the .ssh directory is missing', async function () {
-            ;(await ensureConnectScript(connectScriptPrefix, context)).unwrap()
+            ;(await ensureConnectScript(connectScriptPrefix, globals.context)).unwrap()
         })
 
         it('works if the .ssh directory exists but has different perms', async function () {
             await mkdir(path.join(tmpDir, '.ssh'), 0o777)
-            ;(await ensureConnectScript(connectScriptPrefix, context)).unwrap()
+            ;(await ensureConnectScript(connectScriptPrefix, globals.context)).unwrap()
         })
     })
 })

From 539f0e7d64f3561c587b99593154ddc91b450793 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 21 Oct 2024 17:23:40 -0400
Subject: [PATCH 12/14] remove dead imports

---
 packages/core/src/test/shared/sshConfig.test.ts | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/packages/core/src/test/shared/sshConfig.test.ts b/packages/core/src/test/shared/sshConfig.test.ts
index 054833293b6..96ca450ae14 100644
--- a/packages/core/src/test/shared/sshConfig.test.ts
+++ b/packages/core/src/test/shared/sshConfig.test.ts
@@ -5,13 +5,11 @@
 import assert from 'assert'
 import * as sinon from 'sinon'
 import * as path from 'path'
-import * as vscode from 'vscode'
 import * as http from 'http'
 import { ToolkitError } from '../../shared/errors'
 import { Result } from '../../shared/utilities/result'
 import { ChildProcess, ChildProcessResult } from '../../shared/utilities/processUtils'
 import { SshConfig, ensureConnectScript, sshLogFileLocation } from '../../shared/sshConfig'
-import { FakeExtensionContext } from '../fakeExtensionContext'
 import { fileExists, makeTemporaryToolkitFolder } from '../../shared/filesystemUtilities'
 import {
     DevEnvironmentId,
@@ -20,7 +18,7 @@ import {
     getCodeCatalystSsmEnv,
 } from '../../codecatalyst/model'
 import { StartDevEnvironmentSessionRequest } from 'aws-sdk/clients/codecatalyst'
-import { mkdir, readFile, writeFile } from 'fs/promises'
+import { mkdir, readFile } from 'fs/promises'
 import fs from '../../shared/fs/fs'
 import { globals } from '../../shared'
 

From fdaa72ccc5254785e0c35acaadb7752dcc1887ac Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Mon, 21 Oct 2024 17:46:36 -0400
Subject: [PATCH 13/14] clear global storage instead of deleting

---
 packages/core/src/test/awsService/ec2/sshKeyPair.test.ts | 1 +
 packages/core/src/test/globalSetup.test.ts               | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index 3679a7fdfd0..7039f1e2020 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -45,6 +45,7 @@ describe('SshKeyPair', async function () {
         const content1 = await fs.readFileBytes(keyPair2.getPrivateKeyPath())
         const content2 = await fs.readFileBytes(keyPair.getPrivateKeyPath())
         assert.notStrictEqual(content1, content2)
+        await keyPair2.delete()
     })
 
     it('sets permission of the file to read/write owner', async function () {
diff --git a/packages/core/src/test/globalSetup.test.ts b/packages/core/src/test/globalSetup.test.ts
index 5c71d043ecc..92e2db349ee 100644
--- a/packages/core/src/test/globalSetup.test.ts
+++ b/packages/core/src/test/globalSetup.test.ts
@@ -101,12 +101,13 @@ export const mochaHooks = {
         globals.telemetry.clearRecords()
         globals.telemetry.logger.clear()
         TelemetryDebounceInfo.instance.clear()
-        await fs.delete(globals.context.globalStorageUri.fsPath, { recursive: true, force: true })
 
         // mochaGlobalSetup() set this to a fake, so it's safe to clear it here.
         await globals.globalState.clear()
 
         await testUtil.closeAllEditors()
+        await fs.delete(globals.context.globalStorageUri.fsPath, { recursive: true, force: true })
+        await fs.mkdir(globals.context.globalStorageUri.fsPath)
     },
     async afterEach(this: Mocha.Context) {
         if (openExternalStub.called && openExternalStub.returned(sinon.match.typeOf('undefined'))) {

From 075b050b23c5461452c9a1256ae9f1eeb56c4ff1 Mon Sep 17 00:00:00 2001
From: hkobew <hkobew@amazon.com>
Date: Tue, 22 Oct 2024 10:58:46 -0400
Subject: [PATCH 14/14] delete no longer needed test

---
 packages/core/src/test/awsService/ec2/sshKeyPair.test.ts | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
index 7039f1e2020..144f436e319 100644
--- a/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
+++ b/packages/core/src/test/awsService/ec2/sshKeyPair.test.ts
@@ -126,12 +126,6 @@ describe('SshKeyPair', async function () {
         assert(keyPair.isDeleted())
     })
 
-    it('does not allow writing keys to non-global storage', async function () {
-        await assert.rejects(async () => await SshKeyPair.getSshKeyPair('~/.ssh/someKey', 2000))
-
-        await assert.rejects(async () => await SshKeyPair.getSshKeyPair('/a/path/that/isnt/real/key', 2000))
-    })
-
     describe('isDeleted', async function () {
         it('returns false if key files exist', async function () {
             assert.strictEqual(await keyPair.isDeleted(), false)