Add support for aws-sdk-go-v2 authentication (#202)

Author: njvrzm

go.mod

@@ -7,6 +7,10 @@ toolchain go1.23.6
 require (
 	github.com/aws/aws-sdk-go v1.55.6
 	github.com/aws/aws-sdk-go-v2 v1.36.1
+	github.com/aws/aws-sdk-go-v2/config v1.29.4
+	github.com/aws/aws-sdk-go-v2/credentials v1.17.57
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.27
+	github.com/aws/aws-sdk-go-v2/service/sts v1.33.12
 	github.com/google/go-cmp v0.6.0
 	github.com/grafana/grafana-plugin-sdk-go v0.266.0
 	github.com/grafana/sqlds/v4 v4.1.7
@@ -18,6 +22,13 @@ require (
 require (
 	github.com/BurntSushi/toml v1.4.0 // indirect
 	github.com/apache/arrow-go/v18 v18.0.1-0.20241212180703-82be143d7c30 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.31 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.31 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.12 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.24.14 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.13 // indirect
 	github.com/aws/smithy-go v1.22.2 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect

go.sum

@@ -13,6 +13,28 @@ github.com/aws/aws-sdk-go v1.55.6 h1:cSg4pvZ3m8dgYcgqB97MrcdjUmZ1BeMYKUxMMB89IPk
 github.com/aws/aws-sdk-go v1.55.6/go.mod h1:eRwEWoyTWFMVYVQzKMNHWP5/RV4xIUGMQfXQHfHkpNU=
 github.com/aws/aws-sdk-go-v2 v1.36.1 h1:iTDl5U6oAhkNPba0e1t1hrwAo02ZMqbrGq4k5JBWM5E=
 github.com/aws/aws-sdk-go-v2 v1.36.1/go.mod h1:5PMILGVKiW32oDzjj6RU52yrNrDPUHcbZQYr1sM7qmM=
+github.com/aws/aws-sdk-go-v2/config v1.29.4 h1:ObNqKsDYFGr2WxnoXKOhCvTlf3HhwtoGgc+KmZ4H5yg=
+github.com/aws/aws-sdk-go-v2/config v1.29.4/go.mod h1:j2/AF7j/qxVmsNIChw1tWfsVKOayJoGRDjg1Tgq7NPk=
+github.com/aws/aws-sdk-go-v2/credentials v1.17.57 h1:kFQDsbdBAR3GZsB8xA+51ptEnq9TIj3tS4MuP5b+TcQ=
+github.com/aws/aws-sdk-go-v2/credentials v1.17.57/go.mod h1:2kerxPUUbTagAr/kkaHiqvj/bcYHzi2qiJS/ZinllU0=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.27 h1:7lOW8NUwE9UZekS1DYoiPdVAqZ6A+LheHWb+mHbNOq8=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.27/go.mod h1:w1BASFIPOPUae7AgaH4SbjNbfdkxuggLyGfNFTn8ITY=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.31 h1:lWm9ucLSRFiI4dQQafLrEOmEDGry3Swrz0BIRdiHJqQ=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.31/go.mod h1:Huu6GG0YTfbPphQkDSo4dEGmQRTKb9k9G7RdtyQWxuI=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.31 h1:ACxDklUKKXb48+eg5ROZXi1vDgfMyfIA/WyvqHcHI0o=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.31/go.mod h1:yadnfsDwqXeVaohbGc/RaD287PuyRw2wugkh5ZL2J6k=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.2 h1:Pg9URiobXy85kgFev3og2CuOZ8JZUBENF+dcgWBaYNk=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.2/go.mod h1:FbtygfRFze9usAadmnGJNc8KsP346kEe+y2/oyhGAGc=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.2 h1:D4oz8/CzT9bAEYtVhSBmFj2dNOtaHOtMKc2vHBwYizA=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.2/go.mod h1:Za3IHqTQ+yNcRHxu1OFucBh0ACZT4j4VQFF0BqpZcLY=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.12 h1:O+8vD2rGjfihBewr5bT+QUfYUHIxCVgG61LHoT59shM=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.12/go.mod h1:usVdWJaosa66NMvmCrr08NcWDBRv4E6+YFG2pUdw1Lk=
+github.com/aws/aws-sdk-go-v2/service/sso v1.24.14 h1:c5WJ3iHz7rLIgArznb3JCSQT3uUMiz9DLZhIX+1G8ok=
+github.com/aws/aws-sdk-go-v2/service/sso v1.24.14/go.mod h1:+JJQTxB6N4niArC14YNtxcQtwEqzS3o9Z32n7q33Rfs=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.13 h1:f1L/JtUkVODD+k1+IiSJUUv8A++2qVr+Xvb3xWXETMU=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.13/go.mod h1:tvqlFoja8/s0o+UruA1Nrezo/df0PzdunMDDurUfg6U=
+github.com/aws/aws-sdk-go-v2/service/sts v1.33.12 h1:fqg6c1KVrc3SYWma/egWue5rKI4G2+M4wMQN2JosNAA=
+github.com/aws/aws-sdk-go-v2/service/sts v1.33.12/go.mod h1:7Yn+p66q/jt38qMoVfNvjbm3D89mGBnkwDcijgtih8w=
 github.com/aws/smithy-go v1.22.2 h1:6D9hW43xKFrRx/tXXfAlIZc4JI+yQe6snnWcQyxSyLQ=
 github.com/aws/smithy-go v1.22.2/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=

pkg/awsauth/api_client.go

@@ -0,0 +1,46 @@
+package awsauth
+
+import (
+	"context"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/credentials"
+	"github.com/aws/aws-sdk-go-v2/credentials/ec2rolecreds"
+	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
+)
+
+type LoadOptionsFunc = func(*config.LoadOptions) error
+
+// AWSAPIClient isolates most of our interactions with the AWS SDK to make it easier to mock in tests
+type AWSAPIClient interface {
+	LoadDefaultConfig(ctx context.Context, options ...LoadOptionsFunc) (aws.Config, error)
+	NewStaticCredentialsProvider(key, secret, session string) aws.CredentialsProvider
+	NewSTSClientFromConfig(cfg aws.Config) stscreds.AssumeRoleAPIClient
+	NewAssumeRoleProvider(client stscreds.AssumeRoleAPIClient, roleARN string, optFns ...func(*stscreds.AssumeRoleOptions)) aws.CredentialsProvider
+	NewCredentialsCache(provider aws.CredentialsProvider, optFns ...func(options *aws.CredentialsCacheOptions)) aws.CredentialsProvider
+	NewEC2RoleCreds() aws.CredentialsProvider
+}
+
+type awsAPIClient struct{}
+
+func (c awsAPIClient) NewStaticCredentialsProvider(key, secret, session string) aws.CredentialsProvider {
+	return credentials.NewStaticCredentialsProvider(key, secret, session)
+}
+func (c awsAPIClient) LoadDefaultConfig(ctx context.Context, options ...LoadOptionsFunc) (aws.Config, error) {
+	return config.LoadDefaultConfig(ctx, options...)
+}
+func (c awsAPIClient) NewSTSClientFromConfig(cfg aws.Config) stscreds.AssumeRoleAPIClient {
+	return sts.NewFromConfig(cfg)
+}
+
+func (c awsAPIClient) NewAssumeRoleProvider(client stscreds.AssumeRoleAPIClient, roleARN string, optFns ...func(*stscreds.AssumeRoleOptions)) aws.CredentialsProvider {
+	return stscreds.NewAssumeRoleProvider(client, roleARN, optFns...)
+}
+
+func (c awsAPIClient) NewCredentialsCache(provider aws.CredentialsProvider, optFns ...func(options *aws.CredentialsCacheOptions)) aws.CredentialsProvider {
+	return aws.NewCredentialsCache(provider, optFns...)
+}
+func (c awsAPIClient) NewEC2RoleCreds() aws.CredentialsProvider {
+	return ec2rolecreds.New()
+}

pkg/awsauth/auth.go

@@ -0,0 +1,75 @@
+package awsauth
+
+import (
+	"context"
+	"fmt"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+type ConfigProvider interface {
+	GetConfig(context.Context, Settings) (aws.Config, error)
+}
+
+func NewConfigProvider() ConfigProvider {
+	return newAWSConfigProviderWithClient(awsAPIClient{})
+}
+
+func newAWSConfigProviderWithClient(client AWSAPIClient) *awsConfigProvider {
+	return &awsConfigProvider{client, make(map[uint64]aws.Config)}
+}
+
+type awsConfigProvider struct {
+	client AWSAPIClient
+	cache  map[uint64]aws.Config
+}
+
+func (rcp *awsConfigProvider) GetConfig(ctx context.Context, authSettings Settings) (aws.Config, error) {
+	logger := backend.Logger.FromContext(ctx)
+
+	key := authSettings.Hash()
+	cached, exists := rcp.cache[key]
+	if exists {
+		logger.Debug("returning config from cache")
+		return cached, nil
+	}
+	logger.Debug("creating new config")
+
+	options := authSettings.BaseOptions()
+
+	authType := authSettings.GetAuthType()
+	logger.Debug(fmt.Sprintf("Using auth type: %s", authType))
+	switch authType {
+	case AuthTypeDefault: // nothing else to do here
+	case AuthTypeKeys:
+		options = append(options, authSettings.WithStaticCredentials(rcp.client))
+	case AuthTypeSharedCreds, AuthTypeGrafanaAssumeRole:
+		options = append(options, authSettings.WithSharedCredentials())
+	case AuthTypeEC2IAMRole:
+		// TODO: test this
+		options = append(options, authSettings.WithEC2RoleCredentials(rcp.client))
+	default:
+		return aws.Config{}, fmt.Errorf("unknown auth type: %s", authType)
+	}
+
+	cfg, err := rcp.client.LoadDefaultConfig(ctx, options...)
+	if err != nil {
+		return aws.Config{}, err
+	}
+
+	if authSettings.AssumeRoleARN != "" {
+		options = append(authSettings.BaseOptions(), authSettings.WithAssumeRole(cfg, rcp.client))
+		cfg, err = rcp.client.LoadDefaultConfig(ctx, options...)
+		if err != nil {
+			return aws.Config{}, err
+		}
+	}
+
+	_, err = cfg.Credentials.Retrieve(ctx)
+	if err != nil {
+		return aws.Config{}, fmt.Errorf("error retrieving credentials: %w", err)
+	}
+
+	rcp.cache[key] = cfg
+	return cfg, nil
+}