diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index bb7b36a9d..9d7f6a0b0 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,7 @@ ### Bug Fixes +- Select the Azure managed identity endpoint and request protocol from the host environment instead of always using IMDS. - Recover from concurrent OAuth token cache writes by using the fresh token stored by another process. - Fix Spark runtime selection for major-only runtime versions. - Accept JSON numbers and decimal strings when unmarshalling `int64` API response fields. diff --git a/config/auth_azure_msi.go b/config/auth_azure_msi.go index de51acccc..dae76fef0 100644 --- a/config/auth_azure_msi.go +++ b/config/auth_azure_msi.go @@ -6,6 +6,12 @@ import ( "errors" "fmt" "net/http" + "net/url" + "os" + "path/filepath" + "runtime" + "strconv" + "strings" "time" "github.com/databricks/databricks-sdk-go/config/credentials" @@ -21,12 +27,28 @@ var ( errInvalidTokenExpiry = errors.New("invalid token expiry") ) -// well-known URL for Azure Instance Metadata Service (IMDS) -// https://learn.microsoft.com/en-us/azure-stack/user/instance-metadata-service -var instanceMetadataPrefix = "http://169.254.169.254/metadata" +const ( + azureIMDSAuthority = "http://169.254.169.254" + azureIMDSTokenPath = "/metadata/identity/oauth2/token" -// timeout to wait for IMDS response -const azureMsiTimeout = 10 * time.Second + azureIdentityEndpointEnv = "IDENTITY_ENDPOINT" + azureIdentityHeaderEnv = "IDENTITY_HEADER" + azureIdentityServerThumbprintEnv = "IDENTITY_SERVER_THUMBPRINT" + azureIMDSEndpointEnv = "IMDS_ENDPOINT" + azureMSIEndpointEnv = "MSI_ENDPOINT" + azureMSISecretEnv = "MSI_SECRET" + azurePodIdentityAuthorityHostEnv = "AZURE_POD_IDENTITY_AUTHORITY_HOST" + azureAuthorityHostEnv = "AZURE_AUTHORITY_HOST" + azureClientIDEnv = "AZURE_CLIENT_ID" + azureTenantIDEnv = "AZURE_TENANT_ID" + azureFederatedTokenFileEnv = "AZURE_FEDERATED_TOKEN_FILE" + azureArcLinuxTokenDirectory = "/var/opt/azcmagent/tokens" + azureArcMaximumSecretFileSize int64 = 4096 + + // Managed identity endpoints are local services and should fail promptly + // when the SDK isn't running in the expected Azure environment. + azureMSITimeout = 10 * time.Second +) type AzureMsiCredentials struct{} @@ -55,72 +77,514 @@ func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (creden return newVisitorOAuthCredentials(visitor, inner), nil } -// implementing azureHostResolver for ensureWorkspaceUrl to work +// tokenSourceFor implements azureHostResolver so MSI can resolve a workspace host. func (c AzureMsiCredentials) tokenSourceFor(_ context.Context, cfg *Config, _, resource string) auth.TokenSource { - return &azureMsiTokenSource{ - client: cfg.refreshClient, - resource: resource, - clientId: cfg.AzureClientID, + return newAzureMSITokenSource(cfg.refreshClient, resource, cfg.AzureClientID) +} + +// NewAzureMsiTokenSource returns [oauth2.TokenSource] for passwordless authentication via Azure Managed Identity. +func NewAzureMsiTokenSource(client *httpclient.ApiClient, resource, clientID string) oauth2.TokenSource { + return authconv.OAuth2TokenSource(newAzureMSITokenSource(client, resource, clientID)) +} + +type azureManagedIdentityType string + +const ( + azureManagedIdentityServiceFabric azureManagedIdentityType = "Service Fabric" + azureManagedIdentityAppService azureManagedIdentityType = "App Service" + azureManagedIdentityArc azureManagedIdentityType = "Azure Arc" + azureManagedIdentityML azureManagedIdentityType = "Azure ML" + azureManagedIdentityCloudShell azureManagedIdentityType = "Cloud Shell" + azureManagedIdentityWorkload azureManagedIdentityType = "workload identity" + azureManagedIdentityIMDS azureManagedIdentityType = "IMDS" +) + +type azureManagedIdentityEndpoint struct { + identityType azureManagedIdentityType + url string + secret string + tenantID string + tokenFile string +} + +// azureManagedIdentityEndpointFromEnvironment follows the selection order used +// by Azure Identity's ManagedIdentityCredential. The first matching Azure host +// environment wins, and IMDS is used only when no host-specific environment is +// configured. +func azureManagedIdentityEndpointFromEnvironment(clientID string) (azureManagedIdentityEndpoint, string, error) { + identityEndpoint := getenv(azureIdentityEndpointEnv) + if identityEndpoint != "" { + identityHeader := getenv(azureIdentityHeaderEnv) + if identityHeader != "" { + if getenv(azureIdentityServerThumbprintEnv) != "" { + endpoint := azureManagedIdentityEndpoint{ + identityType: azureManagedIdentityServiceFabric, + url: identityEndpoint, + secret: identityHeader, + } + if clientID != "" { + return endpoint, clientID, fmt.Errorf("azure_client_id is not supported by %q managed identity", endpoint.identityType) + } + return endpoint, clientID, nil + } + return azureManagedIdentityEndpoint{ + identityType: azureManagedIdentityAppService, + url: identityEndpoint, + secret: identityHeader, + }, clientID, nil + } + if getenv(azureIMDSEndpointEnv) != "" { + endpoint := azureManagedIdentityEndpoint{ + identityType: azureManagedIdentityArc, + url: identityEndpoint, + } + if clientID != "" { + return endpoint, clientID, fmt.Errorf("azure_client_id is not supported by %q managed identity", endpoint.identityType) + } + return endpoint, clientID, nil + } + return azureManagedIdentityEndpoint{}, clientID, fmt.Errorf( + "no managed identity endpoint found: %s requires %s or %s", + azureIdentityEndpointEnv, + azureIdentityHeaderEnv, + azureIMDSEndpointEnv, + ) + } + + if msiEndpoint := getenv(azureMSIEndpointEnv); msiEndpoint != "" { + if secret := getenv(azureMSISecretEnv); secret != "" { + return azureManagedIdentityEndpoint{ + identityType: azureManagedIdentityML, + url: msiEndpoint, + secret: secret, + }, clientID, nil + } + endpoint := azureManagedIdentityEndpoint{ + identityType: azureManagedIdentityCloudShell, + url: msiEndpoint, + } + if clientID != "" { + return endpoint, clientID, fmt.Errorf("azure_client_id is not supported by %q managed identity", endpoint.identityType) + } + return endpoint, clientID, nil + } + + authorityHost := getenv(azureAuthorityHostEnv) + tenantID := getenv(azureTenantIDEnv) + tokenFile := getenv(azureFederatedTokenFileEnv) + if authorityHost != "" && tenantID != "" && tokenFile != "" { + if clientID == "" { + clientID = getenv(azureClientIDEnv) + } + endpoint := azureManagedIdentityEndpoint{ + identityType: azureManagedIdentityWorkload, + url: authorityHost, + tenantID: tenantID, + tokenFile: tokenFile, + } + if clientID == "" { + return endpoint, clientID, errors.New("azure_client_id is required for workload identity") + } + return endpoint, clientID, nil + } + + imdsAuthority := getenv(azurePodIdentityAuthorityHostEnv) + if imdsAuthority == "" { + imdsAuthority = azureIMDSAuthority } + return azureManagedIdentityEndpoint{ + identityType: azureManagedIdentityIMDS, + url: strings.TrimRight(imdsAuthority, "/") + azureIMDSTokenPath, + }, clientID, nil } -// NewAzureMsiTokenSource returns [oauth2.TokenSource] for a passwordless authentication via Azure Managed identity -func NewAzureMsiTokenSource(client *httpclient.ApiClient, resource, clientId string) oauth2.TokenSource { - return authconv.OAuth2TokenSource(&azureMsiTokenSource{ - client: client, - resource: resource, - clientId: clientId, - }) +type azureMSITokenSource struct { + client *httpclient.ApiClient + resource string + clientID string + endpoint azureManagedIdentityEndpoint + endpointErr error + now func() time.Time + readAzureArcSecret func(string) (string, error) } -type azureMsiTokenSource struct { - client *httpclient.ApiClient - resource string - clientId string +func newAzureMSITokenSource(client *httpclient.ApiClient, resource, clientID string) *azureMSITokenSource { + endpoint, clientID, err := azureManagedIdentityEndpointFromEnvironment(clientID) + return &azureMSITokenSource{ + client: client, + resource: resource, + clientID: clientID, + endpoint: endpoint, + endpointErr: err, + now: time.Now, + readAzureArcSecret: readAzureArcSecret, + } } -func (s azureMsiTokenSource) Token(ctx context.Context) (*oauth2.Token, error) { - ctx, cancel := context.WithTimeout(ctx, azureMsiTimeout) +func (s *azureMSITokenSource) Token(ctx context.Context) (*oauth2.Token, error) { + if s.endpointErr != nil { + return nil, s.endpointErr + } + ctx, cancel := context.WithTimeout(ctx, azureMSITimeout) defer cancel() - query := map[string]string{ - "api-version": "2018-02-01", - "resource": s.resource, + + now := s.now + if now == nil { + now = time.Now } - if s.clientId != "" { - query["client_id"] = s.clientId + requestTime := now() + + switch s.endpoint.identityType { + case azureManagedIdentityWorkload: + return s.workloadIdentityToken(ctx, requestTime) + case azureManagedIdentityServiceFabric, + azureManagedIdentityAppService, + azureManagedIdentityML, + azureManagedIdentityCloudShell, + azureManagedIdentityIMDS: + request, err := s.endpoint.tokenRequest(s.resource, s.clientID) + if err != nil { + return nil, err + } + inner, err := s.requestToken(ctx, request) + if err != nil { + return nil, err + } + return inner.managedIdentityToken(requestTime, s.endpoint.identityType == azureManagedIdentityML) + case azureManagedIdentityArc: + return s.azureArcToken(ctx, requestTime) + default: + return nil, fmt.Errorf("unknown Azure managed identity type: %q", s.endpoint.identityType) + } +} + +type azureMSITokenRequest struct { + method string + url string + headers map[string]string + data map[string]string + form bool + encodedForm string +} + +func (e azureManagedIdentityEndpoint) tokenRequest(resource, clientID string) (azureMSITokenRequest, error) { + request := azureMSITokenRequest{ + method: http.MethodGet, + url: e.url, + headers: map[string]string{}, + data: map[string]string{ + "resource": resource, + }, + } + + switch e.identityType { + case azureManagedIdentityServiceFabric: + request.headers["Secret"] = e.secret + request.data["api-version"] = "2019-07-01-preview" + case azureManagedIdentityAppService: + request.headers["X-IDENTITY-HEADER"] = e.secret + request.data["api-version"] = "2019-08-01" + if clientID != "" { + request.data["client_id"] = clientID + } + case azureManagedIdentityArc: + request.headers["Metadata"] = "true" + request.data["api-version"] = "2020-06-01" + case azureManagedIdentityML: + request.headers["secret"] = e.secret + request.data["api-version"] = "2017-09-01" + if clientID != "" { + request.data["clientid"] = clientID + } + case azureManagedIdentityCloudShell: + request.method = http.MethodPost + request.headers["Metadata"] = "true" + request.form = true + case azureManagedIdentityIMDS: + request.headers["Metadata"] = "true" + request.data["api-version"] = "2018-02-01" + if clientID != "" { + request.data["client_id"] = clientID + } + case azureManagedIdentityWorkload: + return azureMSITokenRequest{}, errors.New("workload identity requires a token exchange request") + default: + return azureMSITokenRequest{}, fmt.Errorf("unknown Azure managed identity type: %q", e.identityType) } + + return request, nil +} + +func (s *azureMSITokenSource) requestToken(ctx context.Context, request azureMSITokenRequest) (msiToken, error) { var inner msiToken - err := s.client.Do(ctx, http.MethodGet, - fmt.Sprintf("%s/identity/oauth2/token", instanceMetadataPrefix), - httpclient.WithRequestHeader("Metadata", "true"), - httpclient.WithRequestData(query), + requestOptions := []httpclient.DoOption{ + httpclient.WithRequestHeaders(request.headers), httpclient.WithResponseUnmarshal(&inner), - ) + } + if request.encodedForm != "" { + requestOptions = append( + requestOptions, + httpclient.WithRequestHeader("Content-Type", httpclient.UrlEncodedContentType), + httpclient.WithRequestData(strings.NewReader(request.encodedForm)), + ) + } else if request.form { + requestOptions = append(requestOptions, httpclient.WithUrlEncodedData(request.data)) + } else { + requestOptions = append(requestOptions, httpclient.WithRequestData(request.data)) + } + if err := s.client.Do(ctx, request.method, request.url, requestOptions...); err != nil { + return msiToken{}, fmt.Errorf("request managed identity token from %q: %w", request.url, err) + } + return inner, nil +} + +func (s *azureMSITokenSource) workloadIdentityToken(ctx context.Context, requestTime time.Time) (*oauth2.Token, error) { + assertion, err := os.ReadFile(s.endpoint.tokenFile) + if err != nil { + return nil, fmt.Errorf("read federated token file %q: %w", s.endpoint.tokenFile, err) + } + + tokenEndpoint := strings.TrimRight(s.endpoint.url, "/") + "/" + url.PathEscape(s.endpoint.tenantID) + "/oauth2/v2.0/token" + scope := s.resource + if !strings.HasSuffix(scope, "/.default") { + scope = strings.TrimRight(scope, "/") + "/.default" + } + form := url.Values{ + "client_assertion": []string{string(assertion)}, + "client_assertion_type": []string{"urn:ietf:params:oauth:client-assertion-type:jwt-bearer"}, + "client_id": []string{s.clientID}, + "grant_type": []string{"client_credentials"}, + "scope": []string{scope}, + } + request := azureMSITokenRequest{ + method: http.MethodPost, + url: tokenEndpoint, + headers: map[string]string{}, + encodedForm: form.Encode(), + } + inner, err := s.requestToken(ctx, request) + if err != nil { + return nil, err + } + return inner.managedIdentityToken(requestTime, false) +} + +func (s *azureMSITokenSource) azureArcToken(ctx context.Context, requestTime time.Time) (*oauth2.Token, error) { + request, err := s.endpoint.tokenRequest(s.resource, s.clientID) + if err != nil { + return nil, err + } + inner, err := s.requestToken(ctx, request) + if err == nil { + return inner.managedIdentityToken(requestTime, false) + } + + keyFile, challenged, challengeErr := azureArcChallengeKeyFile(err) + if challengeErr != nil { + return nil, fmt.Errorf("handle Azure Arc challenge from %q: %w", request.url, challengeErr) + } + if !challenged { + return nil, err + } + secretReader := s.readAzureArcSecret + if secretReader == nil { + secretReader = readAzureArcSecret + } + secret, err := secretReader(keyFile) + if err != nil { + return nil, fmt.Errorf("read Azure Arc managed identity secret %q: %w", keyFile, err) + } + request.headers["Authorization"] = "Basic " + secret + inner, err = s.requestToken(ctx, request) + if err != nil { + return nil, err + } + return inner.managedIdentityToken(requestTime, false) +} + +func azureArcChallengeKeyFile(err error) (string, bool, error) { + httpError := azureManagedIdentityHTTPError(err) + if httpError == nil || httpError.StatusCode != http.StatusUnauthorized { + return "", false, nil + } + challenge := httpError.Header().Get("WWW-Authenticate") + if challenge == "" { + return "", true, errors.New("Azure Arc managed identity response has no WWW-Authenticate header") + } + _, keyFile, ok := strings.Cut(challenge, "=") + if !ok { + return "", true, fmt.Errorf("invalid Azure Arc WWW-Authenticate header: %q", challenge) + } + keyFile = strings.Trim(strings.TrimSpace(keyFile), `"`) + if keyFile == "" { + return "", true, fmt.Errorf("invalid Azure Arc WWW-Authenticate header: %q", challenge) + } + return keyFile, true, nil +} + +func azureManagedIdentityHTTPError(err error) *httpclient.HttpError { + var azureError *tokenError + if errors.As(err, &azureError) { + return azureError.err + } + var httpError *httpclient.HttpError + if errors.As(err, &httpError) { + return httpError + } + return nil +} + +func readAzureArcSecret(keyFile string) (string, error) { + expectedDirectory, err := azureArcTokenDirectory() + if err != nil { + return "", err + } + return readAzureArcSecretFromDirectory(keyFile, expectedDirectory) +} + +func readAzureArcSecretFromDirectory(keyFile, expectedDirectory string) (string, error) { + info, err := os.Stat(keyFile) + if err != nil { + return "", err + } + if filepath.Clean(filepath.Dir(keyFile)) != filepath.Clean(expectedDirectory) { + return "", fmt.Errorf("unexpected Azure Arc managed identity secret directory %q", filepath.Dir(keyFile)) + } + if filepath.Ext(keyFile) != ".key" { + return "", fmt.Errorf("Azure Arc managed identity secret file %q must have a .key extension", keyFile) + } + if info.Size() > azureArcMaximumSecretFileSize { + return "", fmt.Errorf("Azure Arc managed identity secret file %q exceeds %d bytes", keyFile, azureArcMaximumSecretFileSize) + } + secret, err := os.ReadFile(keyFile) if err != nil { - return nil, fmt.Errorf("token request: %w", err) + return "", err + } + return string(secret), nil +} + +func azureArcTokenDirectory() (string, error) { + switch runtime.GOOS { + case "linux": + return azureArcLinuxTokenDirectory, nil + case "windows": + programData := getenv("PROGRAMDATA") + if programData == "" { + return "", errors.New("PROGRAMDATA is required for Azure Arc managed identity on Windows") + } + return filepath.Join(programData, "AzureConnectedMachineAgent", "Tokens"), nil + default: + return "", fmt.Errorf("Azure Arc managed identity is not supported on %q", runtime.GOOS) } - return inner.Token() } type msiToken struct { - TokenType string `json:"token_type"` - AccessToken string `json:"access_token,omitempty"` - RefreshToken string `json:"refresh_token,omitempty"` - ExpiresOn json.Number `json:"expires_on"` + TokenType string `json:"token_type"` + AccessToken string `json:"access_token,omitempty"` + RefreshToken string `json:"refresh_token,omitempty"` + ExpiresOn azureTokenExpiry `json:"expires_on"` + ExpiresIn azureTokenExpiry `json:"expires_in"` +} + +// azureTokenExpiry accepts the numeric and string expiry representations used +// by the different managed identity endpoints. Azure ML can also return a date +// string, which is parsed only for that environment. +type azureTokenExpiry string + +func (e *azureTokenExpiry) UnmarshalJSON(data []byte) error { + if len(data) > 0 && data[0] == '"' { + var value string + if err := json.Unmarshal(data, &value); err != nil { + return err + } + *e = azureTokenExpiry(value) + return nil + } + var value json.Number + if err := json.Unmarshal(data, &value); err != nil { + return err + } + *e = azureTokenExpiry(value.String()) + return nil +} + +func (e azureTokenExpiry) Int64() (int64, error) { + return strconv.ParseInt(string(e), 10, 64) +} + +func (e azureTokenExpiry) String() string { + return string(e) } func (token msiToken) Token() (*oauth2.Token, error) { + return token.oauth2Token(time.Time{}, false, false) +} + +func (token msiToken) managedIdentityToken(requestTime time.Time, parseAzureMLExpiry bool) (*oauth2.Token, error) { + return token.oauth2Token(requestTime, true, parseAzureMLExpiry) +} + +func (token msiToken) oauth2Token(requestTime time.Time, allowExpiresIn, parseAzureMLExpiry bool) (*oauth2.Token, error) { if token.AccessToken == "" { return nil, fmt.Errorf("token parse: %w", errInvalidToken) } - epoch, err := token.ExpiresOn.Int64() + expiry, err := token.expiry(requestTime, allowExpiresIn, parseAzureMLExpiry) if err != nil { - return nil, fmt.Errorf("%w: %s", errInvalidTokenExpiry, err) + return nil, err + } + tokenType := token.TokenType + if tokenType == "" { + tokenType = "Bearer" } return &oauth2.Token{ - TokenType: token.TokenType, + TokenType: tokenType, AccessToken: token.AccessToken, RefreshToken: token.RefreshToken, - Expiry: time.Unix(epoch, 0), + Expiry: expiry, }, nil } + +func (token msiToken) expiry(requestTime time.Time, allowExpiresIn, parseAzureMLExpiry bool) (time.Time, error) { + if token.ExpiresOn != "" { + epoch, err := token.ExpiresOn.Int64() + if err == nil { + return time.Unix(epoch, 0), nil + } + if parseAzureMLExpiry { + expiry, parseErr := parseAzureMLTokenExpiry(token.ExpiresOn.String()) + if parseErr == nil { + return expiry, nil + } + return time.Time{}, fmt.Errorf("%w: %s", errInvalidTokenExpiry, parseErr) + } + return time.Time{}, fmt.Errorf("%w: %s", errInvalidTokenExpiry, err) + } + if allowExpiresIn && token.ExpiresIn != "" { + seconds, err := token.ExpiresIn.Int64() + if err != nil { + return time.Time{}, fmt.Errorf("%w: %s", errInvalidTokenExpiry, err) + } + return requestTime.Add(time.Duration(seconds) * time.Second), nil + } + return time.Time{}, fmt.Errorf("%w: expires_on is missing", errInvalidTokenExpiry) +} + +func parseAzureMLTokenExpiry(value string) (time.Time, error) { + const utcSuffix = " +00:00" + if !strings.HasSuffix(value, utcSuffix) { + return time.Time{}, fmt.Errorf("unsupported Azure ML token expiry %q", value) + } + value = strings.TrimSuffix(value, utcSuffix) + formats := []string{ + "01/02/2006 15:04:05", + "1/2/2006 15:04:05", + "01/02/2006 03:04:05 PM", + "1/2/2006 3:04:05 PM", + } + for _, format := range formats { + if expiry, err := time.ParseInLocation(format, value, time.UTC); err == nil { + return expiry, nil + } + } + return time.Time{}, fmt.Errorf("unsupported Azure ML token expiry %q", value+utcSuffix) +} diff --git a/config/auth_azure_msi_endpoints_test.go b/config/auth_azure_msi_endpoints_test.go new file mode 100644 index 000000000..5d6a31304 --- /dev/null +++ b/config/auth_azure_msi_endpoints_test.go @@ -0,0 +1,550 @@ +package config + +import ( + "context" + "fmt" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/databricks/databricks-sdk-go/httpclient" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/google/go-cmp/cmp" +) + +const ( + testAzureMSIResource = "https://resource.test/" + testAzureMSIClientID = "test-client-id" +) + +func TestAzureMsiTokenSource_endpointSelection(t *testing.T) { + fixedNow := time.Date(2026, time.August, 17, 12, 0, 0, 0, time.UTC) + fixedExpiry := time.Date(2030, time.January, 1, 0, 0, 0, 0, time.UTC) + federatedTokenFile := filepath.Join(t.TempDir(), "federated-token") + if err := os.WriteFile(federatedTokenFile, []byte("federated-token"), 0o600); err != nil { + t.Fatalf("write federated token: %v", err) + } + + testCases := []struct { + name string + env map[string]string + clientID string + fixture fixtures.HTTPFixture + response any + wantToken string + wantExpiry time.Time + }{ + { + name: "service fabric", + env: map[string]string{ + azureIdentityEndpointEnv: "http://service-fabric.test/service-fabric/token", + azureIdentityHeaderEnv: "service-fabric-secret", + azureIdentityServerThumbprintEnv: "thumbprint", + }, + fixture: fixtures.HTTPFixture{ + Method: http.MethodGet, + Resource: "/service-fabric/token?api-version=2019-07-01-preview&resource=https%3A%2F%2Fresource.test%2F", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Secret": "service-fabric-secret", + }, + }, + response: azureMSITokenResponse("service-fabric-token", fixedExpiry), + wantToken: "service-fabric-token", + wantExpiry: fixedExpiry, + }, + { + name: "app service takes precedence over legacy endpoint", + env: map[string]string{ + azureIdentityEndpointEnv: "http://app-service.test/app-service/token", + azureIdentityHeaderEnv: "app-service-secret", + azureMSIEndpointEnv: "http://azure-ml.test/should-not-be-used", + azureMSISecretEnv: "legacy-secret", + azureAuthorityHostEnv: "http://workload.test/should-not-be-used", + azureTenantIDEnv: "tenant", + azureFederatedTokenFileEnv: federatedTokenFile, + }, + clientID: testAzureMSIClientID, + fixture: fixtures.HTTPFixture{ + Method: http.MethodGet, + Resource: "/app-service/token?api-version=2019-08-01&client_id=test-client-id&resource=https%3A%2F%2Fresource.test%2F", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "X-Identity-Header": "app-service-secret", + }, + }, + response: azureMSITokenResponse("app-service-token", fixedExpiry), + wantToken: "app-service-token", + wantExpiry: fixedExpiry, + }, + { + name: "azure arc", + env: map[string]string{ + azureIdentityEndpointEnv: "http://arc.test/arc/token", + azureIMDSEndpointEnv: "http://arc.test", + }, + fixture: fixtures.HTTPFixture{ + Method: http.MethodGet, + Resource: "/arc/token?api-version=2020-06-01&resource=https%3A%2F%2Fresource.test%2F", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Metadata": "true", + }, + }, + response: azureMSITokenResponse("azure-arc-token", fixedExpiry), + wantToken: "azure-arc-token", + wantExpiry: fixedExpiry, + }, + { + name: "azure ml", + env: map[string]string{ + azureMSIEndpointEnv: "http://azure-ml.test/azure-ml/token", + azureMSISecretEnv: "azure-ml-secret", + azureAuthorityHostEnv: "http://workload.test/should-not-be-used", + azureTenantIDEnv: "tenant", + azureFederatedTokenFileEnv: federatedTokenFile, + azurePodIdentityAuthorityHostEnv: "http://pod.test/should-not-be-used", + }, + clientID: testAzureMSIClientID, + fixture: fixtures.HTTPFixture{ + Method: http.MethodGet, + Resource: "/azure-ml/token?api-version=2017-09-01&clientid=test-client-id&resource=https%3A%2F%2Fresource.test%2F", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Secret": "azure-ml-secret", + }, + }, + response: map[string]any{ + "access_token": "azure-ml-token", + "expires_on": "01/01/2030 00:00:00 +00:00", + "token_type": "Bearer", + }, + wantToken: "azure-ml-token", + wantExpiry: fixedExpiry, + }, + { + name: "cloud shell", + env: map[string]string{ + azureMSIEndpointEnv: "http://cloud-shell.test/cloud-shell/token", + }, + fixture: fixtures.HTTPFixture{ + Method: http.MethodPost, + Resource: "/cloud-shell/token", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + "Metadata": "true", + }, + ExpectedRequest: url.Values{ + "resource": []string{testAzureMSIResource}, + }, + }, + response: azureMSITokenResponse("cloud-shell-token", fixedExpiry), + wantToken: "cloud-shell-token", + wantExpiry: fixedExpiry, + }, + { + name: "workload identity", + env: map[string]string{ + azureAuthorityHostEnv: "http://workload.test/authority/", + azureClientIDEnv: testAzureMSIClientID, + azureTenantIDEnv: "test-tenant", + azureFederatedTokenFileEnv: federatedTokenFile, + }, + fixture: fixtures.HTTPFixture{ + Method: http.MethodPost, + Resource: "/authority/test-tenant/oauth2/v2.0/token", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + }, + ExpectedRequest: url.Values{ + "client_assertion": []string{"federated-token"}, + "client_assertion_type": []string{"urn:ietf:params:oauth:client-assertion-type:jwt-bearer"}, + "client_id": []string{testAzureMSIClientID}, + "grant_type": []string{"client_credentials"}, + "scope": []string{"https://resource.test/.default"}, + }, + }, + response: map[string]any{ + "access_token": "workload-token", + "expires_in": 3600, + }, + wantToken: "workload-token", + wantExpiry: fixedNow.Add(time.Hour), + }, + { + name: "imds fallback", + env: map[string]string{}, + clientID: testAzureMSIClientID, + fixture: fixtures.HTTPFixture{ + Method: http.MethodGet, + Resource: "/metadata/identity/oauth2/token?api-version=2018-02-01&client_id=test-client-id&resource=https%3A%2F%2Fresource.test%2F", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Metadata": "true", + }, + }, + response: azureMSITokenResponse("imds-token", fixedExpiry), + wantToken: "imds-token", + wantExpiry: fixedExpiry, + }, + { + name: "pod identity authority overrides imds authority", + env: map[string]string{ + azureAuthorityHostEnv: "http://incomplete-workload.test", + azureTenantIDEnv: "tenant", + azurePodIdentityAuthorityHostEnv: "http://pod.test/custom/", + }, + fixture: fixtures.HTTPFixture{ + Method: http.MethodGet, + Resource: "/custom/metadata/identity/oauth2/token?api-version=2018-02-01&resource=https%3A%2F%2Fresource.test%2F", + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Metadata": "true", + }, + }, + response: azureMSITokenResponse("pod-identity-token", fixedExpiry), + wantToken: "pod-identity-token", + wantExpiry: fixedExpiry, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + withMockEnv(t, tc.env) + fixture := tc.fixture + fixture.Response = tc.response + httpClient := httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: fixtures.SliceTransport{fixture}, + }) + tokenSource := newAzureMSITokenSource(httpClient, testAzureMSIResource, tc.clientID) + tokenSource.now = func() time.Time { return fixedNow } + + got, gotErr := tokenSource.Token(context.Background()) + if gotErr != nil { + t.Fatalf("Token() failed: %v", gotErr) + } + if got.AccessToken != tc.wantToken { + t.Errorf("Token().AccessToken = %q, want %q", got.AccessToken, tc.wantToken) + } + if got.TokenType != "Bearer" { + t.Errorf("Token().TokenType = %q, want %q", got.TokenType, "Bearer") + } + if diff := cmp.Diff(tc.wantExpiry, got.Expiry); diff != "" { + t.Errorf("Token().Expiry mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestAzureMsiTokenSource_invalidEnvironment(t *testing.T) { + testCases := []struct { + name string + env map[string]string + clientID string + wantErr string + }{ + { + name: "identity endpoint is incomplete", + env: map[string]string{ + azureIdentityEndpointEnv: "http://identity.test/token", + azureMSIEndpointEnv: "http://cloud-shell.test/token", + }, + wantErr: "no managed identity endpoint found: IDENTITY_ENDPOINT requires IDENTITY_HEADER or IMDS_ENDPOINT", + }, + { + name: "service fabric rejects a client ID", + env: map[string]string{ + azureIdentityEndpointEnv: "http://service-fabric.test/token", + azureIdentityHeaderEnv: "secret", + azureIdentityServerThumbprintEnv: "thumbprint", + }, + clientID: testAzureMSIClientID, + wantErr: "azure_client_id is not supported by \"Service Fabric\" managed identity", + }, + { + name: "azure arc rejects a client ID", + env: map[string]string{ + azureIdentityEndpointEnv: "http://arc.test/token", + azureIMDSEndpointEnv: "http://arc.test", + }, + clientID: testAzureMSIClientID, + wantErr: "azure_client_id is not supported by \"Azure Arc\" managed identity", + }, + { + name: "cloud shell rejects a client ID", + env: map[string]string{ + azureMSIEndpointEnv: "http://cloud-shell.test/token", + }, + clientID: testAzureMSIClientID, + wantErr: "azure_client_id is not supported by \"Cloud Shell\" managed identity", + }, + { + name: "workload identity requires a client ID", + env: map[string]string{ + azureAuthorityHostEnv: "http://workload.test", + azureTenantIDEnv: "tenant", + azureFederatedTokenFileEnv: "token-file", + }, + wantErr: "azure_client_id is required for workload identity", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + withMockEnv(t, tc.env) + tokenSource := newAzureMSITokenSource(nil, testAzureMSIResource, tc.clientID) + got, gotErr := tokenSource.Token(context.Background()) + if got != nil { + t.Errorf("Token() = %#v, want nil", got) + } + if gotErr == nil || gotErr.Error() != tc.wantErr { + t.Errorf("Token() error = %q, want %q", gotErr, tc.wantErr) + } + }) + } +} + +func TestAzureMsiTokenSource_azureArcChallenge(t *testing.T) { + withMockEnv(t, map[string]string{ + azureIdentityEndpointEnv: "http://arc.test/token", + azureIMDSEndpointEnv: "http://arc.test", + }) + fixedExpiry := time.Date(2030, time.January, 1, 0, 0, 0, 0, time.UTC) + requestResource := "/token?api-version=2020-06-01&resource=https%3A%2F%2Fresource.test%2F" + httpClient := httpclient.NewApiClient(httpclient.ClientConfig{ + Transport: fixtures.SliceTransport{ + { + Method: http.MethodGet, + Resource: requestResource, + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Metadata": "true", + }, + Status: http.StatusUnauthorized, + ResponseHeaders: map[string][]string{ + "Www-Authenticate": []string{"Basic realm=/var/opt/azcmagent/tokens/identity.key"}, + }, + }, + { + Method: http.MethodGet, + Resource: requestResource, + ExpectedHeaders: map[string]string{ + "Accept": "application/json", + "Authorization": "Basic arc-secret", + "Metadata": "true", + }, + Response: azureMSITokenResponse("azure-arc-token", fixedExpiry), + }, + }, + }) + tokenSource := newAzureMSITokenSource(httpClient, testAzureMSIResource, "") + tokenSource.readAzureArcSecret = func(got string) (string, error) { + want := "/var/opt/azcmagent/tokens/identity.key" + if got != want { + t.Errorf("readAzureArcSecret() path = %q, want %q", got, want) + } + return "arc-secret", nil + } + + got, gotErr := tokenSource.Token(context.Background()) + if gotErr != nil { + t.Fatalf("Token() failed: %v", gotErr) + } + if got.AccessToken != "azure-arc-token" { + t.Errorf("Token().AccessToken = %q, want %q", got.AccessToken, "azure-arc-token") + } +} + +func TestAzureArcChallengeKeyFile(t *testing.T) { + testCases := []struct { + name string + status int + challenge string + wantKeyFile string + wantChallenged bool + wantErr string + }{ + { + name: "non-challenge response", + status: http.StatusInternalServerError, + }, + { + name: "missing header", + status: http.StatusUnauthorized, + wantChallenged: true, + wantErr: "Azure Arc managed identity response has no WWW-Authenticate header", + }, + { + name: "malformed header", + status: http.StatusUnauthorized, + challenge: "Basic realm", + wantChallenged: true, + wantErr: `invalid Azure Arc WWW-Authenticate header: "Basic realm"`, + }, + { + name: "valid header", + status: http.StatusUnauthorized, + challenge: `Basic realm="/var/opt/azcmagent/tokens/identity.key"`, + wantKeyFile: "/var/opt/azcmagent/tokens/identity.key", + wantChallenged: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + header := make(http.Header) + if tc.challenge != "" { + header.Set("WWW-Authenticate", tc.challenge) + } + err := &tokenError{ + err: &httpclient.HttpError{ + Response: &http.Response{ + StatusCode: tc.status, + Header: header, + }, + }, + } + + gotKeyFile, gotChallenged, gotErr := azureArcChallengeKeyFile(err) + if gotKeyFile != tc.wantKeyFile { + t.Errorf("azureArcChallengeKeyFile() key file = %q, want %q", gotKeyFile, tc.wantKeyFile) + } + if gotChallenged != tc.wantChallenged { + t.Errorf("azureArcChallengeKeyFile() challenged = %t, want %t", gotChallenged, tc.wantChallenged) + } + if tc.wantErr == "" { + if gotErr != nil { + t.Errorf("azureArcChallengeKeyFile() failed: %v", gotErr) + } + return + } + if gotErr == nil || gotErr.Error() != tc.wantErr { + t.Errorf("azureArcChallengeKeyFile() error = %q, want %q", gotErr, tc.wantErr) + } + }) + } +} + +func TestParseAzureMLTokenExpiry(t *testing.T) { + testCases := []struct { + name string + value string + want time.Time + wantErr bool + }{ + { + name: "Linux format", + value: "06/20/2030 02:57:58 +00:00", + want: time.Date(2030, time.June, 20, 2, 57, 58, 0, time.UTC), + }, + { + name: "Windows format", + value: "1/16/2030 5:24:12 AM +00:00", + want: time.Date(2030, time.January, 16, 5, 24, 12, 0, time.UTC), + }, + { + name: "invalid format", + value: "2030-01-01T00:00:00Z", + wantErr: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + got, gotErr := parseAzureMLTokenExpiry(tc.value) + if tc.wantErr { + if gotErr == nil { + t.Fatal("parseAzureMLTokenExpiry() succeeded, want error") + } + return + } + if gotErr != nil { + t.Fatalf("parseAzureMLTokenExpiry() failed: %v", gotErr) + } + if diff := cmp.Diff(tc.want, got); diff != "" { + t.Errorf("parseAzureMLTokenExpiry() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestReadAzureArcSecretFromDirectory_validation(t *testing.T) { + expectedDirectory := t.TempDir() + validKeyFile := filepath.Join(expectedDirectory, "identity.key") + if err := os.WriteFile(validKeyFile, []byte("secret"), 0o600); err != nil { + t.Fatalf("write valid key file: %v", err) + } + invalidExtension := filepath.Join(expectedDirectory, "identity.txt") + if err := os.WriteFile(invalidExtension, []byte("secret"), 0o600); err != nil { + t.Fatalf("write invalid-extension key file: %v", err) + } + largeKeyFile := filepath.Join(expectedDirectory, "large.key") + if err := os.WriteFile(largeKeyFile, []byte(strings.Repeat("x", int(azureArcMaximumSecretFileSize+1))), 0o600); err != nil { + t.Fatalf("write large key file: %v", err) + } + otherDirectory := t.TempDir() + unexpectedKeyFile := filepath.Join(otherDirectory, "identity.key") + if err := os.WriteFile(unexpectedKeyFile, []byte("secret"), 0o600); err != nil { + t.Fatalf("write unexpected key file: %v", err) + } + + testCases := []struct { + name string + keyFile string + want string + wantErr string + }{ + { + name: "valid key file", + keyFile: validKeyFile, + want: "secret", + }, + { + name: "unexpected directory", + keyFile: unexpectedKeyFile, + wantErr: fmt.Sprintf("unexpected Azure Arc managed identity secret directory %q", otherDirectory), + }, + { + name: "invalid extension", + keyFile: invalidExtension, + wantErr: fmt.Sprintf("Azure Arc managed identity secret file %q must have a .key extension", invalidExtension), + }, + { + name: "file too large", + keyFile: largeKeyFile, + wantErr: fmt.Sprintf("Azure Arc managed identity secret file %q exceeds 4096 bytes", largeKeyFile), + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + got, gotErr := readAzureArcSecretFromDirectory(tc.keyFile, expectedDirectory) + if tc.wantErr != "" { + if gotErr == nil || gotErr.Error() != tc.wantErr { + t.Errorf("readAzureArcSecretFromDirectory() error = %q, want %q", gotErr, tc.wantErr) + } + return + } + if gotErr != nil { + t.Fatalf("readAzureArcSecretFromDirectory() failed: %v", gotErr) + } + if got != tc.want { + t.Errorf("readAzureArcSecretFromDirectory() = %q, want %q", got, tc.want) + } + }) + } +} + +func azureMSITokenResponse(accessToken string, expiry time.Time) map[string]any { + return map[string]any{ + "access_token": accessToken, + "expires_on": expiry.Unix(), + "token_type": "Bearer", + } +} diff --git a/config/auth_azure_msi_test.go b/config/auth_azure_msi_test.go index d3aeb60e1..d5edf0882 100644 --- a/config/auth_azure_msi_test.go +++ b/config/auth_azure_msi_test.go @@ -44,6 +44,7 @@ func assertHeaders(t *testing.T, cfg *Config, expectedHeaders map[string]string) } func TestMsiHappyFlow(t *testing.T) { + withMockEnv(t, map[string]string{}) assertHeaders(t, &Config{ AzureUseMSI: true, AzureResourceID: "/a/b/c", @@ -84,6 +85,7 @@ func TestMsiHappyFlow(t *testing.T) { } func TestMsiFailsOnResolveWorkspace(t *testing.T) { + withMockEnv(t, map[string]string{}) _, err := authenticateRequest(&Config{ AzureUseMSI: true, AzureResourceID: "/a/b/c", @@ -106,6 +108,7 @@ func TestMsiFailsOnResolveWorkspace(t *testing.T) { } func TestMsiTokenNotFound(t *testing.T) { + withMockEnv(t, map[string]string{}) _, err := authenticateRequest(&Config{ AzureUseMSI: true, AzureClientID: "abc", @@ -122,6 +125,7 @@ func TestMsiTokenNotFound(t *testing.T) { } func TestMsiHappyFlowWithHostAndNoResourceID(t *testing.T) { + withMockEnv(t, map[string]string{}) assertHeaders(t, &Config{ Host: "https://adb-123.4.azuredatabricks.net", AzureUseMSI: true, @@ -149,6 +153,7 @@ func TestMsiHappyFlowWithHostAndNoResourceID(t *testing.T) { } func TestMsiInvalidTokenExpiry(t *testing.T) { + withMockEnv(t, map[string]string{}) _, err := authenticateRequest(&Config{ AzureUseMSI: true, AzureResourceID: "/a/b/c", diff --git a/config/config.go b/config/config.go index 8ba92ab87..39aaf48ca 100644 --- a/config/config.go +++ b/config/config.go @@ -261,7 +261,7 @@ type Config struct { // internal client used in for authentication purposes: // - Databricks Metadata Service: request/refresh tokens from parent processes, like Databricks VSCode extension // - Azure Active Directory (AAD): request/refresh OAuth tokens - // - Azure Instance Metadata Service (IMDS): request/refresh OAuth tokens for Azure Managed Identity + // - Azure managed identity endpoints: request/refresh OAuth tokens for Azure Managed Identity // - Azure Resource Manager (ARM): resolve host if only Azure Databricks Resource ID provided refreshClient *httpclient.ApiClient