Enabled globally disabling instance discovery before token acquisition

Fix #22535

Author: Lei Jin

src/Accounts/Accounts/ChangeLog.md

@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Enabled globally disabling instance discovery before token acquisition [#22535].
 * Added `AsSecureString` to `Get-AzAccessToken` to convert the returned token to SecureString [#24190].
 * Upgraded Azure.Core to 1.37.0.
 

src/Accounts/Authentication/Authentication/Parameters/AuthenticationParameters.cs

@@ -13,6 +13,9 @@
 // ----------------------------------------------------------------------------------
 
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Commands.Common.Exceptions;
+using Microsoft.Azure.Commands.Shared.Config;
+using Microsoft.Azure.PowerShell.Common.Config;
 
 namespace Microsoft.Azure.Commands.Common.Authentication
 {
@@ -28,6 +31,8 @@ public abstract class AuthenticationParameters
 
         public string ResourceId { get; set; }
 
+        public bool? DisableInstanceDiscovery { get; set; } = null;
+
         public AuthenticationParameters(
             PowerShellTokenCacheProvider tokenCacheProvider,
             IAzureEnvironment environment,
@@ -41,6 +46,17 @@ public AuthenticationParameters(
             TokenCache = tokenCache;
             TenantId = tenantId;
             ResourceId = resourceId;
+            try
+            {
+                if (AzureSession.Instance.TryGetComponent<IConfigManager>(nameof(IConfigManager), out var config))
+                {
+                    DisableInstanceDiscovery = config.GetConfigValue<bool>(ConfigKeys.DisableInstanceDiscovery);
+                }
+            }
+            catch(AzPSArgumentException)
+            {
+                DisableInstanceDiscovery = null;
+            }
         }
     }
 }

src/Accounts/Authentication/Config/ConfigInitializer.cs

@@ -232,6 +232,7 @@ private void RegisterConfigs(IConfigManager configManager)
                 false,
                 "AZURE_CLIENTS_SHOW_SECRETS_WARNING",
                 new[] { AppliesTo.Az }));
+            configManager.RegisterConfig(new DisableInstanceDiscoveryConfig());
         }
     }
 }

src/Accounts/Authentication/Config/Definitions/DisableInstanceDiscoveryConfig.cs

@@ -0,0 +1,48 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Commands.Common.Authentication.Properties;
+using Microsoft.Azure.Commands.ResourceManager.Common;
+using Microsoft.Azure.Commands.Shared.Config;
+using Microsoft.Azure.PowerShell.Common.Config;
+using System;
+using System.Collections.Generic;
+
+namespace Microsoft.Azure.Commands.Common.Authentication.Config.Definitions
+{
+    /// <summary>
+    /// Definition of the config to control whether login by WAM (web account manager) or not.
+    /// </summary>
+    internal class DisableInstanceDiscoveryConfig : TypedConfig<bool>
+    {
+        public override object DefaultValue => false; // Opt-in. Will change to opt-out.
+
+        public override string Key => ConfigKeys.DisableInstanceDiscovery;
+
+        public override string HelpMessage => Resources.HelpMessageOfDisableInstanceDiscovery;
+
+        public override IReadOnlyCollection<AppliesTo> CanApplyTo => new[] { AppliesTo.Az };
+
+        protected override void ApplyTyped(bool value)
+        {
+            base.ApplyTyped(value);
+            EventHandler<StreamEventArgs> writeWarningEvent;
+            if (AzureSession.Instance.TryGetComponent(AzureRMCmdlet.WriteWarningKey, out writeWarningEvent) && value)
+            {
+                writeWarningEvent(this, new StreamEventArgs() { Message = Resources.DisableInstanceDiscoveryWarning });
+            }
+        }
+    }
+}

src/Accounts/Authentication/Properties/Resources.Designer.cs

@@ -405,4 +405,10 @@
   <data name="HelpMessageOfDisplaySecretsWarning" xml:space="preserve">
     <value>When enabled, a warning message will be displayed when the cmdlet output contains secrets. Learn more at https://go.microsoft.com/fwlink/?linkid=2258844</value>
   </data>
