auth_wif_test.go

package gosnowflake

import (
	"context"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"strings"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials"
)

type mockWifAttestationCreator struct {
	providerType wifProviderType
	returnError  error
}

func (m *mockWifAttestationCreator) createAttestation() (*wifAttestation, error) {
	if m.returnError != nil {
		return nil, m.returnError
	}
	return &wifAttestation{
		ProviderType: string(m.providerType),
	}, nil
}

func TestGetAttestation(t *testing.T) {
	awsError := errors.New("aws attestation error")
	gcpError := errors.New("gcp attestation error")
	azureError := errors.New("azure attestation error")
	oidcError := errors.New("oidc attestation error")

	provider := &wifAttestationProvider{
		context:      context.Background(),
		awsCreator:   &mockWifAttestationCreator{providerType: awsWif},
		gcpCreator:   &mockWifAttestationCreator{providerType: gcpWif},
		azureCreator: &mockWifAttestationCreator{providerType: azureWif},
		oidcCreator:  &mockWifAttestationCreator{providerType: oidcWif},
	}

	providerWithErrors := &wifAttestationProvider{
		context:      context.Background(),
		awsCreator:   &mockWifAttestationCreator{providerType: awsWif, returnError: awsError},
		gcpCreator:   &mockWifAttestationCreator{providerType: gcpWif, returnError: gcpError},
		azureCreator: &mockWifAttestationCreator{providerType: azureWif, returnError: azureError},
		oidcCreator:  &mockWifAttestationCreator{providerType: oidcWif, returnError: oidcError},
	}

	tests := []struct {
		name             string
		provider         *wifAttestationProvider
		identityProvider string
		expectedResult   *wifAttestation
		expectedError    error
	}{
		{
			name:             "AWS success",
			provider:         provider,
			identityProvider: "AWS",
			expectedResult:   &wifAttestation{ProviderType: string(awsWif)},
			expectedError:    nil,
		},
		{
			name:             "AWS error",
			provider:         providerWithErrors,
			identityProvider: "AWS",
			expectedResult:   nil,
			expectedError:    awsError,
		},
		{
			name:             "GCP success",
			provider:         provider,
			identityProvider: "GCP",
			expectedResult:   &wifAttestation{ProviderType: string(gcpWif)},
			expectedError:    nil,
		},
		{
			name:             "GCP error",
			provider:         providerWithErrors,
			identityProvider: "GCP",
			expectedResult:   nil,
			expectedError:    gcpError,
		},
		{
			name:             "AZURE success",
			provider:         provider,
			identityProvider: "AZURE",
			expectedResult:   &wifAttestation{ProviderType: string(azureWif)},
			expectedError:    nil,
		},
		{
			name:             "AZURE error",
			provider:         providerWithErrors,
			identityProvider: "AZURE",
			expectedResult:   nil,
			expectedError:    azureError,
		},
		{
			name:             "OIDC success",
			provider:         provider,
			identityProvider: "OIDC",
			expectedResult:   &wifAttestation{ProviderType: string(oidcWif)},
			expectedError:    nil,
		},
		{
			name:             "OIDC error",
			provider:         providerWithErrors,
			identityProvider: "OIDC",
			expectedResult:   nil,
			expectedError:    oidcError,
		},
		{
			name:             "Unknown provider",
			provider:         provider,
			identityProvider: "UNKNOWN",
			expectedResult:   nil,
			expectedError:    errors.New("unknown WorkloadIdentityProvider specified: UNKNOWN. Valid values are: AWS, GCP, AZURE, OIDC"),
		},
		{
			name:             "Empty provider",
			provider:         provider,
			identityProvider: "",
			expectedResult:   nil,
			expectedError:    errors.New("unknown WorkloadIdentityProvider specified: . Valid values are: AWS, GCP, AZURE, OIDC"),
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			attestation, err := test.provider.getAttestation(test.identityProvider)
			if test.expectedError != nil {
				assertNilE(t, attestation)
				assertNotNilF(t, err)
				assertEqualE(t, test.expectedError.Error(), err.Error())
			} else if test.expectedResult != nil {
				assertNilE(t, err)
				assertNotNilF(t, attestation)
				assertEqualE(t, test.expectedResult.ProviderType, attestation.ProviderType)
			} else {
				t.Fatal("test case must specify either expectedError or expectedResult")
			}
		})
	}
}

