Use the regions in teleport config instead of ENV for bootstrap (#31701)

gravitational · Sep 11, 2023 · 2933eab · 2933eab
1 parent 4f9603d
commit 2933eab
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 20 deletions.
diff --git a/lib/configurators/aws/aws.go b/lib/configurators/aws/aws.go
@@ -291,8 +291,8 @@ type ConfiguratorConfig struct {
 	AWSSTSClient stsiface.STSAPI
 	// AWSIAMClient AWS IAM client.
 	AWSIAMClient iamiface.IAMAPI
-	// AWSSSMClient AWS SSM Client
-	AWSSSMClient ssmiface.SSMAPI
+	// AWSSSMClient is a mapping of region -> ssm client
+	AWSSSMClients map[string]ssmiface.SSMAPI
 	// Policies instance of the `Policies` that the actions use.
 	Policies awslib.Policies
 	// Identity is the current AWS credentials chain identity.
@@ -333,8 +333,29 @@ func (c *ConfiguratorConfig) CheckAndSetDefaults() error {
 				return trace.Wrap(err)
 			}
 		}
-		if c.AWSSSMClient == nil {
-			c.AWSSSMClient = ssm.New(c.AWSSession)
+		if c.AWSSSMClients == nil {
+			c.AWSSSMClients = make(map[string]ssmiface.SSMAPI)
+			for _, matcher := range c.ServiceConfig.Discovery.AWSMatchers {
+				if !slices.Contains(matcher.Types, services.AWSMatcherEC2) {
+					continue
+				}
+				for _, region := range matcher.Regions {
+					if _, ok := c.AWSSSMClients[region]; ok {
+						continue
+					}
+					session, err := awssession.NewSessionWithOptions(awssession.Options{
+						Config: aws.Config{
+							Region: &region,
+						},
+						SharedConfigState: awssession.SharedConfigEnable,
+					})
+					if err != nil {
+						return trace.Wrap(err)
+					}
+					c.AWSSSMClients[region] = ssm.New(session)
+				}
+			}
+
 		}
 
 		if c.Policies == nil {
@@ -482,7 +503,7 @@ func buildDiscoveryActions(config ConfiguratorConfig, targetCfg targetConfig) ([
 		return nil, err
 	}
 
-	actions = append(actions, buildSSMDocumentCreators(config.AWSSSMClient, targetCfg, proxyAddr)...)
+	actions = append(actions, buildSSMDocumentCreators(config.AWSSSMClients, targetCfg, proxyAddr)...)
 	return actions, nil
 }
 
@@ -744,18 +765,20 @@ func getProxyAddrFromConfig(cfg *servicecfg.Config, flags configurators.Bootstra
 	return "", trace.NotFound("proxy address not found, please provide --proxy, or set either teleport.proxy_server or proxy_service.public_addr in the teleport config")
 }
 
-func buildSSMDocumentCreators(ssm ssmiface.SSMAPI, targetCfg targetConfig, proxyAddr string) []configurators.ConfiguratorAction {
+func buildSSMDocumentCreators(ssm map[string]ssmiface.SSMAPI, targetCfg targetConfig, proxyAddr string) []configurators.ConfiguratorAction {
 	var creators []configurators.ConfiguratorAction
 	for _, matcher := range targetCfg.awsMatchers {
 		if !slices.Contains(matcher.Types, services.AWSMatcherEC2) {
 			continue
 		}
-		ssmCreator := awsSSMDocumentCreator{
-			ssm:      ssm,
-			Name:     matcher.SSM.DocumentName,
-			Contents: EC2DiscoverySSMDocument(proxyAddr),
+		for _, region := range matcher.Regions {
+			ssmCreator := awsSSMDocumentCreator{
+				ssm:      ssm[region],
+				Name:     matcher.SSM.DocumentName,
+				Contents: EC2DiscoverySSMDocument(proxyAddr),
+			}
+			creators = append(creators, &ssmCreator)
 		}
-		creators = append(creators, &ssmCreator)
 	}
 	return creators
 }

diff --git a/lib/configurators/aws/aws_test.go b/lib/configurators/aws/aws_test.go
@@ -1809,14 +1809,15 @@ func TestAWSDocumentConfigurator(t *testing.T) {
 		AWSSession:   &awssession.Session{},
 		AWSIAMClient: &iamMock{},
 		AWSSTSClient: &STSMock{ARN: "arn:aws:iam::1234567:role/example-role"},
-		AWSSSMClient: &SSMMock{
-			t: t,
-			expectedInput: &ssm.CreateDocumentInput{
-				Content:        aws.String(EC2DiscoverySSMDocument("https://proxy.example.org:443")),
-				DocumentType:   aws.String("Command"),
-				DocumentFormat: aws.String("YAML"),
-				Name:           aws.String("document"),
-			},
+		AWSSSMClients: map[string]ssmiface.SSMAPI{
+			"eu-central-1": &SSMMock{
+				t: t,
+				expectedInput: &ssm.CreateDocumentInput{
+					Content:        aws.String(EC2DiscoverySSMDocument("https://proxy.example.org:443")),
+					DocumentType:   aws.String("Command"),
+					DocumentFormat: aws.String("YAML"),
+					Name:           aws.String("document"),
+				}},
 		},
 		ServiceConfig: serviceConfig,
 		Flags: configurators.BootstrapFlags{
@@ -1849,7 +1850,7 @@ func TestAWSConfigurator(t *testing.T) {
 		AWSSession:    &awssession.Session{},
 		AWSIAMClient:  &iamMock{},
 		AWSSTSClient:  &STSMock{ARN: "arn:aws:iam::1234567:role/example-role"},
-		AWSSSMClient:  &SSMMock{},
+		AWSSSMClients: map[string]ssmiface.SSMAPI{"eu-central-1": &SSMMock{}},
 		ServiceConfig: &servicecfg.Config{},
 		Flags: configurators.BootstrapFlags{
 			AttachToUser:        "some-user",