feat(cli): ec2ssm no longer reuses IAM role (#1128)

`ec2ssm` currently will reuse existing IAM roles + profiles from a previous invocation. This was useful in avoiding extra IAM operations in the event of an interrupt or a crash. However, this functionality meant that `ec2ssm` was not idempotent and could not be run concurrently. This property is true now. This commit: * adds a random identifier to the names of new IAM roles and profiles * makes an AWS API call to disassociate instance profiles on exit * removes checks for existing IAM roles and profiles Fixes RAIN-47837 Signed-off-by: nschmeller <[email protected]>
lacework · Jan 26, 2023 · 476d179 · 476d179
1 parent 2f31808
commit 476d179
Show file tree

Hide file tree

Showing 26 changed files with 1,217 additions and 82 deletions.
diff --git a/cli/cmd/agent_aws-install_ec2ssm.go b/cli/cmd/agent_aws-install_ec2ssm.go
@@ -25,6 +25,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/iam/types"
 	"github.com/aws/aws-sdk-go-v2/service/ssm"
@@ -184,7 +185,12 @@ func installAWSSSM(_ *cobra.Command, _ []string) error {
 			cli.StopProgress()
 			err := teardownSSMAccess(cfg, role, instanceProfile, agentCmdState.InstallBYORole) // clean up after ourselves
 			if err != nil {
-				cli.OutputHuman("got an error %v while tearing down IAM role / infra", err)
+				cli.OutputHuman("got an error %v while tearing down IAM role / infra\n", err)
+				cli.Log.Debugw("IAM infra info after error",
+					"role", role,
+					"instance profile", instanceProfile,
+					"error", err,
+				)
 			}
 		}()
 		if err != nil {
@@ -231,7 +237,7 @@ func installAWSSSM(_ *cobra.Command, _ []string) error {
 
 			if !agentCmdState.InstallSkipCreatInfra {
 				// Attach an instance profile with our new role to the instance
-				err = threadRunner.AssociateInstanceProfileWithRunner(cfg, instanceProfile)
+				associationID, err := threadRunner.AssociateInstanceProfileWithRunner(cfg, instanceProfile)
 				if err != nil {
 					cli.OutputHuman(
 						"Failed to attach instance profile %s to instance %s with error %v\n",
@@ -241,6 +247,20 @@ func installAWSSSM(_ *cobra.Command, _ []string) error {
 					)
 					return
 				}
+				defer func(cfg aws.Config, associationID string) {
+					cli.Log.Debugw("disassociating instance profile from runner",
+						"association ID", associationID,
+						"instance_id", threadRunner.InstanceID,
+					)
+					err := threadRunner.DisassociateInstanceProfileFromRunner(cfg, associationID)
+					if err != nil {
+						cli.Log.Debugw("failed to disassociate instance profile from runner",
+							"association ID", associationID,
+							"instance_id", threadRunner.InstanceID,
+							"error", err,
+						)
+					}
+				}(cfg, associationID)
 			}
 
 			// Establish SSM Command connection to the runner

diff --git a/cli/cmd/awsiam.go b/cli/cmd/awsiam.go
@@ -20,12 +20,12 @@ package cmd
 
 import (
 	"context"
-	"fmt"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/iam"
 	"github.com/aws/aws-sdk-go-v2/service/iam/types"
+	"github.com/google/uuid"
 	"github.com/lacework/go-sdk/lwrunner"
 )
 
@@ -193,24 +193,14 @@ func getRoleFromName(c iamGetRoleFromNameAPI, roleName string) (types.Role, erro
 }
 
 type iamCreateSSMRoleAPI interface {
-	GetRole(ctx context.Context, params *iam.GetRoleInput, optFns ...func(*iam.Options)) (*iam.GetRoleOutput, error)
 	CreateRole(ctx context.Context, params *iam.CreateRoleInput, optFns ...func(*iam.Options)) (*iam.CreateRoleOutput, error)
 }
 
 // createSSMRole makes a call to the AWS API to create an IAM role.
 // Returns information about the newly created role and any errors.
 func createSSMRole(c iamCreateSSMRoleAPI) (types.Role, error) {
-	cli.Log.Debug("check if role already exists") // intended for after interrupt or error
-	getOutput, err := c.GetRole(
-		context.Background(),
-		&iam.GetRoleInput{
-			RoleName: aws.String(roleName),
-		},
-	)
-	if err == nil && getOutput.Role != nil {
-		cli.Log.Debug("we previously created the role, use it")
-		return *getOutput.Role, err
-	}
+	const roleNameBase string = "Lacework-Agent-SSM-Install-Role-"
+	roleName := roleNameBase + uuid.New().String()[:5]
 
 	const trustPolicyDocument = `{
 	"Version": "2012-10-17",
@@ -252,8 +242,6 @@ Safe to delete if found`,
 	return *output.Role, nil
 }
 
-const roleName string = "Lacework-Agent-SSM-Install-Role"
-
 // attachSSMPoliciesToRole takes a role, calls the IAM API to attach
 // policies required for SSM to the role, and returns the role along
 // with any errors.
@@ -285,21 +273,9 @@ func setupInstanceProfile(c *iam.Client, role types.Role) (types.InstanceProfile
 }
 
 func createInstanceProfile(c *iam.Client) (types.InstanceProfile, error) {
-	cli.Log.Debug("checking if instance profile already exists") // intended for after interrupt or error
-	getOutput, err := c.GetInstanceProfile(
-		context.Background(),
-		&iam.GetInstanceProfileInput{
-			InstanceProfileName: aws.String(instanceProfileName),
-		},
-	)
-	if err == nil && getOutput.InstanceProfile != nil {
-		cli.Log.Debugw("found existing instance profile",
-			"instance profile", *getOutput.InstanceProfile,
-		)
-		return *getOutput.InstanceProfile, err
-	}
+	const instanceProfileNameBase string = "Lacework-Agent-SSM-Install-Instance-Profile-"
+	instanceProfileName := instanceProfileNameBase + uuid.New().String()[:5]
 
-	cli.Log.Debug("no existing instance profile, creating one now")
 	createOutput, err := c.CreateInstanceProfile(
 		context.Background(),
 		&iam.CreateInstanceProfileInput{
@@ -326,28 +302,7 @@ func createInstanceProfile(c *iam.Client) (types.InstanceProfile, error) {
 	return *createOutput.InstanceProfile, err
 }
 
-const instanceProfileName string = "Lacework-Agent-SSM-Install-Instance-Profile"
-
 func addRoleToInstanceProfile(c *iam.Client, role types.Role, instanceProfile types.InstanceProfile) error {
-	cli.Log.Debug("checking if the role is already associated with the instance profile")
-	if len(instanceProfile.Roles) > 0 {
-		cli.Log.Debugw(
-			"found a role already associated with the instance profile",
-			"found role", instanceProfile.Roles[0],
-			"our role", role,
-		)
-		if *instanceProfile.Roles[0].Arn == *role.Arn {
-			return nil // the correct role is already attached to the instance profile
-		} else { // someone else modified this instance profile. Fail now
-			return fmt.Errorf(
-				"tried to use role %s but pre-existing instance profile %s already has role %s",
-				*role.Arn,
-				instanceProfileName,
-				*instanceProfile.Roles[0].Arn,
-			)
-		}
-	}
-
 	cli.Log.Debugw("adding role to instance profile", "role", role, "instance profile", instanceProfile)
 	_, err := c.AddRoleToInstanceProfile(
 		context.Background(),

diff --git a/cli/cmd/awsiam_test.go b/cli/cmd/awsiam_test.go
@@ -87,14 +87,9 @@ func TestGetRoleFromName(t *testing.T) {
 // createSSMRole tests
 
 type mockCreateSSMRoleClient struct {
-	getRoleMethod    mockGetRoleAPI
 	createRoleMethod mockCreateRoleAPI
 }
 
-func (m mockCreateSSMRoleClient) GetRole(ctx context.Context, params *iam.GetRoleInput, optFns ...func(*iam.Options)) (*iam.GetRoleOutput, error) {
-	return m.getRoleMethod(ctx, params, optFns...)
-}
-
 func (m mockCreateSSMRoleClient) CreateRole(ctx context.Context, params *iam.CreateRoleInput, optFns ...func(*iam.Options)) (*iam.CreateRoleOutput, error) {
 	return m.createRoleMethod(ctx, params, optFns...)
 }
@@ -107,13 +102,6 @@ func TestCreateSSMRole(t *testing.T) {
 	}{
 		{
 			client: mockCreateSSMRoleClient{
-				getRoleMethod: mockGetRoleAPI(func(ctx context.Context, params *iam.GetRoleInput, optFns ...func(*iam.Options)) (*iam.GetRoleOutput, error) {
-					return &iam.GetRoleOutput{
-						Role: &types.Role{
-							RoleName: aws.String(testRoleName),
-						},
-					}, nil
-				}),
 				createRoleMethod: mockCreateRoleAPI(func(ctx context.Context, params *iam.CreateRoleInput, optFns ...func(*iam.Options)) (*iam.CreateRoleOutput, error) {
 					return &iam.CreateRoleOutput{
 						Role: &types.Role{

diff --git a/go.mod b/go.mod
@@ -55,6 +55,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/iam v1.18.23
 	github.com/aws/aws-sdk-go-v2/service/ssm v1.33.1
 	github.com/go-git/go-git/v5 v5.5.2
+	github.com/google/uuid v1.3.0
 	github.com/mattn/go-isatty v0.0.14
 	github.com/spf13/cast v1.5.0
 	golang.org/x/exp v0.0.0-20221208152030-732eee02a75a

diff --git a/go.sum b/go.sum
@@ -259,6 +259,8 @@ github.com/google/pprof v0.0.0-20201203190320-1bf35d6f28c2/go.mod h1:kpwsk12EmLe
 github.com/google/pprof v0.0.0-20201218002935-b9804c9f04c2/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/enterprise-certificate-proxy v0.2.0 h1:y8Yozv7SZtlU//QXbezB6QkpuE6jMD2/gfzk4AftXjs=
 github.com/googleapis/enterprise-certificate-proxy v0.2.0/go.mod h1:8C0jb7/mgJe/9KK8Lm7X9ctZC2t60YyIpYEI16jx0Qg=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=

diff --git a/lwrunner/awsrunner.go b/lwrunner/awsrunner.go
@@ -131,7 +131,8 @@ func (run AWSRunner) SendPublicKey(pubBytes []byte) error {
 // receiving runner. First checks if there are any instance profiles already associated
 // with the runner, and returns an error if so (since a runner can only have one instance
 // profile associated with it). Then associates the instance profile with the runner.
-func (run AWSRunner) AssociateInstanceProfileWithRunner(cfg aws.Config, instanceProfile types.InstanceProfile) error {
+// Returns the association ID or an error.
+func (run AWSRunner) AssociateInstanceProfileWithRunner(cfg aws.Config, instanceProfile types.InstanceProfile) (string, error) {
 	c := ec2.New(ec2.Options{
 		Credentials: cfg.Credentials,
 		Region:      run.Region,
@@ -152,18 +153,18 @@ func (run AWSRunner) AssociateInstanceProfileWithRunner(cfg aws.Config, instance
 		},
 	)
 	if err != nil {
-		return err
+		return "", err
 	}
 
-	alreadyAssociated, err := run.isCorrectInstanceProfileAlreadyAssociated(cfg, describeOutput.IamInstanceProfileAssociations)
+	associationID, err := run.isCorrectInstanceProfileAlreadyAssociated(cfg, describeOutput.IamInstanceProfileAssociations)
 	if err != nil {
-		return err
+		return "", err
 	}
 
-	if alreadyAssociated { // use the existing, correctly configured instance profile
-		return nil
+	if associationID != "" { // use the existing, correctly configured instance profile
+		return associationID, nil
 	} else { // associate our own instance profile
-		_, err = c.AssociateIamInstanceProfile(
+		associateOutput, err := c.AssociateIamInstanceProfile(
 			context.Background(),
 			&ec2.AssociateIamInstanceProfileInput{
 				IamInstanceProfile: &ec2types.IamInstanceProfileSpecification{
@@ -173,22 +174,22 @@ func (run AWSRunner) AssociateInstanceProfileWithRunner(cfg aws.Config, instance
 			},
 		)
 		if err != nil {
-			return err
+			return "", err
 		}
 
-		return nil
+		return *associateOutput.IamInstanceProfileAssociation.AssociationId, nil
 	}
 }
 
 // isCorrectInstanceProfileAlreadyAssociated takes a list of instance profile associations
 // and checks if there is an instance profile associated and if this instance
-// profile has the correct policy for SSM access. Returns `true, nil` if so. Returns
-// `false, nil` if there is no instance profile associated. Returns `false, <error>` if
+// profile has the correct policy for SSM access. Returns `<assoc. id>, nil` if so. Returns
+// `"", nil` if there is no instance profile associated. Returns `"", <error>` if
 // there is an incorrect instance profile associated, or if there was an error in
 // executing this function.
-func (run AWSRunner) isCorrectInstanceProfileAlreadyAssociated(cfg aws.Config, associations []ec2types.IamInstanceProfileAssociation) (bool, error) {
+func (run AWSRunner) isCorrectInstanceProfileAlreadyAssociated(cfg aws.Config, associations []ec2types.IamInstanceProfileAssociation) (string, error) {
 	if len(associations) <= 0 { // no instance profile associated
-		return false, nil
+		return "", nil
 	}
 	instanceProfileName := strings.Split(*associations[0].IamInstanceProfile.Arn, "instance-profile/")[1]
 
@@ -204,13 +205,13 @@ func (run AWSRunner) isCorrectInstanceProfileAlreadyAssociated(cfg aws.Config, a
 		},
 	)
 	if err != nil {
-		return false, err
+		return "", err
 	}
 
 	// Check to see if the instance profile associated with the runner has the correct policy
 
 	if len(getInstanceProfileOutput.InstanceProfile.Roles) <= 0 { // can only have max one role
-		return false, fmt.Errorf(
+		return "", fmt.Errorf(
 			"runner %v already has an instance profile (%v) attached, does not have a role",
 			run,
 			getInstanceProfileOutput.InstanceProfile,
@@ -225,26 +226,42 @@ func (run AWSRunner) isCorrectInstanceProfileAlreadyAssociated(cfg aws.Config, a
 		},
 	)
 	if err != nil {
-		return false, err
+		return "", err
 	}
 
 	for _, policy := range listAttachedRolePoliciesOutput.AttachedPolicies {
 		if *policy.PolicyArn == SSMInstancePolicy {
-			return true, nil // everything is configured correctly, we can return now
+			return *associations[0].AssociationId, nil // everything is configured correctly, we can return now
 		}
 	}
 
 	// The runner has an instance profile attached, the instance profile has a role,
 	// and the role does not have the policy we need for SSM. We can't install on
 	// this instance, return an error
-	return false, fmt.Errorf(
+	return "", fmt.Errorf(
 		"runner %v already has an instance profile (%v) attached, does not have policy %s",
 		run,
 		getInstanceProfileOutput.InstanceProfile,
 		SSMInstancePolicy,
 	)
 }
 
+func (run AWSRunner) DisassociateInstanceProfileFromRunner(cfg aws.Config, associationID string) error {
+	c := ec2.New(ec2.Options{
+		Credentials: cfg.Credentials,
+		Region:      run.Region,
+	})
+
+	_, err := c.DisassociateIamInstanceProfile(
+		context.Background(),
+		&ec2.DisassociateIamInstanceProfileInput{
+			AssociationId: aws.String(associationID),
+		},
+	)
+
+	return err
+}
+
 const SSMInstancePolicy string = "arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore"
 
 // RunSSMCommandOnRemoteHost takes a shell command to install the agent on the runner

diff --git a/vendor/github.com/google/uuid/.travis.yml b/vendor/github.com/google/uuid/.travis.yml
diff --git a/vendor/github.com/google/uuid/CONTRIBUTING.md b/vendor/github.com/google/uuid/CONTRIBUTING.md
diff --git a/vendor/github.com/google/uuid/CONTRIBUTORS b/vendor/github.com/google/uuid/CONTRIBUTORS
diff --git a/vendor/github.com/google/uuid/LICENSE b/vendor/github.com/google/uuid/LICENSE