Fix: use alternate STS endpoint for STS interaction if given (#214)

Author: njvrzm

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 All notable changes to this project will be documented in this file.
 
+## 0.33.1
+- Fix: use alternate STS endpoint for STS interaction if given by @njvrzm in https://github.com/grafana/grafana-aws-sdk/pull/214
+
 ## 0.33.0
 
 - Update CodeBuild metrics and dimensions by @hectorruiz-it in https://github.com/grafana/grafana-aws-sdk/pull/209

pkg/awsauth/auth.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"strings"
 )
 
 type ConfigProvider interface {
@@ -73,3 +74,7 @@ func (rcp *awsConfigProvider) GetConfig(ctx context.Context, authSettings Settin
 	rcp.cache[key] = cfg
 	return cfg, nil
 }
+
+func isStsEndpoint(ep *string) bool {
+	return ep != nil && (strings.HasPrefix(*ep, "sts.") || strings.HasPrefix(*ep, "sts-fips."))
+}

pkg/awsauth/auth_test.go

@@ -52,6 +52,10 @@ func (tc testCase) Run(t *testing.T) {
 		assert.Equal(t, accessKey, creds.AccessKeyID)
 		assert.Equal(t, secret, creds.SecretAccessKey)
 	}
+	if isStsEndpoint(&tc.authSettings.Endpoint) {
+		assert.Equal(t, tc.authSettings.Endpoint, *client.assumeRoleClient.stsConfig.BaseEndpoint)
+		assert.Nil(t, cfg.BaseEndpoint)
+	}
 }
 
 func (tc testCase) assertConfig(t *testing.T, cfg aws.Config) {
@@ -134,6 +138,15 @@ func TestGetAWSConfig_Keys(t *testing.T) {
 				Region:         "ap-south-1",
 			},
 		},
+		{
+			name: "static credentials, sts endpoint",
+			authSettings: Settings{
+				LegacyAuthType: awsds.AuthTypeKeys,
+				AccessKey:      "ubiquitous",
+				SecretKey:      "malevolent",
+				Region:         "ap-south-1",
+			},
+		},
 	}.runAll(t)
 }
 
@@ -155,6 +168,23 @@ func TestGetAWSConfig_Keys_AssumeRule(t *testing.T) {
 				Expiration:      aws.Time(time.Now().Add(time.Hour)),
 			},
 		},
+		{
+			name: "static assume role with sts endpoint - endpoint is nil",
+			authSettings: Settings{
+				AuthType:      AuthTypeKeys,
+				AccessKey:     "tensile",
+				SecretKey:     "diaphanous",
+				Region:        "us-east-1",
+				Endpoint:      "sts.us-east-1.amazonaws.com",
+				AssumeRoleARN: "arn:aws:iam::1234567890:role/aws-service-role",
+			},
+			assumedCredentials: &ststypes.Credentials{
+				AccessKeyId:     aws.String("assumed"),
+				SecretAccessKey: aws.String("role"),
+				SessionToken:    aws.String("session"),
+				Expiration:      aws.Time(time.Now().Add(time.Hour)),
+			},
+		},
 		{
 			name: "static assume role with failure",
 			authSettings: Settings{
@@ -227,25 +257,26 @@ func TestGetAWSConfig_Shared(t *testing.T) {
 func TestGetAWSConfig_UnknownOrMissing(t *testing.T) {
 	testSuite{
 		{
-			name: "shared reads from specified file",
+			name: "unknown auth type fails",
 			authSettings: Settings{
 				AuthType: AuthTypeUnknown,
 			},
 			shouldError: true,
 		},
 		{
-			name: "grafana assume role uses the shared mechanism",
+			name: "random auth type fails",
 			authSettings: Settings{
-				AuthType: AuthTypeMissing,
+				AuthType: "rainbows",
 			},
 			shouldError: true,
 		},
 		{
-			name: "grafana assume role uses the shared mechanism",
-			authSettings: Settings{
-				AuthType: "rainbows",
+			name:         "missing auth type fails back to legacy default (and does not fail)",
+			authSettings: Settings{},
+			environment: map[string]string{
+				"AWS_SHARED_CREDENTIALS_FILE": testDataPath("credentials"),
 			},
-			shouldError: true,
+			shouldError: false,
 		},
 	}.runAll(t)
 }

pkg/awsauth/settings.go

@@ -127,6 +127,9 @@ func (s Settings) WithAssumeRole(cfg aws.Config, client AWSAPIClient) LoadOption
 	cache := client.NewCredentialsCache(provider)
 	return func(options *config.LoadOptions) error {
 		options.Credentials = cache
+		if isStsEndpoint(cfg.BaseEndpoint) {
+			options.BaseEndpoint = ""
+		}
 		return nil
 	}
 }

pkg/awsauth/test_utils.go

@@ -35,7 +35,8 @@ func (m *mockAWSAPIClient) NewStaticCredentialsProvider(key, secret, session str
 	return credentials.NewStaticCredentialsProvider(key, secret, session)
 }
 
-func (m *mockAWSAPIClient) NewSTSClientFromConfig(_ aws.Config) stscreds.AssumeRoleAPIClient {
+func (m *mockAWSAPIClient) NewSTSClientFromConfig(cfg aws.Config) stscreds.AssumeRoleAPIClient {
+	m.assumeRoleClient.stsConfig = cfg
 	return m.assumeRoleClient
 }
 
@@ -54,6 +55,7 @@ func (m *mockAWSAPIClient) NewEC2RoleCreds() aws.CredentialsProvider {
 
 type mockAssumeRoleAPIClient struct {
 	mock.Mock
+	stsConfig aws.Config
 }
 
 func (m *mockAssumeRoleAPIClient) AssumeRole(_ context.Context, params *sts.AssumeRoleInput, _ ...func(*sts.Options)) (*sts.AssumeRoleOutput, error) {