EC2 Auto Discover SSM: add support for debugging custom SSM Docs (#41706)

This PR uses a new AWS API that list the steps of the current
invocation.
After listing them, it will ask for the output of each one.

Previously, we were using a static list of steps: those defined in the
default SSM Document.

However, for custom documts with different list of steps that would
fail.

If the client does not have access to this new API, we will fallback to
the list of steps that exist in the default SSM Document.

If we ask for a status of one of those steps, and we receive a known
error indicating that the step does not exist, instead of failing we
will emit the overall invocation result (which doesnt include
stdout/stderr, but better than nothing)

docs/pages/auto-discovery/servers/ec2-discovery.mdx

@@ -171,6 +171,7 @@ AWS
                 "ec2:DescribeInstances",
                 "ssm:DescribeInstanceInformation",
                 "ssm:GetCommandInvocation",
+                "ssm:ListCommandInvocations",
                 "ssm:SendCommand"
             ],
             "Resource": [
@@ -190,6 +191,7 @@ AWS
                 "ec2:DescribeInstances",
                 "ssm:DescribeInstanceInformation",
                 "ssm:GetCommandInvocation",
+                "ssm:ListCommandInvocations",
                 "ssm:SendCommand"
             ],
             "Resource": [

lib/cloud/aws/policy_statements.go

@@ -147,6 +147,7 @@ func StatementForEC2SSMAutoDiscover() *Statement {
 			"ec2:DescribeInstances",
 			"ssm:DescribeInstanceInformation",
 			"ssm:GetCommandInvocation",
+			"ssm:ListCommandInvocations",
 			"ssm:SendCommand",
 		},
 		Resources: allResources,

lib/cloud/aws/ssm_documents.go

@@ -58,3 +58,10 @@ mainSteps:
 }
 
 const EC2DiscoveryPolicyName = "TeleportEC2Discovery"
+
+// EC2DiscoverySSMDocumentSteps is the list of Steps defined in the default SSM Document for Teleport Discovery.
+// Used to query step results after executing a command using SSM.
+var EC2DiscoverySSMDocumentSteps = []string{
+	"downloadContent",
+	"runShellScript",
+}

lib/configurators/aws/aws_test.go

@@ -584,6 +584,7 @@ func TestAWSIAMDocuments(t *testing.T) {
 						"ec2:DescribeInstances",
 						"ssm:DescribeInstanceInformation",
 						"ssm:GetCommandInvocation",
+						"ssm:ListCommandInvocations",
 						"ssm:SendCommand",
 					},
 					Resources: []string{"*"},
@@ -596,6 +597,7 @@ func TestAWSIAMDocuments(t *testing.T) {
 						"ec2:DescribeInstances",
 						"ssm:DescribeInstanceInformation",
 						"ssm:GetCommandInvocation",
+						"ssm:ListCommandInvocations",
 						"ssm:SendCommand",
 					},
 					Resources: []string{"*"},

lib/integrations/awsoidc/ec2_ssm_iam_config.go

