Skip to content

Fix Endpoint Selection for Azure MSI Authentication - #1813

Open
nickblantz wants to merge 1 commit into
databricks:mainfrom
nickblantz:fix/azure-msi-endpoint-selection
Open

Fix Endpoint Selection for Azure MSI Authentication#1813
nickblantz wants to merge 1 commit into
databricks:mainfrom
nickblantz:fix/azure-msi-endpoint-selection

Conversation

@nickblantz

Copy link
Copy Markdown

Summary

Makes Azure managed identity authentication select the endpoint and request protocol exposed by the host environment instead of always calling IMDS. This enables azure-msi authentication to use host-specific endpoints such as Azure Function App while preserving IMDS as the fallback.

Why

AzureMsiCredentials currently assumes every managed identity is available from the VM IMDS endpoint at 169.254.169.254. Azure exposes different endpoints, headers, parameters, and token response formats on other hosting platforms, so workloads using ARM_USE_MSI=true outside an IMDS environment cannot acquire tokens through the SDK.

Microsoft's ManagedIdentityCredential detects the hosting environment from Azure-provided environment variables and selects the corresponding credential in a defined order. This change follows that selection order and implements each selected endpoint's request contract, using IMDS only when no more specific environment is configured.

What changed

Interface changes

None. The NewAzureMsiTokenSource signature and existing Config fields remain unchanged.

Behavioral changes

  • Selects Service Fabric, App Service, Azure Arc, Azure ML, Cloud Shell, workload identity, or IMDS according to the Azure Identity environment-variable precedence.
  • Uses the endpoint-specific HTTP method, API version, headers, form/query parameters, and client ID parameter name.
  • Honors AZURE_POD_IDENTITY_AUTHORITY_HOST when IMDS is selected.
  • Rejects user-assigned client IDs for Service Fabric, Azure Arc, and Cloud Shell, where Microsoft does not support them, and reports incomplete endpoint configuration instead of silently falling through to another endpoint.
  • Completes the Azure Arc 401 challenge flow by reading the key identified by WWW-Authenticate, validating its platform-specific directory, .key extension, and 4096-byte size limit, and retrying with Basic authorization.
  • Supports numeric or string expires_on values, relative expires_in values, and the Linux and Windows datetime formats returned by Azure ML.
  • Exchanges the projected federated token for an Entra ID access token when workload identity variables are fully configured.

Internal changes

  • Introduces private managed identity endpoint and request representations so selection is separate from request construction and token parsing.
  • Keeps the federated assertion out of HTTP debug output by passing the encoded workload identity form as an io.Reader.
  • Injects the clock and Azure Arc secret reader into the internal token source for deterministic tests without changing the public API.
  • Isolates existing MSI tests from ambient Azure host variables so they continue to exercise the IMDS fallback deliberately.

How is this tested?

Unit tests cover endpoint selection and request construction. End-to-end validation remains for the target Azure Function App environment and an Azure Virtual Machine regression check for the existing IMDS flow.

  • Unit tests matching '^(TestMsi|TestAzureMsi|TestAzureArc|TestParseAzureML|TestReadAzureArc)'
  • Integration test with Pulumi → Databricks Terraform Provider → Databricks Go SDK running from Azure Function App
  • Regression test with Databricks Terraform Provider → Databricks Go SDK running from Azure Virtual Machine

Changes in this pull request were written primarily with Codex and the following prompt:

This implementation includes an incorrect assumption about the usage of the Azure IMDS endpoint in config/auth_azure_msi.go#26. It assumes that this endpoint is used for all MSI authentication, when there are actually multiple endpoints used. The Microsoft implementation of the azure-identity Python module can be found at azure-sdk-for-python/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py. Use the Microsoft reference as a baseline for implementation and mirror its logic. Inspect the test cases included in this repository to create appropriate tests for different logic branches.

@github-actions

Copy link
Copy Markdown

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-go

Inputs:

  • PR number: 1813
  • Commit SHA: a47aaa82e2ba3a4f22334d995ab9e795bd9521cc

Checks will be approved automatically on success.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant