auth-aws: update partition to region lookup for AWS SDK v2

auth/aws/backend.go

@@ -10,10 +10,10 @@ import (
 	"sync"
 	"time"
 
-	"github.com/aws/aws-sdk-go/aws/endpoints"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/iam"
 	"github.com/hashicorp/go-secure-stdlib/awsutil"
+	"github.com/openbao/openbao-plugins/auth/aws/partitions"
 	"github.com/openbao/openbao/sdk/v2/framework"
 	"github.com/openbao/openbao/sdk/v2/helper/consts"
 	"github.com/openbao/openbao/sdk/v2/logical"
@@ -85,11 +85,6 @@ type backend struct {
 	// will be flushed. The empty STS role signifies the master account
 	IAMClientsMap map[string]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
@@ -202,8 +197,6 @@ func Backend(_ *logical.BackendConfig) (*backend, error) {
 		Clean:          b.cleanup,
 	}
 
-	b.partitionToRegionMap = generatePartitionToRegionMap()
-
 	return b, nil
 }
 
@@ -309,11 +302,11 @@ 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 := b.partitionToRegionMap[entity.Partition]
-	if region == nil {
+	region, ok := partitions.GetGlobalRegionForPartition(entity.Partition)
+	if !ok {
 		return "", fmt.Errorf("unable to resolve partition %q to a region", entity.Partition)
 	}
-	iamClient, err := b.clientIAM(ctx, s, region.ID(), entity.AccountNumber)
+	iamClient, err := b.clientIAM(ctx, s, region, entity.AccountNumber)
 	if err != nil {
 		return "", awsutil.AppendAWSError(err)
 	}
@@ -381,39 +374,6 @@ func (b *backend) genDeprecatedPath(path *framework.Path) *framework.Path {
 	return &pathDeprecated
 }
 
-// Adapted from https://docs.aws.amazon.com/sdk-for-go/api/aws/endpoints/
-// the "Enumerating Regions and Endpoint Metadata" section
-func generatePartitionToRegionMap() map[string]*endpoints.Region {
-	partitionToRegion := make(map[string]*endpoints.Region)
-
-	resolver := endpoints.DefaultResolver()
-	partitions := resolver.(endpoints.EnumPartitions).Partitions()
-
-	for _, p := range partitions {
-		// For most partitions, it's fine to choose a single region randomly.
-		// However, there are a few exceptions:
-		//
-		//   For "aws", choose "us-east-1" because it is always enabled (and
-		//   enabled for STS) by default.
-		//
-		//   For "aws-us-gov", choose "us-gov-west-1" because it is the only
-		//   valid region for IAM operations.
-		//   ref: https://github.com/aws/aws-sdk-go/blob/v1.34.25/aws/endpoints/defaults.go#L8176-L8194
-		for _, r := range p.Regions() {
-			if p.ID() == "aws" && r.ID() != "us-east-1" {
-				continue
-			}
-			if p.ID() == "aws-us-gov" && r.ID() != "us-gov-west-1" {
-				continue
-			}
-			partitionToRegion[p.ID()] = &r
-			break
-		}
-	}
-
-	return partitionToRegion
-}
-
 const backendHelp = `
 The aws auth method uses either AWS IAM credentials or AWS-signed EC2 metadata
 to authenticate clients, which are IAM principals or EC2 instances.

auth/aws/backend_test.go

@@ -1901,13 +1901,3 @@ func generateRenewRequest(s logical.Storage, auth *logical.Auth) *logical.Reques
 
 	return renewReq
 }
-
-func TestGeneratePartitionToRegionMap(t *testing.T) {
-	m := generatePartitionToRegionMap()
-	if m["aws"].ID() != "us-east-1" {
-		t.Fatal("expected us-east-1 but received " + m["aws"].ID())
-	}
-	if m["aws-us-gov"].ID() != "us-gov-west-1" {
-		t.Fatal("expected us-gov-west-1 but received " + m["aws-us-gov"].ID())
-	}
-}

auth/aws/partitions/gen.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env sh
+
+# Copyright (c) 2025 OpenBao a Series of LF Projects, LLC
+# SPDX-License-Identifier: MPL-2.0
+
+#################################################################
+# generate a map[string]string mapping from partition -> region #
+#################################################################
+
+# Find the aws module source dir, which contains a json with the data we want
+# We could probably also get the data from some API, but using only content of this repo makes it deterministic
+DIR=$(go list -f '{{.Dir}}' 'github.com/aws/aws-sdk-go-v2/internal/endpoints/awsrulesfn')
+
+TMP_GO=$(mktemp)
+
+# Generate some boilerplate
+echo '// Copyright (c) 2025 OpenBao a Series of LF Projects, LLC' > "$TMP_GO"
+echo '// SPDX-License-Identifier: MPL-2.0' >> "$TMP_GO"
+echo '' >> "$TMP_GO"
+echo "// Code generated by ${0}. DO NOT EDIT." >> "$TMP_GO"
+echo '' >> "$TMP_GO"
+echo 'package partitions' >> "$TMP_GO"
+echo 'var data = map[string]string{' >> "$TMP_GO"
+
+# Generate map entries
+jq -r '
+    .partitions[] |
+        "\(.id | @json): \(.outputs.implicitGlobalRegion | @json),"
+' "$DIR/partitions.json" >> "$TMP_GO"
+
+echo '}' >> "$TMP_GO"
+
+# Format Code
+gofmt "$TMP_GO" > data.go
+rm "$TMP_GO"

auth/aws/partitions/partitions.go

@@ -0,0 +1,11 @@
+// Copyright (c) 2025 OpenBao a Series of LF Projects, LLC
+// SPDX-License-Identifier: MPL-2.0
+
+package partitions
+
+//go:generate ./gen.sh
+
+func GetGlobalRegionForPartition(partition string) (string, bool) {
+	region, ok := data[partition]
+	return region, ok
+}

auth/aws/partitions/partitions_test.go

@@ -0,0 +1,34 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package partitions_test
+
+import (
+	"testing"
+
+	endpoints "github.com/openbao/openbao-plugins/auth/aws/partitions"
+)
+
+func TestGeneratePartitionToRegionMap(t *testing.T) {
+	res, ok := endpoints.GetGlobalRegionForPartition("aws")
+	if !ok {
+		t.Fatal("expected aws to be available")
+	}
+	if res != "us-east-1" {
+		t.Fatal("expected us-east-1 but received " + res)
+	}
+
+	res, ok = endpoints.GetGlobalRegionForPartition("aws-us-gov")
+	if !ok {
+		t.Fatal("expected us-gov-west-1 to be available")
+	}
+	if res != "us-gov-west-1" {
+		t.Fatal("expected us-gov-west-1 but received " + res)
+	}
+
+	res, ok = endpoints.GetGlobalRegionForPartition("gcp")
+	if ok {
+		t.Fatal("expected gcp not to be available, but got" + res)
+	}
+
+}

auth/aws/path_login.go

@@ -33,6 +33,7 @@ import (
 	"github.com/hashicorp/go-secure-stdlib/strutil"
 	uuid "github.com/hashicorp/go-uuid"
 
+	"github.com/openbao/openbao-plugins/auth/aws/partitions"
 	"github.com/openbao/openbao-plugins/auth/aws/pkcs7"
 	"github.com/openbao/openbao/sdk/v2/framework"
 	"github.com/openbao/openbao/sdk/v2/helper/cidrutil"
@@ -1846,12 +1847,12 @@ func (e *iamEntity) canonicalArn() string {
 func (b *backend) fullArn(ctx context.Context, e *iamEntity, s logical.Storage) (string, error) {
 	// Not assuming path is reliable for any entity types
 
-	region := b.partitionToRegionMap[e.Partition]
-	if region == nil {
+	region, ok := partitions.GetGlobalRegionForPartition(e.Partition)
+	if !ok {
 		return "", fmt.Errorf("unable to resolve partition %q to a region", e.Partition)
 	}
 
-	client, err := b.clientIAM(ctx, s, region.ID(), e.AccountNumber)
+	client, err := b.clientIAM(ctx, s, region, e.AccountNumber)
 	if err != nil {
 		return "", fmt.Errorf("error creating IAM client: %w", err)
 	}