From 324b4d4d3f4ac69ba207bfa4db34136b91744246 Mon Sep 17 00:00:00 2001
From: Jim Kalafut <jkalafut@hashicorp.com>
Date: Thu, 13 Feb 2020 22:21:58 -0800
Subject: [PATCH] Stabilize the selection of region from partition in AWS Auth
 (#8161)

AWS client object caches are by region. Some AWS API calls don't care
what region's client they use, but the existing getAnyRegionForAwsPartition
scheme was returning a random region, which in turn triggered maintaining many
more client objects than are necessary (e.g. 18 regions in the main AWS
partition). This can be an issue for heavy STS users bumping up against
STS rate limits, since 18 sets of creds are being cached and renewed per
STS role.
---
 builtin/credential/aws/backend.go    | 24 +++++++++++++++++-------
 builtin/credential/aws/path_login.go |  8 +++++++-
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/builtin/credential/aws/backend.go b/builtin/credential/aws/backend.go
index 352bb14812ef..ce7c39bb6a8e 100644
--- a/builtin/credential/aws/backend.go
+++ b/builtin/credential/aws/backend.go
@@ -65,6 +65,11 @@ type backend struct {
 	// will be flushed. The empty STS role signifies the master account
 	IAMClientsMap map[string]map[string]*iam.IAM
 
+	// Map to associate a partition to a random region in that partition. Users of
+	// this don't care what region in the partition they use, but there is some client
+	// cache efficiency gain if we keep the mapping stable, hence caching a single copy.
+	partitionToRegionMap map[string]*endpoints.Region
+
 	// Map of AWS unique IDs to the full ARN corresponding to that unique ID
 	// This avoids the overhead of an AWS API hit for every login request
 	// using the IAM auth method when bound_iam_principal_arn contains a wildcard
@@ -144,6 +149,8 @@ func Backend(_ *logical.BackendConfig) (*backend, error) {
 		Clean:          b.cleanup,
 	}
 
+	b.partitionToRegionMap = generatePartitionToRegionMap()
+
 	return b, nil
 }
 
@@ -249,7 +256,7 @@ func (b *backend) resolveArnToRealUniqueId(ctx context.Context, s logical.Storag
 	// partition, and passing that region back back to the SDK, so that the SDK can figure out the
 	// proper partition from the arbitrary region we passed in to look up the endpoint.
 	// Sigh
-	region := getAnyRegionForAwsPartition(entity.Partition)
+	region := b.partitionToRegionMap[entity.Partition]
 	if region == nil {
 		return "", fmt.Errorf("unable to resolve partition %q to a region", entity.Partition)
 	}
@@ -293,18 +300,21 @@ func (b *backend) resolveArnToRealUniqueId(ctx context.Context, s logical.Storag
 
 // Adapted from https://docs.aws.amazon.com/sdk-for-go/api/aws/endpoints/
 // the "Enumerating Regions and Endpoint Metadata" section
-func getAnyRegionForAwsPartition(partitionId string) *endpoints.Region {
+func generatePartitionToRegionMap() map[string]*endpoints.Region {
+	partitionToRegion := make(map[string]*endpoints.Region)
+
 	resolver := endpoints.DefaultResolver()
 	partitions := resolver.(endpoints.EnumPartitions).Partitions()
 
 	for _, p := range partitions {
-		if p.ID() == partitionId {
-			for _, r := range p.Regions() {
-				return &r
-			}
+		// Choose a single region randomly from the partition
+		for _, r := range p.Regions() {
+			partitionToRegion[p.ID()] = &r
+			break
 		}
 	}
-	return nil
+
+	return partitionToRegion
 }
 
 const backendHelp = `
diff --git a/builtin/credential/aws/path_login.go b/builtin/credential/aws/path_login.go
index 8cf7552f31a0..6b8fe95d6a0e 100644
--- a/builtin/credential/aws/path_login.go
+++ b/builtin/credential/aws/path_login.go
@@ -1650,7 +1650,13 @@ func (e *iamEntity) canonicalArn() string {
 // This returns the "full" ARN of an iamEntity, how it would be referred to in AWS proper
 func (b *backend) fullArn(ctx context.Context, e *iamEntity, s logical.Storage) (string, error) {
 	// Not assuming path is reliable for any entity types
-	client, err := b.clientIAM(ctx, s, getAnyRegionForAwsPartition(e.Partition).ID(), e.AccountNumber)
+
+	region := b.partitionToRegionMap[e.Partition]
+	if region == nil {
+		return "", fmt.Errorf("unable to resolve partition %q to a region", e.Partition)
+	}
+
+	client, err := b.clientIAM(ctx, s, region.ID(), e.AccountNumber)
 	if err != nil {
 		return "", errwrap.Wrapf("error creating IAM client: {{err}}", err)
 	}