func TestAwsIdentityAttestationCreator(t *testing.T) {
	tests := []struct {
		name             string
		config           Config
		attestationSvc   awsAttestationMetadataProvider
		expectedError    error
		expectedProvider string
		expectedStsHost  string
	}{
		{
			name:           "No attestation service",
			attestationSvc: nil,
			expectedError:  fmt.Errorf("AWS attestation service could not be created"),
		},
		{
			name: "No AWS credentials",
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:  aws.Credentials{},
				region: "us-west-2",
			},
			expectedError: fmt.Errorf("no AWS credentials were found"),
		},
		{
			name: "No AWS region",
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:  mockCreds,
				region: "",
			},
			expectedError: fmt.Errorf("no AWS region was found"),
		},
		{
			name: "Successful attestation",
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:  mockCreds,
				region: "us-west-2",
			},
			expectedProvider: "AWS",
			expectedStsHost:  "sts.us-west-2.amazonaws.com",
		},
		{
			name: "Successful attestation for CN region",
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:  mockCreds,
				region: "cn-northwest-1",
			},
			expectedProvider: "AWS",
			expectedStsHost:  "sts.cn-northwest-1.amazonaws.com.cn",
		},
		{
			// ISO partitions do not use an amazonaws.com suffix at all, so this pins
			// the host that actually reaches the envelope, not just the resolver.
			name: "Successful attestation for ISO region",
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:  mockCreds,
				region: "us-iso-east-1",
			},
			expectedProvider: "AWS",
			expectedStsHost:  "sts.us-iso-east-1.c2s.ic.gov",
		},
		{
			name: "Successful attestation with single role chaining",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{"arn:aws:iam::123456789012:role/test-role"},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds: mockCreds,
				chainingCreds: aws.Credentials{
					AccessKeyID:     "chainedAccessKey",
					SecretAccessKey: "chainedSecretKey",
					SessionToken:    "chainedSessionToken",
				},
				region:          "us-east-1",
				useRoleChaining: true,
			},
			expectedProvider: "AWS",
			expectedStsHost:  "sts.us-east-1.amazonaws.com",
		},
		{
			name: "Successful attestation with multiple role chaining",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{
					"arn:aws:iam::123456789012:role/role1",
					"arn:aws:iam::123456789012:role/role2",
					"arn:aws:iam::123456789012:role/role3",
				},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds: mockCreds,
				chainingCreds: aws.Credentials{
					AccessKeyID:     "finalRoleAccessKey",
					SecretAccessKey: "finalRoleSecretKey",
					SessionToken:    "finalRoleSessionToken",
				},
				region:          "us-west-2",
				useRoleChaining: true,
			},
			expectedProvider: "AWS",
			expectedStsHost:  "sts.us-west-2.amazonaws.com",
		},
		{
			name: "Role chaining with no credentials",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{"arn:aws:iam::123456789012:role/test-role"},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:           aws.Credentials{},
				region:          "us-west-2",
				useRoleChaining: true,
			},
			expectedError: fmt.Errorf("no AWS credentials were found"),
		},
		{
			name: "Role chaining with no region",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{"arn:aws:iam::123456789012:role/test-role"},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds: aws.Credentials{
					AccessKeyID:     "chainedAccessKey",
					SecretAccessKey: "chainedSecretKey",
					SessionToken:    "chainedSessionToken",
				},
				region:          "",
				useRoleChaining: true,
			},
			expectedError: fmt.Errorf("no AWS region was found"),
		},
		{
			name: "Role chaining failure",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{"arn:aws:iam::123456789012:role/test-role"},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:           mockCreds,
				region:          "us-west-2",
				chainingError:   fmt.Errorf("failed to assume role: AccessDenied"),
				useRoleChaining: true,
			},
			expectedError: fmt.Errorf("failed to assume role: AccessDenied"),
		},
		{
			name: "Cross-account role chaining",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{
					"arn:aws:iam::111111111111:role/cross-account-role",
					"arn:aws:iam::222222222222:role/target-role",
				},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds: mockCreds,
				chainingCreds: aws.Credentials{
					AccessKeyID:     "crossAccountAccessKey",
					SecretAccessKey: "crossAccountSecretKey",
					SessionToken:    "crossAccountSessionToken",
				},
				region:          "us-east-1",
				useRoleChaining: true,
			},
			expectedProvider: "AWS",
			expectedStsHost:  "sts.us-east-1.amazonaws.com",
		},
		{
			name: "Role chaining in CN region",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{"arn:aws-cn:iam::123456789012:role/cn-role"},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds: mockCreds,
				chainingCreds: aws.Credentials{
					AccessKeyID:     "cnRoleAccessKey",
					SecretAccessKey: "cnRoleSecretKey",
					SessionToken:    "cnRoleSessionToken",
				},
				region:          "cn-north-1",
				useRoleChaining: true,
			},
			expectedProvider: "AWS",
			expectedStsHost:  "sts.cn-north-1.amazonaws.com.cn",
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			creator := &awsIdentityAttestationCreator{
				attestationServiceFactory: func(ctx context.Context, cfg *Config) awsAttestationMetadataProvider {
					return test.attestationSvc
				},
				ctx: context.Background(),
				cfg: &test.config,
			}
			attestation, err := creator.createAttestation()
			if test.expectedError != nil {
				assertNilF(t, attestation)
				assertNotNilE(t, err)
				assertEqualE(t, test.expectedError.Error(), err.Error())
			} else {
				assertNilE(t, err)
				assertNotNilE(t, attestation)
				assertNotNilE(t, attestation.Credential)
				assertEqualE(t, test.expectedProvider, attestation.ProviderType)
				decoded, err := base64.StdEncoding.DecodeString(attestation.Credential)
				if err != nil {
					t.Fatalf("Failed to decode credential: %v", err)
				}
				var credentialMap map[string]any
				if err := json.Unmarshal(decoded, &credentialMap); err != nil {
					t.Fatalf("Failed to unmarshal credential JSON: %v", err)
				}
				assertEqualE(t, fmt.Sprintf("https://%s?Action=GetCallerIdentity&Version=2011-06-15", test.expectedStsHost), credentialMap["url"])
			}
		})
	}
}

type mockAwsAttestationMetadataProvider struct {
	creds            aws.Credentials
	region           string
	chainingCreds    aws.Credentials
	chainingError    error
	useRoleChaining  bool
	webIdentityToken string
	webIdentityError error
}

var mockCreds = aws.Credentials{
	AccessKeyID:     "mockAccessKey",
	SecretAccessKey: "mockSecretKey",
	SessionToken:    "mockSessionToken",
}

func (m *mockAwsAttestationMetadataProvider) awsCredentials() (aws.Credentials, error) {
	return m.creds, nil
}

func (m *mockAwsAttestationMetadataProvider) awsCredentialsViaRoleChaining() (aws.Credentials, error) {
	if m.chainingError != nil {
		return aws.Credentials{}, m.chainingError
	}
	if m.chainingCreds.AccessKeyID != "" {
		return m.chainingCreds, nil
	}
	return m.creds, nil
}

func (m *mockAwsAttestationMetadataProvider) awsRegion() string {
	return m.region
}

func (m *mockAwsAttestationMetadataProvider) awsWebIdentityToken(_ aws.Credentials, _ string) (string, error) {
	if m.webIdentityError != nil {
		return "", m.webIdentityError
	}
	return m.webIdentityToken, nil
}

func TestAwsIdentityAttestationCreatorOutboundToken(t *testing.T) {
	const jwt = "header.payload.signature"
	tests := []struct {
		name               string
		config             Config
		attestationSvc     awsAttestationMetadataProvider
		expectedError      error
		expectedCredential string
	}{
		{
			name: "Successful outbound attestation",
			config: Config{
				WorkloadIdentityAwsUseOutboundToken: ConfigBoolTrue,
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:            mockCreds,
				region:           "us-west-2",
				webIdentityToken: jwt,
			},
			expectedCredential: jwt,
		},
		{
			name: "Successful outbound attestation with role chaining",
			config: Config{
				WorkloadIdentityAwsUseOutboundToken: ConfigBoolTrue,
				WorkloadIdentityImpersonationPath:   []string{"arn:aws:iam::123456789012:role/test-role"},
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds: mockCreds,
				chainingCreds: aws.Credentials{
					AccessKeyID:     "chainedAccessKey",
					SecretAccessKey: "chainedSecretKey",
					SessionToken:    "chainedSessionToken",
				},
				region:           "us-east-1",
				useRoleChaining:  true,
				webIdentityToken: jwt,
			},
			expectedCredential: jwt,
		},
		{
			name: "STS returns no token",
			config: Config{
				WorkloadIdentityAwsUseOutboundToken: ConfigBoolTrue,
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:            mockCreds,
				region:           "us-west-2",
				webIdentityToken: "",
			},
			expectedError: fmt.Errorf("failed to obtain AWS web identity token from STS"),
		},
		{
			name: "STS call fails",
			config: Config{
				WorkloadIdentityAwsUseOutboundToken: ConfigBoolTrue,
			},
			attestationSvc: &mockAwsAttestationMetadataProvider{
				creds:            mockCreds,
				region:           "us-west-2",
				webIdentityError: fmt.Errorf("AccessDenied"),
			},
			expectedError: fmt.Errorf("AccessDenied"),
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			creator := &awsIdentityAttestationCreator{
				attestationServiceFactory: func(ctx context.Context, cfg *Config) awsAttestationMetadataProvider {
					return test.attestationSvc
				},
				ctx: context.Background(),
				cfg: &test.config,
			}
			attestation, err := creator.createAttestation()
			if test.expectedError != nil {
				assertNilF(t, attestation)
				assertNotNilE(t, err)
				assertEqualE(t, test.expectedError.Error(), err.Error())
			} else {
				assertNilE(t, err)
				assertNotNilE(t, attestation)
				assertEqualE(t, "AWS", attestation.ProviderType)
				// The outbound token is sent verbatim, not base64-encoded.
				assertEqualE(t, test.expectedCredential, attestation.Credential)
			}
		})
	}
}

