Merge pull request #23 from Cabazure/request-scopes

rickykaare · web-flow · commit 96e62e8352a8 · 2025-06-10T16:07:53.000+02:00
Add ability to set auth scopes via request options
diff --git a/samples/AzureRest/AzureRest.Client/AzureRestClientOptions.cs b/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
diff --git a/src/Cabazure.Client.Runtime/Authentication/AzureAuthenticationHandler.cs b/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);
         }
diff --git a/src/Cabazure.Client.Runtime/Authentication/BearerTokenProvider.cs b/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;
     }
 }
diff --git a/src/Cabazure.Client.Runtime/Authentication/IBearerTokenProvider.cs b/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);
     }
 }
diff --git a/src/Cabazure.Client.Runtime/Builder/HttpClientBuilderExtensions.cs b/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));
         }
     }
 }
diff --git a/src/Cabazure.Client.Runtime/HttpRequestMessageExtensions.cs b/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;
+    }
+}
diff --git a/src/Cabazure.Client.Runtime/ICabazureAuthClientOptions.cs b/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.
diff --git a/src/Cabazure.Client/ClientInitializationGenerator.cs b/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));
                             }
                         }
 
diff --git a/test/Cabazure.Client.Runtime.Tests/Authentication/AzureAuthenticationHandlerTests.cs b/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>());
     }
 }
diff --git a/test/Cabazure.Client.Runtime.Tests/Authentication/BearerTokenProviderTests.cs b/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
diff --git a/test/Cabazure.Client.Runtime.Tests/Builder/HttpClientBuilderExtensionsTests.cs b/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);
diff --git a/test/Cabazure.Client.Tests/ClientInitializationGeneratorTests.CanGenerate_Initialization#ClientInitialization.Implementation.g.verified.cs b/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));
                 }
             }
 

Original file line number	Diff line number	Diff line change
`@@ -3,19 +3,24 @@`
`3`	`3`	`public class AzureAuthenticationHandler`
`4`	`4`	`: DelegatingHandler`
`5`	`5`	`{`
	`6`	`+ private readonly string[] defaultScopes;`
`6`	`7`	`private readonly IBearerTokenProvider tokenProvider;`
`7`	`8`
`8`	`9`	`public AzureAuthenticationHandler(`
	`10`	`+ string[] defaultScopes,`
`9`	`11`	`IBearerTokenProvider tokenProvider)`
`10`	`12`	`{`
`11`		`- this.tokenProvider = tokenProvider;`
	`13`	`+ this.defaultScopes = defaultScopes;`
	`14`	`+ this.tokenProvider = tokenProvider;`
`12`	`15`	`}`
`13`	`16`
`14`	`17`	`protected async override Task<HttpResponseMessage> SendAsync(`
`15`	`18`	`HttpRequestMessage request,`
`16`	`19`	`CancellationToken cancellationToken)`
`17`	`20`	`{`
`18`		`- request.Headers.Authorization = await tokenProvider.GetTokenAsync(cancellationToken);`
	`21`	`+ var scopes = request.GetScopes() ?? defaultScopes;`
	`22`	`+`
	`23`	`+ request.Headers.Authorization = await tokenProvider.GetTokenAsync(scopes, cancellationToken);`
`19`	`24`
`20`	`25`	`return await base.SendAsync(request, cancellationToken);`
`21`	`26`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,39 +1,39 @@`
`1`		`-using System.Net.Http.Headers;`
	`1`	`+using System.Collections.Concurrent;`
	`2`	`+using System.Net.Http.Headers;`
`2`	`3`	`using Azure.Core;`
`3`	`4`
`4`	`5`	`namespace Cabazure.Client.Authentication`
`5`	`6`	`{`
`6`	`7`	`public class BearerTokenProvider : IBearerTokenProvider`
`7`	`8`	`{`
`8`		`- private readonly TokenRequestContext context;`
`9`	`9`	`private readonly TokenCredential credential;`
`10`	`10`	`private readonly IDateTimeProvider dateTimeProvider;`
`11`		`- private AccessToken accessToken;`
	`11`	`+ private readonly ConcurrentDictionary<string, AccessToken> accessTokenCache = new();`
`12`	`12`
`13`	`13`	`public BearerTokenProvider(`
`14`		`- TokenRequestContext context,`
`15`	`14`	`TokenCredential credential,`
`16`	`15`	`IDateTimeProvider dateTimeProvider)`
`17`	`16`	`{`
`18`		`- this.context = context;`
`19`	`17`	`this.credential = credential;`
`20`	`18`	`this.dateTimeProvider = dateTimeProvider;`
`21`	`19`	`}`
`22`	`20`
`23`	`21`	`public async Task<AuthenticationHeaderValue> GetTokenAsync(`
	`22`	`+ string[] scopes,`
`24`	`23`	`CancellationToken cancellationToken)`
`25`	`24`	`{`
`26`		`- if (TokenIsExpired())`
	`25`	`+ var key = string.Join(" ", scopes);`
	`26`	`+ if (!accessTokenCache.TryGetValue(key, out var accessToken) \|\| TokenIsExpired(accessToken))`
`27`	`27`	`{`
`28`		`- accessToken = await credential.GetTokenAsync(context, cancellationToken);`
	`28`	`+ accessTokenCache[key] = accessToken = await credential.GetTokenAsync(new TokenRequestContext(scopes), cancellationToken);`
`29`	`29`	`}`
`30`	`30`
`31`	`31`	`return new AuthenticationHeaderValue(`
`32`	`32`	`"Bearer",`
`33`	`33`	`accessToken.Token);`
`34`	`34`	`}`
`35`	`35`
`36`		`- private bool TokenIsExpired()`
	`36`	`+ private bool TokenIsExpired(AccessToken accessToken)`
`37`	`37`	`=> dateTimeProvider.GetDateTime().AddMinutes(1) > accessToken.ExpiresOn;`
`38`	`38`	`}`
`39`	`39`	`}`
Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@ namespace Cabazure.Client.Authentication`
`5`	`5`	`public interface IBearerTokenProvider`
`6`	`6`	`{`
`7`	`7`	`Task<AuthenticationHeaderValue> GetTokenAsync(`
	`8`	`+ string[] scopes,`
