Merged PR 10213: Set MaximumDeflateSize

Brent Schmaltz · Brent Schmaltz · commit e06dc842aef7 · 2023-10-13T00:38:38.000Z
The Decompress method has been adjusted to only process a maximum number of chars.

----
#### AI-Generated Description
This pull request adds support for limiting the size of decompressed tokens in the `DeflateCompressionProvider` and the `JwtTokenUtilities` classes. It also adds unit tests for the new functionality and modifies some existing classes to use the new parameters. The main changes are:

- Added a `MaximumDeflateSize` property to the `DeflateCompressionProvider` class and the `JwtTokenDecryptionParameters` class.
- Added a `maximumDeflateSize` parameter to the `DecompressToken` method in the `JwtTokenUtilities` class and the `DecompressionFunction` delegate.
- Added a `JweDecompressSizeTheoryData` class and a `JWEDecompressionSizeTest` method to test the decompression size limit in both the `JsonWebTokenHandler` and the `JwtSecurityTokenHandler` classes.
- Modified the `CreateCompressionProvider` method in the `CompressionProviderFactory` class to accept a `maximumDeflateSize` parameter and pass it to the `DeflateCompressionProvider` constructor.
- Modified the `DecryptToken` method in the `JsonWebTokenHandler` class and the `DecryptToken` method in the `JwtSecurityTokenHandler` class to pass the `MaximumTokenSizeInBytes` property to the `DecompressToken` method.
diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/JsonWebTokenHandler.cs b/src/Microsoft.IdentityModel.JsonWebTokens/JsonWebTokenHandler.cs
@@ -911,7 +911,8 @@ private string DecryptToken(JsonWebToken jwtToken, TokenValidationParameters val
                 new JwtTokenDecryptionParameters
                 {
                     DecompressionFunction = JwtTokenUtilities.DecompressToken,
-                    Keys = keys
+                    Keys = keys,
+                    MaximumDeflateSize = MaximumTokenSizeInBytes
                 });
         }
 
diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/JwtTokenDecryptionParameters.cs b/src/Microsoft.IdentityModel.JsonWebTokens/JwtTokenDecryptionParameters.cs
@@ -39,7 +39,7 @@ internal class JwtTokenDecryptionParameters
         /// <summary>
         /// Gets or sets the function used to attempt decompression with.
         /// </summary>
-        public Func<byte[], string, string> DecompressionFunction { get; set; }
+        public Func<byte[], string, int, string> DecompressionFunction { get; set; }
 
         /// <summary>
         /// Gets or sets the encryption algorithm (Enc) of the token.
@@ -66,6 +66,15 @@ internal class JwtTokenDecryptionParameters
         /// </summary>
         public IEnumerable<SecurityKey> Keys { get; set; }
 
+        /// <summary>
+        /// Gets and sets the maximum deflate size in chars that will be processed.
+        /// </summary>
+        public int MaximumDeflateSize
+        {
+            get;
+            set;
+        } = TokenValidationParameters.DefaultMaximumTokenSizeInBytes;
+
         /// <summary>
         /// Gets or sets the 'value' of the 'zip' claim.
         /// </summary>
