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

@@ -19,7 +19,7 @@ public class BearerTokenAuthenticationPolicyTests : SyncAsyncPolicyTestBase
         public BearerTokenAuthenticationPolicyTests(bool isAsync) : base(isAsync) { }
 
         [Test]
-        public async Task BearerTokenAuthenticationPolicy_UsesTokenProvidedByCredentials()
+        public async Task BearerTokenAuthenticationPolicy_UsesTokenProvidedByCredentials ()
         {
             var credential = new TokenCredentialStub(
                 (r, c) => r.Scopes.SequenceEqual(new[] { "scope1", "scope2" }) ? new AccessToken("token", DateTimeOffset.MaxValue) : default,
@@ -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()
         {
@@ -804,6 +830,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

@@ -7,6 +7,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

@@ -223,6 +223,14 @@ 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.TokenCredentialOptions 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 SharedTokenCacheCredential : Azure.Core.TokenCredential
     {
         public SharedTokenCacheCredential() { }
@@ -266,6 +274,12 @@ protected UnsafeTokenCacheOptions() { }
         protected internal abstract System.Threading.Tasks.Task<System.ReadOnlyMemory<byte>> RefreshCacheAsync();
         protected internal abstract System.Threading.Tasks.Task TokenCacheUpdatedAsync(Azure.Identity.TokenCacheUpdatedArgs tokenCacheUpdatedArgs);
     }
+    public partial class UserAssertionScope : System.IDisposable
+    {
+        public UserAssertionScope(Azure.Core.AccessToken accessToken) { }
+        public UserAssertionScope(string accessToken) { }
+        public void Dispose() { }
+    }
     public partial class UsernamePasswordCredential : Azure.Core.TokenCredential
     {
         protected UsernamePasswordCredential() { }

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

@@ -121,5 +121,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,95 @@
+// 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>
+    ///
+    /// </summary>
+    public class OnBehalfOfCredential : TokenCredential
+    {
+        private readonly MsalConfidentialClient _client;
+        private readonly string _tenantId;
+        private readonly CredentialPipeline _pipeline;
+        private readonly bool _allowMultiTenantAuthentication;
+
+        /// <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,
+            TokenCredentialOptions options = null)
+            : this(tenantId, clientId, clientSecret, options, null, null)
+        { }
+
+        internal OnBehalfOfCredential(
+            string tenantId,
+            string clientId,
+            string clientSecret,
+            TokenCredentialOptions options,
+            CredentialPipeline pipeline,
+            MsalConfidentialClient client)
+        {
+            Argument.AssertNotNull(clientId, nameof(clientId));
+            Argument.AssertNotNull(clientSecret, nameof(clientSecret));
+
+            options ??= new TokenCredentialOptions();
+            _allowMultiTenantAuthentication = options.AllowMultiTenantAuthentication;
+            _tenantId = Validations.ValidateTenantId(tenantId, nameof(tenantId));
+            _pipeline = pipeline ?? CredentialPipeline.GetInstance(options);
+            _client = client ?? new MsalConfidentialClient(_pipeline, tenantId, clientId, clientSecret, options as ITokenCacheOptions);
+        }
+
+        /// <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);
+
+                AuthenticationResult result = await _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/UserAssertionScope.cs

@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading;
+using Azure.Core;
+using Microsoft.Identity.Client;
+
+namespace Azure.Identity
+{
+    /// <summary>
+    ///
+    /// </summary>
+    public class UserAssertionScope : IDisposable
+    {
+        private static AsyncLocal<UserAssertionScope> _currentAsyncLocal = new();
+        internal UserAssertion UserAssertion { get; }
+
+        /// <summary>
+        /// The current value of the <see cref="AsyncLocal{UserAssertion}"/>.
+        /// </summary>
+        internal static UserAssertionScope Current => _currentAsyncLocal.Value;
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="accessToken"></param>
+        public UserAssertionScope(string accessToken)
+        {
+            UserAssertion = new UserAssertion(accessToken);
+            _currentAsyncLocal.Value = this;
+        }
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="accessToken"></param>
+        public UserAssertionScope(AccessToken accessToken)
+        {
+            UserAssertion = new UserAssertion(accessToken.Token);
+            _currentAsyncLocal.Value = this;
+        }
+
+        /// <summary>
+        ///
+        /// </summary>
+        public void Dispose()
+        {
+            GC.SuppressFinalize(this);
+            _currentAsyncLocal.Value = default;
+        }
+    }
+}

sdk/identity/Azure.Identity/tests/AuthorizationCodeCredentialTests.cs

@@ -26,7 +26,6 @@ public class AuthorizationCodeCredentialTests : ClientTestBase
         private string expectedTenantId;
         private string expectedReplyUri;
         private string clientSecret = Guid.NewGuid().ToString();
-        private Func<string[], string, string, AuthenticationAccount, ValueTask<AuthenticationResult>> silentFactory;
 
         public AuthorizationCodeCredentialTests(bool isAsync) : base(isAsync)
         { }
@@ -52,19 +51,21 @@ public void TestSetup()
                 Guid.NewGuid(),
                 null,
                 "Bearer");
-            silentFactory = (_, _tenantId, _replyUri, _) =>
-            {
-                Assert.AreEqual(expectedTenantId, _tenantId);
-                Assert.AreEqual(expectedReplyUri, _replyUri);
-                return new ValueTask<AuthenticationResult>(result);
-            };
-            mockMsalClient = new MockMsalConfidentialClient(silentFactory);
-            mockMsalClient.AuthcodeFactory = (_, _tenantId, _replyUri, _) =>
-            {
-                Assert.AreEqual(expectedTenantId, _tenantId);
-                Assert.AreEqual(expectedReplyUri, _replyUri);
-                return result;
-            };
+            mockMsalClient = new MockMsalConfidentialClient()
+                .WithSilentFactory(
+                    (_, _tenantId, _replyUri, _) =>
+                    {
+                        Assert.AreEqual(expectedTenantId, _tenantId);
+                        Assert.AreEqual(expectedReplyUri, _replyUri);
+                        return new ValueTask<AuthenticationResult>(result);
+                    })
+                .WithAuthCodeFactory(
+                    (_, _tenantId, _replyUri, _) =>
+                    {
+                        Assert.AreEqual(expectedTenantId, _tenantId);
+                        Assert.AreEqual(expectedReplyUri, _replyUri);
+                        return result;
+                    });
         }
 
         [Test]
@@ -97,7 +98,8 @@ public async Task AuthenticateWithAuthCodeHonorsTenantId([Values(null, TenantIdH
             var context = new TokenRequestContext(new[] { Scope }, tenantId: tenantId);
             expectedTenantId = TenantIdResolver.Resolve(TenantId, context, options.AllowMultiTenantAuthentication);
 
-            AuthorizationCodeCredential cred = InstrumentClient(new AuthorizationCodeCredential(TenantId, ClientId, clientSecret, authCode, options, mockMsalClient));
+            AuthorizationCodeCredential cred = InstrumentClient(
+                new AuthorizationCodeCredential(TenantId, ClientId, clientSecret, authCode, options, mockMsalClient));
 
             AccessToken token = await cred.GetTokenAsync(context);