From 7b968bed4e430ce6253205419b5dd88deac3cf62 Mon Sep 17 00:00:00 2001
From: Peter <34331512+pmaytak@users.noreply.github.com>
Date: Wed, 19 Jun 2024 12:39:17 -0700
Subject: [PATCH] Fixes to claims.

---
 .../Json/JsonClaimSet.cs                      | 33 ++++++++++++++++---
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonClaimSet.cs b/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonClaimSet.cs
index f1f4560973..ba666daf72 100644
--- a/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonClaimSet.cs
+++ b/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonClaimSet.cs
@@ -25,24 +25,33 @@ internal class JsonClaimSet
 
         internal object _claimsLock = new();
         internal readonly Dictionary<string, object> _jsonClaims;
+
 #if NET8_0_OR_GREATER
         internal readonly Dictionary<string, (int startIndex, int length)?> _jsonClaimsBytes;
         internal readonly Memory<byte> _tokenAsMemory;
 #endif
+
         private List<Claim> _claims;
 
         internal JsonClaimSet()
         {
-            _jsonClaims = new();
+            _jsonClaims = [];
+
 #if NET8_0_OR_GREATER
-            _jsonClaimsBytes = new();
+            _jsonClaimsBytes = [];
+            _tokenAsMemory = Memory<byte>.Empty;
 #endif
         }
 
         internal JsonClaimSet(Dictionary<string, object> jsonClaims)
         {
             _jsonClaims = jsonClaims;
+
+#if NET8_0_OR_GREATER
+            _jsonClaimsBytes = [];
+#endif
         }
+
 #if NET8_0_OR_GREATER
         internal JsonClaimSet(
             Dictionary<string, object> jsonClaims,
@@ -54,6 +63,7 @@ internal JsonClaimSet(
             _tokenAsMemory = tokenAsMemory;
         }
 #endif
+
         internal List<Claim> Claims(string issuer)
         {
             if (_claims == null)
@@ -69,6 +79,14 @@ internal List<Claim> CreateClaims(string issuer)
             foreach (KeyValuePair<string, object> kvp in _jsonClaims)
                 CreateClaimFromObject(claims, kvp.Key, kvp.Value, issuer);
 
+#if NET8_0_OR_GREATER
+            // _jsonClaimsBytes is only for string values for known claims, which would not be in _jsonClaims.
+            foreach (KeyValuePair<string, (int StartIndex, int Length)?> kvp in _jsonClaimsBytes)
+            {
+                string value = Encoding.UTF8.GetString(_tokenAsMemory.Slice(kvp.Value.Value.StartIndex, kvp.Value.Value.Length).Span);
+                claims.Add(new Claim(kvp.Key, value, ClaimValueTypes.String, issuer, issuer));
+            }
+#endif
             return claims;
         }
 
@@ -168,7 +186,11 @@ internal Claim GetClaim(string key, string issuer)
             if (key == null)
                 throw LogHelper.LogArgumentNullException(nameof(key));
 
+#if NET8_0_OR_GREATER
+            if (_jsonClaims.TryGetValue(key, out object _) || _jsonClaimsBytes.TryGetValue(key, out _))
+#else
             if (_jsonClaims.TryGetValue(key, out object _))
+#endif
             {
                 foreach (var claim in Claims(issuer))
                     if (claim.Type == key)
@@ -197,7 +219,6 @@ internal string GetStringValue(string key)
                 return obj.ToString();
             }
 #endif
-
             return string.Empty;
         }
 
@@ -486,7 +507,11 @@ internal bool TryGetValue<T>(string key, out T value)
 
         internal bool HasClaim(string claimName)
         {
-            return _jsonClaims.TryGetValue(claimName, out _);
+#if NET8_0_OR_GREATER
+            return _jsonClaims.ContainsKey(claimName) || _jsonClaimsBytes.ContainsKey(claimName);
+#else
+            return _jsonClaims.ContainsKey(claimName);
+#endif
         }
     }
 }