Repository: snowflakedb/gosnowflake
Pull request #1419 (2025-05-21 14:48:26)

SNOW-1902248: AWS WIF support
### Description

SNOW-1902248 Support for Workload Identity Federation for AWS

### Checklist
- [x] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary



Patch:
diff --git a/Jenkinsfile b/Jenkinsfile
index f5d62f0..b8963eb 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -53,6 +53,7 @@ pipeline {
     // remove SCM URL + .git at the end
 
     BASELINE_BRANCH = "${env.CHANGE_TARGET}"
+    SF_ENABLE_EXPERIMENTAL_AUTHENTICATION = true
   }
   stages {
     stage('Checkout') {
diff --git a/auth.go b/auth.go
index 79818d6..9b32f65 100644
--- a/auth.go
+++ b/auth.go
@@ -11,6 +11,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"os"
 	"runtime"
 	"slices"
 	"strconv"
@@ -54,6 +55,8 @@ const (
 	AuthTypeOAuthAuthorizationCode
 	// AuthTypeOAuthClientCredentials is to use non-interactive OAuth2 flow
 	AuthTypeOAuthClientCredentials
+	// AuthTypeWorkloadIdentityFederation is to use CSP identity for authentication
+	AuthTypeWorkloadIdentityFederation
 )
 
 func (authType AuthType) isOauthNativeFlow() bool {
@@ -96,6 +99,9 @@ func determineAuthenticatorType(cfg *Config, value string) error {
 	} else if upperCaseValue == AuthTypeOAuthClientCredentials.String() {
 		cfg.Authenticator = AuthTypeOAuthClientCredentials
 		return nil
+	} else if upperCaseValue == AuthTypeWorkloadIdentityFederation.String() {
+		cfg.Authenticator = AuthTypeWorkloadIdentityFederation
+		return nil
 	} else {
 		// possibly Okta case
 		oktaURLString, err := url.QueryUnescape(lowerCaseValue)
@@ -151,6 +157,8 @@ func (authType AuthType) String() string {
 		return "OAUTH_AUTHORIZATION_CODE"
 	case AuthTypeOAuthClientCredentials:
 		return "OAUTH_CLIENT_CREDENTIALS"
+	case AuthTypeWorkloadIdentityFederation:
+		return "WORKLOAD_IDENTITY"
 	default:
 		return "UNKNOWN"
 	}
@@ -196,6 +204,7 @@ type authRequestData struct {
 	ProofKey                string                       `json:"PROOF_KEY,omitempty"`
 	Token                   string                       `json:"TOKEN,omitempty"`
 	OauthType               string                       `json:"OAUTH_TYPE,omitempty"`
+	Provider                string                       `json:"PROVIDER,omitempty"`
 }
 type authRequest struct {
 	Data authRequestData `json:"data"`
@@ -529,6 +538,19 @@ func createRequestBody(sc *snowflakeConn, sessionParameters map[string]interface
 		requestMain.LoginName = sc.cfg.User
 		requestMain.Token = token
 		requestMain.OauthType = "OAUTH_CLIENT_CREDENTIALS"
+	case AuthTypeWorkloadIdentityFederation:
+		if !experimentalAuthEnabled() {
+			return nil, errors.New("workload identity authentication is not ready to use")
+		}
+		logger.WithContext(sc.ctx).Debug("Workload Identity Federation")
+		wifAttestationProvider := createWifAttestationProvider(sc.ctx)
+		wifAttestation, err := wifAttestationProvider.getAttestation(sc.cfg.WorkloadIdentityProvider)
+		if err != nil {
+			return nil, err
+		}
+		requestMain.Authenticator = AuthTypeWorkloadIdentityFederation.String()
+		requestMain.Token = wifAttestation.Credential
+		requestMain.Provider = wifAttestation.ProviderType
 	}
 
 	authRequest := authRequest{
@@ -541,6 +563,11 @@ func createRequestBody(sc *snowflakeConn, sessionParameters map[string]interface
 	return jsonBody, nil
 }
 
+func experimentalAuthEnabled() bool {
+	val, ok := os.LookupEnv("SF_ENABLE_EXPERIMENTAL_AUTHENTICATION")
+	return ok && strings.EqualFold(val, "true")
+}
+
 // Generate a JWT token in string given the configuration
 func prepareJWTToken(config *Config) (string, error) {
 	if config.PrivateKey == nil {
diff --git a/auth_wif.go b/auth_wif.go
new file mode 100644
index 0000000..4ee54f8
--- /dev/null
+++ b/auth_wif.go
@@ -0,0 +1,250 @@
+package gosnowflake
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/sts"
+	"net/http"
+	"strings"
+	"time"
+)
+
+const (
+	awsWif   wifProviderType = "AWS"
+	gcpWif   wifProviderType = "GCP"
+	azureWif wifProviderType = "AZURE"
+	oidcWif  wifProviderType = "OIDC"
+)
+
+type wifProviderType string
+
+type wifAttestation struct {
+	ProviderType string            `json:"providerType"`
+	Credential   string            `json:"credential"`
+	Metadata     map[string]string `json:"metadata"`
+}
+
+type wifAttestationCreator interface {
+	createAttestation(ctx context.Context) (*wifAttestation, error)
+}
+
+type wifAttestationProvider struct {
+	context      context.Context
+	awsCreator   wifAttestationCreator
+	gcpCreator   wifAttestationCreator
+	azureCreator wifAttestationCreator
+	oidcCreator  wifAttestationCreator
+}
+
+func createWifAttestationProvider(ctx context.Context) *wifAttestationProvider {
+	return &wifAttestationProvider{
+		context:      ctx,
+		awsCreator:   &awsIdentityAttestationCreator{attestationService: createDefaultAwsAttestationService(ctx)},
+		gcpCreator:   nil,
+		azureCreator: nil,
+		oidcCreator:  nil,
+	}
+}
+
+func (p *wifAttestationProvider) getAttestation(identityProvider string) (*wifAttestation, error) {
+	if strings.TrimSpace(identityProvider) == "" {
+		logger.Info("Workload Identity Provider has not been specified. Using autodetect...")
+		return p.createAutodetectAttestation()
+	}
+	creator, err := p.attestationCreator(identityProvider)
+	if err != nil {
+		logger.Errorf("error while creating specified Workload Identity provider %v", err)
+		return nil, err
+	}
+	return creator.createAttestation(p.context)
+}
+
+func (p *wifAttestationProvider) attestationCreator(identityProvider string) (wifAttestationCreator, error) {
+	switch strings.ToUpper(identityProvider) {
+	case string(awsWif):
+		return p.awsCreator, nil
+	case string(gcpWif):
+		return p.gcpCreator, nil
+	case string(azureWif):
+		return p.azureCreator, nil
+	case string(oidcWif):
+		return p.oidcCreator, nil
+	default:
+		return nil, errors.New("unknown Workload Identity provider specified: " + identityProvider)
+	}
+}
+
+func (p *wifAttestationProvider) createAutodetectAttestation() (*wifAttestation, error) {
+	if attestation := p.getAttestationForAutodetect(p.oidcCreator, oidcWif); attestation != nil {
+		return attestation, nil
+	}
+	if attestation := p.getAttestationForAutodetect(p.awsCreator, awsWif); attestation != nil {
+		return attestation, nil
+	}
+	if attestation := p.getAttestationForAutodetect(p.gcpCreator, gcpWif); attestation != nil {
+		return attestation, nil
+	}
+	if attestation := p.getAttestationForAutodetect(p.azureCreator, azureWif); attestation != nil {
+		return attestation, nil
+	}
+	return nil, errors.New("unable to autodetect Workload Identity. None of the supported Workload Identity environments has been identified")
+}
+
+func (p *wifAttestationProvider) getAttestationForAutodetect(
+	creator wifAttestationCreator,
+	providerType wifProviderType,
+) *wifAttestation {
+	attestation, err := creator.createAttestation(p.context)
+	if err != nil {
+		logger.Errorf("Unable to create identity attestation for %s, error: %v", providerType, err)
+		return nil
+	}
+	return attestation
+}
+
+type awsIdentityAttestationCreator struct {
+	attestationService awsAttestationService
+}
+
+type awsAttestationService interface {
+	GetAWSCredentials() aws.Credentials
+	GetAWSRegion() string
+	GetArn() string
+}
+
+type defaultAwsAttestationService struct {
+	ctx context.Context
+	cfg aws.Config
+}
+
+func createDefaultAwsAttestationService(ctx context.Context) *defaultAwsAttestationService {
+	cfg, err := config.LoadDefaultConfig(ctx)
+	if err != nil {
+		logger.Debugf("Unable to load AWS config: %v", err)
+		return nil
+	}
+	return &defaultAwsAttestationService{
+		cfg: cfg,
+		ctx: ctx,
+	}
+}
+
+func (m *defaultAwsAttestationService) GetAWSCredentials() aws.Credentials {
+	creds, err := m.cfg.Credentials.Retrieve(m.ctx)
+	if err != nil {
+		logger.Debugf("Unable to retrieve AWS credentials provider: %v", err)
+		return aws.Credentials{}
+	}
+	return creds
+}
+
+func (m *defaultAwsAttestationService) GetAWSRegion() string {
+	return m.cfg.Region
+}
+
+func (m *defaultAwsAttestationService) GetArn() string {
+	client := sts.NewFromConfig(m.cfg)
+	output, err := client.GetCallerIdentity(m.ctx, &sts.GetCallerIdentityInput{})
+	if err != nil {
+		logger.Debugf("Unable to get caller identity: %v", err)
+		return ""
+	}
+	return aws.ToString(output.Arn)
+}
+
+func (creator *awsIdentityAttestationCreator) createAttestation(ctx context.Context) (*wifAttestation, error) {
+	logger.Debug("Creating AWS identity attestation...")
+
+	if creator.attestationService == nil {
+		logger.Debug("AWS attestation service could not be created.")
+		return nil, nil
+	}
+
+	creds := creator.attestationService.GetAWSCredentials()
+	if creds.AccessKeyID == "" || creds.SecretAccessKey == "" {
+		logger.Debug("No AWS credentials were found.")
+		return nil, nil
+	}
+
+	region := creator.attestationService.GetAWSRegion()
+	if region == "" {
+		logger.Debug("No AWS region was found.")
+		return nil, nil
+	}
+
+	arn := creator.attestationService.GetArn()
+	if arn == "" {
+		logger.Debug("No Caller Identity was found.")
+		return nil, nil
+	}
+
+	stsHostname := fmt.Sprintf("sts.%s.amazonaws.com", region)
+	req, err := creator.createStsRequest(stsHostname)
+	if err != nil {
+		return nil, err
+	}
+
+	err = creator.signRequestWithSigV4(ctx, req, creds, region)
+	if err != nil {
+		return nil, err
+	}
+
+	credential, err := creator.createBase64EncodedRequestCredential(req)
+	if err != nil {
+		return nil, err
+	}
+
+	return &wifAttestation{
+		ProviderType: "AWS",
+		Credential:   credential,
+		Metadata:     map[string]string{"arn": arn},
+	}, nil
+}
+
+func (creator *awsIdentityAttestationCreator) createStsRequest(hostname string) (*http.Request, error) {
+	url := fmt.Sprintf("https://%s/?Action=GetCallerIdentity&Version=2011-06-15", hostname)
+	req, err := http.NewRequest("POST", url, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Host", hostname)
+	req.Header.Set("X-Snowflake-Audience", "snowflakecomputing.com")
+	return req, nil
+}
+
+func (creator *awsIdentityAttestationCreator) signRequestWithSigV4(ctx context.Context, req *http.Request, creds aws.Credentials, region string) error {
+	signer := v4.NewSigner()
+	// as per docs of SignHTTP, the payload hash must be present even if the payload is empty
+	payloadHash := hex.EncodeToString(sha256.New().Sum(nil))
+	err := signer.SignHTTP(ctx, creds, req, payloadHash, "sts", region, time.Now())
+	return err
+}
+
+func (creator *awsIdentityAttestationCreator) createBase64EncodedRequestCredential(req *http.Request) (string, error) {
+	headers := make(map[string]string)
+	for key, values := range req.Header {
+		headers[key] = values[0]
+	}
+
+	assertion := map[string]interface{}{
+		"url":     req.URL.String(),
+		"method":  req.Method,
+		"headers": headers,
+	}
+
+	assertionJSON, err := json.Marshal(assertion)
+	if err != nil {
+		return "", err
+	}
+
+	return base64.StdEncoding.EncodeToString(assertionJSON), nil
+}
diff --git a/connection_configuration.go b/connection_configuration.go
index ee337ce..7836ca6 100644
--- a/connection_configuration.go
+++ b/connection_configuration.go
@@ -190,6 +190,10 @@ func handleSingleParam(cfg *Config, key string, value interface{}) error {
 		cfg.OauthRedirectURI, err = parseString(value)
 	case "oauth_scope":
 		cfg.OauthScope, err = parseString(value)
+	case "workloadIdentityProvider":
+		cfg.WorkloadIdentityProvider, err = parseString(value)
+	case "workloadIdentityEntraResource":
+		cfg.WorkloadIdentityEntraResource, err = parseString(value)
 
 	case "token_file_path":
 		tokenPath, err = parseString(value)
diff --git a/dsn.go b/dsn.go
index 26cae25..d48631f 100644
--- a/dsn.go
+++ b/dsn.go
@@ -120,6 +120,9 @@ type Config struct {
 	DisableConsoleLogin ConfigBool // Indicates whether console login should be disabled
 
 	DisableSamlURLCheck ConfigBool // Indicates whether the SAML URL check should be disabled
+
+	WorkloadIdentityProvider      string // The workload identity provider to use for WIF authentication
+	WorkloadIdentityEntraResource string // The resource to use for WIF authentication on Azure environment
 }
 
 // Validate enables testing if config is correct.
@@ -214,6 +217,12 @@ func DSN(cfg *Config) (dsn string, err error) {
 	if cfg.OauthScope != "" {
 		params.Add("oauthScope", cfg.OauthScope)
 	}
+	if cfg.WorkloadIdentityProvider != "" {
+		params.Add("workloadIdentityProvider", cfg.WorkloadIdentityProvider)
+	}
+	if cfg.WorkloadIdentityEntraResource != "" {
+		params.Add("workloadIdentityEntraResource", cfg.WorkloadIdentityEntraResource)
+	}
 	if cfg.Authenticator != AuthTypeSnowflake {
 		if cfg.Authenticator == AuthTypeOkta {
 			params.Add("authenticator", strings.ToLower(cfg.OktaURL.String()))
@@ -835,6 +844,10 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 
 		case "token":
 			cfg.Token = value
+		case "workloadIdentityProvider":
+			cfg.WorkloadIdentityProvider = value
+		case "workloadIdentityEntraResource":
+			cfg.WorkloadIdentityEntraResource = value
 		case "privateKey":
 			var decodeErr error
 			block, decodeErr := base64.URLEncoding.DecodeString(value)
diff --git a/go.mod b/go.mod
index 6b8d7a6..3cf36f3 100644
--- a/go.mod
+++ b/go.mod
@@ -9,9 +9,11 @@ require (
 	github.com/BurntSushi/toml v1.4.0
 	github.com/apache/arrow-go/v18 v18.0.0
 	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/credentials v1.17.11
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.16.15
 	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/aws/smithy-go v1.20.2
 	github.com/gabriel-vasile/mimetype v1.4.7
 	github.com/golang-jwt/jwt/v5 v5.2.2
@@ -29,13 +31,17 @@ require (
 	github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.2 // indirect
 	github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c // indirect
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.2 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.1 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.5 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.5 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.5 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.7 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.7 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.5 // 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/danieljoos/wincred v1.2.2 // indirect
 	github.com/dvsekhvalnov/jose2go v1.6.0 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect


Test patch:
diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml
index 12243c5..7e16548 100644
--- a/.github/workflows/build-test.yml
+++ b/.github/workflows/build-test.yml
@@ -67,6 +67,7 @@ jobs:
                 GO_TEST_PARAMS: ${{ inputs.goTestParams }}
                 WIREMOCK_PORT: 14335
                 WIREMOCK_HTTPS_PORT: 13567
+                SF_ENABLE_EXPERIMENTAL_AUTHENTICATION: true
               run: ./ci/test.sh
             - name: Upload coverage to Codecov
               uses: codecov/codecov-action@v5
@@ -98,6 +99,7 @@ jobs:
                 GO_TEST_PARAMS: ${{ inputs.goTestParams }}
                 WIREMOCK_PORT: 14335
                 WIREMOCK_HTTPS_PORT: 13567
+                SF_ENABLE_EXPERIMENTAL_AUTHENTICATION: true
               run: ./ci/test.sh
             - name: Upload coverage to Codecov
               uses: codecov/codecov-action@v5
@@ -133,6 +135,7 @@ jobs:
                 GO_TEST_PARAMS: ${{ inputs.goTestParams }}
                 WIREMOCK_PORT: 14335
                 WIREMOCK_HTTPS_PORT: 13567
+                SF_ENABLE_EXPERIMENTAL_AUTHENTICATION: true
               run: ci\\test.bat
             - name: Upload coverage to Codecov
               uses: codecov/codecov-action@v5
diff --git a/auth_wif_test.go b/auth_wif_test.go
new file mode 100644
index 0000000..4036c9b
--- /dev/null
+++ b/auth_wif_test.go
@@ -0,0 +1,243 @@
+package gosnowflake
+
+import (
+	"context"
+	"errors"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"testing"
+)
+
+var awsErrorCreator = &mockWifAttestationCreator{attestation: nil, err: errors.New("AWS error")}
+var gcpErrorCreator = &mockWifAttestationCreator{attestation: nil, err: errors.New("GCP error")}
+var azureErrorCreator = &mockWifAttestationCreator{attestation: nil, err: errors.New("AZURE error")}
+var oidcErrorCreator = &mockWifAttestationCreator{attestation: nil, err: errors.New("OIDC error")}
+
+type mockWifAttestationCreator struct {
+	attestation *wifAttestation
+	err         error
+}
+
+func (m *mockWifAttestationCreator) createAttestation(ctx context.Context) (*wifAttestation, error) {
+	return m.attestation, m.err
+}
+
+func TestAttestationCreator(t *testing.T) {
+	provider := &wifAttestationProvider{
+		context:      context.Background(),
+		awsCreator:   &mockWifAttestationCreator{},
+		gcpCreator:   &mockWifAttestationCreator{},
+		azureCreator: &mockWifAttestationCreator{},
+		oidcCreator:  &mockWifAttestationCreator{},
+	}
+
+	tests := []struct {
+		identityProvider string
+		expectedCreator  wifAttestationCreator
+		expectedError    error
+	}{
+		{"AWS", provider.awsCreator, nil},
+		{"GCP", provider.gcpCreator, nil},
+		{"AZURE", provider.azureCreator, nil},
+		{"OIDC", provider.oidcCreator, nil},
+		{"UNKNOWN", nil, errors.New("unknown Workload Identity provider specified: UNKNOWN")},
+		{"", nil, errors.New("unknown Workload Identity provider specified: ")},
+	}
+
+	for _, test := range tests {
+		t.Run(test.identityProvider, func(t *testing.T) {
+			creator, err := provider.attestationCreator(test.identityProvider)
+			if creator != test.expectedCreator {
+				t.Errorf("expected creator %v, got %v", test.expectedCreator, creator)
+			}
+			if (err != nil && test.expectedError == nil) || (err == nil && test.expectedError != nil) || (err != nil && test.expectedError != nil && err.Error() != test.expectedError.Error()) {
+				t.Errorf("expected error %v, got %v", test.expectedError, err)
+			}
+		})
+	}
+}
+
+func TestCreateAutodetectAttestation(t *testing.T) {
+	tests := []struct {
+		name             string
+		oidcCreator      *mockWifAttestationCreator
+		awsCreator       *mockWifAttestationCreator
+		gcpCreator       *mockWifAttestationCreator
+		azureCreator     *mockWifAttestationCreator
+		expectedProvider string
+		expectedError    string
+	}{
+		{
+			name: "OIDC",
+			oidcCreator: &mockWifAttestationCreator{
+				attestation: &wifAttestation{ProviderType: "OIDC", Credential: "oidc-credential"},
+				err:         nil,
+			},
+			awsCreator:       awsErrorCreator,
+			gcpCreator:       gcpErrorCreator,
+			azureCreator:     azureErrorCreator,
+			expectedProvider: "OIDC",
+			expectedError:    "",
+		},
+		{
+			name:        "AWS",
+			oidcCreator: oidcErrorCreator,
+			awsCreator: &mockWifAttestationCreator{
+				attestation: &wifAttestation{ProviderType: "AWS", Credential: "aws-credential"},
+				err:         nil,
+			},
+			gcpCreator:       gcpErrorCreator,
+			azureCreator:     azureErrorCreator,
+			expectedProvider: "AWS",
+			expectedError:    "",
+		},
+		{
+			name:        "GCP",
+			oidcCreator: oidcErrorCreator,
+			awsCreator:  awsErrorCreator,
+			gcpCreator: &mockWifAttestationCreator{
+				attestation: &wifAttestation{ProviderType: "GCP", Credential: "gcp-credential"},
+				err:         nil,
+			},
+			azureCreator:     azureErrorCreator,
+			expectedProvider: "GCP",
+			expectedError:    "",
+		},
+		{
+			name:        "Azure",
+			oidcCreator: oidcErrorCreator,
+			awsCreator:  awsErrorCreator,
+			gcpCreator:  gcpErrorCreator,
+			azureCreator: &mockWifAttestationCreator{
+				attestation: &wifAttestation{ProviderType: "AZURE", Credential: "azure-credential"},
+				err:         nil,
+			},
+			expectedProvider: "AZURE",
+			expectedError:    "",
+		},
+		{
+			name:             "None",
+			oidcCreator:      oidcErrorCreator,
+			awsCreator:       awsErrorCreator,
+			gcpCreator:       gcpErrorCreator,
+			azureCreator:     azureErrorCreator,
+			expectedProvider: "",
+			expectedError:    "unable to autodetect Workload Identity. None of the supported Workload Identity environments has been identified",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			provider := &wifAttestationProvider{
+				oidcCreator:  test.oidcCreator,
+				awsCreator:   test.awsCreator,
+				gcpCreator:   test.gcpCreator,
+				azureCreator: test.azureCreator,
+			}
+
+			attestation, err := provider.createAutodetectAttestation()
+
+			if test.expectedError == "" {
+				assertNilE(t, err)
+				assertNotNilE(t, attestation)
+				assertEqualE(t, test.expectedProvider, attestation.ProviderType)
+			} else {
+				assertNotNilE(t, err)
+				assertNilE(t, attestation)
+				assertEqualE(t, test.expectedError, err.Error())
+			}
+		})
+	}
+}
+
+func TestAwsIdentityAttestationCreator_CreateAttestation(t *testing.T) {
+	tests := []struct {
+		name                string
+		attestationSvc      awsAttestationService
+		attestationReturned bool
+		expectedProvider    string
+		expectedMetadata    map[string]string
+	}{
+		{
+			name: "No AWS credentials",
+			attestationSvc: &mockAwsAttestationService{
+				creds:  aws.Credentials{},
+				region: "us-west-2",
+				arn:    "arn:aws:iam::123456789012:role/test-role",
+			},
+			attestationReturned: false,
+		},
+		{
+			name: "No AWS region",
+			attestationSvc: &mockAwsAttestationService{
+				creds:  mockCreds,
+				region: "",
+				arn:    "arn:aws:iam::123456789012:role/test-role",
+			},
+			attestationReturned: false,
+		},
+		{
+			name: "No AWS ARN",
+			attestationSvc: &mockAwsAttestationService{
+				creds:  mockCreds,
+				region: "us-west-2",
+				arn:    "",
+			},
+			attestationReturned: false,
+		},
+		{
+			name: "Successful attestation",
+			attestationSvc: &mockAwsAttestationService{
+				creds:  mockCreds,
+				region: "us-west-2",
+				arn:    "arn:aws:iam::123456789012:role/test-role",
+			},
+			attestationReturned: true,
+			expectedProvider:    "AWS",
+			expectedMetadata:    map[string]string{"arn": "arn:aws:iam::123456789012:role/test-role"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			creator := &awsIdentityAttestationCreator{
+				attestationService: test.attestationSvc,
+			}
+
+			ctx := context.Background()
+			attestation, _ := creator.createAttestation(ctx)
+
+			if test.attestationReturned {
+				assertNotNilE(t, attestation)
+				assertNotNilE(t, attestation.Credential)
+				assertEqualE(t, test.expectedProvider, attestation.ProviderType)
+				assertDeepEqualE(t, test.expectedMetadata, attestation.Metadata)
+			} else {
+				assertNilF(t, attestation)
+			}
+		})
+	}
+}
+
+type mockAwsAttestationService struct {
+	creds  aws.Credentials
+	region string
+	arn    string
+}
+
+var mockCreds = aws.Credentials{
+	AccessKeyID:     "mockAccessKey",
+	SecretAccessKey: "mockSecretKey",
+	SessionToken:    "mockSessionToken",
+}
+
+func (m *mockAwsAttestationService) GetAWSCredentials() aws.Credentials {
+	return m.creds
+}
+
+func (m *mockAwsAttestationService) GetAWSRegion() string {
+	return m.region
+}
+
+func (m *mockAwsAttestationService) GetArn() string {
+	return m.arn
+}
diff --git a/ci/container/test_authentication.sh b/ci/container/test_authentication.sh
index 85745d2..3bfb499 100755
--- a/ci/container/test_authentication.sh
+++ b/ci/container/test_authentication.sh
@@ -12,6 +12,7 @@ export RUN_AUTH_TESTS=true
 export AUTHENTICATION_TESTS_ENV="docker"
 
 export RUN_AUTH_TESTS=true
+export SF_ENABLE_EXPERIMENTAL_AUTHENTICATION=true
 export AUTHENTICATION_TESTS_ENV="docker"
 
 go test -run TestExternalBrowser*
diff --git a/ci/test.sh b/ci/test.sh
index b92e447..f94f24f 100755
--- a/ci/test.sh
+++ b/ci/test.sh
@@ -57,6 +57,7 @@ if [[ -n "$JENKINS_HOME" ]]; then
           -e JOB_NAME \
           -e BUILD_NUMBER \
           -e JENKINS_HOME \
+          -e SF_ENABLE_EXPERIMENTAL_AUTHENTICATION=true \
           ${TEST_IMAGE_NAMES[$name]} \
           /mnt/host/ci/container/test_component.sh
           echo "[INFO] Test Results: $WORKSPACE/junit.xml"