Replace overload with delegates to read token values.

AzureAD · Nov 6, 2024 · 3cbc31c · 3cbc31c
1 parent 96b6b36
commit 3cbc31c
Show file tree

Hide file tree

Showing 9 changed files with 124 additions and 90 deletions.
diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/InternalAPI.Unshipped.txt b/src/Microsoft.IdentityModel.JsonWebTokens/InternalAPI.Unshipped.txt
@@ -1,3 +1,7 @@
+Microsoft.IdentityModel.JsonWebTokens.JsonWebToken.JsonWebToken(System.ReadOnlyMemory<char> encodedTokenMemory, Microsoft.IdentityModel.Tokens.ReadTokenPayloadValueDelegate readTokenPayloadValueDelegate) -> void
+Microsoft.IdentityModel.JsonWebTokens.JsonWebToken.ReadTokenPayloadValueDelegate.get -> Microsoft.IdentityModel.Tokens.ReadTokenPayloadValueDelegate
+Microsoft.IdentityModel.JsonWebTokens.JsonWebToken.ReadTokenPayloadValueDelegate.set -> void
+static Microsoft.IdentityModel.JsonWebTokens.JsonWebToken.ReadTokenPayloadValue(ref System.Text.Json.Utf8JsonReader reader, string claimName) -> object
 static Microsoft.IdentityModel.JsonWebTokens.JsonWebTokenHandler.EncryptToken(byte[] innerTokenUtf8Bytes, Microsoft.IdentityModel.Tokens.EncryptingCredentials encryptingCredentials, string compressionAlgorithm, System.Collections.Generic.IDictionary<string, object> additionalHeaderClaims, string tokenType, bool includeKeyIdInHeader) -> string
 static Microsoft.IdentityModel.JsonWebTokens.JsonWebTokenHandler.WriteJweHeader(Microsoft.IdentityModel.Tokens.EncryptingCredentials encryptingCredentials, string compressionAlgorithm, string tokenType, System.Collections.Generic.IDictionary<string, object> jweHeaderClaims, bool includeKeyIdInHeader) -> byte[]
 static Microsoft.IdentityModel.JsonWebTokens.JsonWebTokenHandler.WriteJwsHeader(ref System.Text.Json.Utf8JsonWriter writer, Microsoft.IdentityModel.Tokens.SigningCredentials signingCredentials, Microsoft.IdentityModel.Tokens.EncryptingCredentials encryptingCredentials, System.Collections.Generic.IDictionary<string, object> jweHeaderClaims, System.Collections.Generic.IDictionary<string, object> jwsHeaderClaims, string tokenType, bool includeKeyIdInHeader) -> void
diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonClaimSet.cs b/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonClaimSet.cs
@@ -28,7 +28,7 @@ internal class JsonClaimSet
 
         internal JsonClaimSet()
         {
-            _jsonClaims = new Dictionary<string, object>();
+            _jsonClaims = [];
         }
 
         internal JsonClaimSet(Dictionary<string, object> jsonClaims)

diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonWebToken.PayloadClaimSet.cs b/src/Microsoft.IdentityModel.JsonWebTokens/Json/JsonWebToken.PayloadClaimSet.cs
@@ -5,7 +5,6 @@
 using System.Collections.Generic;
 using System.Text.Json;
 using Microsoft.IdentityModel.Logging;
-using Microsoft.IdentityModel.Tokens;
 using Microsoft.IdentityModel.Tokens.Json;
 
 namespace Microsoft.IdentityModel.JsonWebTokens
