From 5c427a87f76231e71731a7f4dcc24c4c8407e234 Mon Sep 17 00:00:00 2001
From: Gjermund Stensrud <gjermund.stensrud@fhi.no>
Date: Fri, 6 Dec 2024 10:30:52 +0100
Subject: [PATCH] Add support for additional claims in DPoP proof payload

---
 .../AccessTokenHandler.cs                     |  3 +++
 .../AccessTokenManagement/DPoPExtensions.cs   | 10 +++++++
 .../DefaultDPoPProofService.cs                |  8 ++++++
 .../Interfaces/IDPoPProofService.cs           |  5 ++++
 .../ClientTokenManagementApiTests.cs          | 26 +++++++++++++++++++
 5 files changed, 52 insertions(+)

diff --git a/access-token-management/src/AccessTokenManagement/AccessTokenHandler.cs b/access-token-management/src/AccessTokenManagement/AccessTokenHandler.cs
index 9ef4239a..b19d817c 100644
--- a/access-token-management/src/AccessTokenManagement/AccessTokenHandler.cs
+++ b/access-token-management/src/AccessTokenManagement/AccessTokenHandler.cs
@@ -127,6 +127,8 @@ protected virtual async Task<bool> SetDPoPProofTokenAsync(HttpRequestMessage req
 
         if (!string.IsNullOrEmpty(token.DPoPJsonWebKey))
         {
+            request.Options.TryGetValue(new HttpRequestOptionsKey<IReadOnlyDictionary<string, string>>("Duende.AccessTokenManagement.DPoPProofAdditionalPayloadClaims"), out var additionalClaims);
+
             // create proof
             var proofToken = await _dPoPProofService.CreateProofTokenAsync(new DPoPProofRequest
             {
@@ -135,6 +137,7 @@ protected virtual async Task<bool> SetDPoPProofTokenAsync(HttpRequestMessage req
                 Method = request.Method.ToString(),
                 DPoPJsonWebKey = token.DPoPJsonWebKey,
                 DPoPNonce = dpopNonce,
+                AdditionalPayloadClaims = additionalClaims,
             });
 
             if (proofToken != null)
diff --git a/access-token-management/src/AccessTokenManagement/DPoPExtensions.cs b/access-token-management/src/AccessTokenManagement/DPoPExtensions.cs
index 40f1cde9..3c659e76 100644
--- a/access-token-management/src/AccessTokenManagement/DPoPExtensions.cs
+++ b/access-token-management/src/AccessTokenManagement/DPoPExtensions.cs
@@ -77,4 +77,14 @@ public static string GetDPoPUrl(this HttpRequestMessage request)
     {
         return request.RequestUri!.Scheme + "://" + request.RequestUri!.Authority + request.RequestUri!.LocalPath;
     }
+
+    /// <summary>
+    /// Additional claims that will be added to the DPoP proof payload on generation
+    /// </summary>
+    /// <param name="request"></param>
+    /// <param name="customClaims"></param>
+    public static void AddDPoPProofAdditionalPayloadClaims(this HttpRequestMessage request, IDictionary<string, string> customClaims)
+    {
+        request.Options.TryAdd("Duende.AccessTokenManagement.DPoPProofAdditionalPayloadClaims", customClaims.AsReadOnly());
+    }
 }
\ No newline at end of file
diff --git a/access-token-management/src/AccessTokenManagement/DefaultDPoPProofService.cs b/access-token-management/src/AccessTokenManagement/DefaultDPoPProofService.cs
index 0c58012c..5e684849 100644
--- a/access-token-management/src/AccessTokenManagement/DefaultDPoPProofService.cs
+++ b/access-token-management/src/AccessTokenManagement/DefaultDPoPProofService.cs
@@ -119,6 +119,14 @@ await _dPoPNonceStore.StoreNonceAsync(new DPoPNonceContext
             payload.Add(JwtClaimTypes.Nonce, nonce);
         }
 
+        if (request.AdditionalPayloadClaims?.Count > 0)
+        {
+            foreach (var claim in request.AdditionalPayloadClaims)
+            {
+                payload.Add(claim.Key, claim.Value);
+            }
+        }
+
         var handler = new JsonWebTokenHandler() { SetDefaultTimesOnTokenCreation = false };
         var key = new SigningCredentials(jsonWebKey, jsonWebKey.Alg);
         var proofToken = handler.CreateToken(JsonSerializer.Serialize(payload), key, header);
diff --git a/access-token-management/src/AccessTokenManagement/Interfaces/IDPoPProofService.cs b/access-token-management/src/AccessTokenManagement/Interfaces/IDPoPProofService.cs
index 1ba3c4e5..e3873189 100644
--- a/access-token-management/src/AccessTokenManagement/Interfaces/IDPoPProofService.cs
+++ b/access-token-management/src/AccessTokenManagement/Interfaces/IDPoPProofService.cs
@@ -48,6 +48,11 @@ public class DPoPProofRequest
     /// The access token
     /// </summary>
     public string? AccessToken { get; set; }
+
+    /// <summary>
+    /// Additional claims to add to the DPoP proof payload
+    /// </summary>
+    public IReadOnlyDictionary<string, string>? AdditionalPayloadClaims { get; set; }
 }
 
 /// <summary>
diff --git a/access-token-management/test/AccessTokenManagement.Tests/ClientTokenManagementApiTests.cs b/access-token-management/test/AccessTokenManagement.Tests/ClientTokenManagementApiTests.cs
index be19894e..948e268c 100644
--- a/access-token-management/test/AccessTokenManagement.Tests/ClientTokenManagementApiTests.cs
+++ b/access-token-management/test/AccessTokenManagement.Tests/ClientTokenManagementApiTests.cs
@@ -189,6 +189,32 @@ public async Task dpop_tokens_should_be_passed_to_api()
         proofToken.ShouldNotBeNull();
     }
 
+    [Fact]
+    public async Task when_additional_proof_payload_claims_are_defined_they_should_be_included_in_dpop_proof()
+    {
+        string? proofToken = null;
+
+        ApiHost.ApiInvoked += ctx =>
+        {
+            proofToken = ctx.Request.Headers["DPoP"].FirstOrDefault()?.ToString();
+        };
+        var client = _clientFactory.CreateClient("test");
+
+        var requestMessage = new HttpRequestMessage(HttpMethod.Get, ApiHost.Url("/test"));
+        requestMessage.AddDPoPProofAdditionalPayloadClaims(new Dictionary<string, string>() {
+            { "claim_one", "one" },
+            { "claim_two", "two" },
+        });
+
+        var apiResult = await client.SendAsync(requestMessage);
+
+        proofToken.ShouldNotBeNull();
+        var payload = Base64UrlEncoder.Decode(proofToken!.Split('.')[1]);
+        var values = JsonSerializer.Deserialize<Dictionary<string, object>>(payload);
+        values!["claim_one"].ToString().ShouldBe("one");
+        values!["claim_two"].ToString().ShouldBe("two");
+    }
+
     [Fact]
     public async Task when_api_issues_nonce_api_request_should_be_retried_with_new_nonce()
     {