@@ -135,6 +135,7 @@ func NewEC2SSMConfigureClient(ctx context.Context, region string) (EC2SSMConfigu
 // Action: Run a command and get its output.
 //   - ssm:SendCommand
 //   - ssm:GetCommandInvocation
+//   - ssm:ListCommandInvocations
 //
 // Besides setting up the required IAM policies, this method also adds the SSM Document.
 // This SSM Document downloads and runs the Teleport Installer Script, which installs teleport in the target EC2 instance.

lib/srv/server/ssm_install.go

@@ -284,13 +284,6 @@ func skipAWSWaitErr(err error) error {
 	return trace.Wrap(err)
 }
 
-// This list of ssmSteps must match the ones used in awslib.EC2DiscoverySSMDocument
-// These vars are used to obtain specific status of each step when checking a command output.
-var (
-	ssmStepDownloadContent = "downloadContent"
-	ssmStepRunShellScript  = "runShellScript"
-)
-
 func (si *SSMInstaller) checkCommand(ctx context.Context, req SSMRunRequest, commandID, instanceID *string) error {
 	err := req.SSM.WaitUntilCommandExecutedWithContext(ctx, &ssm.GetCommandInvocationInput{
 		CommandId:  commandID,
@@ -301,32 +294,91 @@ func (si *SSMInstaller) checkCommand(ctx context.Context, req SSMRunRequest, com
 		return trace.Wrap(err)
 	}
 
-	// Check 1st step: download Content
-	downloadContentStep, err := si.getCommandStepStatusEvent(ctx, &ssmStepDownloadContent, req, commandID, instanceID)
-	if err != nil {
+	invocationSteps, err := si.getInvocationSteps(ctx, req, commandID, instanceID)
+	switch {
+	case trace.IsAccessDenied(err):
+		// getInvocationSteps uses `ssm:ListCommandInvocations` to gather all the executed steps.
+		// Using `ssm:ListCommandInvocations` is not always possible because previous Docs versions (pre-v16) did not ask for that permission.
+		// If the IAM role does not have access to that action, an Access Denied is returned here.
+		// The process continues but the user is warned that they should add that permission to get better diagnostics.
+		si.Logger.WarnContext(ctx,
+			"Add ssm:ListCommandInvocations action to IAM Role to improve diagnostics of EC2 Teleport installation failures",
+			"error", err)
+
+		invocationSteps = awslib.EC2DiscoverySSMDocumentSteps
+
+	case err != nil:
 		return trace.Wrap(err)
 	}
 
-	// Only check runShellScript step if downloadContent was a success.
-	if downloadContentStep.Metadata.Code != libevents.SSMRunSuccessCode {
-		return trace.Wrap(si.Emitter.EmitAuditEvent(ctx, downloadContentStep))
+	for i, step := range invocationSteps {
+		stepResultEvent, err := si.getCommandStepStatusEvent(ctx, step, req, commandID, instanceID)
+		if err != nil {
+			var invalidPluginNameErr *ssm.InvalidPluginName
+			if errors.As(err, &invalidPluginNameErr) {
+				// If using a custom SSM Document and the client does not have access to ssm:ListCommandInvocations
+				// the list of invocationSteps (ie plugin name) might be wrong.
+				// If that's the case, emit an event with the overall invocation result (ignoring specific steps' stdout and stderr).
+				invocationResultEvent, err := si.getCommandStepStatusEvent(ctx, "" /*no step*/, req, commandID, instanceID)
+				if err != nil {
+					return trace.Wrap(err)
+				}
+
+				return trace.Wrap(si.Emitter.EmitAuditEvent(ctx, invocationResultEvent))
+			}
+
+			return trace.Wrap(err)
+		}
+
+		// Emit an event for the first failed step or for the latest step.
+		lastStep := i+1 == len(invocationSteps)
+		if stepResultEvent.Metadata.Code != libevents.SSMRunSuccessCode || lastStep {
+			return trace.Wrap(si.Emitter.EmitAuditEvent(ctx, stepResultEvent))
+		}
 	}
 
-	// Check 2nd step: run shell script
-	runShellScriptStep, err := si.getCommandStepStatusEvent(ctx, &ssmStepRunShellScript, req, commandID, instanceID)
-	if err != nil {
-		return trace.Wrap(err)
-	}
-
-	return trace.Wrap(si.Emitter.EmitAuditEvent(ctx, runShellScriptStep))
+	return nil
 }
 
-func (si *SSMInstaller) getCommandStepStatusEvent(ctx context.Context, step *string, req SSMRunRequest, commandID, instanceID *string) (*apievents.SSMRun, error) {
-	stepResult, err := req.SSM.GetCommandInvocationWithContext(ctx, &ssm.GetCommandInvocationInput{
+func (si *SSMInstaller) getInvocationSteps(ctx context.Context, req SSMRunRequest, commandID, instanceID *string) ([]string, error) {
+	// ssm:ListCommandInvocations is used to list the actual steps because users might be using a custom SSM Document.
+	listCommandInvocationResp, err := req.SSM.ListCommandInvocationsWithContext(ctx, &ssm.ListCommandInvocationsInput{
 		CommandId:  commandID,
 		InstanceId: instanceID,
-		PluginName: step,
+		Details:    aws.Bool(true),
 	})
+	if err != nil {
+		return nil, trace.Wrap(awslib.ConvertRequestFailureError(err))
+	}
+
+	// We only expect a single invocation because we are sending both the CommandID and the InstanceID.
+	// This call happens after WaitUntilCommandExecuted, so there's no reason for this to ever return 0 elements.
+	if len(listCommandInvocationResp.CommandInvocations) == 0 {
+		si.Logger.WarnContext(ctx,
+			"No command invocation was found.",
+			"command_id", aws.StringValue(commandID),
+			"instance_id", aws.StringValue(instanceID),
+		)
+		return nil, trace.BadParameter("no command invocation was found")
+	}
+	commandInvocation := listCommandInvocationResp.CommandInvocations[0]
+
+	documentSteps := make([]string, 0, len(commandInvocation.CommandPlugins))
+	for _, step := range commandInvocation.CommandPlugins {
+		documentSteps = append(documentSteps, aws.StringValue(step.Name))
+	}
+	return documentSteps, nil
+}
+
+func (si *SSMInstaller) getCommandStepStatusEvent(ctx context.Context, step string, req SSMRunRequest, commandID, instanceID *string) (*apievents.SSMRun, error) {
+	getCommandInvocationReq := &ssm.GetCommandInvocationInput{
+		CommandId:  commandID,
+		InstanceId: instanceID,
+	}
+	if step != "" {
+		getCommandInvocationReq.PluginName = aws.String(step)
+	}
+	stepResult, err := req.SSM.GetCommandInvocationWithContext(ctx, getCommandInvocationReq)
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}

lib/srv/server/ssm_install_test.go

@@ -30,7 +30,6 @@ import (
 	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
 	"github.com/google/uuid"
-	"github.com/gravitational/trace"
 	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport/api/types/events"
@@ -39,10 +38,10 @@ import (
 
 type mockSSMClient struct {
 	ssmiface.SSMAPI
-	commandOutput            *ssm.SendCommandOutput
-	invokeOutputStepDownload *ssm.GetCommandInvocationOutput
-	invokeOutputStepScript   *ssm.GetCommandInvocationOutput
-	describeOutput           *ssm.DescribeInstanceInformationOutput
+	commandOutput          *ssm.SendCommandOutput
+	commandInvokeOutput    map[string]*ssm.GetCommandInvocationOutput
+	describeOutput         *ssm.DescribeInstanceInformationOutput
+	listCommandInvocations *ssm.ListCommandInvocationsOutput
 }
 
 const docWithoutSSHDConfigPathParam = "ssmdocument-without-sshdConfigPath-param"
@@ -55,13 +54,10 @@ func (sm *mockSSMClient) SendCommandWithContext(_ context.Context, input *ssm.Se
 }
 
 func (sm *mockSSMClient) GetCommandInvocationWithContext(_ context.Context, input *ssm.GetCommandInvocationInput, _ ...request.Option) (*ssm.GetCommandInvocationOutput, error) {
-	if aws.StringValue(input.PluginName) == "downloadContent" {
-		return sm.invokeOutputStepDownload, nil
+	if stepResult, found := sm.commandInvokeOutput[aws.StringValue(input.PluginName)]; found {
+		return stepResult, nil
 	}
-	if aws.StringValue(input.PluginName) == "runShellScript" {
-		return sm.invokeOutputStepScript, nil
-	}
-	return nil, trace.NotFound("plugin name is required")
+	return nil, &ssm.InvalidPluginName{}
 }
 
 func (sm *mockSSMClient) DescribeInstanceInformationWithContext(_ context.Context, input *ssm.DescribeInstanceInformationInput, _ ...request.Option) (*ssm.DescribeInstanceInformationOutput, error) {
@@ -71,6 +67,13 @@ func (sm *mockSSMClient) DescribeInstanceInformationWithContext(_ context.Contex
 	return sm.describeOutput, nil
 }
 
+func (sm *mockSSMClient) ListCommandInvocationsWithContext(aws.Context, *ssm.ListCommandInvocationsInput, ...request.Option) (*ssm.ListCommandInvocationsOutput, error) {
+	if sm.listCommandInvocations == nil {
+		return nil, awserr.NewRequestFailure(awserr.New("AccessDeniedException", "message", nil), http.StatusBadRequest, uuid.NewString())
+	}
+	return sm.listCommandInvocations, nil
+}
+
 func (sm *mockSSMClient) WaitUntilCommandExecutedWithContext(aws.Context, *ssm.GetCommandInvocationInput, ...request.WaiterOption) error {
 	if aws.StringValue(sm.commandOutput.Command.Status) == ssm.CommandStatusFailed {
 		return awserr.New(request.WaiterResourceNotReadyErrorCode, "err", nil)
@@ -110,13 +113,15 @@ func TestSSMInstaller(t *testing.T) {
 							CommandId: aws.String("command-id-1"),
 						},
 					},
-					invokeOutputStepDownload: &ssm.GetCommandInvocationOutput{
-						Status:       aws.String(ssm.CommandStatusSuccess),
-						ResponseCode: aws.Int64(0),
-					},
-					invokeOutputStepScript: &ssm.GetCommandInvocationOutput{
-						Status:       aws.String(ssm.CommandStatusSuccess),
-						ResponseCode: aws.Int64(0),
+					commandInvokeOutput: map[string]*ssm.GetCommandInvocationOutput{
+						"downloadContent": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
+						"runShellScript": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
 					},
 				},
 				Region:    "eu-central-1",
@@ -155,13 +160,15 @@ func TestSSMInstaller(t *testing.T) {
 							CommandId: aws.String("command-id-1"),
 						},
 					},
-					invokeOutputStepDownload: &ssm.GetCommandInvocationOutput{
-						Status:       aws.String(ssm.CommandStatusSuccess),
-						ResponseCode: aws.Int64(0),
-					},
-					invokeOutputStepScript: &ssm.GetCommandInvocationOutput{
-						Status:       aws.String(ssm.CommandStatusSuccess),
-						ResponseCode: aws.Int64(0),
+					commandInvokeOutput: map[string]*ssm.GetCommandInvocationOutput{
+						"downloadContent": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
+						"runShellScript": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
 					},
 				},
 				Region:    "eu-central-1",
@@ -200,11 +207,13 @@ func TestSSMInstaller(t *testing.T) {
 							CommandId: aws.String("command-id-1"),
 						},
 					},
-					invokeOutputStepDownload: &ssm.GetCommandInvocationOutput{
-						Status:                aws.String(ssm.CommandStatusFailed),
-						ResponseCode:          aws.Int64(1),
-						StandardErrorContent:  aws.String("timeout error"),
-						StandardOutputContent: aws.String(""),
+					commandInvokeOutput: map[string]*ssm.GetCommandInvocationOutput{
+						"downloadContent": {
+							Status:                aws.String(ssm.CommandStatusFailed),
+							ResponseCode:          aws.Int64(1),
+							StandardErrorContent:  aws.String("timeout error"),
+							StandardOutputContent: aws.String(""),
+						},
 					},
 				},
 				Region:    "eu-central-1",
@@ -245,17 +254,19 @@ func TestSSMInstaller(t *testing.T) {
 							CommandId: aws.String("command-id-1"),
 						},
 					},
-					invokeOutputStepDownload: &ssm.GetCommandInvocationOutput{
-						Status:                aws.String(ssm.CommandStatusSuccess),
-						ResponseCode:          aws.Int64(0),
-						StandardErrorContent:  aws.String("no error"),
-						StandardOutputContent: aws.String(""),
-					},
-					invokeOutputStepScript: &ssm.GetCommandInvocationOutput{
-						Status:                aws.String(ssm.CommandStatusFailed),
-						ResponseCode:          aws.Int64(1),
-						StandardErrorContent:  aws.String("timeout error"),
-						StandardOutputContent: aws.String(""),
+					commandInvokeOutput: map[string]*ssm.GetCommandInvocationOutput{
+						"downloadContent": {
+							Status:                aws.String(ssm.CommandStatusSuccess),
+							ResponseCode:          aws.Int64(0),
+							StandardErrorContent:  aws.String("no error"),
+							StandardOutputContent: aws.String(""),
+						},
+						"runShellScript": {
+							Status:                aws.String(ssm.CommandStatusFailed),
+							ResponseCode:          aws.Int64(1),
+							StandardErrorContent:  aws.String("timeout error"),
+							StandardOutputContent: aws.String(""),
+						},
 					},
 				},
 				Region:    "eu-central-1",
@@ -299,13 +310,15 @@ func TestSSMInstaller(t *testing.T) {
 							CommandId: aws.String("command-id-1"),
 						},
 					},
-					invokeOutputStepDownload: &ssm.GetCommandInvocationOutput{
-						Status:       aws.String(ssm.CommandStatusSuccess),
-						ResponseCode: aws.Int64(0),
-					},
-					invokeOutputStepScript: &ssm.GetCommandInvocationOutput{
-						Status:       aws.String(ssm.CommandStatusSuccess),
-						ResponseCode: aws.Int64(0),
+					commandInvokeOutput: map[string]*ssm.GetCommandInvocationOutput{
+						"downloadContent": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
+						"runShellScript": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
 					},
 					describeOutput: &ssm.DescribeInstanceInformationOutput{
 						InstanceInformationList: []*ssm.InstanceInformation{
@@ -385,6 +398,106 @@ func TestSSMInstaller(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "ssm with custom steps",
+			req: SSMRunRequest{
+				Instances: []EC2Instance{
+					{InstanceID: "instance-id-1"},
+				},
+				DocumentName: document,
+				Params:       map[string]string{"token": "abcdefg"},
+				SSM: &mockSSMClient{
+					commandOutput: &ssm.SendCommandOutput{
+						Command: &ssm.Command{
+							CommandId: aws.String("command-id-1"),
+						},
+					},
+					commandInvokeOutput: map[string]*ssm.GetCommandInvocationOutput{
+						"downloadContentCustom": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
+						"runShellScriptCustom": {
+							Status:                aws.String(ssm.CommandStatusSuccess),
+							ResponseCode:          aws.Int64(0),
+							StandardOutputContent: aws.String("custom output"),
+						},
+					},
+					listCommandInvocations: &ssm.ListCommandInvocationsOutput{
+						CommandInvocations: []*ssm.CommandInvocation{{
+							CommandPlugins: []*ssm.CommandPlugin{
+								{Name: aws.String("downloadContentCustom")},
+								{Name: aws.String("runShellScriptCustom")},
+							},
+						}},
+					},
+				},
+				Region:    "eu-central-1",
+				AccountID: "account-id",
+			},
+			conf: SSMInstallerConfig{
+				Emitter: &mockEmitter{},
+			},
+			expectedEvents: []events.AuditEvent{
+				&events.SSMRun{
+					Metadata: events.Metadata{
+						Type: libevent.SSMRunEvent,
+						Code: libevent.SSMRunSuccessCode,
+					},
+					CommandID:      "command-id-1",
+					InstanceID:     "instance-id-1",
+					AccountID:      "account-id",
+					Region:         "eu-central-1",
+					ExitCode:       0,
+					Status:         ssm.CommandStatusSuccess,
+					StandardOutput: "custom output",
+					InvocationURL:  "https://eu-central-1.console.aws.amazon.com/systems-manager/run-command/command-id-1/instance-id-1",
+				},
+			},
+		},
+		{
+			name: "ssm with custom steps but without listing permissions only returns the overall result",
+			req: SSMRunRequest{
+				Instances: []EC2Instance{
+					{InstanceID: "instance-id-1"},
+				},
+				DocumentName: document,
+				Params:       map[string]string{"token": "abcdefg"},
+				SSM: &mockSSMClient{
+					commandOutput: &ssm.SendCommandOutput{
+						Command: &ssm.Command{
+							CommandId: aws.String("command-id-1"),
+						},
+					},
+					commandInvokeOutput: map[string]*ssm.GetCommandInvocationOutput{
+						"": {
+							Status:       aws.String(ssm.CommandStatusSuccess),
+							ResponseCode: aws.Int64(0),
+						},
+					},
+				},
+				Region:    "eu-central-1",
+				AccountID: "account-id",
+			},
+			conf: SSMInstallerConfig{
+				Emitter: &mockEmitter{},
+			},
+			expectedEvents: []events.AuditEvent{
+				&events.SSMRun{
+					Metadata: events.Metadata{
+						Type: libevent.SSMRunEvent,
+						Code: libevent.SSMRunSuccessCode,
+					},
+					CommandID:     "command-id-1",
+					InstanceID:    "instance-id-1",
+					AccountID:     "account-id",
+					Region:        "eu-central-1",
+					ExitCode:      0,
+					Status:        ssm.CommandStatusSuccess,
+					InvocationURL: "https://eu-central-1.console.aws.amazon.com/systems-manager/run-command/command-id-1/instance-id-1",
+				},
+			},
+		},
 		// todo(amk): test that incomplete commands eventually return
 		// an event once completed
 	} {