Add ability to set auth scopes via request options

samples/AzureRest/AzureRest.Client/AzureRestClientOptions.cs

@@ -13,9 +13,9 @@ Uri ICabazureClientOptions.GetBaseAddress()
             // for the client to work correctly with relative paths.
             => new("https://management.azure.com/");
 
-        string ICabazureAuthClientOptions.GetScope()
+        string[] ICabazureAuthClientOptions.GetScopes()
             // The default scope for Microsoft Graph.
-            => "https://management.azure.com/.default";
+            => new string[] { "https://management.azure.com/.default" };
 
         TokenCredential ICabazureAuthClientOptions.GetCredential()
             => Credential

src/Cabazure.Client.Runtime/Authentication/AzureAuthenticationHandler.cs

@@ -3,19 +3,24 @@
     public class AzureAuthenticationHandler
         : DelegatingHandler
     {
+        private readonly string[] defaultScopes;
         private readonly IBearerTokenProvider tokenProvider;
 
         public AzureAuthenticationHandler(
+            string[] defaultScopes,
             IBearerTokenProvider tokenProvider)
         {
-            this.tokenProvider = tokenProvider;     
+            this.defaultScopes = defaultScopes;
+            this.tokenProvider = tokenProvider;
         }
 
         protected async override Task<HttpResponseMessage> SendAsync(
             HttpRequestMessage request,
             CancellationToken cancellationToken)
         {
-            request.Headers.Authorization = await tokenProvider.GetTokenAsync(cancellationToken);
+            var scopes = request.GetScopes() ?? defaultScopes;
+
+            request.Headers.Authorization = await tokenProvider.GetTokenAsync(scopes, cancellationToken);
 
             return await base.SendAsync(request, cancellationToken);
         }

src/Cabazure.Client.Runtime/Authentication/BearerTokenProvider.cs

@@ -1,39 +1,39 @@
-using System.Net.Http.Headers;
+using System.Collections.Concurrent;
+using System.Net.Http.Headers;
 using Azure.Core;
 
 namespace Cabazure.Client.Authentication
 {
     public class BearerTokenProvider : IBearerTokenProvider
     {
-        private readonly TokenRequestContext context;
         private readonly TokenCredential credential;
         private readonly IDateTimeProvider dateTimeProvider;
-        private AccessToken accessToken;
+        private readonly ConcurrentDictionary<string, AccessToken> accessTokenCache = new();
 
         public BearerTokenProvider(
-            TokenRequestContext context,
             TokenCredential credential,
             IDateTimeProvider dateTimeProvider)
         {
-            this.context = context;
             this.credential = credential;
             this.dateTimeProvider = dateTimeProvider;
         }
 
         public async Task<AuthenticationHeaderValue> GetTokenAsync(
+            string[] scopes,
             CancellationToken cancellationToken)
         {
-            if (TokenIsExpired())
+            var key = string.Join(" ", scopes);
+            if (!accessTokenCache.TryGetValue(key, out var accessToken) || TokenIsExpired(accessToken))
             {
-                accessToken = await credential.GetTokenAsync(context, cancellationToken);
+                accessTokenCache[key] = accessToken = await credential.GetTokenAsync(new TokenRequestContext(scopes), cancellationToken);
             }
 
             return new AuthenticationHeaderValue(
                 "Bearer",
                 accessToken.Token);
         }
 
-        private bool TokenIsExpired()
+        private bool TokenIsExpired(AccessToken accessToken)
             => dateTimeProvider.GetDateTime().AddMinutes(1) > accessToken.ExpiresOn;
     }
 }

src/Cabazure.Client.Runtime/Authentication/IBearerTokenProvider.cs

@@ -5,6 +5,7 @@ namespace Cabazure.Client.Authentication
     public interface IBearerTokenProvider
     {
         Task<AuthenticationHeaderValue> GetTokenAsync(
+            string[] scopes,
             CancellationToken cancellationToken);
     }
 }

src/Cabazure.Client.Runtime/Builder/HttpClientBuilderExtensions.cs

@@ -16,22 +16,21 @@ public static IHttpClientBuilder AddAuthentication(
             TokenCredential credential)
             => AddAuthentication(
                 builder,
-                new TokenRequestContext(new[] { scope }),
+                new[] { scope },
                 credential);
 
         public static IHttpClientBuilder AddAuthentication(
             this IHttpClientBuilder builder,
-            TokenRequestContext context,
+            string[] scopes,
             TokenCredential credential)
         {
             var tokenProvider = new BearerTokenProvider(
-                context,
                 credential,
                 new DateTimeProvider());
 
             return builder
                 .AddHttpMessageHandler(
-                    () => new AzureAuthenticationHandler(tokenProvider));
+                    () => new AzureAuthenticationHandler(scopes, tokenProvider));
         }
     }
 }

