Support Tenant Id Challenges / Hints

sdk/core/Azure.Core/src/Pipeline/BearerTokenAuthenticationPolicy.cs

@@ -136,7 +136,7 @@ private async ValueTask ProcessAsync(HttpMessage message, ReadOnlyMemory<HttpPip
                 // If we succeed in getting the context, authenticate the request and pass it up the policy chain.
                 if (async)
                 {
-                    if (await AuthorizeRequestOnChallengeAsync((message)).ConfigureAwait(false))
+                    if (await AuthorizeRequestOnChallengeAsync(message).ConfigureAwait(false))
                     {
                         await ProcessNextAsync(message, pipeline).ConfigureAwait(false);
                     }

sdk/core/Azure.Core/src/Shared/IBearerTokenChallengeOptions.cs

@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+namespace Azure.Core
+{
+    /// <summary>
+    /// Options to be exposed in client options classes related to bearer token authorization challenge scenarios.
+    /// </summary>
+    internal interface IBearerTokenChallengeOptions
+    {
+        /// <summary>
+        ///  Disables tenant discovery through the authorization challenge when the client is configured to use a <see cref="TokenCredential"/>.
+        /// When disabled, the client will not attempt an initial un-authroized request to prompt a challenge in order to discover the correct tenant for the resource.
+        /// </summary>
+        public bool DisableTenantDiscovery { get; }
+    }
+}

sdk/core/Azure.Core/src/TokenRequestContext.cs

@@ -18,6 +18,7 @@ public TokenRequestContext(string[] scopes, string? parentRequestId)
             Scopes = scopes;
             ParentRequestId = parentRequestId;
             Claims = default;
+            TenantIdHint = default;
         }
 
         /// <summary>
@@ -31,6 +32,19 @@ public TokenRequestContext(string[] scopes, string? parentRequestId = default, s
             Scopes = scopes;
             ParentRequestId = parentRequestId;
             Claims = claims;
+            TenantIdHint = default;
+        }
+
+        /// <summary>
+        /// Creates a new TokenRequest with the specified scopes.
+        /// </summary>
+        /// <param name="options"> The options to initialized the <see cref="TokenRequestContext"/> </param>
+        public TokenRequestContext(TokenRequestContextOptions options)
+        {
+            Scopes = options.Scopes;
+            ParentRequestId = options.ParentRequestId;
+            Claims = options.Claims;
+            TenantIdHint = options.TenantIdHint;
         }
 
         /// <summary>
@@ -47,5 +61,10 @@ public TokenRequestContext(string[] scopes, string? parentRequestId = default, s
         /// Additional claims to be included in the token. See <see href="https://openid.net/specs/openid-connect-core-1_0-final.html#ClaimsParameter">https://openid.net/specs/openid-connect-core-1_0-final.html#ClaimsParameter</see> for more information on format and content.
         /// </summary>
         public string? Claims { get; }
+
+        /// <summary>
+        /// A hint to indicate which tenantId is preferred.
+        /// </summary>
+        public string? TenantIdHint { get; }
     }
 }

sdk/core/Azure.Core/src/TokenRequestContextOptions.cs

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+namespace Azure.Core
+{
+    /// <summary>
+    /// Options for <see cref="TokenRequestContext"/> initialization.
+    /// </summary>
+    public struct TokenRequestContextOptions
+    {
+        /// <summary>
+        /// The scopes required for the token.
+        /// </summary>
+        public string[] Scopes { get; set; }
+
+        /// <summary>
+        /// The <see cref="Request.ClientRequestId"/> of the request requiring a token for authentication, if applicable.
+        /// </summary>
+        public string? ParentRequestId { get; set; }
+
+        /// <summary>
+        /// Additional claims to be included in the token. See <see href="https://openid.net/specs/openid-connect-core-1_0-final.html#ClaimsParameter">https://openid.net/specs/openid-connect-core-1_0-final.html#ClaimsParameter</see> for more information on format and content.
+        /// </summary>
+        public string? Claims { get; set; }
+
+        /// <summary>
+        /// A hint to indicate which tenantId is preferred.
+        /// </summary>
+        public string? TenantIdHint { get; set; }
+    }
+}

sdk/identity/Azure.Identity/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Fixes and improvements
 
 - Added `LoginHint` property to `InteractiveBrowserCredentialOptions` which allows a user name to be pre-selected for interactive logins. Setting this option skips the account selection prompt and immediately attempts to login with the specified account.
