feat: add support for web identity in auth token generation

signer/msk_auth_token_provider.go

@@ -6,7 +6,6 @@ import (
 	"encoding/base64"
 	"encoding/hex"
 	"fmt"
-
 	"log"
 	"net/http"
 	"net/url"
@@ -77,6 +76,23 @@ func GenerateAuthTokenFromRole(
 	return constructAuthToken(ctx, region, credentials)
 }
 
+// GenerateAuthTokenFromWebIdentity generates base64 encoded signed url as auth token by loading IAM credentials from a web identity role Arn.
+func GenerateAuthTokenFromWebIdentity(
+	ctx context.Context, region string, roleArn string, webIdentityToken string, stsSessionName string,
+) (string, int64, error) {
+	if stsSessionName == "" {
+		stsSessionName = DefaultSessionName
+	}
+
+	credentials, err := loadCredentialsFromWebIdentityParameters(ctx, region, roleArn, webIdentityToken, stsSessionName)
+
+	if err != nil {
+		return "", 0, fmt.Errorf("failed to load credentials: %w", err)
+	}
+
+	return constructAuthToken(ctx, region, credentials)
+}
+
 // GenerateAuthTokenFromCredentialsProvider generates base64 encoded signed url as auth token by loading IAM credentials
 // from an aws credentials provider
 func GenerateAuthTokenFromCredentialsProvider(
@@ -150,6 +166,41 @@ func loadCredentialsFromRoleArn(
 	return &creds, nil
 }
 
+// Loads credentials from a named by assuming the passed web identity role and id token.
+// This implementation creates a new sts client for every call to get or refresh token. In order to avoid this, please
+// use your own credentials' provider.
+// If you wish to use regional endpoint, please pass your own credentials' provider.
+func loadCredentialsFromWebIdentityParameters(
+	ctx context.Context, region, roleArn, webIdentityToken, stsSessionName string,
+) (*aws.Credentials, error) {
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
+
+	if err != nil {
+		return nil, fmt.Errorf("unable to load SDK config: %w", err)
+	}
+
+	stsClient := sts.NewFromConfig(cfg)
+
+	assumeRoleWithWebIdentityInput := &sts.AssumeRoleWithWebIdentityInput{
+		RoleArn:          aws.String(roleArn),
+		RoleSessionName:  aws.String(stsSessionName),
+		WebIdentityToken: aws.String(webIdentityToken),
+	}
+	assumeRoleWithWebIdentityOutput, err := stsClient.AssumeRoleWithWebIdentity(ctx, assumeRoleWithWebIdentityInput)
+	if err != nil {
+		return nil, fmt.Errorf("unable to assume role with web identity, %s: %w", roleArn, err)
+	}
+
+	//Create new aws.Credentials instance using the credentials from AssumeRoleWithWebIdentityOutput.Credentials
+	creds := aws.Credentials{
+		AccessKeyID:     *assumeRoleWithWebIdentityOutput.Credentials.AccessKeyId,
+		SecretAccessKey: *assumeRoleWithWebIdentityOutput.Credentials.SecretAccessKey,
+		SessionToken:    *assumeRoleWithWebIdentityOutput.Credentials.SessionToken,
+	}
+
+	return &creds, nil
+}
+
 // Loads credentials from the credentials provider
 func loadCredentialsFromCredentialsProvider(
 	ctx context.Context, credentialsProvider aws.CredentialsProvider,