Add DisableInstanceDiscovery option for private cloud auth (Azure#19872)

Author: jonnguyen1

sdk/azidentity/azidentity.go

@@ -15,6 +15,7 @@ import (
 	"net/url"
 	"os"
 	"regexp"
+	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/cloud"
@@ -56,20 +57,31 @@ var getConfidentialClient = func(clientID, tenantID string, cred confidential.Cr
 		confidential.WithHTTPClient(newPipelineAdapter(co)),
 	}
 	o = append(o, additionalOpts...)
+	if strings.ToLower(tenantID) == "adfs" {
+		o = append(o, confidential.WithInstanceDiscovery(false))
+	}
 	return confidential.New(clientID, cred, o...)
 }
 
-var getPublicClient = func(clientID, tenantID string, co *azcore.ClientOptions) (public.Client, error) {
+var getPublicClient = func(clientID, tenantID string, co *azcore.ClientOptions, additionalOpts ...public.Option) (public.Client, error) {
 	if !validTenantID(tenantID) {
 		return public.Client{}, errors.New(tenantIDValidationErr)
 	}
 	authorityHost, err := setAuthorityHost(co.Cloud)
 	if err != nil {
 		return public.Client{}, err
 	}
-	return public.New(clientID,
+
+	o := []public.Option{
 		public.WithAuthority(runtime.JoinPaths(authorityHost, tenantID)),
 		public.WithHTTPClient(newPipelineAdapter(co)),
+	}
+	o = append(o, additionalOpts...)
+	if strings.ToLower(tenantID) == "adfs" {
+		o = append(o, public.WithInstanceDiscovery(false))
+	}
+	return public.New(clientID,
+		o...,
 	)
 }
 

sdk/azidentity/client_assertion_credential.go

@@ -32,6 +32,9 @@ type ClientAssertionCredential struct {
 // ClientAssertionCredentialOptions contains optional parameters for ClientAssertionCredential.
 type ClientAssertionCredentialOptions struct {
 	azcore.ClientOptions
+
+	// DisableInstanceDiscovery allows disconnected cloud solutions to skip instance discovery for unknown authority hosts.
+	DisableInstanceDiscovery bool
 }
 
 // NewClientAssertionCredential constructs a ClientAssertionCredential. The getAssertion function must be thread safe. Pass nil for options to accept defaults.
@@ -47,7 +50,7 @@ func NewClientAssertionCredential(tenantID, clientID string, getAssertion func(c
 			return getAssertion(ctx)
 		},
 	)
-	c, err := getConfidentialClient(clientID, tenantID, cred, &options.ClientOptions)
+	c, err := getConfidentialClient(clientID, tenantID, cred, &options.ClientOptions, confidential.WithInstanceDiscovery(!options.DisableInstanceDiscovery))
 	if err != nil {
 		return nil, err
 	}

sdk/azidentity/client_certificate_credential.go

@@ -25,6 +25,9 @@ const credNameCert = "ClientCertificateCredential"
 type ClientCertificateCredentialOptions struct {
 	azcore.ClientOptions
 
+	// DisableInstanceDiscovery allows disconnected cloud solutions to skip instance discovery for unknown authority hosts.
+	DisableInstanceDiscovery bool
+
 	// SendCertificateChain controls whether the credential sends the public certificate chain in the x5c
 	// header of each token request's JWT. This is required for Subject Name/Issuer (SNI) authentication.
 	// Defaults to False.
@@ -52,6 +55,7 @@ func NewClientCertificateCredential(tenantID string, clientID string, certs []*x
 	if options.SendCertificateChain {
 		o = append(o, confidential.WithX5C())
 	}
+	o = append(o, confidential.WithInstanceDiscovery(!options.DisableInstanceDiscovery))
 	c, err := getConfidentialClient(clientID, tenantID, cred, &options.ClientOptions, o...)
 	if err != nil {
 		return nil, err

sdk/azidentity/client_certificate_credential_test.go

@@ -19,6 +19,7 @@ import (
 	"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/internal/mock"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/recording"
 )
 
 type certTest struct {
@@ -224,6 +225,52 @@ func TestClientCertificateCredential_Live(t *testing.T) {
 			testGetTokenSuccess(t, cred)
 		})
 	}
+	t.Run("instance discovery disabled", func(t *testing.T) {
+		if liveSP.pemPath == "" {
+			t.Skip("no certificate file specified")
+		}
+		certData, err := os.ReadFile(liveSP.pemPath)
+		if err != nil {
+			t.Fatalf(`failed to read cert: %v`, err)
+		}
+		certs, key, err := ParseCertificates(certData, nil)
+		if err != nil {
+			t.Fatalf(`failed to parse cert: %v`, err)
+		}
+		o, stop := initRecording(t)
+		defer stop()
+		opts := &ClientCertificateCredentialOptions{ClientOptions: o, DisableInstanceDiscovery: true}
+		cred, err := NewClientCertificateCredential(liveSP.tenantID, liveSP.clientID, certs, key, opts)
+		if err != nil {
+			t.Fatalf("failed to construct credential: %v", err)
+		}
+		testGetTokenSuccess(t, cred)
+	})
+}
+
+func TestClientCertificateCredentialADFS_Live(t *testing.T) {
+	if recording.GetRecordMode() != recording.PlaybackMode {
+		if adfsLiveSP.clientID == "" || adfsLiveSP.certPath == "" || adfsScope == "" {
+			t.Skip("set ADFS_SP_* to run this test live")
+		}
+	}
+	certData, err := os.ReadFile(adfsLiveSP.certPath)
+	if err != nil {
+		t.Fatalf(`failed to read cert: %v`, err)
+	}
+	certs, key, err := ParseCertificates(certData, nil)
+	if err != nil {
+		t.Fatalf(`failed to parse cert: %v`, err)
+	}
+	o, stop := initRecording(t)
+	defer stop()
+	o.Cloud.ActiveDirectoryAuthorityHost = adfsAuthority
+	opts := &ClientCertificateCredentialOptions{ClientOptions: o, DisableInstanceDiscovery: true}
+	cred, err := NewClientCertificateCredential("adfs", adfsLiveSP.clientID, certs, key, opts)
+	if err != nil {
+		t.Fatalf("failed to construct credential: %v", err)
+	}
+	testGetTokenSuccess(t, cred, adfsScope)
 }
 
 func TestClientCertificateCredential_InvalidCertLive(t *testing.T) {

sdk/azidentity/client_secret_credential.go

@@ -20,6 +20,9 @@ const credNameSecret = "ClientSecretCredential"
 // ClientSecretCredentialOptions contains optional parameters for ClientSecretCredential.
 type ClientSecretCredentialOptions struct {
 	azcore.ClientOptions
+
+	// DisableInstanceDiscovery allows disconnected cloud solutions to skip instance discovery for unknown authority hosts.
+	DisableInstanceDiscovery bool
 }
 
 // ClientSecretCredential authenticates an application with a client secret.
@@ -36,7 +39,7 @@ func NewClientSecretCredential(tenantID string, clientID string, clientSecret st
 	if err != nil {
 		return nil, err
 	}
-	c, err := getConfidentialClient(clientID, tenantID, cred, &options.ClientOptions)
+	c, err := getConfidentialClient(clientID, tenantID, cred, &options.ClientOptions, confidential.WithInstanceDiscovery(!options.DisableInstanceDiscovery))
 	if err != nil {
 		return nil, err
 	}

sdk/azidentity/client_secret_credential_test.go

@@ -13,6 +13,7 @@ import (
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/recording"
 )
 
 const secret = "secret"
@@ -40,14 +41,39 @@ func TestClientSecretCredential_GetTokenSuccess(t *testing.T) {
 }
 
 func TestClientSecretCredential_Live(t *testing.T) {
+	for _, disabledID := range []bool{true, false} {
+		name := "default options"
+		if disabledID {
+			name = "instance discovery disabled"
+		}
+		t.Run(name, func(t *testing.T) {
+			opts, stop := initRecording(t)
+			defer stop()
+			o := ClientSecretCredentialOptions{ClientOptions: opts, DisableInstanceDiscovery: disabledID}
+			cred, err := NewClientSecretCredential(liveSP.tenantID, liveSP.clientID, liveSP.secret, &o)
+			if err != nil {
+				t.Fatalf("failed to construct credential: %v", err)
+			}
+			testGetTokenSuccess(t, cred)
+		})
+	}
+}
+
+func TestClientSecretCredentialADFS_Live(t *testing.T) {
+	if recording.GetRecordMode() != recording.PlaybackMode {
+		if adfsLiveSP.clientID == "" || adfsLiveSP.secret == "" || adfsScope == "" {
+			t.Skip("set ADFS_SP_* environment variables to run this test live")
+		}
+	}
 	opts, stop := initRecording(t)
 	defer stop()
-	o := ClientSecretCredentialOptions{ClientOptions: opts}
-	cred, err := NewClientSecretCredential(liveSP.tenantID, liveSP.clientID, liveSP.secret, &o)
+	opts.Cloud.ActiveDirectoryAuthorityHost = adfsAuthority
+	o := ClientSecretCredentialOptions{ClientOptions: opts, DisableInstanceDiscovery: true}
+	cred, err := NewClientSecretCredential("adfs", adfsLiveSP.clientID, adfsLiveSP.secret, &o)
 	if err != nil {
 		t.Fatalf("failed to construct credential: %v", err)
 	}
-	testGetTokenSuccess(t, cred)
+	testGetTokenSuccess(t, cred, adfsScope)
 }
 
 func TestClientSecretCredential_InvalidSecretLive(t *testing.T) {

sdk/azidentity/default_azure_credential.go

@@ -23,6 +23,9 @@ import (
 type DefaultAzureCredentialOptions struct {
 	azcore.ClientOptions
 
+	// DisableInstanceDiscovery allows disconnected cloud solutions to skip instance discovery for unknown authority hosts.
+	DisableInstanceDiscovery bool
+
 	// TenantID identifies the tenant the Azure CLI should authenticate in.
 	// Defaults to the CLI's default tenant, which is typically the home tenant of the user logged in to the CLI.
 	TenantID string
@@ -56,7 +59,7 @@ func NewDefaultAzureCredential(options *DefaultAzureCredentialOptions) (*Default
 		options = &DefaultAzureCredentialOptions{}
 	}
 
-	envCred, err := NewEnvironmentCredential(&EnvironmentCredentialOptions{ClientOptions: options.ClientOptions})
+	envCred, err := NewEnvironmentCredential(&EnvironmentCredentialOptions{ClientOptions: options.ClientOptions, DisableInstanceDiscovery: options.DisableInstanceDiscovery})
 	if err == nil {
 		creds = append(creds, envCred)
 	} else {

sdk/azidentity/device_code_credential.go

@@ -22,13 +22,18 @@ const credNameDeviceCode = "DeviceCodeCredential"
 type DeviceCodeCredentialOptions struct {
 	azcore.ClientOptions
 
+	// ClientID is the ID of the application users will authenticate to.
+	// Defaults to the ID of an Azure development application.
+	ClientID string
+
+	// DisableInstanceDiscovery allows disconnected cloud solutions to skip instance discovery for unknown authority hosts.
+	DisableInstanceDiscovery bool
+
 	// TenantID is the Azure Active Directory tenant the credential authenticates in. Defaults to the
 	// "organizations" tenant, which can authenticate work and school accounts. Required for single-tenant
 	// applications.
 	TenantID string
-	// ClientID is the ID of the application users will authenticate to.
-	// Defaults to the ID of an Azure development application.
-	ClientID string
+
 	// UserPrompt controls how the credential presents authentication instructions. The credential calls
 	// this function with authentication details when it receives a device code. By default, the credential
 	// prints these details to stdout.
@@ -78,7 +83,7 @@ func NewDeviceCodeCredential(options *DeviceCodeCredentialOptions) (*DeviceCodeC
 		cp = *options
 	}
 	cp.init()
-	c, err := getPublicClient(cp.ClientID, cp.TenantID, &cp.ClientOptions)
+	c, err := getPublicClient(cp.ClientID, cp.TenantID, &cp.ClientOptions, public.WithInstanceDiscovery(!cp.DisableInstanceDiscovery))
 	if err != nil {
 		return nil, err
 	}

sdk/azidentity/device_code_credential_test.go

@@ -89,15 +89,44 @@ func TestDeviceCodeCredential_Live(t *testing.T) {
 	if recording.GetRecordMode() != recording.PlaybackMode {
 		t.Skip("this test requires manual recording and can't pass live in CI")
 	}
+	for _, disabledID := range []bool{true, false} {
+		name := "default options"
+		if disabledID {
+			name = "instance discovery disabled"
+		}
+		t.Run(name, func(t *testing.T) {
+			o, stop := initRecording(t)
+			defer stop()
+			opts := DeviceCodeCredentialOptions{TenantID: liveUser.tenantID, ClientOptions: o, DisableInstanceDiscovery: disabledID}
+			if recording.GetRecordMode() == recording.PlaybackMode {
+				opts.UserPrompt = func(ctx context.Context, m DeviceCodeMessage) error { return nil }
+			}
+			cred, err := NewDeviceCodeCredential(&opts)
+			if err != nil {
+				t.Fatal(err)
+			}
+			testGetTokenSuccess(t, cred)
+		})
+	}
+}
+
+func TestDeviceCodeCredentialADFS_Live(t *testing.T) {
+	if recording.GetRecordMode() != recording.PlaybackMode {
+		t.Skip("this test requires manual recording and can't pass live in CI")
+	}
+	if adfsLiveSP.clientID == "" {
+		t.Skip("set ADFS_SP_* environment variables to run this test")
+	}
 	o, stop := initRecording(t)
 	defer stop()
-	opts := DeviceCodeCredentialOptions{TenantID: liveUser.tenantID, ClientOptions: o}
+	o.Cloud.ActiveDirectoryAuthorityHost = adfsAuthority
+	opts := DeviceCodeCredentialOptions{TenantID: "adfs", ClientID: adfsLiveUser.clientID, ClientOptions: o, DisableInstanceDiscovery: true}
 	if recording.GetRecordMode() == recording.PlaybackMode {
 		opts.UserPrompt = func(ctx context.Context, m DeviceCodeMessage) error { return nil }
 	}
 	cred, err := NewDeviceCodeCredential(&opts)
 	if err != nil {
 		t.Fatal(err)
 	}
-	testGetTokenSuccess(t, cred)
+	testGetTokenSuccess(t, cred, adfsScope)
 }

sdk/azidentity/environment_credential.go

@@ -23,6 +23,9 @@ const envVarSendCertChain = "AZURE_CLIENT_SEND_CERTIFICATE_CHAIN"
 // EnvironmentCredentialOptions contains optional parameters for EnvironmentCredential
 type EnvironmentCredentialOptions struct {
 	azcore.ClientOptions
+
+	// DisableInstanceDiscovery allows disconnected cloud solutions to skip instance discovery for unknown authority hosts.
+	DisableInstanceDiscovery bool
 }
 
 // EnvironmentCredential authenticates a service principal with a secret or certificate, or a user with a password, depending
@@ -74,7 +77,7 @@ func NewEnvironmentCredential(options *EnvironmentCredentialOptions) (*Environme
 	}
 	if clientSecret := os.Getenv(azureClientSecret); clientSecret != "" {
 		log.Write(EventAuthentication, "EnvironmentCredential will authenticate with ClientSecretCredential")
-		o := &ClientSecretCredentialOptions{ClientOptions: options.ClientOptions}
+		o := &ClientSecretCredentialOptions{ClientOptions: options.ClientOptions, DisableInstanceDiscovery: options.DisableInstanceDiscovery}
 		cred, err := NewClientSecretCredential(tenantID, clientID, clientSecret, o)
 		if err != nil {
 			return nil, err
@@ -95,7 +98,7 @@ func NewEnvironmentCredential(options *EnvironmentCredentialOptions) (*Environme
 		if err != nil {
 			return nil, fmt.Errorf(`failed to load certificate from "%s": %v`, certPath, err)
 		}
-		o := &ClientCertificateCredentialOptions{ClientOptions: options.ClientOptions}
+		o := &ClientCertificateCredentialOptions{ClientOptions: options.ClientOptions, DisableInstanceDiscovery: options.DisableInstanceDiscovery}
 		if v, ok := os.LookupEnv(envVarSendCertChain); ok {
 			o.SendCertificateChain = v == "1" || strings.ToLower(v) == "true"
 		}
@@ -108,7 +111,7 @@ func NewEnvironmentCredential(options *EnvironmentCredentialOptions) (*Environme
 	if username := os.Getenv(azureUsername); username != "" {
 		if password := os.Getenv(azurePassword); password != "" {
 			log.Write(EventAuthentication, "EnvironmentCredential will authenticate with UsernamePasswordCredential")
-			o := &UsernamePasswordCredentialOptions{ClientOptions: options.ClientOptions}
+			o := &UsernamePasswordCredentialOptions{ClientOptions: options.ClientOptions, DisableInstanceDiscovery: options.DisableInstanceDiscovery}
 			cred, err := NewUsernamePasswordCredential(tenantID, clientID, username, password, o)
 			if err != nil {
 				return nil, err