ec2: explicitly check if SSM agent is running (#3617)

Problem:
If the error is not related to permissions or instance status, we throw a general error about the SSM agent not running. However, this may not always be the case and we shouldn't assume other errors couldn't occur.

Solution:
We can leverage DescribeInstanceInformation from the SSM SDK to check the ping status of the SSM agent of target EC2 instance. If this call returns an error, or the ping status is outdated, we can conclude the SSM agent isn't running.

Now, there is a specific error message for when an instance is not able to connect. Therefore, the default error message has been changed so that it does not mention the SSM agent (since we would have already checked that).

Author: Hweinstock

src/ec2/model.ts

@@ -12,7 +12,7 @@ import { ToolkitError } from '../shared/errors'
 import { SsmClient } from '../shared/clients/ssmClient'
 import { Ec2Client } from '../shared/clients/ec2Client'
 
-export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect'
+export type Ec2ConnectErrorCode = 'EC2SSMStatus' | 'EC2SSMPermission' | 'EC2SSMConnect' | 'EC2SSMAgentStatus'
 
 import { openRemoteTerminal } from '../shared/remoteSession'
 import { DefaultIamClient } from '../shared/clients/iamClient'
@@ -71,6 +71,7 @@ export class Ec2ConnectionManager {
     public async checkForStartSessionError(selection: Ec2Selection): Promise<void> {
         const isInstanceRunning = await this.isInstanceRunning(selection.instanceId)
         const hasProperPolicies = await this.hasProperPolicies(selection.instanceId)
+        const isSsmAgentRunning = (await this.ssmClient.getInstanceAgentPingStatus(selection.instanceId)) == 'Online'
 
         if (!isInstanceRunning) {
             const message = 'Ensure the target instance is running and not currently starting, stopping, or stopped.'
@@ -87,6 +88,15 @@ export class Ec2ConnectionManager {
                 documentationUri: documentationUri,
             })
         }
+
+        if (!isSsmAgentRunning) {
+            this.throwConnectionError('Is SSM Agent running on the target instance?', selection, {
+                code: 'EC2SSMAgentStatus',
+                documentationUri: vscode.Uri.parse(
+                    'https://docs.aws.amazon.com/systems-manager/latest/userguide/ssm-agent.html'
+                ),
+            })
+        }
     }
 
     private async openSessionInTerminal(session: Session, selection: Ec2Selection) {
@@ -110,8 +120,8 @@ export class Ec2ConnectionManager {
             await this.openSessionInTerminal(response, selection)
         } catch (err: unknown) {
             // Default error if pre-check fails.
-            this.throwConnectionError('Check that the SSM Agent is running on target instance', selection, {
-                code: 'EC2SSMConnect',
+            this.throwConnectionError('Unable to connect to target instance. ', selection, {
+                code: 'EC2SSMAgentStatus',
                 cause: err as Error,
             })
         }

src/shared/clients/ssmClient.ts

@@ -6,6 +6,7 @@
 import { SSM } from 'aws-sdk'
 import { getLogger } from '../logger/logger'
 import globals from '../extensionGlobals'
+import { pageableToCollection } from '../utilities/collectionUtils'
 
 export class SsmClient {
     public constructor(public readonly regionCode: string) {}
@@ -32,4 +33,30 @@ export class SsmClient {
         const response = await client.startSession({ Target: target }).promise()
         return response
     }
+
+    public async describeInstance(target: string): Promise<SSM.InstanceInformation> {
+        const client = await this.createSdkClient()
+        const requester = async (req: SSM.DescribeInstanceInformationRequest) =>
+            client.describeInstanceInformation(req).promise()
+        const request: SSM.DescribeInstanceInformationRequest = {
+            InstanceInformationFilterList: [
+                {
+                    key: 'InstanceIds',
+                    valueSet: [target],
+                },
+            ],
+        }
+
+        const response = await pageableToCollection(requester, request, 'NextToken', 'InstanceInformationList')
+            .flatten()
+            .flatten()
+            .promise()
+
+        return response[0]!
+    }
+
+    public async getInstanceAgentPingStatus(target: string): Promise<string> {
+        const instanceInformation = await this.describeInstance(target)
+        return instanceInformation ? instanceInformation.PingStatus! : 'Inactive'
+    }
 }

src/test/ec2/model.test.ts

@@ -17,6 +17,10 @@ describe('Ec2ConnectClient', function () {
         public constructor() {
             super('test-region')
         }
+
+        public override async getInstanceAgentPingStatus(target: string): Promise<string> {
+            return target.split(':')[2]
+        }
     }
 
     class MockEc2Client extends Ec2Client {
@@ -82,36 +86,44 @@ describe('Ec2ConnectClient', function () {
 
             await assertThrowsErrorCode(
                 {
-                    instanceId: 'pending:noPolicies',
+                    instanceId: 'pending:noPolicies:Online',
                     region: 'test-region',
                 },
                 'EC2SSMStatus'
             )
 
             await assertThrowsErrorCode(
                 {
-                    instanceId: 'shutting-down:noPolicies',
+                    instanceId: 'shutting-down:noPolicies:Online',
                     region: 'test-region',
                 },
                 'EC2SSMStatus'
             )
 
             await assertThrowsErrorCode(
                 {
-                    instanceId: 'running:noPolicies',
+                    instanceId: 'running:noPolicies:Online',
                     region: 'test-region',
                 },
                 'EC2SSMPermission'
             )
 
             await assertThrowsErrorCode(
                 {
-                    instanceId: 'running:hasPolicies',
+                    instanceId: 'running:hasPolicies:Offline',
                     region: 'test-region',
                 },
-                'EC2SSMConnect'
+                'EC2SSMAgentStatus'
             )
         })
+
+        it('does not throw an error if all checks pass', async function () {
+            const passingInstance = {
+                instanceId: 'running:hasPolicies:Online',
+                region: 'test-region',
+            }
+            assert.doesNotThrow(async () => await client.checkForStartSessionError(passingInstance))
+        })
     })
 
     describe('hasProperPolicies', async function () {