rest: Make Amazon AWS STS domain configurable for AWS rest plugin

Fixes: #5915

Signed-off-by: Johan Fylling <johan.dev@fylling.se>

docs/content/configuration.md

@@ -537,6 +537,7 @@ If specifying `web_identity_credentials`, OPA will expect to find environment va
 | --- | --- | --- | --- |
 | `services[_].credentials.s3_signing.web_identity_credentials.aws_region` | `string` | Yes | The AWS region to use for the sts regional endpoint. Uses the global endpoint by default |
 | `services[_].credentials.s3_signing.web_identity_credentials.session_name` | `string` | No | The session name used to identify the assumed role session. Default: `open-policy-agent` |
+| `services[_].credentials.s3_signing.web_identity_credentials.aws_domain` | `string` | No | The AWS domain name to use. Default: `amazonaws.com`. Can also be set via the `AWS_DOMAIN` environment variable (config takes precedence) |
 
 #### GCP Metadata Token
 

plugins/rest/aws.go

@@ -34,8 +34,9 @@ const (
 	ecsRelativePathEnvVar     = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"
 
 	// ref. https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp_enable-regions.html
-	stsDefaultPath = "https://sts.amazonaws.com"
-	stsRegionPath  = "https://sts.%s.amazonaws.com"
+	stsDefaultDomain = "amazonaws.com"
+	stsDefaultPath   = "https://sts.%s"
+	stsRegionPath    = "https://sts.%s.%s"
 
 	// ref. https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html
 	accessKeyEnvVar               = "AWS_ACCESS_KEY_ID"
@@ -43,6 +44,7 @@ const (
 	securityTokenEnvVar           = "AWS_SECURITY_TOKEN"
 	sessionTokenEnvVar            = "AWS_SESSION_TOKEN"
 	awsRegionEnvVar               = "AWS_REGION"
+	awsDomainEnvVar               = "AWS_DOMAIN"
 	awsRoleArnEnvVar              = "AWS_ROLE_ARN"
 	awsWebIdentityTokenFileEnvVar = "AWS_WEB_IDENTITY_TOKEN_FILE"
 	awsCredentialsFileEnvVar      = "AWS_SHARED_CREDENTIALS_FILE"
@@ -322,6 +324,7 @@ type awsWebIdentityCredentialService struct {
 	WebIdentityTokenFile string
 	RegionName           string `json:"aws_region"`
 	SessionName          string `json:"session_name"`
+	Domain               string `json:"aws_domain"`
 	stsURL               string
 	creds                aws.Credentials
 	expiration           time.Time
@@ -338,6 +341,10 @@ func (cs *awsWebIdentityCredentialService) populateFromEnv() error {
 		return errors.New("no " + awsWebIdentityTokenFileEnvVar + " set in environment")
 	}
 
+	if cs.Domain == "" {
+		cs.Domain = os.Getenv(awsDomainEnvVar)
+	}
+
 	if cs.RegionName == "" {
 		if cs.RegionName = os.Getenv(awsRegionEnvVar); cs.RegionName == "" {
 			return errors.New("no " + awsRegionEnvVar + " set in environment or configuration")
@@ -347,14 +354,21 @@ func (cs *awsWebIdentityCredentialService) populateFromEnv() error {
 }
 
 func (cs *awsWebIdentityCredentialService) stsPath() string {
+	var domain string
+	if cs.Domain != "" {
+		domain = strings.ToLower(cs.Domain)
+	} else {
+		domain = stsDefaultDomain
+	}
+
 	var stsPath string
 	switch {
 	case cs.stsURL != "":
 		stsPath = cs.stsURL
 	case cs.RegionName != "":
-		stsPath = fmt.Sprintf(stsRegionPath, strings.ToLower(cs.RegionName))
+		stsPath = fmt.Sprintf(stsRegionPath, strings.ToLower(cs.RegionName), domain)
 	default:
-		stsPath = stsDefaultPath
+		stsPath = fmt.Sprintf(stsDefaultPath, domain)
 	}
 	return stsPath
 }

plugins/rest/aws_test.go

@@ -1016,13 +1016,95 @@ func TestWebIdentityCredentialService(t *testing.T) {
 func TestStsPath(t *testing.T) {
 	cs := awsWebIdentityCredentialService{}
 
-	assertEq(cs.stsPath(), stsDefaultPath, t)
+	defaultPath := fmt.Sprintf(stsDefaultPath, stsDefaultDomain)
+	assertEq(defaultPath, cs.stsPath(), t)
 
 	cs.RegionName = "us-east-2"
-	assertEq(cs.stsPath(), "https://sts.us-east-2.amazonaws.com", t)
+	assertEq("https://sts.us-east-2.amazonaws.com", cs.stsPath(), t)
+
+	cs.Domain = "example.com"
+	assertEq("https://sts.us-east-2.example.com", cs.stsPath(), t)
 
 	cs.stsURL = "http://test.com"
-	assertEq(cs.stsPath(), "http://test.com", t)
+	assertEq("http://test.com", cs.stsPath(), t)
+}
+
+func TestStsPathFromEnv(t *testing.T) {
+	t.Setenv(awsRoleArnEnvVar, "role:arn")
+	t.Setenv(awsWebIdentityTokenFileEnvVar, "/nonsense")
+
+	tests := []struct {
+		note string
+		env  map[string]string
+		cs   awsWebIdentityCredentialService
+		want string
+	}{
+		{
+			note: "region set in config",
+			cs: awsWebIdentityCredentialService{
+				RegionName: "us-east-2",
+			},
+			want: "https://sts.us-east-2.amazonaws.com",
+		},
+		{
+			note: "region set in env",
+			env: map[string]string{
+				awsRegionEnvVar: "us-east-1",
+			},
+			want: "https://sts.us-east-1.amazonaws.com",
+		},
+		{
+			note: "region set in env and config (config wins)",
+			env: map[string]string{
+				awsRegionEnvVar: "us-east-1",
+			},
+			cs: awsWebIdentityCredentialService{
+				RegionName: "us-east-2",
+			},
+			want: "https://sts.us-east-2.amazonaws.com",
+		},
+		{
+			note: "domain set in config",
+			cs: awsWebIdentityCredentialService{
+				RegionName: "us-east-2",
+				Domain:     "foo.example.com",
+			},
+			want: "https://sts.us-east-2.foo.example.com",
+		},
+		{
+			note: "domain set in env",
+			env: map[string]string{
+				awsDomainEnvVar: "bar.example.com",
+			},
+			cs: awsWebIdentityCredentialService{
+				RegionName: "us-east-2", // Region must always be set
+			},
+			want: "https://sts.us-east-2.bar.example.com",
+		},
+		{
+			note: "domain set in env and config (config wins)",
+			env: map[string]string{
+				awsDomainEnvVar: "bar.example.com",
+			},
+			cs: awsWebIdentityCredentialService{
+				RegionName: "us-east-2", // Region must always be set
+				Domain:     "foo.example.com",
+			},
+			want: "https://sts.us-east-2.foo.example.com",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			for k, v := range tc.env {
+				t.Setenv(k, v)
+			}
+			if err := tc.cs.populateFromEnv(); err != nil {
+				t.Fatalf("Unexpected err: %s", err)
+			}
+			assertEq(tc.want, tc.cs.stsPath(), t)
+		})
+	}
 }
 
 // simulate EC2 metadata service