feat: Use AWS version 2 to get the AWS credentials for AWS MSK

Author: Neurostep

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/aws/aws-sdk-go-v2 v1.26.1
 	github.com/aws/aws-sdk-go-v2/config v1.27.11
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.53.1
+	github.com/aws/aws-sdk-go-v2/service/sts v1.28.6
 	github.com/getsentry/sentry-go v0.12.0
 	github.com/go-kit/kit v0.9.0
 	github.com/google/uuid v1.6.0
@@ -61,7 +62,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sqs v1.24.4 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.20.5 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.4 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sts v1.28.6 // indirect
 	github.com/aws/smithy-go v1.20.2 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect

pkg/pubsub/kafka/config.go

@@ -10,6 +10,10 @@ import (
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/sts"
 	"github.com/twmb/franz-go/pkg/kgo"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	stsv2 "github.com/aws/aws-sdk-go-v2/service/sts"
+
 	awssasl "github.com/twmb/franz-go/pkg/sasl/aws"
 	"github.com/twmb/franz-go/pkg/sasl/plain"
 
@@ -24,10 +28,14 @@ type Config struct {
 	// Kafka configuration provided by go-sdk
 	KafkaConfig pubsub.Kafka
 	// AWS session reference, it will be used in case AWS MSK IAM authentication mechanism is used
+	//
+	// Deprecated: Use AwsConfig instead
 	AwsSession *session.Session
 	// MsgHandler is a function that will be called when a message is received
 	MsgHandler MsgHandler
-	Logger     sdklogger.Logger
+	// AWS configuration reference, it will be used in case AWS MSK IAM authentication mechanism is used
+	AwsConfig *aws.Config
+	Logger    sdklogger.Logger
 }
 
 const tlsConnectionTimeout = 10 * time.Second
@@ -43,7 +51,7 @@ func newConfig(c Config, opts ...kgo.Opt) ([]kgo.Opt, error) {
 		case pubsub.Plain:
 			options = append(options, getPlainSaslOption(c.KafkaConfig.SASL))
 		case pubsub.AWSMskIam:
-			options = append(options, getAwsMskIamSaslOption(c.KafkaConfig.SASL.AWSMskIam, c.AwsSession))
+			options = append(options, getAwsMskIamSaslOption(c.KafkaConfig.SASL.AWSMskIam, c.AwsSession, c.AwsConfig))
 		}
 	}
 
@@ -101,11 +109,11 @@ func getPlainSaslOption(saslConf pubsub.SASL) kgo.Opt {
 	}.AsMechanism())
 }
 
-func getAwsMskIamSaslOption(iamConf pubsub.SASLAwsMskIam, s *session.Session) kgo.Opt {
+func getAwsMskIamSaslOption(iamConf pubsub.SASLAwsMskIam, s *session.Session, awsCfg *aws.Config) kgo.Opt {
 	var opt kgo.Opt
 
-	// no AWS session provided
-	if s == nil {
+	// no AWS session and AWS config provided
+	if s == nil && awsCfg == nil {
 		opt = kgo.SASL(awssasl.Auth{
 			AccessKey:    iamConf.AccessKey,
 			SecretKey:    iamConf.SecretKey,
@@ -115,40 +123,85 @@ func getAwsMskIamSaslOption(iamConf pubsub.SASLAwsMskIam, s *session.Session) kg
 	} else {
 		opt = kgo.SASL(
 			awssasl.ManagedStreamingIAM(func(ctx context.Context) (awssasl.Auth, error) {
-				// If assumable role is not provided, we try to get credentials from the provided AWS session
-				if iamConf.AssumableRole == "" {
-					val, err := s.Config.Credentials.Get()
-					if err != nil {
-						return awssasl.Auth{}, err
-					}
-
-					return awssasl.Auth{
-						AccessKey:    val.AccessKeyID,
-						SecretKey:    val.SecretAccessKey,
-						SessionToken: val.SessionToken,
-						UserAgent:    iamConf.UserAgent,
-					}, nil
+				if s != nil {
+					return getAwsSaslAuthFromSession(iamConf, s)
 				}
 
-				svc := sts.New(s)
-
-				res, stsErr := svc.AssumeRole(&sts.AssumeRoleInput{
-					RoleArn:         &iamConf.AssumableRole,
-					RoleSessionName: &iamConf.SessionName,
-				})
-				if stsErr != nil {
-					return awssasl.Auth{}, stsErr
-				}
-
-				return awssasl.Auth{
-					AccessKey:    *res.Credentials.AccessKeyId,
-					SecretKey:    *res.Credentials.SecretAccessKey,
-					SessionToken: *res.Credentials.SessionToken,
-					UserAgent:    iamConf.UserAgent,
-				}, nil
+				return getAwsSaslAuthFromConfig(ctx, iamConf, awsCfg)
 			}),
 		)
 	}
 
 	return opt
 }
+
+func getAwsSaslAuthFromSession(iamConf pubsub.SASLAwsMskIam, s *session.Session) (awssasl.Auth, error) {
+	// If assumable role is not provided, we try to get credentials from the provided AWS session
+	if iamConf.AssumableRole == "" {
+		val, err := s.Config.Credentials.Get()
+		if err != nil {
+			return awssasl.Auth{}, err
+		}
+
+		return awssasl.Auth{
+			AccessKey:    val.AccessKeyID,
+			SecretKey:    val.SecretAccessKey,
+			SessionToken: val.SessionToken,
+			UserAgent:    iamConf.UserAgent,
+		}, nil
+	}
+
+	svc := sts.New(s)
+
+	res, stsErr := svc.AssumeRole(&sts.AssumeRoleInput{
+		RoleArn:         &iamConf.AssumableRole,
+		RoleSessionName: &iamConf.SessionName,
+	})
+	if stsErr != nil {
+		return awssasl.Auth{}, stsErr
+	}
+
+	return awssasl.Auth{
+		AccessKey:    *res.Credentials.AccessKeyId,
+		SecretKey:    *res.Credentials.SecretAccessKey,
+		SessionToken: *res.Credentials.SessionToken,
+		UserAgent:    iamConf.UserAgent,
+	}, nil
+}
+
+func getAwsSaslAuthFromConfig(
+	ctx context.Context,
+	iamConf pubsub.SASLAwsMskIam,
+	awsCfg *aws.Config) (awssasl.Auth, error) {
+	// If assumable role is not provided, we try to get credentials from the provided AWS config
+	if iamConf.AssumableRole == "" {
+		val, err := awsCfg.Credentials.Retrieve(ctx)
+		if err != nil {
+			return awssasl.Auth{}, err
+		}
+
+		return awssasl.Auth{
+			AccessKey:    val.AccessKeyID,
+			SecretKey:    val.SecretAccessKey,
+			SessionToken: val.SessionToken,
+			UserAgent:    iamConf.UserAgent,
+		}, nil
+	}
+
+	client := stsv2.NewFromConfig(*awsCfg)
+
+	res, stsErr := client.AssumeRole(ctx, &stsv2.AssumeRoleInput{
+		RoleArn:         &iamConf.AssumableRole,
+		RoleSessionName: &iamConf.SessionName,
+	})
+	if stsErr != nil {
+		return awssasl.Auth{}, stsErr
+	}
+
+	return awssasl.Auth{
+		AccessKey:    *res.Credentials.AccessKeyId,
+		SecretKey:    *res.Credentials.SecretAccessKey,
+		SessionToken: *res.Credentials.SessionToken,
+		UserAgent:    iamConf.UserAgent,
+	}, nil
+}