Implemement OnBehalfOfCredential

sdk/core/Azure.Core/api/Azure.Core.net461.cs

@@ -454,6 +454,7 @@ internal RetryOptions() { }
     public abstract partial class TokenCredential
     {
         protected TokenCredential() { }
+        public virtual bool SupportsCaching { get { throw null; } }
         public abstract Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
         public abstract System.Threading.Tasks.ValueTask<Azure.Core.AccessToken> GetTokenAsync(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
     }

sdk/core/Azure.Core/api/Azure.Core.net5.0.cs

@@ -454,6 +454,7 @@ internal RetryOptions() { }
     public abstract partial class TokenCredential
     {
         protected TokenCredential() { }
+        public virtual bool SupportsCaching { get { throw null; } }
         public abstract Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
         public abstract System.Threading.Tasks.ValueTask<Azure.Core.AccessToken> GetTokenAsync(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
     }

sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs

@@ -454,6 +454,7 @@ internal RetryOptions() { }
     public abstract partial class TokenCredential
     {
         protected TokenCredential() { }
+        public virtual bool SupportsCaching { get { throw null; } }
         public abstract Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
         public abstract System.Threading.Tasks.ValueTask<Azure.Core.AccessToken> GetTokenAsync(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
     }

sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs

@@ -454,6 +454,7 @@ internal RetryOptions() { }
     public abstract partial class TokenCredential
     {
         protected TokenCredential() { }
+        public virtual bool SupportsCaching { get { throw null; } }
         public abstract Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
         public abstract System.Threading.Tasks.ValueTask<Azure.Core.AccessToken> GetTokenAsync(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken);
     }

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

@@ -197,12 +197,18 @@ public async ValueTask<string> GetHeaderValueAsync(HttpMessage message, TokenReq
                 bool getTokenFromCredential;
                 TaskCompletionSource<HeaderValueInfo> headerValueTcs;
                 TaskCompletionSource<HeaderValueInfo>? backgroundUpdateTcs;
+                HeaderValueInfo info;
                 int maxCancellationRetries = 3;
 
+                if (_credential.SupportsCaching)
+                {
+                    info = await GetHeaderValueFromCredentialAsync(context, async, message.CancellationToken).ConfigureAwait(false);
+                    return info.HeaderValue;
+                }
+
                 while (true)
                 {
                     (headerValueTcs, backgroundUpdateTcs, getTokenFromCredential) = GetTaskCompletionSources(context);
-                    HeaderValueInfo info;
                     if (getTokenFromCredential)
                     {
                         if (backgroundUpdateTcs != null)

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

@@ -1,9 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Azure.Core.Pipeline;
 
 namespace Azure.Core
 {
@@ -27,5 +27,11 @@ public abstract class TokenCredential
         /// <param name="cancellationToken">The <see cref="CancellationToken"/> to use.</param>
         /// <returns>A valid <see cref="AccessToken"/>.</returns>
         public abstract AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Indicates if a <see cref="TokenCredential"/> supports token caching.
+        /// If <c>true</c>, upstream components such as a <see cref="HttpPipelinePolicy"/> should rely on the credential itself to manage any token caching.
+        /// </summary>
+        public virtual bool SupportsCaching => false;
     }
 }

sdk/core/Azure.Core/tests/BearerTokenAuthenticationPolicyTests.cs

@@ -33,6 +33,32 @@ public async Task BearerTokenAuthenticationPolicy_UsesTokenProvidedByCredentials
             Assert.AreEqual("Bearer token", authValue);
         }
 
+        [Test]
+        public async Task BearerTokenAuthenticationPolicy_SupportsCaching_True()
+        {
+            int expectedCalls = 3;
+            int tokenCalls = 0;
+            var credential = new TokenCredentialStub(
+                (r, c) =>
+                {
+                    tokenCalls++;
+                    return r.Scopes.SequenceEqual(new[] { "scope1", "scope2" }) ? new AccessToken("token", DateTimeOffset.MaxValue) : default;
+                },
+                IsAsync);
+            credential.SetSupportsCaching(true);
+            var policy = new BearerTokenAuthenticationPolicy(credential, new[] { "scope1", "scope2" });
+
+            MockTransport transport = CreateMockTransport(_ => new MockResponse(200));
+
+            for (int i = 0; i < expectedCalls; i++)
+            {
+                await SendGetRequest(transport, policy, uri: new Uri("https://example.com"));
+            }
+
+            Assert.That(transport.Requests.Select(r => r.Headers.Single()), Has.All.Property("Name").EqualTo("Authorization"));
+            Assert.AreEqual(expectedCalls, tokenCalls);
+        }
+
         [Test]
         public async Task BearerTokenAuthenticationPolicy_RequestsTokenEveryRequest()
         {
@@ -807,6 +833,10 @@ public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext request
 
             public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken)
                 => _getTokenHandler(requestContext, cancellationToken);
+
+            private bool _supportsCaching = false;
+            public void SetSupportsCaching(bool supportsCaching) => _supportsCaching = supportsCaching;
+            public override bool SupportsCaching => _supportsCaching;
         }
     }
 }

sdk/identity/Azure.Identity/CHANGELOG.md

@@ -17,6 +17,7 @@
 - Added support to `ManagedIdentityCredential` for Bridge to Kubernetes local development authentication.
 - TenantId values returned from service challenge responses can now be used to request tokens from the correct tenantId. To support this feature, there is a new `AllowMultiTenantAuthentication` option on `TokenCredentialOptions`.
   - By default, `AllowMultiTenantAuthentication` is false. When this option property is false and the tenant Id configured in the credential options differs from the tenant Id set in the `TokenRequestContext` sent to a credential, an `AuthorizationFailedException` will be thrown. This is potentially breaking change as it could be a different exception than what was thrown previously. This exception behavior can be overridden by either setting an `AppContext` switch named "Azure.Identity.EnableLegacyTenantSelection" to `true` or by setting the environment variable "AZURE_IDENTITY_ENABLE_LEGACY_TENANT_SELECTION" to "true". Note: AppContext switches can also be configured via configuration like below:
+- Added `OnBehalfOfFlowCredential` which enables support for AAD On-Behalf-Of (OBO) flow. See the [Azure Active Directory documentation](https://docs.microsoft.com/azure/active-directory/develop/v2-oauth2-on-behalf-of-flow) to learn more about OBO flow scenarios.
 
 ```xml  
 <ItemGroup>

sdk/identity/Azure.Identity/api/Azure.Identity.netstandard2.0.cs

@@ -225,6 +225,19 @@ public ManagedIdentityCredential(string clientId = null, Azure.Identity.TokenCre
         public override Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public override System.Threading.Tasks.ValueTask<Azure.Core.AccessToken> GetTokenAsync(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }
+    public partial class OnBehalfOfCredential : Azure.Core.TokenCredential
+    {
+        protected OnBehalfOfCredential() { }
+        public OnBehalfOfCredential(string tenantId, string clientId, string clientSecret, Azure.Identity.OnBehalfOfCredentialOptions options = null) { }
+        public override bool SupportsCaching { get { throw null; } }
+        public override Azure.Core.AccessToken GetToken(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken) { throw null; }
+        public override System.Threading.Tasks.ValueTask<Azure.Core.AccessToken> GetTokenAsync(Azure.Core.TokenRequestContext requestContext, System.Threading.CancellationToken cancellationToken) { throw null; }
+    }
+    public partial class OnBehalfOfCredentialOptions : Azure.Identity.TokenCredentialOptions
+    {
+        public OnBehalfOfCredentialOptions() { }
+        public Azure.Identity.TokenCachePersistenceOptions TokenCachePersistenceOptions { get { throw null; } set { } }
+    }
     [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]
     public readonly partial struct RegionalAuthority : System.IEquatable<Azure.Identity.RegionalAuthority>
     {
@@ -314,6 +327,11 @@ public SharedTokenCacheCredentialOptions(Azure.Identity.TokenCachePersistenceOpt
         public Azure.Identity.TokenCachePersistenceOptions TokenCachePersistenceOptions { get { throw null; } }
         public string Username { get { throw null; } set { } }
     }
+    public partial class TokenCacheNotificationDetails
+    {
+        internal TokenCacheNotificationDetails() { }
+        public string SuggestedCacheKey { get { throw null; } }
+    }
     public partial class TokenCachePersistenceOptions
     {
         public TokenCachePersistenceOptions() { }
@@ -335,8 +353,27 @@ public abstract partial class UnsafeTokenCacheOptions : Azure.Identity.TokenCach
     {
         protected UnsafeTokenCacheOptions() { }
         protected internal abstract System.Threading.Tasks.Task<System.ReadOnlyMemory<byte>> RefreshCacheAsync();
+        protected internal virtual System.Threading.Tasks.Task<Azure.Identity.UserAssertionCacheDetails> RefreshCacheAsync(Azure.Identity.TokenCacheNotificationDetails details) { throw null; }
         protected internal abstract System.Threading.Tasks.Task TokenCacheUpdatedAsync(Azure.Identity.TokenCacheUpdatedArgs tokenCacheUpdatedArgs);
     }
+    [System.Runtime.InteropServices.StructLayoutAttribute(System.Runtime.InteropServices.LayoutKind.Sequential)]
+    public partial struct UserAssertionCacheDetails
+    {
+        private object _dummy;
+        private int _dummyPrimitive;
+        public System.ReadOnlyMemory<byte> CacheBytes { get { throw null; } set { } }
+    }
+    public partial class UserAssertionScope : System.IDisposable
+    {
+        public UserAssertionScope(string accessToken, Azure.Identity.UserAssertionScopeOptions options = null) { }
+        public void Dispose() { }
+    }
+    public partial class UserAssertionScopeOptions
+    {
+        public UserAssertionScopeOptions() { }
+        public System.Func<System.Threading.Tasks.Task<Azure.Identity.UserAssertionCacheDetails>> HydrateCache { get { throw null; } set { } }
+        public System.Func<Azure.Identity.UserAssertionCacheDetails, System.Threading.Tasks.Task> PersistCache { get { throw null; } set { } }
+    }
     public partial class UsernamePasswordCredential : Azure.Core.TokenCredential
     {
         protected UsernamePasswordCredential() { }

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

@@ -12,7 +12,10 @@
     <EnableBannedApiAnalyzers>false</EnableBannedApiAnalyzers>
   </PropertyGroup>
   <ItemGroup>
-    <PackageReference Include="Azure.Core" />
+    <!-- TODO: Revert after TokenRequestContext changes ship in Azure.Core -->
+    <!-- <PackageReference Include="Azure.Core" /> -->
+    <ProjectReference Include="$(MSBuildThisFileDirectory)..\..\..\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/MsalConfidentialClient.cs

@@ -125,5 +125,25 @@ public virtual async ValueTask<AuthenticationResult> AcquireTokenByAuthorization
                 .ExecuteAsync(async, cancellationToken)
                 .ConfigureAwait(false);
         }
+
+        public virtual async ValueTask<AuthenticationResult> AcquireTokenOnBehalfOf(
+            string[] scopes,
+            string tenantId,
+            UserAssertion userAssertionValue,
+            bool async,
+            CancellationToken cancellationToken)
+        {
+            IConfidentialClientApplication client = await GetClientAsync(async, cancellationToken).ConfigureAwait(false);
+
+            var builder = client.AcquireTokenOnBehalfOf(scopes, userAssertionValue);
+
+            if (!string.IsNullOrEmpty(tenantId))
+            {
+                builder.WithAuthority(Pipeline.AuthorityHost.AbsoluteUri, tenantId);
+            }
+            return await builder
+                .ExecuteAsync(async, cancellationToken)
+                .ConfigureAwait(false);
+        }
     }
 }

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

@@ -0,0 +1,100 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core;
+using Azure.Core.Pipeline;
+using Microsoft.Identity.Client;
+
+namespace Azure.Identity
+{
+    /// <summary>
+    /// Enables authentication to Azure Active Directory using an On-Behalf-Of flow.``
+    /// </summary>
+    public class OnBehalfOfCredential : TokenCredential
+    {
+        private readonly MsalConfidentialClient _client;
+        private readonly string _tenantId;
+        private readonly CredentialPipeline _pipeline;
+        private readonly bool _allowMultiTenantAuthentication;
+        private readonly string _clientId;
+        private readonly string _clientSecret;
+
+        /// <inheritdoc />
+        public override bool SupportsCaching => true;
+
+        /// <summary>
+        /// Protected constructor for mocking.
+        /// </summary>
+        protected OnBehalfOfCredential()
+        { }
+
+        /// <summary>
+        /// Creates an instance of the <see cref="OnBehalfOfCredential"/> with the details needed to authenticate with Azure Active Directory.
+        /// Calls to <see cref="GetToken"/> and <see cref="GetTokenAsync"/> should be wrapped with a using block that constructs an instance of
+        /// <see cref="UserAssertionScope"/> with the user's <see cref="AccessToken"/> to be used in the On-Behalf-Of flow.
+        /// </summary>
+        /// <param name="tenantId">The Azure Active Directory tenant (directory) Id of the service principal.</param>
+        /// <param name="clientId">The client (application) ID of the service principal</param>
+        /// <param name="clientSecret">A client secret that was generated for the App Registration used to authenticate the client.</param>
+        /// <param name="options">Options that allow to configure the management of the requests sent to the Azure Active Directory service.</param>
+        public OnBehalfOfCredential(
+            string tenantId,
+            string clientId,
+            string clientSecret,
+            OnBehalfOfCredentialOptions options = null)
+            : this(tenantId, clientId, clientSecret, options, null, null)
+        { }
+
+        internal OnBehalfOfCredential(
+            string tenantId,
+            string clientId,
+            string clientSecret,
+            OnBehalfOfCredentialOptions options,
+            CredentialPipeline pipeline,
+            MsalConfidentialClient client)
+        {
+            Argument.AssertNotNull(clientId, nameof(clientId));
+            Argument.AssertNotNull(clientSecret, nameof(clientSecret));
+
+            options ??= new OnBehalfOfCredentialOptions();
+            _pipeline = pipeline ?? CredentialPipeline.GetInstance(options);
+            _allowMultiTenantAuthentication = options.AllowMultiTenantAuthentication;
+            _tenantId = Validations.ValidateTenantId(tenantId, nameof(tenantId));
+            _clientId = clientId;
+            _clientSecret = clientSecret;
+            _client = client;
+        }
+
+        /// <inheritdoc />
+        public override AccessToken GetToken(TokenRequestContext requestContext, CancellationToken cancellationToken) =>
+            GetTokenInternalAsync(requestContext, false, cancellationToken).EnsureCompleted();
+
+        /// <inheritdoc />
+        public override ValueTask<AccessToken> GetTokenAsync(TokenRequestContext requestContext, CancellationToken cancellationToken) =>
+            GetTokenInternalAsync(requestContext, true, cancellationToken);
+
+        internal async ValueTask<AccessToken> GetTokenInternalAsync(TokenRequestContext requestContext, bool async, CancellationToken cancellationToken)
+        {
+            using CredentialDiagnosticScope scope = _pipeline.StartGetTokenScope("OnBehalfOfCredential.GetToken", requestContext);
+
+            try
+            {
+                var tenantId = TenantIdResolver.Resolve(_tenantId, requestContext, _allowMultiTenantAuthentication);
+                UserAssertionScope.Current.Client = _client ?? new MsalConfidentialClient(_pipeline, tenantId, _clientId, _clientSecret, UserAssertionScope.Current.CacheOptions, default);
+
+                AuthenticationResult result = await UserAssertionScope.Current.Client
+                    .AcquireTokenOnBehalfOf(requestContext.Scopes, tenantId, UserAssertionScope.Current.UserAssertion, async, cancellationToken)
+                    .ConfigureAwait(false);
+
+                return new AccessToken(result.AccessToken, result.ExpiresOn);
+            }
+            catch (Exception e)
+            {
+                throw scope.FailWrapAndThrow(e);
+            }
+        }
+    }
+}

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

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+namespace Azure.Identity
+{
+    /// <summary>
+    ///
+    /// </summary>
+    public class OnBehalfOfCredentialOptions : TokenCredentialOptions, ITokenCacheOptions
+    {
+        /// <summary>
+        /// The <see cref="TokenCachePersistenceOptions"/>.
+        /// </summary>
+        public TokenCachePersistenceOptions TokenCachePersistenceOptions { get; set; }
+    }
+}

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

@@ -6,7 +6,6 @@
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
-using Azure.Core.Pipeline;
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.Extensions.Msal;
 
@@ -88,7 +87,7 @@ internal TokenCache(TokenCachePersistenceOptions options, MsalCacheHelperWrapper
         /// <summary>
         /// A delegate that will be called before the cache is accessed. The data returned will be used to set the current state of the cache.
         /// </summary>
-        internal Func<Task<ReadOnlyMemory<byte>>> RefreshCacheFromOptionsAsync;
+        internal Func<TokenCacheNotificationDetails, Task<UserAssertionCacheDetails>> RefreshCacheFromOptionsAsync;
 
         internal virtual async Task RegisterCache(bool async, ITokenCache tokenCache, CancellationToken cancellationToken)
         {
@@ -143,7 +142,8 @@ private async Task OnBeforeCacheAccessAsync(TokenCacheNotificationArgs args)
             {
                 if (RefreshCacheFromOptionsAsync != null)
                 {
-                    Data = (await RefreshCacheFromOptionsAsync().ConfigureAwait(false)).ToArray();
+                    Data = (await RefreshCacheFromOptionsAsync(new TokenCacheNotificationDetails(args)).ConfigureAwait(false))
+                        .CacheBytes.ToArray();
                 }
                 args.TokenCache.DeserializeMsalV3(Data, true);