func TestGcpIdentityAttestationCreator(t *testing.T) {
	tests := []struct {
		name                string
		wiremockMappingPath string
		config              Config
		expectedError       error
		expectedSub         string
	}{
		{
			name:                "Successful flow",
			wiremockMappingPath: "auth/wif/gcp/successful_flow.json",
			expectedError:       nil,
			expectedSub:         "some-subject",
		},
		{
			name:                "Successful impersonation flow",
			wiremockMappingPath: "auth/wif/gcp/successful_impersionation_flow.json",
			config: Config{
				WorkloadIdentityImpersonationPath: []string{
					"delegate1",
					"delegate2",
					"targetServiceAccount",
				},
			},
			expectedError: nil,
			expectedSub:   "some-impersonated-subject",
		},
		{
			name:                "No GCP credential - http error",
			wiremockMappingPath: "auth/wif/gcp/http_error.json",
			expectedError:       fmt.Errorf("no GCP token was found"),
			expectedSub:         "",
		},
		{
			name:                "missing issuer claim",
			wiremockMappingPath: "auth/wif/gcp/missing_issuer_claim.json",
			expectedError:       fmt.Errorf("could not extract claims from token: missing issuer claim in JWT token"),
			expectedSub:         "",
		},
		{
			name:                "missing sub claim",
			wiremockMappingPath: "auth/wif/gcp/missing_sub_claim.json",
			expectedError:       fmt.Errorf("could not extract claims from token: missing sub claim in JWT token"),
			expectedSub:         "",
		},
		{
			name:                "unparsable token",
			wiremockMappingPath: "auth/wif/gcp/unparsable_token.json",
			expectedError:       fmt.Errorf("could not extract claims from token: unable to extract JWT claims from token: token is malformed: token contains an invalid number of segments"),
			expectedSub:         "",
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			creator := &gcpIdentityAttestationCreator{
				cfg:                    &test.config,
				metadataServiceBaseURL: wiremock.baseURL(),
				iamCredentialsURL:      wiremock.baseURL(),
			}
			wiremock.registerMappings(t, wiremockMapping{filePath: test.wiremockMappingPath})
			attestation, err := creator.createAttestation()

			if test.expectedError != nil {
				assertNilF(t, attestation)
				assertNotNilF(t, err)
				assertEqualE(t, test.expectedError.Error(), err.Error())
			} else {
				assertNilF(t, err)
				assertNotNilF(t, attestation)
				assertEqualE(t, string(gcpWif), attestation.ProviderType)
				assertEqualE(t, test.expectedSub, attestation.Metadata["sub"])
			}
		})
	}
}

func TestOidcIdentityAttestationCreator(t *testing.T) {
	const (
		/*
		 * {
		 *   "sub": "some-subject",
		 *   "iat": 1743761213,
		 *   "exp": 1743764813,
		 *   "aud": "www.example.com"
		 * }
		 */
		missingIssuerClaimToken = "eyJ0eXAiOiJhdCtqd3QiLCJhbGciOiJFUzI1NiIsImtpZCI6ImU2M2I5NzA1OTRiY2NmZTAxMDlkOTg4OWM2MDk3OWEwIn0.eyJzdWIiOiJzb21lLXN1YmplY3QiLCJpYXQiOjE3NDM3NjEyMTMsImV4cCI6MTc0Mzc2NDgxMywiYXVkIjoid3d3LmV4YW1wbGUuY29tIn0.H6sN6kjA82EuijFcv-yCJTqau5qvVTCsk0ZQ4gvFQMkB7c71XPs4lkwTa7ZlNNlx9e6TpN1CVGnpCIRDDAZaDw" // pragma: allowlist secret
		/*
		 * {
		 *   "iss": "https://accounts.google.com",
		 *   "iat": 1743761213,
		 *   "exp": 1743764813,
		 *   "aud": "www.example.com"
		 * }
		 */
		missingSubClaimToken = "eyJ0eXAiOiJhdCtqd3QiLCJhbGciOiJFUzI1NiIsImtpZCI6ImU2M2I5NzA1OTRiY2NmZTAxMDlkOTg4OWM2MDk3OWEwIn0.eyJpc3MiOiJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20iLCJpYXQiOjE3NDM3NjEyMTMsImV4cCI6MTc0Mzc2NDgxMywiYXVkIjoid3d3LmV4YW1wbGUuY29tIn0.w0njdpfWFETVK8Ktq9GdvuKRQJjvhOplcSyvQ_zHHwBUSMapqO1bjEWBx5VhGkdECZIGS1VY7db_IOqT45yOMA" // pragma: allowlist secret
		/*
		 * {
		 *     "iss": "https://oidc.eks.us-east-2.amazonaws.com/id/3B869BC5D12CEB5515358621D8085D58",
		 *     "iat": 1743692017,
		 *     "exp": 1775228014,
		 *     "aud": "www.example.com",
		 *     "sub": "system:serviceaccount:poc-namespace:oidc-sa"
		 * }
		 */
		validToken      = "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJodHRwczovL29pZGMuZWtzLnVzLWVhc3QtMi5hbWF6b25hd3MuY29tL2lkLzNCODY5QkM1RDEyQ0VCNTUxNTM1ODYyMUQ4MDg1RDU4IiwiaWF0IjoxNzQ0Mjg3ODc4LCJleHAiOjE3NzU4MjM4NzgsImF1ZCI6Ind3dy5leGFtcGxlLmNvbSIsInN1YiI6InN5c3RlbTpzZXJ2aWNlYWNjb3VudDpwb2MtbmFtZXNwYWNlOm9pZGMtc2EifQ.a8H6KRIF1XmM8lkqL6kR8ccInr7wAzQrbKd3ZHFgiEg" // pragma: allowlist secret
		unparsableToken = "unparsable_token"
		emptyToken      = ""
	)

	type testCase struct {
		name          string
		token         string
		expectedError error
		expectedSub   string
	}

	tests := []testCase{
		{
			name:          "no token input",
			token:         emptyToken,
			expectedError: fmt.Errorf("no OIDC token was specified"),
		},
		{
			name:          "valid token returns proper attestation",
			token:         validToken,
			expectedError: nil,
			expectedSub:   "system:serviceaccount:poc-namespace:oidc-sa",
		},
		{
			name:          "missing issuer returns error",
			token:         missingIssuerClaimToken,
			expectedError: errors.New("missing issuer claim in JWT token"),
		},
		{
			name:          "missing sub returns error",
			token:         missingSubClaimToken,
			expectedError: errors.New("missing sub claim in JWT token"),
		},
		{
			name:          "unparsable token returns error",
			token:         unparsableToken,
			expectedError: errors.New("unable to extract JWT claims from token: token is malformed: token contains an invalid number of segments"),
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			creator := &oidcIdentityAttestationCreator{token: func() (string, error) {
				return test.token, nil
			}}
			attestation, err := creator.createAttestation()

			if test.expectedError != nil {
				assertNotNilE(t, err)
				assertNilF(t, attestation)
				assertEqualE(t, test.expectedError.Error(), err.Error())
			} else {
				assertNilE(t, err)
				assertNotNilE(t, attestation)
				assertEqualE(t, string(oidcWif), attestation.ProviderType)
				assertEqualE(t, test.expectedSub, attestation.Metadata["sub"])
			}
		})
	}
}

