From 001a0772088174199fe96123b87dd3b908a245e4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andreas=20N=C3=A4geli?= <andr.naegeli@gmail.com>
Date: Tue, 30 Dec 2025 15:34:06 +0100
Subject: [PATCH] Allow signing keys to be fetched via a custom logic

---
 .../Bearer/BearerAuthenticationConcern.cs         | 15 +++++++++++++--
 .../Bearer/BearerAuthenticationConcernBuilder.cs  | 13 +++++++++++++
 .../Bearer/TokenValidationOptions.cs              |  6 ++++++
 .../Authentication/BearerAuthenticationTests.cs   | 14 ++++++++++++++
 4 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/Modules/Authentication/Bearer/BearerAuthenticationConcern.cs b/Modules/Authentication/Bearer/BearerAuthenticationConcern.cs
index 94922ee6f..2e761bdd6 100644
--- a/Modules/Authentication/Bearer/BearerAuthenticationConcern.cs
+++ b/Modules/Authentication/Bearer/BearerAuthenticationConcern.cs
@@ -1,8 +1,10 @@
 using System.IdentityModel.Tokens.Jwt;
 using System.Text.Json;
 using System.Text.Json.Serialization;
+
 using GenHTTP.Api.Content;
 using GenHTTP.Api.Protocol;
+
 using Microsoft.IdentityModel.Logging;
 using Microsoft.IdentityModel.Tokens;
 
@@ -75,7 +77,16 @@ internal BearerAuthenticationConcern(IHandler content, TokenValidationOptions va
 
         if (issuer != null && _issuerKeys == null)
         {
-            _issuerKeys = await FetchSigningKeys(issuer);
+            if (ValidationOptions.CustomKeyResolver != null)
+            {
+                var unvalidatedToken = tokenHandler.ReadJwtToken(tokenString);
+                
+                _issuerKeys = await ValidationOptions.CustomKeyResolver(unvalidatedToken);
+            }
+            else
+            {
+                _issuerKeys = await FetchSigningKeysAsync(issuer);
+            }
         }
 
         var validationParameters = new TokenValidationParameters
@@ -122,7 +133,7 @@ internal BearerAuthenticationConcern(IHandler content, TokenValidationOptions va
         }
     }
 
-    private static async Task<ICollection<SecurityKey>> FetchSigningKeys(string issuer)
+    private static async ValueTask<ICollection<SecurityKey>> FetchSigningKeysAsync(string issuer)
     {
         try
         {
diff --git a/Modules/Authentication/Bearer/BearerAuthenticationConcernBuilder.cs b/Modules/Authentication/Bearer/BearerAuthenticationConcernBuilder.cs
index 30b8e29dd..9b4c46b12 100644
--- a/Modules/Authentication/Bearer/BearerAuthenticationConcernBuilder.cs
+++ b/Modules/Authentication/Bearer/BearerAuthenticationConcernBuilder.cs
@@ -2,6 +2,7 @@
 using GenHTTP.Api.Content;
 using GenHTTP.Api.Content.Authentication;
 using GenHTTP.Api.Protocol;
+using Microsoft.IdentityModel.Tokens;
 
 namespace GenHTTP.Modules.Authentication.Bearer;
 
@@ -53,6 +54,18 @@ public BearerAuthenticationConcernBuilder Validation(Func<JwtSecurityToken, Task
         _options.CustomValidator = validator;
         return this;
     }
+     
+    /// <summary>
+    /// Adds a custom logic to fetch the signing keys. By default, the keys will be
+    /// downloaded from the URL returned by the .well-known information returned by
+    /// the issuer.
+    /// </summary>
+    /// <param name="keyResolver">The logic used to resolve the signing keys for a given incoming token</param>
+    public BearerAuthenticationConcernBuilder KeyResolver(Func<JwtSecurityToken, ValueTask<ICollection<SecurityKey>>> keyResolver)
+    {
+        _options.CustomKeyResolver = keyResolver;
+        return this;
+    }
 
     /// <summary>
     /// Optionally register a function that will compute the user that
diff --git a/Modules/Authentication/Bearer/TokenValidationOptions.cs b/Modules/Authentication/Bearer/TokenValidationOptions.cs
index 5a2db0e2e..8b89eb347 100644
--- a/Modules/Authentication/Bearer/TokenValidationOptions.cs
+++ b/Modules/Authentication/Bearer/TokenValidationOptions.cs
@@ -1,7 +1,10 @@
 using System.IdentityModel.Tokens.Jwt;
+
 using GenHTTP.Api.Content.Authentication;
 using GenHTTP.Api.Protocol;
 
+using Microsoft.IdentityModel.Tokens;
+
 namespace GenHTTP.Modules.Authentication.Bearer;
 
 internal sealed class TokenValidationOptions
@@ -16,4 +19,7 @@ internal sealed class TokenValidationOptions
     internal Func<JwtSecurityToken, Task>? CustomValidator { get; set; }
 
     internal Func<IRequest, JwtSecurityToken, ValueTask<IUser?>>? UserMapping { get; set; }
+    
+    internal Func<JwtSecurityToken, ValueTask<ICollection<SecurityKey>>>? CustomKeyResolver { get; set; }
+    
 }
diff --git a/Testing/Acceptance/Modules/Authentication/BearerAuthenticationTests.cs b/Testing/Acceptance/Modules/Authentication/BearerAuthenticationTests.cs
index 4280357e4..d39794b05 100644
--- a/Testing/Acceptance/Modules/Authentication/BearerAuthenticationTests.cs
+++ b/Testing/Acceptance/Modules/Authentication/BearerAuthenticationTests.cs
@@ -40,6 +40,20 @@ public async Task TestCustomValidator(TestEngine engine)
 
         await response.AssertStatusAsync(HttpStatusCode.Forbidden);
     }
+    
+    [TestMethod]
+    [MultiEngineTest]
+    public async Task TestCustomKeyResolver(TestEngine engine)
+    {
+        var auth = BearerAuthentication.Create()
+                                       .Issuer("https://facebook.com")
+                                       .KeyResolver(_ => throw new ProviderException(ResponseStatus.Forbidden, "Nah"))
+                                       .AllowExpired();
+
+        using var response = await Execute(auth, engine, ValidToken);
+
+        await response.AssertStatusAsync(HttpStatusCode.Forbidden);
+    }
 
     [TestMethod]
     [MultiEngineTest]