@@ -40,7 +39,8 @@ internal JsonClaimSet CreatePayloadClaimSet(ReadOnlySpan<byte> byteSpan)
             {
                 if (reader.TokenType == JsonTokenType.PropertyName)
                 {
-                    ReadPayloadValue(ref reader, claims);
+                    string claimName = reader.GetString();
+                    claims[claimName] = ReadTokenPayloadValueDelegate(ref reader, claimName);
                 }
                 // We read a JsonTokenType.StartObject above, exiting and positioning reader at next token.
                 else if (JsonSerializerPrimitives.IsReaderAtTokenType(ref reader, JsonTokenType.EndObject, false))
@@ -52,72 +52,63 @@ internal JsonClaimSet CreatePayloadClaimSet(ReadOnlySpan<byte> byteSpan)
             return new JsonClaimSet(claims);
         }
 
-        private protected virtual void ReadPayloadValue(ref Utf8JsonReader reader, IDictionary<string, object> claims)
+        /// <summary>
+        /// Reads and saves the value of the payload claim from the reader.
+        /// </summary>
+        /// <param name="reader">The reader over the JWT.</param>
+        /// <param name="claimName">The claim at the current position of the reader.</param>
+        /// <returns>A claim that was read.</returns>
+        internal static object ReadTokenPayloadValue(ref Utf8JsonReader reader, string claimName)
         {
             if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Aud))
             {
-                _audiences = [];
+                List<string> _audiences = [];
                 reader.Read();
                 if (reader.TokenType == JsonTokenType.StartArray)
                 {
                     JsonSerializerPrimitives.ReadStringsSkipNulls(ref reader, _audiences, JwtRegisteredClaimNames.Aud, ClassName);
-                    claims[JwtRegisteredClaimNames.Aud] = _audiences;
                 }
                 else
                 {
                     if (reader.TokenType != JsonTokenType.Null)
                     {
                         _audiences.Add(JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.Aud, ClassName));
-                        claims[JwtRegisteredClaimNames.Aud] = _audiences[0];
-                    }
-                    else
-                    {
-                        claims[JwtRegisteredClaimNames.Aud] = _audiences;
+                        return _audiences[0];
                     }
                 }
+                return _audiences;
             }
             else if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Azp))
             {
-                _azp = JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.Azp, ClassName, true);
-                claims[JwtRegisteredClaimNames.Azp] = _azp;
+                return JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.Azp, ClassName, true);
             }
             else if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Exp))
             {
-                _exp = JsonSerializerPrimitives.ReadLong(ref reader, JwtRegisteredClaimNames.Exp, ClassName, true);
-                _expDateTime = EpochTime.DateTime(_exp.Value);
-                claims[JwtRegisteredClaimNames.Exp] = _exp;
+                return JsonSerializerPrimitives.ReadLong(ref reader, JwtRegisteredClaimNames.Exp, ClassName, true);
             }
             else if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Iat))
             {
-                _iat = JsonSerializerPrimitives.ReadLong(ref reader, JwtRegisteredClaimNames.Iat, ClassName, true);
-                _iatDateTime = EpochTime.DateTime(_iat.Value);
-                claims[JwtRegisteredClaimNames.Iat] = _iat;
+                return JsonSerializerPrimitives.ReadLong(ref reader, JwtRegisteredClaimNames.Iat, ClassName, true);
             }
             else if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Iss))
             {
-                _iss = JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.Iss, ClassName, true);
-                claims[JwtRegisteredClaimNames.Iss] = _iss;
+                return JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.Iss, ClassName, true);
             }
             else if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Jti))
             {
-                _jti = JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.Jti, ClassName, true);
-                claims[JwtRegisteredClaimNames.Jti] = _jti;
+                return JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.Jti, ClassName, true);
             }
             else if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Nbf))
             {
-                _nbf = JsonSerializerPrimitives.ReadLong(ref reader, JwtRegisteredClaimNames.Nbf, ClassName, true);
-                _nbfDateTime = EpochTime.DateTime(_nbf.Value);
-                claims[JwtRegisteredClaimNames.Nbf] = _nbf;
+                return JsonSerializerPrimitives.ReadLong(ref reader, JwtRegisteredClaimNames.Nbf, ClassName, true);
             }
             else if (reader.ValueTextEquals(JwtPayloadUtf8Bytes.Sub))
             {
-                _sub = JsonSerializerPrimitives.ReadStringOrNumberAsString(ref reader, JwtRegisteredClaimNames.Sub, ClassName, true);
-                claims[JwtRegisteredClaimNames.Sub] = _sub;
+                return JsonSerializerPrimitives.ReadStringOrNumberAsString(ref reader, JwtRegisteredClaimNames.Sub, ClassName, true);
             }
             else
             {
-                string propertyName = reader.GetString();
-                claims[propertyName] = JsonSerializerPrimitives.ReadPropertyValueAsObject(ref reader, propertyName, JsonClaimSet.ClassName, true);
+                return JsonSerializerPrimitives.ReadPropertyValueAsObject(ref reader, claimName, JsonClaimSet.ClassName, true);
             }
         }
     }

diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/JsonWebToken.cs b/src/Microsoft.IdentityModel.JsonWebTokens/JsonWebToken.cs
@@ -86,6 +86,25 @@ public JsonWebToken(string jwtEncodedString)
             _encodedToken = jwtEncodedString;
         }
 
+        /// <summary>
+        /// Initializes a new instance of <see cref="JsonWebToken"/> from a ReadOnlyMemory{char} in JWS or JWE Compact serialized format.
+        /// </summary>
+        /// <param name="encodedTokenMemory">A ReadOnlyMemory{char} containing the JSON Web Token serialized in JWS or JWE Compact format.</param>
+        /// <param name="readTokenPayloadValueDelegate">A custom delegate to be called when each payload claim is being read. If null, default implementation is called.</param>
+        internal JsonWebToken(
+            ReadOnlyMemory<char> encodedTokenMemory,
+            ReadTokenPayloadValueDelegate readTokenPayloadValueDelegate)
+        {
+            if (encodedTokenMemory.IsEmpty)
+                throw LogHelper.LogExceptionMessage(new ArgumentNullException(nameof(encodedTokenMemory)));
+
+            ReadTokenPayloadValueDelegate = readTokenPayloadValueDelegate ?? ReadTokenPayloadValue;
+
+            ReadToken(encodedTokenMemory);
+
+            _encodedTokenMemory = encodedTokenMemory;
+        }
+
         /// <summary>
         /// Initializes a new instance of <see cref="JsonWebToken"/> from a ReadOnlyMemory{char} in JWS or JWE Compact serialized format.
         /// </summary>
@@ -141,6 +160,24 @@ public JsonWebToken(string header, string payload)
             _encodedToken = encodedToken;
         }
 
+        /// <summary>
+        /// Called for each claim when token payload is being read.
+        /// </summary>
+        /// <remarks>
+        /// An example implementation:
+        /// <code>
+        /// object ReadPayloadValueDelegate(ref Utf8JsonReader reader, string claimName) =>
+        /// {
+        ///     if (reader.ValueTextEquals("CustomProp"))
+        ///     {
+        ///         return JsonSerializerPrimitives.ReadString(ref reader, JwtRegisteredClaimNames.CustomProp, ClassName, true);    
+        ///     }
+        ///     return JsonWebToken.ReadTokenPayloadValue(ref reader, claimName);
+        /// }
+        /// </code>
+        /// </remarks>
+        internal ReadTokenPayloadValueDelegate ReadTokenPayloadValueDelegate { get; set; } = ReadTokenPayloadValue;
+
         internal string ActualIssuer { get; set; }
 
         internal ClaimsIdentity ActorClaimsIdentity { get; set; }

diff --git a/src/Microsoft.IdentityModel.Tokens/Delegates.cs b/src/Microsoft.IdentityModel.Tokens/Delegates.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Text.Json;
 using System.Threading.Tasks;
 
 namespace Microsoft.IdentityModel.Tokens
@@ -206,4 +207,13 @@ namespace Microsoft.IdentityModel.Tokens
     /// <returns>The validated <see cref="SecurityToken"/>.</returns>
     internal delegate ValidationResult<SecurityKey> SignatureValidationDelegate(SecurityToken token, ValidationParameters validationParameters, BaseConfiguration? configuration, CallContext? callContext);
 #nullable restore