func TestAzureIdentityAttestationCreator(t *testing.T) {
	tests := []struct {
		name                string
		wiremockMappingPath string
		metadataProvider    *mockAzureAttestationMetadataProvider
		cfg                 *Config
		expectedIss         string
		expectedError       error
	}{
		/*
		 * {
		 *     "iss": "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
		 *     "sub": "77213E30-E8CB-4595-B1B6-5F050E8308FD"
		 * }
		 */
		{
			name:                "Successful flow",
			wiremockMappingPath: "auth/wif/azure/successful_flow_basic.json",
			metadataProvider:    azureVMMetadataProvider(),
			expectedIss:         "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
			expectedError:       nil,
		},
		/*
		 * {
		 *     "iss": "https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/",
		 *     "sub": "77213E30-E8CB-4595-B1B6-5F050E8308FD"
		 * }
		 */
		{
			name:                "Successful flow v2 issuer",
			wiremockMappingPath: "auth/wif/azure/successful_flow_v2_issuer.json",
			metadataProvider:    azureVMMetadataProvider(),
			expectedIss:         "https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/",
			expectedError:       nil,
		},
		/*
		 * {
		 *     "iss": "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
		 *     "sub": "77213E30-E8CB-4595-B1B6-5F050E8308FD"
		 * }
		 */
		{
			name:                "Successful flow azure functions",
			wiremockMappingPath: "auth/wif/azure/successful_flow_azure_functions.json",
			metadataProvider:    azureFunctionsMetadataProvider(),
			expectedIss:         "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
			expectedError:       nil,
		},
		/*
		 * {
		 *     "iss": "https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/",
		 *     "sub": "77213E30-E8CB-4595-B1B6-5F050E8308FD"
		 * }
		 */
		{
			name:                "Successful flow azure functions v2 issuer",
			wiremockMappingPath: "auth/wif/azure/successful_flow_azure_functions_v2_issuer.json",
			metadataProvider:    azureFunctionsMetadataProvider(),
			expectedIss:         "https://login.microsoftonline.com/fa15d692-e9c7-4460-a743-29f29522229/",
			expectedError:       nil,
		},
		/*
		 * {
		 *     "iss": "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
		 *     "sub": "77213E30-E8CB-4595-B1B6-5F050E8308FD"
		 * }
		 */
		{
			name:                "Successful flow azure functions no client ID",
			wiremockMappingPath: "auth/wif/azure/successful_flow_azure_functions_no_client_id.json",
			metadataProvider: &mockAzureAttestationMetadataProvider{
				identityEndpointValue: wiremock.baseURL() + "/metadata/identity/endpoint/from/env",
				identityHeaderValue:   "some-identity-header-from-env",
				clientIDValue:         "",
			},
			expectedIss:   "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
			expectedError: nil,
		},
		/*
		 * {
		 *     "iss": "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
		 *     "sub": "77213E30-E8CB-4595-B1B6-5F050E8308FD"
		 * }
		 */
		{
			name:                "Successful flow azure functions custom Entra resource",
			wiremockMappingPath: "auth/wif/azure/successful_flow_azure_functions_custom_entra_resource.json",
			metadataProvider:    azureFunctionsMetadataProvider(),
			cfg:                 &Config{WorkloadIdentityEntraResource: "api://1111111-2222-3333-44444-55555555"},
			expectedIss:         "https://sts.windows.net/fa15d692-e9c7-4460-a743-29f29522229/",
			expectedError:       nil,
		},
		{
			name:                "Non-json response",
			wiremockMappingPath: "auth/wif/azure/non_json_response.json",
			metadataProvider:    azureVMMetadataProvider(),
			expectedError:       fmt.Errorf("failed to extract token from JSON: invalid character 'o' in literal null (expecting 'u')"),
		},
		{
			name: "Identity endpoint but no identity header",
			metadataProvider: &mockAzureAttestationMetadataProvider{
				identityEndpointValue: wiremock.baseURL() + "/metadata/identity/endpoint/from/env",
				identityHeaderValue:   "",
				clientIDValue:         "managed-client-id-from-env",
			},
			expectedError: fmt.Errorf("managed identity is not enabled on this Azure function"),
		},
		{
			name:                "Unparsable token",
			wiremockMappingPath: "auth/wif/azure/unparsable_token.json",
			metadataProvider:    azureVMMetadataProvider(),
			expectedError:       fmt.Errorf("failed to extract sub and iss claims from token: unable to extract JWT claims from token: token is malformed: token contains an invalid number of segments"),
		},
		{
			name:                "HTTP error",
			metadataProvider:    azureVMMetadataProvider(),
			wiremockMappingPath: "auth/wif/azure/http_error.json",
			expectedError:       fmt.Errorf("could not fetch Azure token"),
		},
		{
			name:                "Missing sub or iss claim",
			wiremockMappingPath: "auth/wif/azure/missing_issuer_claim.json",
			metadataProvider:    azureVMMetadataProvider(),
			expectedError:       fmt.Errorf("failed to extract sub and iss claims from token: missing issuer claim in JWT token"),
		},
		{
			name:                "Missing sub claim",
			wiremockMappingPath: "auth/wif/azure/missing_sub_claim.json",
			metadataProvider:    azureVMMetadataProvider(),
			expectedError:       fmt.Errorf("failed to extract sub and iss claims from token: missing sub claim in JWT token"),
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			if test.wiremockMappingPath != "" {
				wiremock.registerMappings(t, wiremockMapping{filePath: test.wiremockMappingPath})
			}
			creator := &azureIdentityAttestationCreator{
				cfg:                              test.cfg,
				azureMetadataServiceBaseURL:      wiremock.baseURL(),
				azureAttestationMetadataProvider: test.metadataProvider,
				workloadIdentityEntraResource:    determineEntraResource(test.cfg),
			}
			attestation, err := creator.createAttestation()

			if test.expectedError != nil {
				assertNilF(t, attestation)
				assertNotNilE(t, err)
				assertEqualE(t, test.expectedError.Error(), err.Error())
			} else {
				assertNilF(t, err)
				assertNotNilF(t, attestation)
				assertEqualE(t, string(azureWif), attestation.ProviderType)
				assertEqualE(t, test.expectedIss, attestation.Metadata["iss"])
				assertEqualE(t, "77213E30-E8CB-4595-B1B6-5F050E8308FD", attestation.Metadata["sub"])
			}
		})
	}
}

