fix(ec2): check if instance connectable before connecting. (#3613)

Hweinstock · web-flow · commit 0fe859487da0 · 2023-07-03T13:22:36.000-07:00
Problem:
If the user changes the IAM role associated with the instance, the connect command fails whether it be adding or removing the necessary policies.

The root of this problem is ultimately a bug in the usage of the startSession SDK. That is, there are ways this command can fail such that it doesn't throw an error, but the terminal will crash on entry.

Solution:
We check for ways the startSession command can fail before running it. If any are detected, we throw a relevant error. This way, we only call this command when we are confident it will succeed.
diff --git a/src/ec2/model.ts b/src/ec2/model.ts
@@ -4,18 +4,19 @@
  */
 import * as vscode from 'vscode'
 import { Session } from 'aws-sdk/clients/ssm'
-import { AWSError, IAM } from 'aws-sdk'
+import { IAM } from 'aws-sdk'
 import { Ec2Selection } from './utils'
 import { getOrInstallCli } from '../shared/utilities/cliUtils'
 import { isCloud9 } from '../shared/extensionUtilities'
-import { ToolkitError, isAwsError } from '../shared/errors'
+import { ToolkitError } from '../shared/errors'
 import { SsmClient } from '../shared/clients/ssmClient'
 import { Ec2Client } from '../shared/clients/ec2Client'
 
 export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect'
 
 import { openRemoteTerminal } from '../shared/remoteSession'
 import { DefaultIamClient } from '../shared/clients/iamClient'
+import { ErrorInformation } from '../shared/errors'
 
 export class Ec2ConnectionManager {
     private ssmClient: SsmClient
@@ -41,56 +42,51 @@ export class Ec2ConnectionManager {
     }
 
     protected async getAttachedPolicies(instanceId: string): Promise<IAM.attachedPoliciesListType> {
-        try {
-            const IamRole = await this.ec2Client.getAttachedIamRole(instanceId)
-            const iamResponse = await this.iamClient.listAttachedRolePolicies(IamRole!.Arn!)
-            return iamResponse.AttachedPolicies!
-        } catch (err) {
+        const IamRole = await this.ec2Client.getAttachedIamRole(instanceId)
+        if (!IamRole) {
             return []
         }
+        const iamResponse = await this.iamClient.listAttachedRolePolicies(IamRole!.Arn!)
+
+        return iamResponse.AttachedPolicies ?? []
     }
 
     public async hasProperPolicies(instanceId: string): Promise<boolean> {
         const attachedPolicies = (await this.getAttachedPolicies(instanceId)).map(policy => policy.PolicyName!)
         const requiredPolicies = ['AmazonSSMManagedInstanceCore', 'AmazonSSMManagedEC2InstanceDefaultPolicy']
 
-        return requiredPolicies.every(policy => attachedPolicies.includes(policy))
+        return requiredPolicies.length !== 0 && requiredPolicies.every(policy => attachedPolicies.includes(policy))
     }
 
-    public async handleStartSessionError(err: AWSError, selection: Ec2Selection): Promise<Error> {
-        const isInstanceRunning = (await this.ec2Client.getInstanceStatus(selection.instanceId)) == 'running'
+    public async isInstanceRunning(instanceId: string): Promise<boolean> {
+        const instanceStatus = await this.ec2Client.getInstanceStatus(instanceId)
+        return instanceStatus == 'running'
+    }
+
+    private throwConnectionError(message: string, selection: Ec2Selection, errorInfo: ErrorInformation) {
         const generalErrorMessage = `Unable to connect to target instance ${selection.instanceId} on region ${selection.region}. `
+        throw new ToolkitError(generalErrorMessage + message, errorInfo)
+    }
+
+    public async checkForStartSessionError(selection: Ec2Selection): Promise<void> {
+        const isInstanceRunning = await this.isInstanceRunning(selection.instanceId)
         const hasProperPolicies = await this.hasProperPolicies(selection.instanceId)
 
         if (!isInstanceRunning) {
-            throw new ToolkitError(
-                generalErrorMessage +
-                    'Ensure the target instance is running and not currently starting, stopping, or stopped.',
-                { code: 'EC2SSMStatus' }
-            )
+            const message = 'Ensure the target instance is running and not currently starting, stopping, or stopped.'
+            this.throwConnectionError(message, selection, { code: 'EC2SSMStatus' })
         }
 
         if (!hasProperPolicies) {
-            throw new ToolkitError(
-                generalErrorMessage + 'Ensure the IAM role attached to the instance has the required policies.',
-                {
-                    code: 'EC2SSMPermission',
-                    documentationUri: vscode.Uri.parse(
-                        'https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-getting-started-instance-profile.html'
-                    ),
-                }
+            const message = 'Ensure the IAM role attached to the instance has the required policies.'
+            const documentationUri = vscode.Uri.parse(
+                'https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-getting-started-instance-profile.html'
             )
+            this.throwConnectionError(message, selection, {
+                code: 'EC2SSMPermission',
+                documentationUri: documentationUri,
+            })
         }
-
-        throw new ToolkitError(
-            'Ensure SSM is running on target instance. For more information see the documentation.',
-            {
-                code: 'EC2SSMConnect',
-                documentationUri: vscode.Uri.parse(
-                    'https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-getting-started.html'
-                ),
-            }
-        )
     }
 
     private async openSessionInTerminal(session: Session, selection: Ec2Selection) {
@@ -108,15 +104,16 @@ export class Ec2ConnectionManager {
     }
 
     public async attemptEc2Connection(selection: Ec2Selection): Promise<void> {
+        await this.checkForStartSessionError(selection)
         try {
             const response = await this.ssmClient.startSession(selection.instanceId)
             await this.openSessionInTerminal(response, selection)
-        } catch (err) {
-            if (isAwsError(err)) {
-                await this.handleStartSessionError(err, selection)
-            } else {
-                throw err
-            }
+        } catch (err: unknown) {
+            // Default error if pre-check fails.
+            this.throwConnectionError('Check that the SSM Agent is running on target instance', selection, {
+                code: 'EC2SSMConnect',
+                cause: err as Error,
+            })
         }
     }
 }
diff --git a/src/shared/clients/ec2Client.ts b/src/shared/clients/ec2Client.ts
@@ -19,28 +19,31 @@ export class Ec2Client {
     private async createSdkClient(): Promise<EC2> {
         return await globals.sdkClientBuilder.createAwsService(EC2, undefined, this.regionCode)
     }
-    public async getInstances(): Promise<AsyncCollection<Ec2Instance>> {
+
+    public async getInstances(filters?: EC2.Filter[]): Promise<AsyncCollection<EC2.Instance>> {
         const client = await this.createSdkClient()
+
         const requester = async (request: EC2.DescribeInstancesRequest) => client.describeInstances(request).promise()
-        const collection = pageableToCollection(requester, {}, 'NextToken', 'Reservations')
+        const collection = pageableToCollection(
+            requester,
+            filters ? { Filters: filters } : {},
+            'NextToken',
+            'Reservations'
+        )
         const instances = this.getInstancesFromReservations(collection)
         return instances
     }
 
-    private lookupTagKey(tags: EC2.Tag[], targetKey: string): string | undefined {
-        return tags.filter(tag => tag.Key == targetKey)[0].Value
-    }
-
     public getInstancesFromReservations(
         reservations: AsyncCollection<EC2.ReservationList | undefined>
-    ): AsyncCollection<Ec2Instance> {
+    ): AsyncCollection<EC2.Instance> {
         return reservations
             .flatten()
             .map(instanceList => instanceList?.Instances)
             .flatten()
             .filter(instance => instance!.InstanceId !== undefined)
             .map(instance => {
-                return instance!.Tags ? { ...instance, name: this.lookupTagKey(instance!.Tags!, 'Name') } : instance!
+                return instance!.Tags ? { ...instance, name: lookupTagKey(instance!.Tags!, 'Name') } : instance!
             })
     }
 
@@ -62,19 +65,23 @@ export class Ec2Client {
         return response[0]
     }
 
-    /**
-     * Retrieve IAM role attached to given EC2 instance.
-     * @param instanceId target EC2 instance ID
-     * @returns IAM role associated with instance, or undefined if none exists.
-     */
-    public async getAttachedIamRole(instanceId: string): Promise<IamInstanceProfile | undefined> {
-        const client = await this.createSdkClient()
-        const instanceFilter: EC2.Filter[] = [
+    public getInstancesFilter(instanceIds: string[]): EC2.Filter[] {
+        return [
             {
                 Name: 'instance-id',
-                Values: [instanceId],
+                Values: instanceIds,
             },
         ]
+    }
+
+    /**
+     * Retrieve IAM Association for a given EC2 instance.
+     * @param instanceId target EC2 instance ID
+     * @returns IAM Association for instance
+     */
+    private async getIamInstanceProfileAssociation(instanceId: string): Promise<EC2.IamInstanceProfileAssociation> {
+        const client = await this.createSdkClient()
+        const instanceFilter = this.getInstancesFilter([instanceId])
         const requester = async (request: EC2.DescribeIamInstanceProfileAssociationsRequest) =>
             client.describeIamInstanceProfileAssociations(request).promise()
         const response = await pageableToCollection(
@@ -84,9 +91,27 @@ export class Ec2Client {
             'IamInstanceProfileAssociations'
         )
             .flatten()
-            .map(val => val?.IamInstanceProfile)
+            .filter(association => association !== undefined)
             .promise()
 
-        return response && response.length ? response[0] : undefined
+        return response[0]!
     }
+
+    /**
+     * Retrieve IAM role attached to given EC2 instance.
+     * @param instanceId target EC2 instance ID
+     * @returns IAM role associated with instance or undefined if none exists.
+     */
+    public async getAttachedIamRole(instanceId: string): Promise<IamInstanceProfile | undefined> {
+        const association = await this.getIamInstanceProfileAssociation(instanceId)
+        return association ? association.IamInstanceProfile : undefined
+    }
+}
+
+export function getNameOfInstance(instance: EC2.Instance): string | undefined {
+    return instance.Tags ? lookupTagKey(instance.Tags, 'Name') : undefined
+}
+
+function lookupTagKey(tags: EC2.Tag[], targetKey: string) {
+    return tags.filter(tag => tag.Key == targetKey)[0].Value
 }
diff --git a/src/test/ec2/model.test.ts b/src/test/ec2/model.test.ts
@@ -10,7 +10,7 @@ import { Ec2Client } from '../../shared/clients/ec2Client'
 import { attachedPoliciesListType } from 'aws-sdk/clients/iam'
 import { Ec2Selection } from '../../ec2/utils'
 import { ToolkitError } from '../../shared/errors'
-import { AWSError, EC2 } from 'aws-sdk'
+import { EC2 } from 'aws-sdk'
 
 describe('Ec2ConnectClient', function () {
     class MockSsmClient extends SsmClient {
@@ -42,9 +42,25 @@ describe('Ec2ConnectClient', function () {
             return new MockEc2Client()
         }
     }
+
+    describe('isInstanceRunning', async function () {
+        let client: MockEc2ConnectClient
+
+        before(function () {
+            client = new MockEc2ConnectClient()
+        })
+
+        it('only returns true with the instance is running', async function () {
+            const actualFirstResult = await client.isInstanceRunning('running:noPolicies')
+            const actualSecondResult = await client.isInstanceRunning('stopped:noPolicies')
+
+            assert.strictEqual(true, actualFirstResult)
+            assert.strictEqual(false, actualSecondResult)
+        })
+    })
+
     describe('handleStartSessionError', async function () {
         let client: MockEc2ConnectClientForError
-        const dummyError: AWSError = { name: 'testName', message: 'testMessage', code: 'testCode', time: new Date() }
 
         class MockEc2ConnectClientForError extends MockEc2ConnectClient {
             public override async hasProperPolicies(instanceId: string): Promise<boolean> {
@@ -58,7 +74,7 @@ describe('Ec2ConnectClient', function () {
         it('determines which error to throw based on if instance is running', async function () {
             async function assertThrowsErrorCode(testInstance: Ec2Selection, errCode: Ec2ConnectErrorCode) {
                 try {
-                    await client.handleStartSessionError(dummyError, testInstance)
+                    await client.checkForStartSessionError(testInstance)
                 } catch (err: unknown) {
                     assert.strictEqual((err as ToolkitError).code, errCode)
                 }
diff --git a/src/test/shared/clients/defaultEc2Client.test.ts b/src/test/shared/clients/defaultEc2Client.test.ts
@@ -94,3 +94,31 @@ describe('extractInstancesFromReservations', function () {
             )
         })
 })
+
+describe('getSingleInstanceFilter', function () {
+    const client = new Ec2Client('')
+
+    it('returns proper filter when given instanceId', function () {
+        const testInstanceId1 = 'test'
+        const actualFilters1 = client.getInstancesFilter([testInstanceId1])
+        const expectedFilters1: EC2.Filter[] = [
+            {
+                Name: 'instance-id',
+                Values: [testInstanceId1],
+            },
+        ]
+
+        assert.deepStrictEqual(expectedFilters1, actualFilters1)
+
+        const testInstanceId2 = 'test2'
+        const actualFilters2 = client.getInstancesFilter([testInstanceId1, testInstanceId2])
+        const expectedFilters2: EC2.Filter[] = [
+            {
+                Name: 'instance-id',
+                Values: [testInstanceId1, testInstanceId2],
+            },
+        ]
+
+        assert.deepStrictEqual(expectedFilters2, actualFilters2)
+    })
+})