-</root>
+  <data name="DisableInstanceDiscoveryWarning" xml:space="preserve">
+    <value>Set it to true to disable both instance discovery and authority validation. This functionality is intended for use in scenarios where the metadata endpoint cannot be reached, such as in private clouds or Azure Stack. It is crucial to ensure that the configured authority host is valid and trustworthy.</value>
+  </data>
+  <data name="HelpMessageOfDisableInstanceDiscovery" xml:space="preserve">
+    <value>Set it to true to disable both instance discovery and authority validation. This functionality is intended for use in scenarios where the metadata endpoint cannot be reached, such as in private clouds or Azure Stack. The process of instance discovery entails retrieving authority metadata from https://login.microsoft.com/ to validate the authority. By setting this to true, the validation of the authority is disabled. As a result, it is crucial to ensure that the configured authority host is valid and trustworthy.</value>
+  </data>
+</root>

src/Accounts/Authentication/Properties/Resources.resx

@@ -45,6 +45,7 @@ public override Task<IAccessToken> Authenticate(AuthenticationParameters paramet
                 AuthorityHost = new Uri(authority),
                 TokenCachePersistenceOptions = spParameters.TokenCacheProvider.GetTokenCachePersistenceOptions()
             };
+            options.DisableInstanceDiscovery = spParameters.DisableInstanceDiscovery ?? options.DisableInstanceDiscovery;
             options.Diagnostics.IsTelemetryEnabled = false; // disable telemetry to avoid error thrown from Azure.Core that AssemblyInformationalVersion is null
             TokenCredential tokenCredential = new ClientAssertionCredential(tenantId, spParameters.ClientId, () => GetClientAssertion(spParameters), options);
             string parametersLog = $"- ClientId:'{spParameters.ClientId}', TenantId:'{tenantId}', ClientAssertion:'***' Scopes:'{string.Join(",", scopes)}'";

src/Accounts/Authenticators/ClientAssertionAuthenticator.cs

@@ -53,6 +53,7 @@ public override Task<IAccessToken> Authenticate(AuthenticationParameters paramet
                 TenantId = tenantId,
                 TokenCachePersistenceOptions = tokenCacheProvider.GetTokenCachePersistenceOptions(),
             };
+            options.DisableInstanceDiscovery = deviceCodeParameters.DisableInstanceDiscovery ?? options.DisableInstanceDiscovery;
             var codeCredential = new DeviceCodeCredential(options);
 
             TracingAdapter.Information($"{DateTime.Now:T} - [DeviceCodeAuthenticator] Calling DeviceCodeCredential.AuthenticateAsync - TenantId:'{options.TenantId}', Scopes:'{string.Join(",", scopes)}', AuthorityHost:'{options.AuthorityHost}'");

src/Accounts/Authenticators/DeviceCodeAuthenticator.cs

@@ -12,23 +12,20 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using System;
-using System.Diagnostics;
-using System.Net;
-using System.Net.Sockets;
-using System.Runtime.InteropServices;
-using System.Threading;
-using System.Threading.Tasks;
-
 using Azure.Core;
 using Azure.Identity;