type mockAzureAttestationMetadataProvider struct {
	identityEndpointValue string
	identityHeaderValue   string
	clientIDValue         string
}

func (m *mockAzureAttestationMetadataProvider) identityEndpoint() string {
	return m.identityEndpointValue
}

func (m *mockAzureAttestationMetadataProvider) identityHeader() string {
	return m.identityHeaderValue
}

func (m *mockAzureAttestationMetadataProvider) clientID() string {
	return m.clientIDValue
}

func azureFunctionsMetadataProvider() *mockAzureAttestationMetadataProvider {
	return &mockAzureAttestationMetadataProvider{
		identityEndpointValue: wiremock.baseURL() + "/metadata/identity/endpoint/from/env",
		identityHeaderValue:   "some-identity-header-from-env",
		clientIDValue:         "managed-client-id-from-env",
	}
}

func azureVMMetadataProvider() *mockAzureAttestationMetadataProvider {
	return &mockAzureAttestationMetadataProvider{
		identityEndpointValue: "",
		identityHeaderValue:   "",
		clientIDValue:         "",
	}
}

// Running this test locally:
// * Push branch to repository
// * Set PARAMETERS_SECRET
// * Run ci/test_wif.sh
func TestWorkloadIdentityAuthOnCloudVM(t *testing.T) {
	account := os.Getenv("SNOWFLAKE_TEST_WIF_ACCOUNT")
	host := os.Getenv("SNOWFLAKE_TEST_WIF_HOST")
	provider := os.Getenv("SNOWFLAKE_TEST_WIF_PROVIDER")
	println("provider = " + provider)
	if account == "" || host == "" || provider == "" {
		t.Skip("Test can run only on cloud VM with env variables set")
	}
	testCases := []struct {
		name             string
		skip             func() (bool, string)
		setupCfg         func(*testing.T, *Config)
		expectedUsername string
	}{
		{
			name: "provider=" + provider,
			setupCfg: func(_ *testing.T, config *Config) {
				if provider != "GCP+OIDC" {
					config.WorkloadIdentityProvider = provider
				} else {
					config.WorkloadIdentityProvider = "OIDC"
					config.Token = func() string {
						metadataURL := defaultGcpMetadataServiceBase + "/computeMetadata/v1/instance/service-accounts/default/identity?audience=snowflakecomputing.com"
						cmd := exec.Command("wget", "-O", "-", "--header=Metadata-Flavor: Google", metadataURL)
						output, err := cmd.Output()
						if err != nil {
							t.Fatalf("error executing GCP metadata request: %v", err)
						}
						token := strings.TrimSpace(string(output))
						if token == "" {
							t.Fatal("failed to retrieve GCP access token: empty response")
						}
						return token
					}()
				}
			},
			expectedUsername: os.Getenv("SNOWFLAKE_TEST_WIF_USERNAME"),
		},
		{
			name: "provider=" + provider + ",impersonation",
			skip: func() (bool, string) {
				if provider != "AWS" && provider != "GCP" {
					return true, "Impersonation is supported only on AWS and GCP"
				}
				return false, ""
			},
			setupCfg: func(t *testing.T, config *Config) {
				config.WorkloadIdentityProvider = provider
				impersonationPath := os.Getenv("SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH")
				assertNotEqualF(t, impersonationPath, "", "SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH is not set")
				config.WorkloadIdentityImpersonationPath = strings.Split(impersonationPath, ",")
				assertNotEqualF(t, os.Getenv("SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION"), "", "SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION is not set")
			},
			expectedUsername: os.Getenv("SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION"),
		},
		{
			// AWS WIF supports two attestation methods: GetCallerIdentity (default) and
			// GetWebIdentityToken (enabled via WorkloadIdentityAwsUseOutboundToken). This case
			// covers the GetWebIdentityToken method against the dedicated TEST_WIF_E2E_AWS_WITH_ISSUER
			// Snowflake user (configured with an ISSUER). Since an EC2 instance can only have one IAM
			// role, the test impersonates a dedicated role that maps to that user.
			name: "provider=" + provider + ",outboundToken",
			skip: func() (bool, string) {
				if provider != "AWS" {
					return true, "GetWebIdentityToken outbound token is supported only on AWS"
				}
				return false, ""
			},
			setupCfg: func(_ *testing.T, config *Config) {
				config.WorkloadIdentityProvider = provider
				config.WorkloadIdentityAwsUseOutboundToken = ConfigBoolTrue
				config.WorkloadIdentityImpersonationPath = []string{"arn:aws:iam::376129840140:role/drivers-wif-automated-tests-with-issuer"}
			},
			expectedUsername: "TEST_WIF_E2E_AWS_WITH_ISSUER",
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			if tc.skip != nil {
				if skip, msg := tc.skip(); skip {
					t.Skip(msg)
				}
			}
			config := &Config{
				Account:       account,
				Host:          host,
				Authenticator: AuthTypeWorkloadIdentityFederation,
			}
			tc.setupCfg(t, config)
			connector := NewConnector(SnowflakeDriver{}, *config)
			db := sql.OpenDB(connector)
			defer db.Close()
			currentUser := runSelectCurrentUser(t, db)
			assertEqualE(t, currentUser, tc.expectedUsername)
		})
	}
}

