How to use AWS WebIdentity to assume a role using GCP Workload Idenity in Go

aws_helpers.go

package awshelpers

import (
	"context"
	"errors"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/idtoken"
	"google.golang.org/api/impersonate"
	"google.golang.org/api/option"
)

var _ stscreds.IdentityTokenRetriever = &tokenRetriever{}

// tokenRetriever exists to satisfy the stscreds.IdentityTokenRetriever interface
type tokenRetriever struct {
	tokenSource oauth2.TokenSource
}

func newTokenRetriever(tokenSource oauth2.TokenSource) *tokenRetriever {
	return &tokenRetriever{
		tokenSource: tokenSource,
	}
}

// GetIdentityToken returns the identity token from the token source
func (tr tokenRetriever) GetIdentityToken() ([]byte, error) {
	t, err := tr.tokenSource.Token()
	if err != nil {
		return nil, err
	}

	return []byte(t.AccessToken), nil
}

type config struct {
	tokenSource oauth2.TokenSource
}

// ConfigOption is a function that modifies a config
type ConfigOption func(context.Context, *config) error

// NewConfigWithAssumedRole creates a new AWS config. First it assumes a role with a web identity token,
// then it creates a new AWS config with the assumed role. The returned config has a credentials provider
// that will refresh the GCP and AWS assumed role's credentials when they expire.
//
// The JWT sent to AWS is provided by Google with minimal configuration options. It is valid for 1 hour,
// not configurable. The only value that can be set is the audience, which can be any string. We have used
// the name of the service account that created the JWT in the past, but it can be anything.
// Here's an example of what the JWT claims look like:
//
//	{
//	  "aud": "[email protected]",
//	  "azp": "12345678910",
//	  "exp": 1711984401,
//	  "iat": 1711980801,
//	  "iss": "https://accounts.google.com",
//	  "sub": "12345678910"
//	}
//
// There's a surprise around the audience mapping once that JWT is translated into an AWS token to write
// policy against. For some reason, unless the audience is set to the same value as the sub, AWS will
// rewrite the 'azp' claim to the 'aud' claim, and the 'aud' claim to the 'oaud' claim.
// https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_iam-condition-keys.html
//
// Here's an example of what the trust policy needs to look like on the attached role:
//
//	{
//	    "Version": "2012-10-17",
//	    "Statement": [
//	        {
//	            "Effect": "Allow",
//	            "Principal": {
//	                "Federated": "accounts.google.com"
//	            },
//	            "Action": "sts:AssumeRoleWithWebIdentity",
//	            "Condition": {
//	                "StringEquals": {
//	                    "accounts.google.com:sub": "12345678910"
//	                    "accounts.google.com:aud": "12345678910",
//	                    "accounts.google.com:oaud": "[email protected]",
//	                }
//	            }
//	        }
//	    ]
//	}
func NewConfigWithAssumedRole(c context.Context, role, region string, opts ...ConfigOption) (aws.Config, error) {
	cfg := &config{}
	for _, opt := range opts {
		if err := opt(c, cfg); err != nil {
			return aws.Config{}, err
		}
	}

	// one of the ConfigOptions must have set the token source
	if cfg.tokenSource == nil {
		return aws.Config{}, errors.New("token source is required")
	}

	// we have to create a separate AWS config to assume the role with the web identity token. This is
	// an intermediary step, so we can create the credentials provider to pass into the final AWS config.
	stsCfg, err := awsconfig.LoadDefaultConfig(c, awsconfig.WithRegion(region))
	if err != nil {
		return aws.Config{}, err
	}

	// with the token source, create the AWS config to assume that role
	awsCfg, err := awsconfig.LoadDefaultConfig(c,
		awsconfig.WithRegion(region),
		awsconfig.WithCredentialsProvider(stscreds.NewWebIdentityRoleProvider(
			sts.NewFromConfig(stsCfg),
			role,
			newTokenRetriever(cfg.tokenSource),
		)),
	)
	if err != nil {
		return aws.Config{}, err
	}

	return awsCfg, nil
}

// WithImpersonatedTokenSource uses GCP Application Default Credentials, generally using workload identity
// when running in GCP, to impersonate the target principal and generate an ID token for the audience.
// This is useful when you need to assume an AWS role using a GCP service account.
// The base account must have roles/iam.serviceAccountTokenCreator granted on the target service account.
func WithImpersonatedTokenSource(targetPrincipal, audience string) ConfigOption {
	return func(c context.Context, cfg *config) error {
		tokenSource, err := impersonate.IDTokenSource(c, impersonate.IDTokenConfig{
			TargetPrincipal: targetPrincipal,
			Audience:        audience,
		})
		if err != nil {
			return err
		}

		cfg.tokenSource = tokenSource

		return nil
	}
}

// WithDefaultTokenSource uses GCP Application Default Credentials to generate a token source for the audience.
// Scopes default to https://www.googleapis.com/auth/cloud-platform
//
// This will not work locally with the gcloud default user account. You must use a service account.
func WithDefaultTokenSource(audience string) ConfigOption {
	return func(c context.Context, cfg *config) error {
		creds, err := google.FindDefaultCredentials(c, "https://www.googleapis.com/auth/cloud-platform")
		if err != nil {
			return err
		}

		cfg.tokenSource, err = idtoken.NewTokenSource(c, audience, option.WithCredentials(creds))
		if err != nil {
			return err
		}

		return nil
	}
}

main.go

package main

func main() {
 	awsCfg, err := awshelpers.NewConfigWithAssumedRole(c, "arn:aws:iam::your-role", "us-west-2",
		awshelpers.WithImpersonatedTokenSource("[email protected]", "audience"),
	)

	// Create an Amazon S3 service client
	s3Client := s3.NewFromConfig(awsCfg)
  
    ...
}

This was surprisingly difficult to find any documentation on, but turned out to be fairly simple in practice. Let me know if there is anything that can be improved, but it has been tested to work.