From 8fe46a8d1862b2b37907adb7f38c9e332a0a29d6 Mon Sep 17 00:00:00 2001
From: Travis Walker <travis.walker@microsoft.com>
Date: Fri, 23 Aug 2024 00:20:47 -0700
Subject: [PATCH] Changing the handling of client claims to use JSON (#4886)

* Changing the handling of client claims to use JSON

* Updating tests to account for JSON formatting

---------

Co-authored-by: trwalke <trwalke@microsoft.com>
---
 .../Internal/JsonWebToken.cs                  | 23 +++++++++++--------
 .../TestConstants.cs                          |  2 +-
 .../ClientCredentialWithCertTest.cs           |  4 +++-
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/src/client/Microsoft.Identity.Client/Internal/JsonWebToken.cs b/src/client/Microsoft.Identity.Client/Internal/JsonWebToken.cs
index 12c8bce83b..35cf767c4e 100644
--- a/src/client/Microsoft.Identity.Client/Internal/JsonWebToken.cs
+++ b/src/client/Microsoft.Identity.Client/Internal/JsonWebToken.cs
@@ -8,6 +8,11 @@
 using Microsoft.Identity.Client.PlatformsCommon.Interfaces;
 using Microsoft.Identity.Client.Utils;
 using System.Security.Cryptography;
+#if SUPPORTS_SYSTEM_TEXT_JSON
+using JObject = System.Text.Json.Nodes.JsonObject;
+#else
+using Microsoft.Identity.Json.Linq;
+#endif
 
 namespace Microsoft.Identity.Client.Internal
 {
@@ -65,18 +70,18 @@ private string CreateJsonPayload()
                 payload.Append('{');
             }
 
-            int i = 0;
-            foreach (var kvp in _claimsToSign)
+            var json = new JObject();
+
+            foreach (var claim in _claimsToSign)
             {
-                payload.Append($"\"{kvp.Key}\":\"{kvp.Value}\"");
-                
-                if (i!= _claimsToSign.Count-1)
-                {
-                    i++;
-                    payload.Append(',');
-                }
+                json[claim.Key] = claim.Value;
             }
 
+            var jsonClaims = JsonHelper.JsonObjectToString(json);
+
+            //Remove extra brackets from JSON result
+            payload.Append(jsonClaims.Substring(1, jsonClaims.Length - 2));
+
             payload.Append('}');
 
             return payload.ToString();
diff --git a/tests/Microsoft.Identity.Test.Common/TestConstants.cs b/tests/Microsoft.Identity.Test.Common/TestConstants.cs
index 307c643957..a02a197c4c 100644
--- a/tests/Microsoft.Identity.Test.Common/TestConstants.cs
+++ b/tests/Microsoft.Identity.Test.Common/TestConstants.cs
@@ -35,7 +35,7 @@ public static HashSet<string> s_scope
         public const string PublicCloudConfidentialClientID = "88f91eac-c606-4c67-a0e2-a5e8a186854f";
         public const string AutomationTestCertName = "LabAuth.MSIDLab.com";
         public static Dictionary<string, string> AdditionalAssertionClaims =>
-            new Dictionary<string, string>() { { "Key1", "Val1" }, { "Key2", "Val2" } };
+            new Dictionary<string, string>() { { "Key1", "Val1" }, { "Key2", "Val2" }, { "customClaims", "{\"xms_az_claim\": [\"GUID\", \"GUID2\", \"GUID3\"]}" } };
 
         public static readonly SortedSet<string> s_scopeForAnotherResource = new SortedSet<string>(new[] { "r2/scope1", "r2/scope2" }, StringComparer.OrdinalIgnoreCase);
         public static readonly SortedSet<string> s_cacheMissScope = new SortedSet<string>(new[] { "r3/scope1", "r3/scope2" }, StringComparer.OrdinalIgnoreCase);
diff --git a/tests/Microsoft.Identity.Test.Unit/PublicApiTests/ClientCredentialWithCertTest.cs b/tests/Microsoft.Identity.Test.Unit/PublicApiTests/ClientCredentialWithCertTest.cs
index e184fbd22c..453a33e361 100644
--- a/tests/Microsoft.Identity.Test.Unit/PublicApiTests/ClientCredentialWithCertTest.cs
+++ b/tests/Microsoft.Identity.Test.Unit/PublicApiTests/ClientCredentialWithCertTest.cs
@@ -480,7 +480,7 @@ public void ClientAssertionTests(bool sendX5C, bool useSha2AndPss, bool addExtra
             if (appendDefaultClaims == false && addExtraClaims == false)
                 appendDefaultClaims = true;
 
-            int expectedPayloadClaimsCount = (appendDefaultClaims ? 6 : 0) + (addExtraClaims ? 2 : 0);
+            int expectedPayloadClaimsCount = (appendDefaultClaims ? 6 : 0) + (addExtraClaims ? 3 : 0);
             Assert.AreEqual(expectedPayloadClaimsCount, decodedToken.Payload.Count);
             if (appendDefaultClaims)
             {
@@ -506,6 +506,8 @@ public void ClientAssertionTests(bool sendX5C, bool useSha2AndPss, bool addExtra
             {
                 Assert.AreEqual("Val1", decodedToken.Payload["Key1"]);
                 Assert.AreEqual("Val2", decodedToken.Payload["Key2"]);
+                //Ensure JSON formatting is preserved
+                Assert.AreEqual("{\"xms_az_claim\": [\"GUID\", \"GUID2\", \"GUID3\"]}", decodedToken.Payload["customClaims"]);
             }
 
             if (useSha2AndPss)