func TestAwsStsEndpointFor(t *testing.T) {
	testcases := []struct {
		name              string
		wifHost           string
		region            string
		expectedAuthority string
		expectedBaseURL   string
		expectedOverride  bool
		expectedError     string
	}{
		// Regional defaults come from the SDK's endpoint rules. The commercial and
		// China cases pin that delegating there did not move the endpoint anyone
		// already depends on; the rest are partitions whose DNS suffix is not
		// amazonaws.com and which a hand-formatted hostname gets wrong.
		{name: "regional default us-east-1", wifHost: "", region: "us-east-1", expectedAuthority: "sts.us-east-1.amazonaws.com", expectedBaseURL: "https://sts.us-east-1.amazonaws.com", expectedOverride: false},
		{name: "regional default cn-north-1", wifHost: "", region: "cn-north-1", expectedAuthority: "sts.cn-north-1.amazonaws.com.cn", expectedBaseURL: "https://sts.cn-north-1.amazonaws.com.cn", expectedOverride: false},
		{name: "regional default us-gov-west-1", wifHost: "", region: "us-gov-west-1", expectedAuthority: "sts.us-gov-west-1.amazonaws.com", expectedBaseURL: "https://sts.us-gov-west-1.amazonaws.com", expectedOverride: false},
		{name: "regional default us-iso-east-1", wifHost: "", region: "us-iso-east-1", expectedAuthority: "sts.us-iso-east-1.c2s.ic.gov", expectedBaseURL: "https://sts.us-iso-east-1.c2s.ic.gov", expectedOverride: false},
		{name: "regional default us-isob-east-1", wifHost: "", region: "us-isob-east-1", expectedAuthority: "sts.us-isob-east-1.sc2s.sgov.gov", expectedBaseURL: "https://sts.us-isob-east-1.sc2s.sgov.gov", expectedOverride: false},
		{name: "regional default eu-isoe-west-1", wifHost: "", region: "eu-isoe-west-1", expectedAuthority: "sts.eu-isoe-west-1.cloud.adc-e.uk", expectedBaseURL: "https://sts.eu-isoe-west-1.cloud.adc-e.uk", expectedOverride: false},
		{name: "regional default us-isof-south-1", wifHost: "", region: "us-isof-south-1", expectedAuthority: "sts.us-isof-south-1.csp.hci.ic.gov", expectedBaseURL: "https://sts.us-isof-south-1.csp.hci.ic.gov", expectedOverride: false},
		{name: "regional default eusc-de-east-1", wifHost: "", region: "eusc-de-east-1", expectedAuthority: "sts.eusc-de-east-1.amazonaws.eu", expectedBaseURL: "https://sts.eusc-de-east-1.amazonaws.eu", expectedOverride: false},
		{name: "bare host", wifHost: "sts.custom.example.com", region: "us-custom-1", expectedAuthority: "sts.custom.example.com", expectedBaseURL: "https://sts.custom.example.com", expectedOverride: true},
		{name: "host with port", wifHost: "sts.custom.example.com:8443", region: "us-custom-1", expectedAuthority: "sts.custom.example.com:8443", expectedBaseURL: "https://sts.custom.example.com:8443", expectedOverride: true},
		{name: "full URL", wifHost: "https://sts.custom.example.com", region: "us-custom-1", expectedAuthority: "sts.custom.example.com", expectedBaseURL: "https://sts.custom.example.com", expectedOverride: true},
		{name: "trailing slashes", wifHost: "https://sts.custom.example.com///", region: "us-custom-1", expectedAuthority: "sts.custom.example.com", expectedBaseURL: "https://sts.custom.example.com", expectedOverride: true},
		{name: "http scheme", wifHost: "http://sts.custom.example.com", region: "us-custom-1", expectedAuthority: "sts.custom.example.com", expectedBaseURL: "http://sts.custom.example.com", expectedOverride: true},
		{name: "whitespace", wifHost: "  sts.custom.example.com  ", region: "us-custom-1", expectedAuthority: "sts.custom.example.com", expectedBaseURL: "https://sts.custom.example.com", expectedOverride: true},
		{name: "invalid scheme", wifHost: "ftp://sts.custom.example.com", region: "us-custom-1", expectedError: `workloadIdentityHost "ftp://sts.custom.example.com" must use https or http, got scheme "ftp"`},
		{name: "with query", wifHost: "https://sts.custom.example.com?Action=Foo", region: "us-custom-1", expectedError: `workloadIdentityHost "https://sts.custom.example.com?Action=Foo" must not contain user info, a query or a fragment`},
		{name: "no hostname", wifHost: "https:///sts", region: "us-custom-2", expectedError: `workloadIdentityHost "https:///sts" does not contain a hostname`},
		{name: "fragment", wifHost: "https://sts.custom.example.com#frag", region: "us-custom-2", expectedError: `workloadIdentityHost "https://sts.custom.example.com#frag" must not contain user info, a query or a fragment`},
		{name: "user info", wifHost: "https://user:pass@sts.custom.example.com", region: "us-custom-2", expectedError: `workloadIdentityHost "https://user:pass@sts.custom.example.com" must not contain user info, a query or a fragment`}, // pragma: allowlist secret
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			cfg := &Config{WorkloadIdentityHost: tc.wifHost}
			endpoint, err := awsStsEndpointFor(context.Background(), cfg, tc.region)
			if tc.expectedError != "" {
				assertNotNilE(t, err)
				assertEqualE(t, err.Error(), tc.expectedError)
			} else {
				assertNilE(t, err)
				assertEqualE(t, endpoint.authority, tc.expectedAuthority)
				assertEqualE(t, endpoint.baseURL, tc.expectedBaseURL)
				assertEqualE(t, endpoint.overridden, tc.expectedOverride)
			}
		})
	}
}

