refactor key class

Author: Hweinstock

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
     }

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

@@ -20,26 +20,27 @@ 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 () => {
             await this.delete()
         })
     }
 
-    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 {

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 })
         })
     })
 

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,44 +13,38 @@ 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 () {
         await keyPair.delete()
     })
 
     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,29 +66,25 @@ 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)
     })
 
     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)