diff --git a/src/Microsoft.IdentityModel.JsonWebTokens/JwtTokenUtilities.cs b/src/Microsoft.IdentityModel.JsonWebTokens/JwtTokenUtilities.cs
@@ -112,14 +112,15 @@ public static string CreateEncodedSignature(string input, SigningCredentials sig
         /// <summary>
         /// Decompress JWT token bytes.
         /// </summary>
-        /// <param name="tokenBytes"></param>
-        /// <param name="algorithm"></param>
+        /// <param name="tokenBytes">the bytes to be decompressed.</param>
+        /// <param name="algorithm">the decompress algorithm.</param>
+        /// <param name="maximumDeflateSize">maximum number of chars that will be decompressed.</param>
         /// <exception cref="ArgumentNullException">if <paramref name="tokenBytes"/> is null.</exception>
         /// <exception cref="ArgumentNullException">if <paramref name="algorithm"/> is null.</exception>
         /// <exception cref="NotSupportedException">if the decompression <paramref name="algorithm"/> is not supported.</exception>
         /// <exception cref="SecurityTokenDecompressionFailedException">if decompression using <paramref name="algorithm"/> fails.</exception>
         /// <returns>Decompressed JWT token</returns>
-        internal static string DecompressToken(byte[] tokenBytes, string algorithm)
+        internal static string DecompressToken(byte[] tokenBytes, string algorithm, int maximumDeflateSize)
         {
             if (tokenBytes == null)
                 throw LogHelper.LogArgumentNullException(nameof(tokenBytes));
@@ -130,7 +131,7 @@ internal static string DecompressToken(byte[] tokenBytes, string algorithm)
             if (!CompressionProviderFactory.Default.IsSupportedAlgorithm(algorithm))
                 throw LogHelper.LogExceptionMessage(new NotSupportedException(LogHelper.FormatInvariant(TokenLogMessages.IDX10682, LogHelper.MarkAsNonPII(algorithm))));
 
-            var compressionProvider = CompressionProviderFactory.Default.CreateCompressionProvider(algorithm);
+            var compressionProvider = CompressionProviderFactory.Default.CreateCompressionProvider(algorithm, maximumDeflateSize);
 
             var decompressedBytes = compressionProvider.Decompress(tokenBytes);
 
@@ -241,7 +242,7 @@ internal static string DecryptJwtToken(
                 if (string.IsNullOrEmpty(zipAlgorithm))
                     return Encoding.UTF8.GetString(decryptedTokenBytes);
 
-                return decryptionParameters.DecompressionFunction(decryptedTokenBytes, zipAlgorithm);
+                return decryptionParameters.DecompressionFunction(decryptedTokenBytes, zipAlgorithm, decryptionParameters.MaximumDeflateSize);
             }
             catch (Exception ex)
             {
diff --git a/src/Microsoft.IdentityModel.Tokens/CompressionProviderFactory.cs b/src/Microsoft.IdentityModel.Tokens/CompressionProviderFactory.cs
@@ -78,6 +78,17 @@ private static bool IsSupportedCompressionAlgorithm(string algorithm)
         /// <param name="algorithm">the decompression algorithm.</param>
         /// <returns>a <see cref="ICompressionProvider"/>.</returns>
         public ICompressionProvider CreateCompressionProvider(string algorithm)
+        {
+            return CreateCompressionProvider(algorithm, TokenValidationParameters.DefaultMaximumTokenSizeInBytes);
+        }
+
+        /// <summary>
+        /// Returns a <see cref="ICompressionProvider"/> for a specific algorithm.
+        /// </summary>
+        /// <param name="algorithm">the decompression algorithm.</param>
+        /// <param name="maximumDeflateSize">the maximum deflate size in chars that will be processed.</param>
+        /// <returns>a <see cref="ICompressionProvider"/>.</returns>
+        public ICompressionProvider CreateCompressionProvider(string algorithm, int maximumDeflateSize)
         {
             if (string.IsNullOrEmpty(algorithm))
                 throw LogHelper.LogArgumentNullException(nameof(algorithm));
@@ -86,7 +97,7 @@ public ICompressionProvider CreateCompressionProvider(string algorithm)
                 return CustomCompressionProvider;
 
             if (algorithm.Equals(CompressionAlgorithms.Deflate))
-                return new DeflateCompressionProvider();
+                return new DeflateCompressionProvider { MaximumDeflateSize = maximumDeflateSize };
 
             throw LogHelper.LogExceptionMessage(new NotSupportedException(LogHelper.FormatInvariant(LogMessages.IDX10652, LogHelper.MarkAsNonPII(algorithm))));
         }
diff --git a/src/Microsoft.IdentityModel.Tokens/DeflateCompressionProvider.cs b/src/Microsoft.IdentityModel.Tokens/DeflateCompressionProvider.cs
@@ -3,6 +3,10 @@
 
 using Microsoft.IdentityModel.Logging;
 using System;
+#if NET461_OR_GREATER
+using System.Buffers;
+#endif
+
 using System.IO;
 using System.IO.Compression;
 using System.Text;
@@ -14,6 +18,8 @@ namespace Microsoft.IdentityModel.Tokens
     /// </summary>
     public class DeflateCompressionProvider : ICompressionProvider
     {
+        private int _maximumTokenSizeInBytes = TokenValidationParameters.DefaultMaximumTokenSizeInBytes;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="DeflateCompressionProvider"/> class used to compress and decompress used the <see cref="CompressionAlgorithms.Deflate"/> algorithm.
         /// </summary>
@@ -41,6 +47,16 @@ public DeflateCompressionProvider(CompressionLevel compressionLevel)
         /// </summary>
         public CompressionLevel CompressionLevel { get; private set; } = CompressionLevel.Optimal;
 
+        /// <summary>
+        /// Gets and sets the maximum deflate size in chars that will be processed.
+        /// </summary>
+        /// <exception cref="ArgumentOutOfRangeException">'value' less than 1.</exception>
+        public int MaximumDeflateSize
+        {
+            get => _maximumTokenSizeInBytes;
+            set => _maximumTokenSizeInBytes = (value < 1) ? throw LogHelper.LogExceptionMessage(new ArgumentOutOfRangeException(nameof(value), LogHelper.FormatInvariant(LogMessages.IDX10101, LogHelper.MarkAsNonPII(value)))) : value;
+        }
+
         /// <summary>
         /// Decompress the value using DEFLATE algorithm.
         /// </summary>
@@ -51,16 +67,43 @@ public byte[] Decompress(byte[] value)
             if (value == null)
                 throw LogHelper.LogArgumentNullException(nameof(value));
 
-            using (var inputStream = new MemoryStream(value))
+            char[] chars = null;
+            try
             {
-                using (var deflateStream = new DeflateStream(inputStream, CompressionMode.Decompress))
+#if NET461_OR_GREATER
+                chars = ArrayPool<char>.Shared.Rent(MaximumDeflateSize);
+#else
+                chars = new char[MaximumDeflateSize];
+#endif
+                using (var inputStream = new MemoryStream(value))
                 {
-                    using (var reader = new StreamReader(deflateStream, Encoding.UTF8))
+                    using (var deflateStream = new DeflateStream(inputStream, CompressionMode.Decompress))
                     {
-                        return Encoding.UTF8.GetBytes(reader.ReadToEnd());
+                        using (var reader = new StreamReader(deflateStream, Encoding.UTF8))
+                        {
+                            // if there is one more char to read, then the token is too large.
+                            int bytesRead = reader.Read(chars, 0, MaximumDeflateSize);
+                            if (reader.Peek() != -1)
+                            {
+                                throw LogHelper.LogExceptionMessage(
+                                    new SecurityTokenDecompressionFailedException(
+                                        LogHelper.FormatInvariant(
+                                            LogMessages.IDX10816,
+                                            LogHelper.MarkAsNonPII(MaximumDeflateSize))));
+                            }
+
+                            return Encoding.UTF8.GetBytes(chars, 0, bytesRead);
+                        }
                     }
                 }
             }
+            finally
+            {
+#if NET461_OR_GREATER
+                if (chars != null)
+                    ArrayPool<char>.Shared.Return(chars);
+#endif
+            }
         }
 
         /// <summary>
diff --git a/src/Microsoft.IdentityModel.Tokens/LogMessages.cs b/src/Microsoft.IdentityModel.Tokens/LogMessages.cs
@@ -239,6 +239,8 @@ internal static class LogMessages
         public const string IDX10812 = "IDX10812: Unable to create a {0} from the properties found in the JsonWebKey: '{1}'.";
         public const string IDX10813 = "IDX10813: Unable to create a {0} from the properties found in the JsonWebKey: '{1}', Exception '{2}'.";
         public const string IDX10814 = "IDX10814: Unable to create a {0} from the properties found in the JsonWebKey: '{1}'. Missing: '{2}'.";
+        public const string IDX10815 = "IDX10815: Depth of JSON: '{0}' exceeds max depth of '{1}'.";
+        public const string IDX10816 = "IDX10816: Decompressing would result in a token with a size greater than allowed. Maximum size allowed: '{0}'.";
 
         // Base64UrlEncoding
         public const string IDX10820 = "IDX10820: Invalid character found in Base64UrlEncoding. Character: '{0}', Encoding: '{1}'.";
diff --git a/src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs b/src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs
@@ -1760,6 +1760,7 @@ protected string DecryptToken(JwtSecurityToken jwtToken, TokenValidationParamete
                 EncodedToken = jwtToken.RawData,
                 HeaderAsciiBytes = Encoding.ASCII.GetBytes(jwtToken.EncodedHeader),
                 InitializationVectorBytes = Base64UrlEncoder.DecodeBytes(jwtToken.RawInitializationVector),
+                MaximumDeflateSize = MaximumTokenSizeInBytes,
                 Keys = keys,
                 Zip = jwtToken.Header.Zip,
             });
diff --git a/test/Microsoft.IdentityModel.JsonWebTokens.Tests/JsonWebTokenHandlerTests.cs b/test/Microsoft.IdentityModel.JsonWebTokens.Tests/JsonWebTokenHandlerTests.cs
@@ -8,10 +8,8 @@
 using System.IdentityModel.Tokens.Jwt.Tests;
 using System.IO;
 using System.Linq;
-using System.Runtime.InteropServices;
 using System.Security.Claims;
 using System.Security.Cryptography;
-using System.Security.Policy;
 using System.Text;
 using System.Threading.Tasks;
 using Microsoft.IdentityModel.Json;
@@ -3400,6 +3398,70 @@ public static TheoryData<CreateTokenTheoryData> JWECompressionTheoryData
             }
         }
 
