Address code review feedback: remove host checks, early return, extract identity binding configuration

Co-authored-by: bcho <[email protected]>

Author: Copilot

sdk/azidentity/workload_identity.go

@@ -105,44 +105,9 @@ func NewWorkloadIdentityCredential(options *WorkloadIdentityCredentialOptions) (
 
 	w := WorkloadIdentityCredential{file: file, mtx: &sync.RWMutex{}}
 
-	// Check for identity binding mode environment variables
-	kubernetesTokenEndpointStr := os.Getenv(azureKubernetesTokenEndpoint)
-	kubernetesSNIName := os.Getenv(azureKubernetesSNIName)
-	kubernetesCAFile := os.Getenv(azureKubernetesCAFile)
-
-	// If any of the identity binding environment variables are present, enable identity binding mode
-	if kubernetesTokenEndpointStr != "" || kubernetesSNIName != "" || kubernetesCAFile != "" {
-		// All three variables must be present for identity binding mode
-		if kubernetesTokenEndpointStr == "" || kubernetesSNIName == "" || kubernetesCAFile == "" {
-			return nil, errors.New("identity binding mode requires all three environment variables: AZURE_KUBERNETES_TOKEN_ENDPOINT, AZURE_KUBERNETES_SNI_NAME, and AZURE_KUBERNETES_CA_FILE")
-		}
-		
-		// Parse the Kubernetes token endpoint URL
-		kubernetesTokenEndpoint, err := url.Parse(kubernetesTokenEndpointStr)
-		if err != nil {
-			return nil, fmt.Errorf("failed to parse Kubernetes token endpoint URL: %w", err)
-		}
-		
-		w.identityBinding = true
-		w.kubernetesTokenEndpoint = kubernetesTokenEndpoint
-		w.kubernetesSNIName = kubernetesSNIName
-		w.kubernetesCAFile = kubernetesCAFile
-		
-		// Validate CA file exists and is readable during construction
-		if _, err := os.Stat(kubernetesCAFile); err != nil {
-			return nil, fmt.Errorf("failed to read Kubernetes CA file: %w", err)
-		}
-		
-		// Validate CA certificate is parseable during construction
-		caCert, err := os.ReadFile(kubernetesCAFile)
-		if err != nil {
-			return nil, fmt.Errorf("failed to read Kubernetes CA file: %w", err)
-		}
-		
-		caCertPool := x509.NewCertPool()
-		if !caCertPool.AppendCertsFromPEM(caCert) {
-			return nil, errors.New("failed to parse Kubernetes CA certificate")
-		}
+	// Configure identity binding if environment variables are present
+	if err := w.configureIdentityBinding(); err != nil {
+		return nil, err
 	}
 
 	caco := ClientAssertionCredentialOptions{
@@ -207,24 +172,69 @@ 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() error {
+	// Check for identity binding mode environment variables
+	kubernetesTokenEndpointStr := os.Getenv(azureKubernetesTokenEndpoint)
+	kubernetesSNIName := os.Getenv(azureKubernetesSNIName)
+	kubernetesCAFile := os.Getenv(azureKubernetesCAFile)
+
+	// If any of the identity binding environment variables are present, enable identity binding mode
+	if kubernetesTokenEndpointStr != "" || kubernetesSNIName != "" || kubernetesCAFile != "" {
+		// 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")
+		}
+
+		// Parse the Kubernetes token endpoint URL
+		kubernetesTokenEndpoint, err := url.Parse(kubernetesTokenEndpointStr)
+		if err != nil {
+			return fmt.Errorf("failed to parse Kubernetes token endpoint URL: %w", err)
+		}
+
+		w.identityBinding = true
+		w.kubernetesTokenEndpoint = kubernetesTokenEndpoint
+		w.kubernetesSNIName = kubernetesSNIName
+		w.kubernetesCAFile = kubernetesCAFile
+
+		// Validate CA file exists and is readable during construction
+		if _, err := os.Stat(kubernetesCAFile); err != nil {
+			return fmt.Errorf("failed to read Kubernetes CA file: %w", err)
+		}
+
+		// Validate CA certificate is parseable during construction
+		caCert, err := os.ReadFile(kubernetesCAFile)
+		if err != nil {
+			return fmt.Errorf("failed to read Kubernetes CA file: %w", err)
+		}
+
+		caCertPool := x509.NewCertPool()
+		if !caCertPool.AppendCertsFromPEM(caCert) {
+			return errors.New("failed to parse Kubernetes CA certificate")
+		}
+	}
+
+	return nil
+}
+
 // loadKubernetesCA loads and caches the Kubernetes CA certificate
 func (w *WorkloadIdentityCredential) loadKubernetesCA() (*x509.CertPool, error) {
 	w.mtx.RLock()
 	if w.caExpires.After(time.Now()) && w.caCertPool != nil {
 		defer w.mtx.RUnlock()
 		return w.caCertPool, nil
 	}
-	
+
 	// ensure only one goroutine at a time updates the CA cert
 	w.mtx.RUnlock()
 	w.mtx.Lock()
 	defer w.mtx.Unlock()
-	
+
 	// double check because another goroutine may have acquired the write lock first and done the update
 	if now := time.Now(); w.caExpires.After(now) && w.caCertPool != nil {
 		return w.caCertPool, nil
 	}
-	
+
 	// Load the CA certificate for the Kubernetes endpoint
 	caCert, err := os.ReadFile(w.kubernetesCAFile)
 	if err != nil {
@@ -235,12 +245,12 @@ func (w *WorkloadIdentityCredential) loadKubernetesCA() (*x509.CertPool, error)
 	if !caCertPool.AppendCertsFromPEM(caCert) {
 		return nil, errors.New("failed to parse Kubernetes CA certificate")
 	}
-	
+
 	// Cache the CA certificate for 10 minutes (same as token assertion)
 	w.caCert = caCert
 	w.caCertPool = caCertPool
 	w.caExpires = time.Now().Add(10 * time.Minute)
-	
+
 	return caCertPool, nil
 }
 
@@ -252,42 +262,38 @@ type identityBindingTransport struct {
 }
 
 func (t *identityBindingTransport) Do(req *http.Request) (*http.Response, error) {
-	// Check if this is a token request to the Azure authority host
-	if strings.HasSuffix(req.URL.Path, "/oauth2/v2.0/token") && (req.URL.Host == "login.microsoftonline.com" || 
-		req.URL.Host == "login.microsoftonline.us" || 
-		req.URL.Host == "login.partner.microsoftonline.cn" ||
-		req.URL.Host == "login.microsoftonline.de") {
-		// This is a token request, redirect to Kubernetes endpoint
-		
-		// Load the CA certificate (this will use cached version if still valid)
-		caCertPool, err := t.credential.loadKubernetesCA()
-		if err != nil {
-			return nil, err
-		}
-		
-		// Create custom transport with the CA and SNI configuration
-		transport := &http.Transport{
-			TLSClientConfig: &tls.Config{
-				RootCAs:    caCertPool,
-				ServerName: t.kubernetesSNIName,
-			},
-		}
-		
-		// Clone the request to avoid modifying the original
-		newReq := req.Clone(req.Context())
-		
-		// Update the URL to point to the Kubernetes endpoint
-		newReq.URL.Scheme = t.credential.kubernetesTokenEndpoint.Scheme
-		newReq.URL.Host = t.credential.kubernetesTokenEndpoint.Host
-		newReq.Host = t.credential.kubernetesTokenEndpoint.Host
-		
-		// Preserve the original path (contains tenant ID and token endpoint path)
-		// The path should be something like "/tenant-id/oauth2/v2.0/token"
-		// Keep the original path to maintain the token request structure
-		
-		return transport.RoundTrip(newReq)
+	// Return early if this is not a token request
+	if !strings.HasSuffix(req.URL.Path, "/oauth2/v2.0/token") {
+		return http.DefaultTransport.RoundTrip(req)
 	}
-	
-	// For non-token requests, use the default transport
-	return http.DefaultTransport.RoundTrip(req)
+
+	// This is a token request, redirect to Kubernetes endpoint
+
+	// Load the CA certificate (this will use cached version if still valid)
+	caCertPool, err := t.credential.loadKubernetesCA()
+	if err != nil {
+		return nil, err
+	}
+
+	// Create custom transport with the CA and SNI configuration
+	transport := &http.Transport{
+		TLSClientConfig: &tls.Config{
+			RootCAs:    caCertPool,
+			ServerName: t.kubernetesSNIName,
+		},
+	}
+
+	// Clone the request to avoid modifying the original
+	newReq := req.Clone(req.Context())
+
+	// Update the URL to point to the Kubernetes endpoint
+	newReq.URL.Scheme = t.credential.kubernetesTokenEndpoint.Scheme
+	newReq.URL.Host = t.credential.kubernetesTokenEndpoint.Host
+	newReq.Host = t.credential.kubernetesTokenEndpoint.Host
+
+	// Preserve the original path (contains tenant ID and token endpoint path)
+	// The path should be something like "/tenant-id/oauth2/v2.0/token"
+	// Keep the original path to maintain the token request structure
+
+	return transport.RoundTrip(newReq)
 }

sdk/azidentity/workload_identity_test.go

@@ -328,7 +328,7 @@ kAF/AQQQAAQBQAAkCNADCBADI==
 	if err != nil {
 		t.Fatal(err)
 	}
-	
+
 	mockCredential := &WorkloadIdentityCredential{
 		identityBinding:         true,
 		kubernetesTokenEndpoint: kubernetesURL,
@@ -343,7 +343,7 @@ kAF/AQQQAAQBQAAkCNADCBADI==
 		kubernetesSNIName: "kubernetes.default.svc",
 	}
 
-	// Test with Azure authority host request (should be redirected)
+	// Test with token endpoint request (should be redirected)
 	req := &http.Request{
 		Method: "POST",
 		Header: make(http.Header),
@@ -362,7 +362,7 @@ kAF/AQQQAAQBQAAkCNADCBADI==
 		t.Error("expected an error due to network connectivity")
 	}
 
-	// Test with non-Azure host request (should not be redirected)
+	// Test with non-token request (should not be redirected)
 	req2 := &http.Request{
 		Method: "GET",
 		Header: make(http.Header),
@@ -405,7 +405,7 @@ func TestWorkloadIdentityCredential_IdentityBinding_Success(t *testing.T) {
 	if err != nil {
 		t.Skipf("test certificate not found, run: cd /tmp/test-certs && openssl req -x509 -newkey rsa:2048 -keyout key.pem -out cert.pem -days 1 -nodes -subj \"/C=US/ST=Test/L=Test/O=Test/OU=Test/CN=kubernetes.default.svc\"")
 	}
-	
+
 	if err := os.WriteFile(caFile, caCert, os.ModePerm); err != nil {
 		t.Fatalf("failed to write CA file: %v", err)
 	}
@@ -479,7 +479,7 @@ klmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890abcdefghij
 klmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890abcdefghij
 klmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890abcdefghij
 -----END CERTIFICATE-----`
-	
+
 	if err := os.WriteFile(caFile, []byte(caCert), os.ModePerm); err != nil {
 		t.Fatalf("failed to write CA file: %v", err)
 	}