-using Azure.Identity.Broker;
 
 using Hyak.Common;
 
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 
+using System;
+using System.Net;
+using System.Net.Sockets;
+using System.Threading;
+using System.Threading.Tasks;
+
 namespace Microsoft.Azure.PowerShell.Authenticators
 {
     /// <summary>
@@ -68,6 +65,7 @@ public override Task<IAccessToken> Authenticate(AuthenticationParameters paramet
                 RedirectUri = GetReplyUrl(onPremise, interactiveParameters.PromptAction),
                 LoginHint = interactiveParameters.UserId,
             };
+            options.DisableInstanceDiscovery = interactiveParameters.DisableInstanceDiscovery ?? options.DisableInstanceDiscovery;
             var browserCredential = new InteractiveBrowserCredential(options);
 
             TracingAdapter.Information($"{DateTime.Now:T} - [InteractiveUserAuthenticator] Calling InteractiveBrowserCredential.AuthenticateAsync with TenantId:'{options.TenantId}', Scopes:'{string.Join(",", scopes)}', AuthorityHost:'{options.AuthorityHost}', RedirectUri:'{options.RedirectUri}'");

src/Accounts/Authenticators/InteractiveUserAuthenticator.cs

@@ -60,6 +60,7 @@ public override Task<IAccessToken> Authenticate(AuthenticationParameters paramet
                 RedirectUri = GetReplyUrl(onPremise, interactiveParameters.PromptAction),
                 LoginHint = interactiveParameters.UserId
             };
+            options.DisableInstanceDiscovery = interactiveParameters.DisableInstanceDiscovery ?? options.DisableInstanceDiscovery;
             var browserCredential = new InteractiveBrowserCredential(options);
 
             TracingAdapter.Information($"{DateTime.Now:T} - [InteractiveWamAuthenticator] Calling InteractiveBrowserCredential.AuthenticateAsync with TenantId:'{options.TenantId}', Scopes:'{string.Join(",", scopes)}', AuthorityHost:'{options.AuthorityHost}', RedirectUri:'{options.RedirectUri}'");

src/Accounts/Authenticators/InteractiveWamAuthenticator.cs

@@ -12,20 +12,16 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using System;
-using System.Text.RegularExpressions;
-using System.Threading;
-using System.Threading.Tasks;
-
 using Azure.Core;
-using Azure.Identity;
-
-using Hyak.Common;
 
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Azure.PowerShell.Authenticators.Factories;
 
+using System.Text.RegularExpressions;
+using System.Threading;
+using System.Threading.Tasks;
+
 namespace Microsoft.Azure.PowerShell.Authenticators
 {
     public class ManagedServiceIdentityAuthenticator : DelegatingAuthenticator

src/Accounts/Authenticators/ManagedServiceIdentityAuthenticator.cs

@@ -55,6 +55,7 @@ public override Task<IAccessToken> Authenticate(AuthenticationParameters paramet
                 SendCertificateChain = spParameters.SendCertificateChain ?? default(bool)
             };
 
+            options.DisableInstanceDiscovery = spParameters.DisableInstanceDiscovery ?? options.DisableInstanceDiscovery;
             TokenCredential tokenCredential = null;
             string parametersLog = null;
             if (!string.IsNullOrEmpty(spParameters.Thumbprint))

src/Accounts/Authenticators/ServicePrincipalAuthenticator.cs

@@ -68,6 +68,7 @@ private static SharedTokenCacheCredentialOptions GetTokenCredentialOptions(Silen
             options.Username = silentParameters.UserId;
             options.AuthorityHost = new Uri(authority);
             options.TenantId = tenantId;
+            options.DisableInstanceDiscovery = silentParameters.DisableInstanceDiscovery ?? options.DisableInstanceDiscovery;
             return options;
         }
 

src/Accounts/Authenticators/SilentAuthenticator.cs

@@ -59,6 +59,7 @@ public override Task<IAccessToken> Authenticate(AuthenticationParameters paramet
                 AuthorityHost = new Uri(authority),
                 TokenCachePersistenceOptions = tokenCacheProvider.GetTokenCachePersistenceOptions()
             };
+            credentialOptions.DisableInstanceDiscovery = upParameters.DisableInstanceDiscovery ?? credentialOptions.DisableInstanceDiscovery;
             if (upParameters.Password != null)
             {
                 passwordCredential = new UsernamePasswordCredential(upParameters.UserId, upParameters.Password.ConvertToString(), tenantId, clientId, credentialOptions);

src/Accounts/Authenticators/UsernamePasswordAuthenticator.cs

@@ -35,6 +35,7 @@ internal static class ConfigKeys
         //Use DisableErrorRecordsPersistence as opt-out for now, will replace it with EnableErrorRecordsPersistence as opt-in at next major release (November 2023)
         public const string DisableErrorRecordsPersistence = "DisableErrorRecordsPersistence";
         public const string EnableErrorRecordsPersistence = "EnableErrorRecordsPersistence";
+        public const string DisableInstanceDiscovery = "DisableInstanceDiscovery";
         public const string CheckForUpgrade = "CheckForUpgrade";
         public const string EnvCheckForUpgrade = "AZUREPS_CHECK_FOR_UPGRADE";
         public const string DisplaySecretsWarning = "DisplaySecretsWarning";