+
+    /// <summary>
+    /// Definition for ReadTokenPayloadValueDelegate.
+    /// Called for each claim when token payload is being read.
+    /// </summary>
+    /// <param name="reader">Reader for the underlying token bytes.</param>
+    /// <param name="claimName">The name of the claim being read.</param>
+    /// <returns></returns>
+    internal delegate object ReadTokenPayloadValueDelegate(ref Utf8JsonReader reader, string claimName);
 }
diff --git a/src/Microsoft.IdentityModel.Tokens/InternalAPI.Unshipped.txt b/src/Microsoft.IdentityModel.Tokens/InternalAPI.Unshipped.txt
@@ -10,9 +10,12 @@ Microsoft.IdentityModel.Tokens.IssuerValidationSource.IssuerMatchedConfiguration
 Microsoft.IdentityModel.Tokens.IssuerValidationSource.IssuerMatchedValidationParameters = 2 -> Microsoft.IdentityModel.Tokens.IssuerValidationSource
 Microsoft.IdentityModel.Tokens.LifetimeValidationError._expires -> System.DateTime
 Microsoft.IdentityModel.Tokens.LifetimeValidationError._notBefore -> System.DateTime
+Microsoft.IdentityModel.Tokens.ReadTokenPayloadValueDelegate
 Microsoft.IdentityModel.Tokens.TokenTypeValidationError
 Microsoft.IdentityModel.Tokens.TokenTypeValidationError.TokenTypeValidationError(Microsoft.IdentityModel.Tokens.MessageDetail messageDetail, System.Type exceptionType, System.Diagnostics.StackFrame stackFrame, string invalidTokenType) -> void
 Microsoft.IdentityModel.Tokens.TokenTypeValidationError._invalidTokenType -> string
+Microsoft.IdentityModel.Tokens.TokenValidationParameters.ReadTokenPayloadValue.get -> Microsoft.IdentityModel.Tokens.ReadTokenPayloadValueDelegate
+Microsoft.IdentityModel.Tokens.TokenValidationParameters.ReadTokenPayloadValue.set -> void
 Microsoft.IdentityModel.Tokens.TokenValidationParameters.TimeProvider.get -> System.TimeProvider
 Microsoft.IdentityModel.Tokens.TokenValidationParameters.TimeProvider.set -> void
 Microsoft.IdentityModel.Tokens.ValidationError.GetException(System.Type exceptionType, System.Exception innerException) -> System.Exception

diff --git a/src/Microsoft.IdentityModel.Tokens/TokenValidationParameters.cs b/src/Microsoft.IdentityModel.Tokens/TokenValidationParameters.cs
@@ -449,6 +449,11 @@ public string NameClaimType
         /// </summary>
         public IDictionary<string, object> PropertyBag { get; set; }
 
+        /// <summary>
+        /// Gets or sets a delegate that will be called when reading token payload claims.
+        /// </summary>
+        internal ReadTokenPayloadValueDelegate ReadTokenPayloadValue { get; set; }
+
         /// <summary>
         /// Gets or sets a boolean to control if configuration required to be refreshed before token validation.
         /// </summary>

diff --git a/test/Microsoft.IdentityModel.JsonWebTokens.Tests/CustomJsonWebToken.cs b/test/Microsoft.IdentityModel.JsonWebTokens.Tests/CustomJsonWebToken.cs
diff --git a/test/Microsoft.IdentityModel.JsonWebTokens.Tests/JsonWebTokenTests.cs b/test/Microsoft.IdentityModel.JsonWebTokens.Tests/JsonWebTokenTests.cs
@@ -12,9 +12,11 @@
 using System.Reflection;
 using System.Security.Claims;
 using System.Text;
