feat: use the cloud endpoints from the credentials for sovereign clouds

README.md

@@ -30,6 +30,30 @@ Example pipeline step:
 The Azure Credentials Secret structure is fully compatible with the standard
 [Azure Official Provider][azop]
 
+### Sovereign clouds
+
+The function targets the Azure public cloud unless the credentials say
+otherwise. For Azure Government, Azure China or another cloud, add the endpoint
+URLs that `az ad sp create-for-rbac --sdk-auth` prints to the credentials JSON;
+they apply to service principal and workload identity credentials alike:
+
+```json
+{
+  "clientId": "your-client-id",
+  "clientSecret": "your-client-secret",
+  "tenantId": "your-tenant-id",
+  "subscriptionId": "your-subscription-id",
+  "activeDirectoryEndpointUrl": "https://login.microsoftonline.us",
+  "resourceManagerEndpointUrl": "https://management.usgovcloudapi.net/",
+  "managementEndpointUrl": "https://management.core.usgovcloudapi.net/"
+}
+```
+
+`activeDirectoryEndpointUrl` and `resourceManagerEndpointUrl` select the cloud.
+The public, Government and China endpoints map to the Azure SDK's built-in
+configurations; any other pair is used as given, with `managementEndpointUrl`
+as the token audience when present.
+
 Example XR status after e2e query:
 
 ```yaml

fn.go

@@ -4,11 +4,15 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"net/url"
 	"regexp"
 	"strings"
 	"sync/atomic"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph"
@@ -36,6 +40,12 @@ const (
 	ClientID = "clientId"
 	// ClientSecret defines the azure credentials key for client secret
 	ClientSecret = "clientSecret"
+	// ActiveDirectoryEndpointURL defines the azure credentials key for the Entra ID authority of the cloud
+	ActiveDirectoryEndpointURL = "activeDirectoryEndpointUrl"
+	// ResourceManagerEndpointURL defines the azure credentials key for the Resource Manager endpoint of the cloud
+	ResourceManagerEndpointURL = "resourceManagerEndpointUrl"
+	// ManagementEndpointURL defines the azure credentials key for the Resource Manager token audience of the cloud
+	ManagementEndpointURL = "managementEndpointUrl"
 	// WorkloadIdentityCredentialPath defines the azure credentials key for federated token file path
 	WorkloadIdentityCredentialPath = "federatedTokenFile"
 )
@@ -504,7 +514,13 @@ func (a *AzureQuery) azQuery(ctx context.Context, azureCreds interface{}, in *v1
 func (a *AzureQuery) initializeWorkloadIdentityProvider(azureCreds map[string]string, log logging.Logger) (*armresourcegraph.Client, error) {
 	tokenFilePath := azureCreds[WorkloadIdentityCredentialPath]
 
+	cloudConfig, err := cloudConfiguration(azureCreds)
+	if err != nil {
+		return nil, err
+	}
+
 	options := &azidentity.WorkloadIdentityCredentialOptions{
+		ClientOptions: azcore.ClientOptions{Cloud: cloudConfig},
 		TokenFilePath: tokenFilePath,
 	}
 
@@ -521,15 +537,15 @@ func (a *AzureQuery) initializeWorkloadIdentityProvider(azureCreds map[string]st
 	}
 
 	// Create Azure credential
-	log.Info("Initializing workload identity provider", "tokenFile", tokenFilePath)
+	log.Info("Initializing workload identity provider", "tokenFile", tokenFilePath, "authority", cloudConfig.ActiveDirectoryAuthorityHost)
 
 	cred, err := azidentity.NewWorkloadIdentityCredential(options)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to obtain workloadidentity credentials")
 	}
 
 	// Create and authorize a ResourceGraph client
-	client, err := armresourcegraph.NewClient(cred, nil)
+	client, err := armresourcegraph.NewClient(cred, &arm.ClientOptions{ClientOptions: azcore.ClientOptions{Cloud: cloudConfig}})
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create client")
 	}
@@ -544,15 +560,20 @@ func (a *AzureQuery) initializeClientSecretProvider(azureCreds map[string]string
 	clientID := azureCreds[ClientID]
 	clientSecret := azureCreds[ClientSecret]
 
+	cloudConfig, err := cloudConfiguration(azureCreds)
+	if err != nil {
+		return nil, err
+	}
+
 	// To configure DefaultAzureCredential to authenticate a user-assigned managed identity,
 	// set the environment variable AZURE_CLIENT_ID to the identity's client ID.
-	cred, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, nil)
+	cred, err := azidentity.NewClientSecretCredential(tenantID, clientID, clientSecret, &azidentity.ClientSecretCredentialOptions{ClientOptions: azcore.ClientOptions{Cloud: cloudConfig}})
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to obtain clientsecret credentials")
 	}
 
 	// Create and authorize a ResourceGraph client
-	client, err := armresourcegraph.NewClient(cred, nil)
+	client, err := armresourcegraph.NewClient(cred, &arm.ClientOptions{ClientOptions: azcore.ClientOptions{Cloud: cloudConfig}})
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create client")
 	}
@@ -986,3 +1007,52 @@ func (f *Function) checkContextTargetHasData(req *fnv1.RunFunctionRequest, in *v
 	}
 	return false
 }
+
+// cloudConfiguration selects the Azure cloud from the endpoint URLs that
+// `az ad sp create-for-rbac --sdk-auth` writes into the credentials. Without
+// them the public cloud is used. The public, Government and China endpoints
+// map to the SDK's own configurations; any other pair is used as given.
+func cloudConfiguration(azureCreds map[string]string) (cloud.Configuration, error) {
+	authority := azureCreds[ActiveDirectoryEndpointURL]
+	manager := azureCreds[ResourceManagerEndpointURL]
+	if authority == "" && manager == "" {
+		return cloud.AzurePublic, nil
+	}
+
+	for _, known := range []cloud.Configuration{cloud.AzurePublic, cloud.AzureGovernment, cloud.AzureChina} {
+		if sameHost(authority, known.ActiveDirectoryAuthorityHost) && sameHost(manager, known.Services[cloud.ResourceManager].Endpoint) {
+			return known, nil
+		}
+	}
+
+	if authority == "" || manager == "" {
+		return cloud.Configuration{}, errors.Errorf("both %s and %s are required for a cloud the SDK does not know", ActiveDirectoryEndpointURL, ResourceManagerEndpointURL)
+	}
+	audience := azureCreds[ManagementEndpointURL]
+	if audience == "" {
+		audience = manager
+	}
+	return cloud.Configuration{
+		ActiveDirectoryAuthorityHost: authority,
+		Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
+			cloud.ResourceManager: {Audience: audience, Endpoint: strings.TrimSuffix(manager, "/")},
+		},
+	}, nil
+}
+
+// sameHost reports whether the given URL names the same host as the known
+// one. An unset URL matches, so that one endpoint is enough to pick a cloud.
+func sameHost(given, known string) bool {
+	if given == "" {
+		return true
+	}
+	g, err := url.Parse(given)
+	if err != nil {
+		return false
+	}
+	k, err := url.Parse(known)
+	if err != nil {
+		return false
+	}
+	return strings.EqualFold(g.Host, k.Host)
+}

fn_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph"
 	"github.com/google/go-cmp/cmp"
@@ -3162,3 +3163,96 @@ func validateBackwardsCompatibility(rsp *fnv1.RunFunctionResponse) error {
 
 	return nil
 }
+
+func TestCloudConfiguration(t *testing.T) {
+	custom := cloud.Configuration{
+		ActiveDirectoryAuthorityHost: "https://login.example.test/",
+		Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
+			cloud.ResourceManager: {Audience: "https://management.core.example.test/", Endpoint: "https://management.example.test"},
+		},
+	}
+	customWithoutAudience := cloud.Configuration{
+		ActiveDirectoryAuthorityHost: "https://login.example.test/",
+		Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
+			cloud.ResourceManager: {Audience: "https://management.example.test/", Endpoint: "https://management.example.test"},
+		},
+	}
+
+	cases := map[string]struct {
+		reason  string
+		creds   map[string]string
+		want    cloud.Configuration
+		wantErr bool
+	}{
+		"NoEndpoints": {
+			reason: "Credentials without endpoint URLs keep targeting the public cloud",
+			creds:  map[string]string{ClientID: "client", ClientSecret: "secret", TenantID: "tenant"},
+			want:   cloud.AzurePublic,
+		},
+		"PublicEndpoints": {
+			reason: "The public endpoints map to the SDK's public cloud configuration",
+			creds: map[string]string{
+				ActiveDirectoryEndpointURL: "https://login.microsoftonline.com",
+				ResourceManagerEndpointURL: "https://management.azure.com/",
+				ManagementEndpointURL:      "https://management.core.windows.net/",
+			},
+			want: cloud.AzurePublic,
+		},
+		"Government": {
+			reason: "The Azure Government endpoints map to the SDK's Government configuration",
+			creds: map[string]string{
+				ActiveDirectoryEndpointURL: "https://login.microsoftonline.us",
+				ResourceManagerEndpointURL: "https://management.usgovcloudapi.net/",
+				ManagementEndpointURL:      "https://management.core.usgovcloudapi.net/",
+			},
+			want: cloud.AzureGovernment,
+		},
+		"GovernmentAuthorityOnly": {
+			reason: "One known endpoint is enough to pick the cloud",
+			creds:  map[string]string{ActiveDirectoryEndpointURL: "https://login.microsoftonline.us/"},
+			want:   cloud.AzureGovernment,
+		},
+		"China": {
+			reason: "The Azure China endpoints map to the SDK's China configuration",
+			creds: map[string]string{
+				ActiveDirectoryEndpointURL: "https://login.chinacloudapi.cn",
+				ResourceManagerEndpointURL: "https://management.chinacloudapi.cn/",
+			},
+			want: cloud.AzureChina,
+		},
+		"Custom": {
+			reason: "Unknown endpoints are used as given, with managementEndpointUrl as the token audience",
+			creds: map[string]string{
+				ActiveDirectoryEndpointURL: "https://login.example.test/",
+				ResourceManagerEndpointURL: "https://management.example.test/",
+				ManagementEndpointURL:      "https://management.core.example.test/",
+			},
+			want: custom,
+		},
+		"CustomWithoutAudience": {
+			reason: "Without managementEndpointUrl the Resource Manager endpoint is the audience",
+			creds: map[string]string{
+				ActiveDirectoryEndpointURL: "https://login.example.test/",
+				ResourceManagerEndpointURL: "https://management.example.test/",
+			},
+			want: customWithoutAudience,
+		},
+		"CustomMissingAuthority": {
+			reason:  "An unknown Resource Manager endpoint without its authority is an error rather than a silent public cloud",
+			creds:   map[string]string{ResourceManagerEndpointURL: "https://management.example.test/"},
+			wantErr: true,
+		},
+	}
+
+	for name, tc := range cases {
+		t.Run(name, func(t *testing.T) {
+			got, err := cloudConfiguration(tc.creds)
+			if (err != nil) != tc.wantErr {
+				t.Fatalf("%s\ncloudConfiguration(...): error = %v, wantErr %v", tc.reason, err, tc.wantErr)
+			}
+			if diff := cmp.Diff(tc.want, got); diff != "" {
+				t.Errorf("%s\ncloudConfiguration(...): -want, +got:\n%s", tc.reason, diff)
+			}
+		})
+	}
+}