// TestAzureEntraResourceQueryEncoding is a regression test for SNOW-3649876.
// WorkloadIdentityEntraResource (and the managed identity client id) must be
// percent-encoded when building the Azure IMDS / App Service identity query
// string, so a value containing "&" or "=" stays confined to its own parameter
// instead of adding or replacing others such as client_id.
func TestAzureEntraResourceQueryEncoding(t *testing.T) {
	const resourceWithSeparators = "api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad&client_id=other-uami-guid"
	const legitResource = "api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad"

	t.Run("VM path keeps the value in one parameter", func(t *testing.T) {
		creator := &azureIdentityAttestationCreator{
			azureMetadataServiceBaseURL:   defaultAzureMetadataServiceBase,
			workloadIdentityEntraResource: resourceWithSeparators,
		}
		req, err := creator.azureVMIdentityRequest("")
		assertNilF(t, err)
		q := req.URL.Query()
		// The embedded client_id must NOT appear as its own query parameter.
		assertEqualE(t, "", q.Get("client_id"))
		// The whole value must round-trip as the single resource value.
		assertEqualE(t, resourceWithSeparators, q.Get("resource"))
		assertEqualE(t, "2018-02-01", q.Get("api-version"))
	})

	t.Run("Functions path keeps the value in one parameter", func(t *testing.T) {
		creator := &azureIdentityAttestationCreator{
			workloadIdentityEntraResource: resourceWithSeparators,
		}
		req, err := creator.azureFunctionsIdentityRequest("http://127.0.0.1:41812/msi/token", "secret-header", "")
		assertNilF(t, err)
		q := req.URL.Query()
		assertEqualE(t, "", q.Get("client_id"))
		assertEqualE(t, resourceWithSeparators, q.Get("resource"))
		assertEqualE(t, "2019-08-01", q.Get("api-version"))
	})

	t.Run("legitimate resource round-trips", func(t *testing.T) {
		creator := &azureIdentityAttestationCreator{
			azureMetadataServiceBaseURL:   defaultAzureMetadataServiceBase,
			workloadIdentityEntraResource: legitResource,
		}
		req, err := creator.azureVMIdentityRequest("")
		assertNilF(t, err)
		q := req.URL.Query()
		assertEqualE(t, legitResource, q.Get("resource"))
		assertEqualE(t, "2018-02-01", q.Get("api-version"))
		assertEqualE(t, "", q.Get("client_id"))
	})
}

// TestAzureVMIdentityRequestClientID checks that the VM IMDS and App Service
// token requests include client_id when a user-assigned managed identity is
// selected, and omit it when neither WorkloadIdentityAzureClientID nor
// MANAGED_IDENTITY_CLIENT_ID is set. The config value wins when both are set.
func TestAzureVMIdentityRequestClientID(t *testing.T) {
	clientIDs := func(t *testing.T, cfg *Config) (vm, functions string) {
		t.Helper()
		creator := &azureIdentityAttestationCreator{
			cfg:                              cfg,
			azureMetadataServiceBaseURL:      defaultAzureMetadataServiceBase,
			workloadIdentityEntraResource:    determineEntraResource(cfg),
			azureAttestationMetadataProvider: &defaultAzureAttestationMetadataProvider{},
		}
		clientID := creator.managedIdentityClientID()
		vmReq, err := creator.azureVMIdentityRequest(clientID)
		assertNilF(t, err)
		functionsReq, err := creator.azureFunctionsIdentityRequest("http://127.0.0.1:41812/msi/token", "identity-header", clientID)
		assertNilF(t, err)
		return vmReq.URL.Query().Get("client_id"), functionsReq.URL.Query().Get("client_id")
	}

	t.Run("omits client_id when unset", func(t *testing.T) {
		t.Setenv("MANAGED_IDENTITY_CLIENT_ID", "")
		vm, functions := clientIDs(t, nil)
		assertEqualE(t, "", vm)
		assertEqualE(t, "", functions)
	})

	t.Run("includes MANAGED_IDENTITY_CLIENT_ID", func(t *testing.T) {
		t.Setenv("MANAGED_IDENTITY_CLIENT_ID", "uami-from-env")
		vm, functions := clientIDs(t, nil)
		assertEqualE(t, "uami-from-env", vm)
		assertEqualE(t, "uami-from-env", functions)
	})

	t.Run("includes WorkloadIdentityAzureClientID", func(t *testing.T) {
		t.Setenv("MANAGED_IDENTITY_CLIENT_ID", "")
		cfg := &Config{WorkloadIdentityAzureClientID: "uami-from-config"}
		vm, functions := clientIDs(t, cfg)
		assertEqualE(t, "uami-from-config", vm)
		assertEqualE(t, "uami-from-config", functions)
	})

	t.Run("config overrides the environment variable", func(t *testing.T) {
		t.Setenv("MANAGED_IDENTITY_CLIENT_ID", "uami-from-env")
		cfg := &Config{WorkloadIdentityAzureClientID: "uami-from-config"}
		vm, functions := clientIDs(t, cfg)
		assertEqualE(t, "uami-from-config", vm)
		assertEqualE(t, "uami-from-config", functions)
	})

	t.Run("percent-encodes client_id", func(t *testing.T) {
		const clientID = "uami&object_id=other"
		creator := &azureIdentityAttestationCreator{
			azureMetadataServiceBaseURL:   defaultAzureMetadataServiceBase,
			workloadIdentityEntraResource: determineEntraResource(nil),
		}
		vmReq, err := creator.azureVMIdentityRequest(clientID)
		assertNilF(t, err)
		functionsReq, err := creator.azureFunctionsIdentityRequest("http://127.0.0.1:41812/msi/token", "identity-header", clientID)
		assertNilF(t, err)
		for _, req := range []*http.Request{vmReq, functionsReq} {
			q := req.URL.Query()
			assertEqualE(t, clientID, q.Get("client_id"))
			assertEqualE(t, "", q.Get("object_id"))
		}
	})
}

// TestValidateLocalMetadataEndpoint is a unit test for the host validation
// added for SNOW-3649837: only loopback / link-local identity endpoints are
// permitted, so IDENTITY_ENDPOINT cannot direct the request, or the
// IDENTITY_HEADER value, to a remote host.
func TestValidateLocalMetadataEndpoint(t *testing.T) {
	tests := []struct {
		name     string
		endpoint string
		wantErr  bool
	}{
		{"loopback ipv4", "http://127.0.0.1:41812/msi/token", false},
		{"loopback ipv4 non-localhost octet", "http://127.5.6.7/msi/token", false},
		{"localhost hostname", "http://localhost:41812/msi/token", false},
		{"loopback ipv6", "http://[::1]:41812/msi/token", false},
		{"link-local ipv4 (IMDS)", "http://169.254.169.254/metadata/identity", false},
		{"https loopback", "https://127.0.0.1/msi", false},
		{"remote https host", "https://alice.example/msi", true},
		{"public ip", "http://8.8.8.8/msi", true},
		{"private routable ip", "http://10.0.0.5/msi", true},
		{"non-http scheme", "file:///etc/passwd", true},
		{"empty endpoint", "", true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := validateLocalMetadataEndpoint(tt.endpoint)
			if tt.wantErr {
				assertNotNilE(t, err)
			} else {
				assertNilE(t, err)
			}
		})
	}
}