+using System.Text.Json;
 using System.Threading;
 using Microsoft.IdentityModel.TestUtils;
 using Microsoft.IdentityModel.Tokens;
+using Microsoft.IdentityModel.Tokens.Json;
 using Microsoft.IdentityModel.Tokens.Json.Tests;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
@@ -1741,35 +1743,59 @@ public void StringAndMemoryConstructors_CreateEquivalentTokens(JwtTheoryData the
         }
 
         [Fact]
-        public void DerivedJsonWebToken_IsCreatedCorrectly()
+        public void CreateTokenWithoutKeyIdentifiersInHeader()
         {
-            var expectedCustomClaim = "customclaim";
-            var tokenStr = new JsonWebTokenHandler().CreateToken(new SecurityTokenDescriptor
+            string rawToken = new JsonWebTokenHandler().CreateToken(new SecurityTokenDescriptor
+            {
+                SigningCredentials = KeyingMaterial.JsonWebKeyRsa256SigningCredentials,
+                IncludeKeyIdInHeader = false
+            });
+            var token = new JsonWebToken(rawToken);
+            Assert.False(token.TryGetHeaderValue(JwtHeaderParameterNames.Kid, out string _));
+            Assert.False(token.TryGetHeaderValue(JwtHeaderParameterNames.X5t, out string _));
+        }
+
+        [Fact]
+        public void ReadTokenDelegates_CalledCorrectly()
+        {
+            var tokenSpan = new JsonWebTokenHandler().CreateToken(new SecurityTokenDescriptor
             {
                 Issuer = Default.Issuer,
                 Claims = new Dictionary<string, object>
                 {
-                    { "CustomClaim", expectedCustomClaim },
+                    { "CustomPayload", "custom_payload" },
+                },
+                AdditionalHeaderClaims = new Dictionary<string, object>
+                {
+                    { "CustomHeader", "custom_header" }
                 }
-            });
+            }).AsMemory();
+
+            object ReadPayloadValue(ref Utf8JsonReader reader, string claimName)
+            {
+                if (reader.ValueTextEquals("CustomPayload"u8))
+                {
+                    return new CustomPayloadClaim(JsonSerializerPrimitives.ReadString(ref reader, "CustomPayload", string.Empty, true));
+                }
+                return JsonWebToken.ReadTokenPayloadValue(ref reader, claimName);
+            }
+
+            var jwt = new JsonWebToken(tokenSpan, ReadPayloadValue);
 
-            var derivedToken = new CustomJsonWebToken(tokenStr);
+            Assert.True(jwt.TryGetHeaderValue<string>("CustomHeader", out var actualHeaderClaim));
+            Assert.True(jwt.TryGetPayloadValue<CustomPayloadClaim>("CustomPayload", out var actualPayloadClaim));
 
-            Assert.Equal(expectedCustomClaim, derivedToken.CustomClaim);
-            Assert.Equal(Default.Issuer, derivedToken.Issuer);
+            Assert.Equal("custom_header", actualHeaderClaim);
+            Assert.Equal("custom_payload", actualPayloadClaim.CustomValue);
         }
 
-        [Fact]
-        public void CreateTokenWithoutKeyIdentifiersInHeader()
+        private class CustomHeaderClaim(string customValue)
         {
-            string rawToken = new JsonWebTokenHandler().CreateToken(new SecurityTokenDescriptor
-            {
-                SigningCredentials = KeyingMaterial.JsonWebKeyRsa256SigningCredentials,
-                IncludeKeyIdInHeader = false
-            });
-            var token = new JsonWebToken(rawToken);
-            Assert.False(token.TryGetHeaderValue(JwtHeaderParameterNames.Kid, out string _));
-            Assert.False(token.TryGetHeaderValue(JwtHeaderParameterNames.X5t, out string _));
+            public string CustomValue { get; set; } = customValue;
+        }
+        private class CustomPayloadClaim(string customValue)
+        {
+            public string CustomValue { get; set; } = customValue;
         }
     }