Repository: envoyproxy/ai-gateway
Pull request #517 (2025-03-23 18:45:02)

feat: add proxy URL option for Azure authentication 
**Commit Message**

Add proxy URL option for Azure authentication.


**Related Issues/PRs (if applicable)**
Related PR: #445 

xiaolin593: Forced-push for DCO

Patch:
diff --git a/internal/controller/tokenprovider/azure_token_provider.go b/internal/controller/tokenprovider/azure_token_provider.go
index 9072f03..1d5908e 100644
--- a/internal/controller/tokenprovider/azure_token_provider.go
+++ b/internal/controller/tokenprovider/azure_token_provider.go
@@ -7,7 +7,11 @@ package tokenprovider
 
 import (
 	"context"
+	"net/http"
+	"net/url"
+	"os"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 )
@@ -20,7 +24,8 @@ type azureTokenProvider struct {
 
 // NewAzureTokenProvider creates a new TokenProvider with the given tenant ID, client ID, client secret, and token request options.
 func NewAzureTokenProvider(tenantID, clientID, clientSecret string, tokenOption policy.TokenRequestOptions) (TokenProvider, error) {
-	credential, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, nil)
+	clientOptions := GetClientSecretCredentialOptions()
+	credential, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, clientOptions)
 	if err != nil {
 		return nil, err
 	}
@@ -35,3 +40,19 @@ func (a *azureTokenProvider) GetToken(ctx context.Context) (TokenExpiry, error)
 	}
 	return TokenExpiry{Token: azureToken.Token, ExpiresAt: azureToken.ExpiresOn}, nil
 }
+
+func GetClientSecretCredentialOptions() *azidentity.ClientSecretCredentialOptions {
+	if azureProxyURL := os.Getenv("AI_GATEWAY_AZURE_PROXY_URL"); azureProxyURL != "" {
+		proxyURL, err := url.Parse(azureProxyURL)
+		if err == nil {
+			customTransport := &http.Transport{Proxy: http.ProxyURL(proxyURL)}
+			customHTTPClient := &http.Client{Transport: customTransport}
+			return &azidentity.ClientSecretCredentialOptions{
+				ClientOptions: azcore.ClientOptions{
+					Transport: customHTTPClient,
+				},
+			}
+		}
+	}
+	return nil
+}
diff --git a/manifests/charts/ai-gateway-helm/values.yaml b/manifests/charts/ai-gateway-helm/values.yaml
index 66bbaf1..fceee46 100644
--- a/manifests/charts/ai-gateway-helm/values.yaml
+++ b/manifests/charts/ai-gateway-helm/values.yaml
@@ -46,6 +46,9 @@ controller:
   # AWS STS request when rotating OIDC credentials will be configured to use AI_GATEWAY_STS_PROXY_URL proxy if set.
   #  - key: AI_GATEWAY_STS_PROXY_URL
   #    value: some-proxy-placeholder
+  # Azure authentication request will be configured to use AI_GATEWAY_AZURE_PROXY_URL proxy if set.
+  #  - key: AI_GATEWAY_AZURE_PROXY_URL
+  #    value: some-proxy-placeholder
   podEnv: {}
   # Example of volumes
   #  - mountPath: /placeholder/path


Test patch:
diff --git a/internal/controller/tokenprovider/azure_token_provider_test.go b/internal/controller/tokenprovider/azure_token_provider_test.go
index c42b1c7..38003e8 100644
--- a/internal/controller/tokenprovider/azure_token_provider_test.go
+++ b/internal/controller/tokenprovider/azure_token_provider_test.go
@@ -7,6 +7,8 @@ package tokenprovider
 
 import (
 	"context"
+	"net/http"
+	"os"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
@@ -40,4 +42,32 @@ func TestAzureTokenProvider_GetToken(t *testing.T) {
 		require.Error(t, err)
 		require.Contains(t, err.Error(), "Tenant 'invalidtenantid' not found. Check to make sure you have the correct tenant ID and are signing into the correct cloud.")
 	})
+
+	t.Run("azure proxy url", func(t *testing.T) {
+		// Set environment variable for the test
+		mockProxyURL := "http://localhost:8888"
+		os.Setenv("AI_GATEWAY_AZURE_PROXY_URL", mockProxyURL)
+		defer os.Unsetenv("AI_GATEWAY_AZURE_PROXY_URL")
+
+		opts := GetClientSecretCredentialOptions()
+
+		require.NotNil(t, opts)
+		require.NotNil(t, opts.ClientOptions.Transport)
+
+		// Assert that the transport has a proxy set
+		transport, ok := opts.ClientOptions.Transport.(*http.Client)
+		require.True(t, ok)
+		require.NotNil(t, transport.Transport)
+
+		// Check the proxy URL (optional, deeper inspection)
+		innerTransport, ok := transport.Transport.(*http.Transport)
+		require.True(t, ok)
+		require.NotNil(t, innerTransport.Proxy)
+
+		req, _ := http.NewRequest("GET", "http://example.com", nil)
+		proxyFunc := innerTransport.Proxy
+		proxyURL, err := proxyFunc(req)
+		require.NoError(t, err)
+		require.Equal(t, "http://localhost:8888", proxyURL.String())
+	})
 }