src/Cabazure.Client.Runtime/HttpRequestMessageExtensions.cs

@@ -0,0 +1,21 @@
+namespace Cabazure.Client;
+
+public static class HttpRequestMessageExtensions
+{
+    private const string ScopesKey = "scopes";
+
+    public static void SetScopes(
+        this HttpRequestMessage request,
+        params string[] scopes)
+        => request.Properties[ScopesKey] = scopes;
+
+    public static string[]? GetScopes(this HttpRequestMessage request)
+    {
+        if (request.Properties.TryGetValue(ScopesKey, out var value) && value is string[] scopes)
+        {
+            return scopes;
+        }
+
+        return null;
+    }
+}

src/Cabazure.Client.Runtime/ICabazureAuthClientOptions.cs

@@ -8,10 +8,10 @@ namespace Cabazure.Client
     public interface ICabazureAuthClientOptions : ICabazureClientOptions
     {
         /// <summary>
-        /// The scope to use for authorization token.
+        /// The scopes to use for authorization token.
         /// </summary>
-        /// <returns>The authorization scope.</returns>
-        string GetScope();
+        /// <returns>The authorization scopes.</returns>
+        string[] GetScopes();
 
         /// <summary>
         /// The credential to use for authentication.

src/Cabazure.Client/ClientInitializationGenerator.cs

@@ -135,15 +135,14 @@ void ConfigureAuthHandler(IList<DelegatingHandler> handlers, IServiceProvider se
 
                             if (options is ICabazureAuthClientOptions authOptions)
                             {
-                                var scope = authOptions.GetScope();
+                                var scopes = authOptions.GetScopes();
                                 var credential = authOptions.GetCredential();
 
                                 var tokenProvider = new BearerTokenProvider(
-                                    new TokenRequestContext(new [] { scope }),
                                     credential,
                                     new DateTimeProvider());
 
-                                handlers.Add(new AzureAuthenticationHandler(tokenProvider));
+                                handlers.Add(new AzureAuthenticationHandler(scopes, tokenProvider));
                             }
                         }
 

test/Cabazure.Client.Runtime.Tests/Authentication/AzureAuthenticationHandlerTests.cs

@@ -17,7 +17,7 @@ internal async Task Should_Get_Token_On_Send(
         CancellationToken cancellationToken)
     {
         tokenProvider
-            .GetTokenAsync(default)
+            .GetTokenAsync(default, default)
             .ReturnsForAnyArgs(authenticationHeader);
         handler
             .InvokeProtectedMethod<Task<HttpResponseMessage>>("SendAsync", request, cancellationToken)
@@ -29,6 +29,6 @@ internal async Task Should_Get_Token_On_Send(
 
         _ = tokenProvider
             .Received(1)
-            .GetTokenAsync(Arg.Any<CancellationToken>());
+            .GetTokenAsync(Arg.Any<string[]>(), Arg.Any<CancellationToken>());
     }
 }

test/Cabazure.Client.Runtime.Tests/Authentication/BearerTokenProviderTests.cs

@@ -9,8 +9,8 @@ public class BearerTokenProviderTests
     internal async Task Should_Return_Token_From_TokenCredential(
         [Frozen] TokenCredential credential,
         [Frozen] IDateTimeProvider dateTimeProvider,
-        [Frozen] TokenRequestContext context,
         BearerTokenProvider sut,
+        string[] scopes,
         DateTimeOffset timestamp,
         AccessToken accessToken,
         CancellationToken cancellationToken)
@@ -23,7 +23,7 @@ internal async Task Should_Return_Token_From_TokenCredential(
             .GetTokenAsync(default, default)
             .ReturnsForAnyArgs(accessToken);
 
-        var response = await sut.GetTokenAsync(cancellationToken);
+        var response = await sut.GetTokenAsync(scopes, cancellationToken);
 
         response
             .Parameter
@@ -35,8 +35,8 @@ internal async Task Should_Return_Token_From_TokenCredential(
     internal async Task Should_Use_Cached_Token_From_TokenCredential(
         [Frozen] TokenCredential credential,
         [Frozen] IDateTimeProvider dateTimeProvider,
-        [Frozen] TokenRequestContext context,
         BearerTokenProvider sut,
+        string[] scopes,
         DateTimeOffset timestamp,
         CancellationToken cancellationToken)
     {
@@ -52,8 +52,8 @@ internal async Task Should_Use_Cached_Token_From_TokenCredential(
             .GetTokenAsync(default, default)
             .ReturnsForAnyArgs(accessToken);
 
-        var response1 = await sut.GetTokenAsync(cancellationToken);
-        var response2 = await sut.GetTokenAsync(cancellationToken);
+        var response1 = await sut.GetTokenAsync(scopes, cancellationToken);
+        var response2 = await sut.GetTokenAsync(scopes, cancellationToken);
 
         response1
             .Should()
@@ -70,8 +70,8 @@ internal async Task Should_Use_Cached_Token_From_TokenCredential(
     internal async Task Should_Renew_Expired_Token(
         [Frozen] TokenCredential credential,
         [Frozen] IDateTimeProvider dateTimeProvider,
-        [Frozen] TokenRequestContext context,
         BearerTokenProvider sut,
+        string[] scopes,
         DateTimeOffset timestamp,
         CancellationToken cancellationToken)
     {
@@ -87,9 +87,9 @@ internal async Task Should_Renew_Expired_Token(
             .GetTokenAsync(default, default)
             .ReturnsForAnyArgs(accessToken);
 
-        await sut.GetTokenAsync(cancellationToken);
-        await sut.GetTokenAsync(cancellationToken);
-        await sut.GetTokenAsync(cancellationToken);
+        await sut.GetTokenAsync(scopes, cancellationToken);
+        await sut.GetTokenAsync(scopes, cancellationToken);
+        await sut.GetTokenAsync(scopes, cancellationToken);
 
         _ = credential
             .Received(3)
@@ -102,8 +102,8 @@ internal async Task Should_Renew_Expired_Token(
     internal async Task Should_Use_Bearer_AuthorizationToken(
         [Frozen] TokenCredential credential,
         [Frozen] IDateTimeProvider dateTimeProvider,
-        [Frozen] TokenRequestContext context,
         BearerTokenProvider sut,
+        string[] scopes,
         DateTimeOffset timestamp,
         CancellationToken cancellationToken)
     {
@@ -119,7 +119,7 @@ internal async Task Should_Use_Bearer_AuthorizationToken(
             .GetTokenAsync(default, default)
             .ReturnsForAnyArgs(accessToken);
 
-        var response = await sut.GetTokenAsync(cancellationToken);
+        var response = await sut.GetTokenAsync(scopes, cancellationToken);
 
         response
             .Scheme

test/Cabazure.Client.Runtime.Tests/Builder/HttpClientBuilderExtensionsTests.cs

@@ -29,18 +29,18 @@ internal void SetBaseAddress_Should_Set_BaseAddress_On_HttpClient(
     }
 
     [Theory, AutoNSubstituteData]
-    internal void AddAuthentication_With_TokenContext_Should_Configure_AzureAuthenticationHandler(
+    internal void AddAuthentication_With_Scopes_Should_Configure_AzureAuthenticationHandler(
         ServiceCollection services,
         IHttpClientBuilder builder,
         string name,
-        TokenRequestContext context,
+        string[] context,
         TokenCredential credential)
     {
         builder.Name.Returns(name);
         builder.Services.Returns(services);
 
         builder.AddAuthentication(
-            context, 
+            context,
             credential);
 
         services.AddHttpClient(name);

test/Cabazure.Client.Tests/ClientInitializationGeneratorTests.CanGenerate_Initialization#ClientInitialization.Implementation.g.verified.cs

@@ -46,15 +46,14 @@ void ConfigureAuthHandler(IList<DelegatingHandler> handlers, IServiceProvider se
 
                 if (options is ICabazureAuthClientOptions authOptions)
                 {
-                    var scope = authOptions.GetScope();
+                    var scopes = authOptions.GetScopes();
                     var credential = authOptions.GetCredential();
 
                     var tokenProvider = new BearerTokenProvider(
-                        new TokenRequestContext(new [] { scope }),
                         credential,
                         new DateTimeProvider());
 
-                    handlers.Add(new AzureAuthenticationHandler(tokenProvider));
+                    handlers.Add(new AzureAuthenticationHandler(scopes, tokenProvider));
                 }
             }