`8`	`9`	`CancellationToken cancellationToken);`
`9`	`10`	`}`
`10`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,22 +16,21 @@ public static IHttpClientBuilder AddAuthentication(`
`16`	`16`	`TokenCredential credential)`
`17`	`17`	`=> AddAuthentication(`
`18`	`18`	`builder,`
`19`		`- new TokenRequestContext(new[] { scope }),`
	`19`	`+ new[] { scope },`
`20`	`20`	`credential);`
`21`	`21`
`22`	`22`	`public static IHttpClientBuilder AddAuthentication(`
`23`	`23`	`this IHttpClientBuilder builder,`
`24`		`- TokenRequestContext context,`
	`24`	`+ string[] scopes,`
`25`	`25`	`TokenCredential credential)`
`26`	`26`	`{`
`27`	`27`	`var tokenProvider = new BearerTokenProvider(`
`28`		`- context,`
`29`	`28`	`credential,`
`30`	`29`	`new DateTimeProvider());`
`31`	`30`
`32`	`31`	`return builder`
`33`	`32`	`.AddHttpMessageHandler(`
`34`		`- () => new AzureAuthenticationHandler(tokenProvider));`
	`33`	`+ () => new AzureAuthenticationHandler(scopes, tokenProvider));`
`35`	`34`	`}`
`36`	`35`	`}`
`37`	`36`	`}`
Original file line number	Diff line number	Diff line change
`@@ -135,15 +135,14 @@ void ConfigureAuthHandler(IList<DelegatingHandler> handlers, IServiceProvider se`
`135`	`135`
`136`	`136`	`if (options is ICabazureAuthClientOptions authOptions)`
`137`	`137`	`{`
`138`		`- var scope = authOptions.GetScope();`
	`138`	`+ var scopes = authOptions.GetScopes();`
`139`	`139`	`var credential = authOptions.GetCredential();`
`140`	`140`
`141`	`141`	`var tokenProvider = new BearerTokenProvider(`
`142`		`- new TokenRequestContext(new [] { scope }),`
`143`	`142`	`credential,`
`144`	`143`	`new DateTimeProvider());`
`145`	`144`
`146`		`- handlers.Add(new AzureAuthenticationHandler(tokenProvider));`
	`145`	`+ handlers.Add(new AzureAuthenticationHandler(scopes, tokenProvider));`
`147`	`146`	`}`
`148`	`147`	`}`
`149`	`148`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@ internal async Task Should_Get_Token_On_Send(`
`17`	`17`	`CancellationToken cancellationToken)`
`18`	`18`	`{`
`19`	`19`	`tokenProvider`
`20`		`- .GetTokenAsync(default)`
	`20`	`+ .GetTokenAsync(default, default)`
`21`	`21`	`.ReturnsForAnyArgs(authenticationHeader);`
`22`	`22`	`handler`
`23`	`23`	`.InvokeProtectedMethod<Task<HttpResponseMessage>>("SendAsync", request, cancellationToken)`
`@@ -29,6 +29,6 @@ internal async Task Should_Get_Token_On_Send(`
`29`	`29`
`30`	`30`	`_ = tokenProvider`
`31`	`31`	`.Received(1)`
`32`		`- .GetTokenAsync(Arg.Any<CancellationToken>());`
	`32`	`+ .GetTokenAsync(Arg.Any<string[]>(), Arg.Any<CancellationToken>());`
`33`	`33`	`}`
`34`	`34`	`}`