+- TenantId values returned from service challenge responses can now be used to request tokens from the correct tenantId.
 
 ## 1.4.0 (2021-05-12)
 

sdk/identity/Azure.Identity/src/AuthorizationCodeCredential.cs

@@ -5,7 +5,6 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
-using Azure.Core.Diagnostics;
 using Azure.Core.Pipeline;
 using Microsoft.Identity.Client;
 
@@ -18,12 +17,13 @@ namespace Azure.Identity
     /// </summary>
     public class AuthorizationCodeCredential : TokenCredential
     {
-        private readonly IConfidentialClientApplication _confidentialClient;
-        private readonly ClientDiagnostics _clientDiagnostics;
         private readonly string _authCode;
         private readonly string _clientId;
         private readonly CredentialPipeline _pipeline;
         private AuthenticationRecord _record;
+        private readonly string _tenantId;
+        private readonly TokenCredentialOptions _options;
+        private readonly MsalConfidentialClient _client;
 
         /// <summary>
         /// Protected constructor for mocking.
@@ -55,26 +55,26 @@ public AuthorizationCodeCredential(string tenantId, string clientId, string clie
         /// See https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-auth-code-flow for more information.</param>
         /// <param name="options">Options that allow to configure the management of the requests sent to the Azure Active Directory service.</param>
         public AuthorizationCodeCredential(string tenantId, string clientId, string clientSecret, string authorizationCode, TokenCredentialOptions options)
-        {
-            Validations.ValidateTenantId(tenantId, nameof(tenantId));
-
-            if (clientSecret is null) throw new ArgumentNullException(nameof(clientSecret));
-
-            _clientId = clientId ?? throw new ArgumentNullException(nameof(clientId));
-
-            _authCode = authorizationCode ?? throw new ArgumentNullException(nameof(authorizationCode));
+            : this(tenantId, clientId, clientSecret, authorizationCode, options, null)
+        {}
 
-            options ??= new TokenCredentialOptions();
-
-            _pipeline = CredentialPipeline.GetInstance(options);
-
-            _confidentialClient = ConfidentialClientApplicationBuilder.Create(clientId).WithHttpClientFactory(new HttpPipelineClientFactory(_pipeline.HttpPipeline)).WithTenantId(tenantId).WithClientSecret(clientSecret).Build();
-
-            _clientDiagnostics = new ClientDiagnostics(options);
+        internal AuthorizationCodeCredential(string tenantId, string clientId, string clientSecret, string authorizationCode, TokenCredentialOptions options, MsalConfidentialClient client)
+        {
+            Argument.AssertNotNull(clientSecret, nameof(clientSecret));
+            Argument.AssertNotNull(clientId, nameof(clientId));
+            Argument.AssertNotNull(authorizationCode, nameof(authorizationCode));
+            _tenantId = Validations.ValidateTenantId(tenantId, nameof(tenantId));
+            _clientId = clientId;
+            _authCode = authorizationCode ;
+            _options = options ?? new TokenCredentialOptions();
+            _pipeline = CredentialPipeline.GetInstance(_options);
+
+            _client = client ?? new MsalConfidentialClient(_pipeline, tenantId, clientId, clientSecret, options as ITokenCacheOptions);
         }
 
         /// <summary>
-        /// Obtains a token from the Azure Active Directory service, using the specified authorization code authenticate. This method is called automatically by Azure SDK client libraries. You may call this method directly, but you must also handle token caching and token refreshing.
+        /// Obtains a token from the Azure Active Directory service, using the specified authorization code authenticate. This method is called automatically by
+        /// Azure SDK client libraries. You may call this method directly, but you must also handle token caching and token refreshing.
         /// </summary>
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
@@ -85,7 +85,8 @@ public override AccessToken GetToken(TokenRequestContext requestContext, Cancell
         }
 
         /// <summary>
-        /// Obtains a token from the Azure Active Directory service, using the specified authorization code authenticate. This method is called automatically by Azure SDK client libraries. You may call this method directly, but you must also handle token caching and token refreshing.
+        /// Obtains a token from the Azure Active Directory service, using the specified authorization code authenticate. This method is called automatically by
+        /// Azure SDK client libraries. You may call this method directly, but you must also handle token caching and token refreshing.
         /// </summary>
         /// <param name="requestContext">The details of the authentication request.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> controlling the request lifetime.</param>
@@ -101,19 +102,24 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
 
             try
             {
-                AccessToken token = default;
+                AccessToken token;
+                var tenantId = TenantIdResolver.Resolve(_tenantId, requestContext, _options);
 
                 if (_record is null)
                 {
-                    AuthenticationResult result = await _confidentialClient.AcquireTokenByAuthorizationCode(requestContext.Scopes, _authCode).ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+                    AuthenticationResult result = await _client
+                        .AcquireTokenByAuthorizationCodeAsync(requestContext.Scopes, _authCode, tenantId, async, cancellationToken)
+                        .ConfigureAwait(false);
 
                     _record = new AuthenticationRecord(result, _clientId);
 
                     token = new AccessToken(result.AccessToken, result.ExpiresOn);
                 }
                 else
                 {
-                    AuthenticationResult result = await _confidentialClient.AcquireTokenSilent(requestContext.Scopes, (AuthenticationAccount)_record).ExecuteAsync(async, cancellationToken).ConfigureAwait(false);
+                    AuthenticationResult result = await _client
+                        .AcquireTokenSilentAsync(requestContext.Scopes, (AuthenticationAccount)_record, tenantId, async, cancellationToken)
+                        .ConfigureAwait(false);
 
                     token = new AccessToken(result.AccessToken, result.ExpiresOn);
                 }

sdk/identity/Azure.Identity/src/Azure.Identity.csproj

@@ -11,7 +11,12 @@
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
   </PropertyGroup>
   <ItemGroup>
+    <!-- TODO: uncomment when Azure.Core ships
     <PackageReference Include="Azure.Core" />
+    end TODO -->
+    <!-- TODO: revert when Azure.Core ships -->
+    <ProjectReference Include="..\..\..\core\Azure.Core\src\Azure.Core.csproj" />
+    <!--  end TODO-->
     <PackageReference Include="System.Memory" />
     <PackageReference Include="System.Text.Json" />
     <PackageReference Include="System.Threading.Tasks.Extensions" />

sdk/identity/Azure.Identity/src/AzureCliCredential.cs

@@ -86,7 +86,7 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
 
             try
             {
-                AccessToken token = await RequestCliAccessTokenAsync(async, requestContext.Scopes, cancellationToken).ConfigureAwait(false);
+                AccessToken token = await RequestCliAccessTokenAsync(async, requestContext, cancellationToken).ConfigureAwait(false);
                 return scope.Succeeded(token);
             }
             catch (Exception e)
@@ -95,13 +95,14 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
             }
         }
 
-        private async ValueTask<AccessToken> RequestCliAccessTokenAsync(bool async, string[] scopes, CancellationToken cancellationToken)
+        private async ValueTask<AccessToken> RequestCliAccessTokenAsync(bool async, TokenRequestContext context, CancellationToken cancellationToken)
         {
-            string resource = ScopeUtilities.ScopesToResource(scopes);
+            string resource = ScopeUtilities.ScopesToResource(context.Scopes);
+            string tenantId = TenantIdResolver.Resolve(null, context, null);
 
             ScopeUtilities.ValidateScope(resource);
 
-            GetFileNameAndArguments(resource, out string fileName, out string argument);
+            GetFileNameAndArguments(resource, tenantId, out string fileName, out string argument);
             ProcessStartInfo processStartInfo = GetAzureCliProcessStartInfo(fileName, argument);
             using var processRunner = new ProcessRunner(_processService.Create(processStartInfo), TimeSpan.FromMilliseconds(CliProcessTimeoutMs), cancellationToken);
 
@@ -157,9 +158,13 @@ private ProcessStartInfo GetAzureCliProcessStartInfo(string fileName, string arg
                 Environment = { { "PATH", _path } }
             };
 
-        private static void GetFileNameAndArguments(string resource, out string fileName, out string argument)
+        private static void GetFileNameAndArguments(string resource, string tenantId, out string fileName, out string argument)
         {
-            string command = $"az account get-access-token --output json --resource {resource}";
+            string command = tenantId switch
+            {
+                null => $"az account get-access-token --output json --resource {resource}",
+                _ => $"az account get-access-token --output json --resource {resource} -tenant {tenantId}"
+            };
 
             if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
             {

sdk/identity/Azure.Identity/src/AzurePowerShellCredential.cs

@@ -40,6 +40,8 @@ public class AzurePowerShellCredential : TokenCredential
         private static readonly string DefaultWorkingDir =
             RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? DefaultWorkingDirWindows : DefaultWorkingDirNonWindows;
 
+        private readonly AzurePowerShellCredentialOptions _options;
+
         private const int ERROR_FILE_NOT_FOUND = 2;
 
         /// <summary>
@@ -59,6 +61,7 @@ public AzurePowerShellCredential(AzurePowerShellCredentialOptions options) : thi
         internal AzurePowerShellCredential(AzurePowerShellCredentialOptions options, CredentialPipeline pipeline, IProcessService processService)
         {
             UseLegacyPowerShell = false;
+            _options = options;
             _pipeline = pipeline ?? CredentialPipeline.GetInstance(options);
             _processService = processService ?? ProcessService.Default;
         }
@@ -91,15 +94,15 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
 
             try
             {
-                AccessToken token = await RequestAzurePowerShellAccessTokenAsync(async, requestContext.Scopes, cancellationToken).ConfigureAwait(false);
+                AccessToken token = await RequestAzurePowerShellAccessTokenAsync(async, requestContext, cancellationToken).ConfigureAwait(false);
                 return scope.Succeeded(token);
             }
             catch (Win32Exception ex) when (ex.NativeErrorCode == ERROR_FILE_NOT_FOUND && RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
             {
                 UseLegacyPowerShell = true;
                 try
                 {
-                    AccessToken token = await RequestAzurePowerShellAccessTokenAsync(async, requestContext.Scopes, cancellationToken).ConfigureAwait(false);
+                    AccessToken token = await RequestAzurePowerShellAccessTokenAsync(async, requestContext, cancellationToken).ConfigureAwait(false);
                     return scope.Succeeded(token);
                 }
                 catch (Exception e)
@@ -113,13 +116,14 @@ private async ValueTask<AccessToken> GetTokenImplAsync(bool async, TokenRequestC
             }
         }
 
-        private async ValueTask<AccessToken> RequestAzurePowerShellAccessTokenAsync(bool async, string[] scopes, CancellationToken cancellationToken)
+        private async ValueTask<AccessToken> RequestAzurePowerShellAccessTokenAsync(bool async, TokenRequestContext context, CancellationToken cancellationToken)
         {
-            string resource = ScopeUtilities.ScopesToResource(scopes);
+            string resource = ScopeUtilities.ScopesToResource(context.Scopes);
 
             ScopeUtilities.ValidateScope(resource);
+            var tenantId = TenantIdResolver.Resolve(null, context, _options);
 
-            GetFileNameAndArguments(resource, out string fileName, out string argument);
+            GetFileNameAndArguments(resource, tenantId, out string fileName, out string argument);
             ProcessStartInfo processStartInfo = GetAzurePowerShellProcessStartInfo(fileName, argument);
             using var processRunner = new ProcessRunner(
                 _processService.Create(processStartInfo),
@@ -185,7 +189,7 @@ private static ProcessStartInfo GetAzurePowerShellProcessStartInfo(string fileNa
                 WorkingDirectory = DefaultWorkingDir
             };
 
-        private void GetFileNameAndArguments(string resource, out string fileName, out string argument)
+        private void GetFileNameAndArguments(string resource, string tenantId, out string fileName, out string argument)
         {
             string powershellExe = "pwsh -NonInteractive -EncodedCommand";
 
@@ -194,6 +198,8 @@ private void GetFileNameAndArguments(string resource, out string fileName, out s
                 powershellExe = "powershell -NonInteractive -EncodedCommand";
             }
 
+            var tenantIdArg = tenantId == null ? string.Empty : $" -TenantId {tenantId}";
+
             string command = @$"
 $ErrorActionPreference = 'Stop'
 [version]$minimumVersion = '2.2.0'
@@ -205,7 +211,7 @@ private void GetFileNameAndArguments(string resource, out string fileName, out s
     exit
 }}
 
-$token = Get-AzAccessToken -ResourceUrl '{resource}'
+$token = Get-AzAccessToken -ResourceUrl '{resource}'{tenantIdArg}
 
 $x = $token | ConvertTo-Xml
 return $x.Objects.FirstChild.OuterXml