+        [Theory, MemberData(nameof(JweDecompressSizeTheoryData))]
+        public void JWEDecompressionSizeTest(JWEDecompressionTheoryData theoryData)
+        {
+            var context = TestUtilities.WriteHeader($"{this}.JWEDecompressionTest", theoryData);
+
+            try
+            {
+                var handler = new JsonWebTokenHandler();
+                CompressionProviderFactory.Default = theoryData.CompressionProviderFactory;
+                var validationResult = handler.ValidateTokenAsync(theoryData.JWECompressionString, theoryData.ValidationParameters).Result;
+                theoryData.ExpectedException.ProcessException(validationResult.Exception, context);
+            }
+            catch (Exception ex)
+            {
+                theoryData.ExpectedException.ProcessException(ex, context);
+            }
+
+            TestUtilities.AssertFailIfErrors(context);
+        }
+
+        public static TheoryData<JWEDecompressionTheoryData> JweDecompressSizeTheoryData()
+        {
+            // The character 'U' compresses better because UUU in base 64, repeated characters compress best.
+            JsonWebTokenHandler jwth = new JsonWebTokenHandler();
+            SecurityKey key = new SymmetricSecurityKey(new byte[256 / 8]);
+            EncryptingCredentials encryptingCredentials = new EncryptingCredentials(key, "dir", "A128CBC-HS256");
+            TokenValidationParameters validationParameters = new TokenValidationParameters { TokenDecryptionKey = key };
+
+            TheoryData<JWEDecompressionTheoryData> theoryData = new TheoryData<JWEDecompressionTheoryData>();
+            string strU = new string('U', 100_000_000);
+            string strUU = new string('U', 40_000_000);
+            string payload = $@"{{""U"":""{strU}"", ""UU"":""{strUU}""}}";
+
+            string token = jwth.CreateToken(payload, encryptingCredentials, "DEF");
+            theoryData.Add(new JWEDecompressionTheoryData
+            {
+                CompressionProviderFactory = new CompressionProviderFactory(),
+                ValidationParameters = validationParameters,
+                JWECompressionString = token,
+                TestId = "DeflateSizeExceeded",
+                ExpectedException = new ExpectedException(
+                    typeof(SecurityTokenDecompressionFailedException),
+                    "IDX10679:",
+                    typeof(SecurityTokenDecompressionFailedException))
+            });
+
+            strUU = new string('U', 50_000_000);
+            payload = $@"{{""U"":""{strU}"", ""UU"":""{strUU}""}}";
+
+            token = jwth.CreateToken(payload, encryptingCredentials, "DEF");
+            theoryData.Add(new JWEDecompressionTheoryData
+            {
+                CompressionProviderFactory = new CompressionProviderFactory(),
+                ValidationParameters = validationParameters,
+                JWECompressionString = token,
+                TestId = "TokenSizeExceeded",
+                ExpectedException = new ExpectedException(
+                    typeof(ArgumentException),
+                    "IDX10209:")
+            });
+
+            return theoryData;
+        }
+
         [Theory, MemberData(nameof(JWEDecompressionTheoryData))]
         public void JWEDecompressionTest(JWEDecompressionTheoryData theoryData)
         {
@@ -3408,7 +3470,7 @@ public void JWEDecompressionTest(JWEDecompressionTheoryData theoryData)
             try
             {
                 var handler = new JsonWebTokenHandler();
-                CompressionProviderFactory.Default = theoryData.CompressionProviderFactory;
+                //CompressionProviderFactory.Default = theoryData.CompressionProviderFactory;
                 var validationResult = handler.ValidateToken(theoryData.JWECompressionString, theoryData.ValidationParameters);
                 var validatedToken = validationResult.SecurityToken as JsonWebToken;
                 if (validationResult.Exception != null)
@@ -3493,29 +3555,32 @@ public static TheoryData<JWEDecompressionTheoryData> JWEDecompressionTheoryData(
                     TestId = "InvalidToken",
                     ExpectedException = new ExpectedException(typeof(SecurityTokenDecompressionFailedException), "IDX10679:", typeof(InvalidDataException))
                 },
-                new JWEDecompressionTheoryData
-                {
-                    ValidationParameters = Default.JWECompressionTokenValidationParameters,
-                    JWECompressionString = ReferenceTokens.JWECompressionTokenWithDEF,
-                    CompressionProviderFactory = null,
-                    TestId = "NullCompressionProviderFactory",
-                    ExpectedException = ExpectedException.ArgumentNullException("IDX10000:")
-                },
-                new JWEDecompressionTheoryData
-                {
-                    ValidationParameters = Default.JWECompressionTokenValidationParameters,
-                    CompressionProviderFactory = compressionProviderFactoryForCustom,
-                    JWECompressionString = ReferenceTokens.JWECompressionTokenWithCustomAlgorithm,
-                    TestId = "CustomCompressionProviderSucceeds"
-                },
-                new JWEDecompressionTheoryData
-                {
-                    ValidationParameters = Default.JWECompressionTokenValidationParameters,
-                    JWECompressionString = ReferenceTokens.JWECompressionTokenWithDEF,
-                    CompressionProviderFactory = compressionProviderFactoryForCustom2,
-                    TestId = "CustomCompressionProviderFails",
-                    ExpectedException = new ExpectedException(typeof(SecurityTokenDecompressionFailedException), "IDX10679:", typeof(SecurityTokenDecompressionFailedException))
-                }
+                // Skip these tests as they set a static
+                // We need to have a replacement model for custom compression
+                // https://identitydivision.visualstudio.com/Engineering/_workitems/edit/2719954
+                //new JWEDecompressionTheoryData
+                //{
+                //    ValidationParameters = Default.JWECompressionTokenValidationParameters,
+                //    JWECompressionString = ReferenceTokens.JWECompressionTokenWithDEF,
+                //    CompressionProviderFactory = null,
+                //    TestId = "NullCompressionProviderFactory",
+                //    ExpectedException = ExpectedException.ArgumentNullException("IDX10000:")
+                //},
+                //new JWEDecompressionTheoryData
+                //{
+                //    ValidationParameters = Default.JWECompressionTokenValidationParameters,
+                //    CompressionProviderFactory = compressionProviderFactoryForCustom,
+                //    JWECompressionString = ReferenceTokens.JWECompressionTokenWithCustomAlgorithm,
+                //    TestId = "CustomCompressionProviderSucceeds"
+                //},
+                //new JWEDecompressionTheoryData
+                //{
+                //    ValidationParameters = Default.JWECompressionTokenValidationParameters,
+                //    JWECompressionString = ReferenceTokens.JWECompressionTokenWithDEF,
+                //    CompressionProviderFactory = compressionProviderFactoryForCustom2,
+                //    TestId = "CustomCompressionProviderFails",
+                //    ExpectedException = new ExpectedException(typeof(SecurityTokenDecompressionFailedException), "IDX10679:", typeof(SecurityTokenDecompressionFailedException))
+                //}
             };
         }
 
diff --git a/test/System.IdentityModel.Tokens.Jwt.Tests/JwtSecurityTokenHandlerTests.cs b/test/System.IdentityModel.Tokens.Jwt.Tests/JwtSecurityTokenHandlerTests.cs

Original file line number	Diff line number	Diff line change
`@@ -911,7 +911,8 @@ private string DecryptToken(JsonWebToken jwtToken, TokenValidationParameters val`
`911`	`911`	`new JwtTokenDecryptionParameters`
`912`	`912`	`{`
`913`	`913`	`DecompressionFunction = JwtTokenUtilities.DecompressToken,`
`914`		`- Keys = keys`
	`914`	`+ Keys = keys,`
	`915`	`+ MaximumDeflateSize = MaximumTokenSizeInBytes`
`915`	`916`	`});`
`916`	`917`	`}`
`917`	`918`