// TestAzureFunctionsIdentityRequestEndpointValidation is the behavioral
// regression test for SNOW-3649837: a non-loopback IDENTITY_ENDPOINT must be
// rejected before the request is built, so the platform-managed IDENTITY_HEADER
// secret is never attached or sent; a loopback endpoint is accepted and the
// secret is attached as expected.
func TestAzureFunctionsIdentityRequestEndpointValidation(t *testing.T) {
	const secret = "platform-managed-identity-header-secret" // pragma: allowlist secret

	t.Run("non-loopback endpoint rejected and secret not sent", func(t *testing.T) {
		creator := &azureIdentityAttestationCreator{
			workloadIdentityEntraResource: determineEntraResource(nil),
		}
		req, err := creator.azureFunctionsIdentityRequest("https://attacker.example/msi", secret, "")
		// No request must be constructed, hence the secret cannot leave the process.
		assertNotNilE(t, err)
		assertNilF(t, req)
	})

	t.Run("loopback endpoint accepted and secret attached", func(t *testing.T) {
		creator := &azureIdentityAttestationCreator{
			workloadIdentityEntraResource: determineEntraResource(nil),
		}
		req, err := creator.azureFunctionsIdentityRequest("http://127.0.0.1:41812/msi/token", secret, "")
		assertNilF(t, err)
		assertNotNilF(t, req)
		assertEqualE(t, secret, req.Header.Get("X-IDENTITY-HEADER"))
	})
}

// TestWorkloadIdentityHostRoutesOutboundStsCalls drives the real assumeRole and
// awsWebIdentityToken against a local server. These are the only AWS WIF paths
// that open a socket to STS from the driver, so they are the paths that hard-fail
// in a partition where the regional default does not resolve. The mock metadata
// provider used by the other AWS tests replaces both methods wholesale, so
// nothing else in the suite would catch withStsBaseEndpoint being dropped here.
func TestWorkloadIdentityHostRoutesOutboundStsCalls(t *testing.T) {
	const assumeRoleResponse = `<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
  <AssumeRoleResult>
    <Credentials>
      <AccessKeyId>CHAINED_KEY</AccessKeyId>
      <SecretAccessKey>CHAINED_SECRET</SecretAccessKey>
      <SessionToken>CHAINED_TOKEN</SessionToken>
      <Expiration>2035-01-01T00:00:00Z</Expiration>
    </Credentials>
  </AssumeRoleResult>
</AssumeRoleResponse>`

	const webIdentityResponse = `<GetWebIdentityTokenResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
  <GetWebIdentityTokenResult>
    <WebIdentityToken>header.payload.signature</WebIdentityToken>
  </GetWebIdentityTokenResult>
</GetWebIdentityTokenResponse>`

	// newStsStub returns a server that records the Action it was asked for, so the
	// assertions prove the call actually landed here rather than at the regional default.
	newStsStub := func(t *testing.T, body string) (*httptest.Server, *[]string) {
		t.Helper()
		var actions []string
		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if err := r.ParseForm(); err == nil {
				actions = append(actions, r.Form.Get("Action"))
			}
			w.Header().Set("Content-Type", "text/xml")
			_, _ = w.Write([]byte(body))
		}))
		t.Cleanup(srv.Close)
		return srv, &actions
	}

	baseCreds := aws.Credentials{AccessKeyID: "k", SecretAccessKey: "s", SessionToken: "t"}

	t.Run("role chaining reaches the configured host", func(t *testing.T) {
		srv, actions := newStsStub(t, assumeRoleResponse)
		provider := &defaultAwsAttestationMetadataProvider{
			ctx:    context.Background(),
			cfg:    &Config{WorkloadIdentityHost: srv.URL},
			awsCfg: aws.Config{Region: "us-custom-1"},
		}

		creds, err := provider.assumeRole(baseCreds, "arn:aws:iam::123456789012:role/target")
		assertNilF(t, err)
		assertEqualE(t, creds.AccessKeyID, "CHAINED_KEY")
		assertEqualE(t, creds.SessionToken, "CHAINED_TOKEN")
		assertEqualE(t, len(*actions), 1)
		assertEqualE(t, (*actions)[0], "AssumeRole")
	})

	t.Run("every hop of a multi-role path reaches the configured host", func(t *testing.T) {
		srv, actions := newStsStub(t, assumeRoleResponse)
		provider := &defaultAwsAttestationMetadataProvider{
			ctx: context.Background(),
			cfg: &Config{
				WorkloadIdentityHost: srv.URL,
				WorkloadIdentityImpersonationPath: []string{
					"arn:aws:iam::111111111111:role/first",
					"arn:aws:iam::222222222222:role/second",
				},
			},
			awsCfg: aws.Config{
				Region:      "us-custom-1",
				Credentials: credentials.StaticCredentialsProvider{Value: baseCreds},
			},
		}

		creds, err := provider.awsCredentialsViaRoleChaining()
		assertNilF(t, err)
		assertEqualE(t, creds.AccessKeyID, "CHAINED_KEY")
		// One AssumeRole per role in the path - a hop that ignored the override
		// would have gone to the regional endpoint and never shown up here.
		assertEqualE(t, len(*actions), 2)
	})

	t.Run("outbound token reaches the configured host", func(t *testing.T) {
		srv, actions := newStsStub(t, webIdentityResponse)
		provider := &defaultAwsAttestationMetadataProvider{
			ctx:    context.Background(),
			cfg:    &Config{WorkloadIdentityHost: srv.URL},
			awsCfg: aws.Config{Region: "us-custom-1"},
		}

		token, err := provider.awsWebIdentityToken(baseCreds, "us-custom-1")
		assertNilF(t, err)
		assertEqualE(t, token, "header.payload.signature")
		assertEqualE(t, len(*actions), 1)
		assertEqualE(t, (*actions)[0], "GetWebIdentityToken")
	})

	t.Run("a malformed host fails before any request is made", func(t *testing.T) {
		_, actions := newStsStub(t, assumeRoleResponse)
		provider := &defaultAwsAttestationMetadataProvider{
			ctx:    context.Background(),
			cfg:    &Config{WorkloadIdentityHost: "ftp://sts.custom.example.com"},
			awsCfg: aws.Config{Region: "us-custom-1"},
		}

		_, err := provider.assumeRole(baseCreds, "arn:aws:iam::123456789012:role/target")
		assertNotNilF(t, err)
		assertStringContainsE(t, err.Error(), "must use https or http")
		// The stub stays silent - validation rejects the host before the SDK is built,
		// so role chaining never falls back to the regional endpoint.
		assertEqualE(t, len(*actions), 0)
	})
}