[azidentity] Implement support for custom token endpoint mode in WorkloadIdentityCredential with CA data support

sdk/azidentity/azidentity.go

@@ -34,6 +34,9 @@ const (
 	azureClientID                   = "AZURE_CLIENT_ID"
 	azureClientSecret               = "AZURE_CLIENT_SECRET"
 	azureFederatedTokenFile         = "AZURE_FEDERATED_TOKEN_FILE"
+	azureKubernetesCAFile           = "AZURE_KUBERNETES_CA_FILE"
+	azureKubernetesSNIName          = "AZURE_KUBERNETES_SNI_NAME"
+	azureKubernetesTokenEndpoint    = "AZURE_KUBERNETES_TOKEN_ENDPOINT"
 	azurePassword                   = "AZURE_PASSWORD"
 	azureRegionalAuthorityName      = "AZURE_REGIONAL_AUTHORITY_NAME"
 	azureTenantID                   = "AZURE_TENANT_ID"

sdk/azidentity/workload_identity.go

@@ -7,9 +7,16 @@
 package azidentity
 
 import (
+	"bytes"
 	"context"
+	"crypto/tls"
+	"crypto/x509"
 	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
 	"os"
+	"strings"
 	"sync"
 	"time"
 
@@ -87,14 +94,22 @@ func NewWorkloadIdentityCredential(options *WorkloadIdentityCredentialOptions) (
 			return nil, errors.New("no tenant ID specified. Check pod configuration or set TenantID in the options")
 		}
 	}
+
 	w := WorkloadIdentityCredential{file: file, mtx: &sync.RWMutex{}}
-	caco := ClientAssertionCredentialOptions{
+	caco := &ClientAssertionCredentialOptions{
 		AdditionallyAllowedTenants: options.AdditionallyAllowedTenants,
 		Cache:                      options.Cache,
 		ClientOptions:              options.ClientOptions,
 		DisableInstanceDiscovery:   options.DisableInstanceDiscovery,
 	}
-	cred, err := NewClientAssertionCredential(tenantID, clientID, w.getAssertion, &caco)
+
+	// configure identity binding if environment variables are present.
+	// In identity binding enabled mode, a dedicated transport will be used for proxying token requests to a dedicated endpoint.
+	if err := w.configureIdentityBinding(caco); err != nil {
+		return nil, err
+	}
+
+	cred, err := NewClientAssertionCredential(tenantID, clientID, w.getAssertion, caco)
 	if err != nil {
 		return nil, err
 	}
@@ -139,3 +154,205 @@ func (w *WorkloadIdentityCredential) getAssertion(context.Context) (string, erro
 	}
 	return w.assertion, nil
 }
+
+// configureIdentityBinding configures identity binding mode if the required environment variables are present
+func (w *WorkloadIdentityCredential) configureIdentityBinding(caco *ClientAssertionCredentialOptions) error {
+	// check for identity binding mode environment variables
+	kubernetesTokenEndpointStr := os.Getenv(azureKubernetesTokenEndpoint)
+	kubernetesSNIName := os.Getenv(azureKubernetesSNIName)
+	kubernetesCAFile := os.Getenv(azureKubernetesCAFile)
+
+	if kubernetesTokenEndpointStr == "" && kubernetesSNIName == "" && kubernetesCAFile == "" {
+		// identity binding is not set
+		return nil
+	}
+
+	// All three variables must be present for identity binding mode
+	if kubernetesTokenEndpointStr == "" || kubernetesSNIName == "" || kubernetesCAFile == "" {
+		return errors.New("identity binding mode requires all three environment variables: AZURE_KUBERNETES_TOKEN_ENDPOINT, AZURE_KUBERNETES_SNI_NAME, and AZURE_KUBERNETES_CA_FILE")
+	}
+
+	transporter, err := newIdentityBindingTransport(
+		kubernetesCAFile, kubernetesSNIName, kubernetesTokenEndpointStr,
+		caco.Transport,
+	)
+	if err != nil {
+		return err
+	}
+	caco.Transport = transporter
+	return nil
+}
+
+const (
+	tokenEndpointSuffix = "/oauth2/v2.0/token"
+	caReloadInterval    = 10 * time.Minute
+)
+
+// identityBindingTransport is a custom HTTP transport that redirects token requests
+// to the Kubernetes token endpoint when in identity binding mode
+type identityBindingTransport struct {
+	caFile              string
+	sniName             string
+	tokenEndpoint       *url.URL
+	fallbackTransporter policy.Transporter
+
+	mtx *sync.RWMutex
+
+	nextRead  time.Time
+	currentCA []byte
+	transport *http.Transport
+}
+
+func newIdentityBindingTransport(
+	caFile, sniName, tokenEndpointStr string,
+	fallbackTransporter policy.Transporter,
+) (*identityBindingTransport, error) {
+	tokenEndpoint, err := url.Parse(tokenEndpointStr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse token endpoint URL %q: %w", tokenEndpointStr, err)
+	}
+
+	if fallbackTransporter == nil {
+		// FIXME: can we callback to the defaultHTTPClient from azcore/runtime?
+		fallbackTransporter = http.DefaultClient
+	}
+
+	initialTransport := func() *http.Transport {
+		// try reusing the user provided transport if available
+		if httpClient, ok := fallbackTransporter.(*http.Client); ok {
+			if transport, ok := httpClient.Transport.(*http.Transport); ok {
+				return transport.Clone()
+			}
+		}
+
+		// if the user did not provide a policy.Transporter or it's not a *http.Client,
+		// we fall back to the default one.
+		// FIXME: can we callback to the defaultHTTPClient from azcore/runtime?
+		if transport, ok := http.DefaultTransport.(*http.Transport); ok {
+			return transport.Clone()
+		}
+
+		// this should not happen, but if the user mutates the net/http.DefaultTransport
+		// to something else, we fall back to a sane default
+		return &http.Transport{
+			ForceAttemptHTTP2:   true,
+			MaxIdleConns:        100,
+			IdleConnTimeout:     90 * time.Second,
+			TLSHandshakeTimeout: 10 * time.Second,
+		}
+	}()
+
+	tr := &identityBindingTransport{
+		caFile:              caFile,
+		sniName:             sniName,
+		tokenEndpoint:       tokenEndpoint,
+		fallbackTransporter: fallbackTransporter,
+		mtx:                 &sync.RWMutex{},
+		transport:           initialTransport,
+	}
+
+	// perform an initial load to surface any issues with the CA file and transport settings.
+	// Lock is not held here as this is called in the constructor
+	if err := tr.reloadCA(); err != nil {
+		return nil, err
+	}
+
+	return tr, nil
+}
+
+func (i *identityBindingTransport) Do(req *http.Request) (*http.Response, error) {
+	if !strings.HasSuffix(req.URL.Path, tokenEndpointSuffix) {
+		// not a token request, fallback to the original transporter
+		return i.fallbackTransporter.Do(req)
+	}
+
+	tr, err := i.getTokenTransporter()
+	if err != nil {
+		return nil, err
+	}
+
+	newReq := req.Clone(req.Context())
+	newReq.URL.Scheme = i.tokenEndpoint.Scheme // this will always be https
+	newReq.URL.Host = i.tokenEndpoint.Host
+	newReq.URL.Path = ""
+	newReq.Host = i.tokenEndpoint.Host
+
+	return tr.RoundTrip(newReq)
+}
+
+func (i *identityBindingTransport) getTokenTransporter() (*http.Transport, error) {
+	i.mtx.RLock()
+	if i.nextRead.Before(time.Now()) {
+		i.mtx.RUnlock()
+		i.mtx.Lock()
+		defer i.mtx.Unlock()
+		// double check on the read time
+		if now := time.Now(); i.nextRead.Before(now) {
+			if err := i.reloadCA(); err != nil {
+				// we return error if any attempt of reloading CA fails
+				// This should surface in the token calls and we expect the caller to
+				// have proper error handling / rate limit so we don't fall into deadloop here
+				// due to scenario like broken CA file.
+				return nil, err
+			}
+		}
+	} else {
+		defer i.mtx.RUnlock()
+	}
+	return i.transport, nil
+}
+
+func (i *identityBindingTransport) createTransportWithCAPool(
+	fromTransport *http.Transport,
+	caPool *x509.CertPool,
+) *http.Transport {
+	transport := fromTransport.Clone()
+	if transport.TLSClientConfig == nil {
+		transport.TLSClientConfig = &tls.Config{}
+	}
+	transport.TLSClientConfig.ServerName = i.sniName
+	transport.TLSClientConfig.RootCAs = caPool
+	return transport
+}
+
+// reloadCA attempts to read the latest CA from the CA file and updates the transport if the content has changed.
+// If a new CA is discovered, the existing transport will be replaced with a new one that uses the new CA.
+// It expects the caller to hold the write lock on i.mtx to ensure thread safety.
+func (i *identityBindingTransport) reloadCA() error {
+	newCA, err := os.ReadFile(i.caFile)
+	if err != nil {
+		return fmt.Errorf("read CA file %q: %w", i.caFile, err)
+	}
+
+	if len(newCA) == 0 {
+		// the CA file might be in the middle of rotation without the content written.
+		// We return nil and rely on next check.
+		return nil
+	}
+
+	if bytes.Equal(i.currentCA, newCA) {
+		// no change in CA content, no need to replace
+		i.nextRead = time.Now().Add(caReloadInterval)
+		return nil
+	}
+
+	newCAPool := x509.NewCertPool()
+	if !newCAPool.AppendCertsFromPEM(newCA) {
+		return fmt.Errorf("parse CA file %q: no valid certificates found", i.caFile)
+	}
+
+	newTransport := i.createTransportWithCAPool(i.transport, newCAPool)
+	oldTransport := i.transport
+
+	i.transport = newTransport
+	i.currentCA = newCA
+	i.nextRead = time.Now().Add(caReloadInterval)
+
+	if oldTransport != nil {
+		// drop any idle connections from previous transport so new requests can be
+		// moved to the new transport
+		oldTransport.CloseIdleConnections()
+	}
+
+	return nil
+}