Create exception type for capacity exceeded

dvsekhvalnov · Dec 19, 2024 · 8cf59d9 · 8cf59d9
1 parent 88d4a56
commit 8cf59d9
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 27 deletions.
diff --git a/UnitTests/SecurityVulnerabilitiesTest.cs b/UnitTests/SecurityVulnerabilitiesTest.cs
@@ -197,8 +197,7 @@ public void DeflateBomb()
             string bomb = Jose.JWT.Encode(payload, publicKey, JweAlgorithm.RSA_OAEP, JweEncryption.A256GCM, JweCompression.DEF);
 
             // when
-            Exception thrownException = Assert.Throws<JoseException>(() => Jose.JWT.Decode(bomb, privateKey));
-            Assert.IsType<NotSupportedException>(thrownException.InnerException);
+            Exception thrownException = Assert.Throws<CapacityExceededException>(() => Jose.JWT.Decode(bomb, privateKey));
         }
 
         [Fact]
@@ -219,6 +218,5 @@ public void TruncatedGcmAuthTag()
                 Console.Out.WriteLine(e.ToString());
             }
         }
-
     }
 }
diff --git a/jose-jwt/JWT.cs b/jose-jwt/JWT.cs
@@ -1,7 +1,5 @@
-using Jose;
 using System;
 using System.Collections.Generic;
-using System.Linq;
 using System.Text;
 
 namespace Jose
@@ -499,7 +497,7 @@ public static byte[] DecryptBytes(string token, object key, JweAlgorithm? alg =
 
         private static byte[] DecodeBytes(Compact.Iterator parts, object key = null, JwsAlgorithm? expectedJwsAlg = null, JweAlgorithm? expectedJweAlg = null, JweEncryption? expectedJweEnc = null, JwtSettings settings = null, byte[] payload = null)
         {
-            Ensure.IsNotEmpty(parts.Token, "Incoming token expected to be in compact serialization form, not empty, whitespace or null.");            
+            Ensure.IsNotEmpty(parts.Token, "Incoming token expected to be in compact serialization form, not empty, whitespace or null.");
 
             if (parts.Count == 5) //encrypted JWT
             {
@@ -512,7 +510,7 @@ private static byte[] DecodeBytes(Compact.Iterator parts, object key = null, Jws
             }
             else if (parts.Count == 3) // signed JWT
             {
-                if (expectedJweAlg != null || expectedJweEnc !=null)
+                if (expectedJweAlg != null || expectedJweEnc != null)
                 {
                     throw new InvalidAlgorithmException("Signed tokens can't assert encryption type.");
                 }
@@ -610,4 +608,10 @@ public class InvalidAlgorithmException : JoseException
         public InvalidAlgorithmException(string message) : base(message) { }
         public InvalidAlgorithmException(string message, Exception innerException) : base(message, innerException) { }
     }
+
+    public class CapacityExceededException : JoseException
+    {
+        public CapacityExceededException(string message) : base(message) { }
+        public CapacityExceededException(string message, Exception innerException) : base(message, innerException) { }
+    }
 }
diff --git a/jose-jwt/compression/CappedMemoryStream.cs b/jose-jwt/compression/CappedMemoryStream.cs
@@ -16,7 +16,7 @@ public override void Write(byte[] buffer, int offset, int count)
     {
         if (Length + Math.Min(count, buffer.Length - offset) > maxCapacity)
         {
-            throw new NotSupportedException("Exceeding maximum memory stream size.");
+            throw new CapacityExceededException("Exceeding maximum memory stream size.");
         }
 
         base.Write(buffer, offset, count);
@@ -26,7 +26,7 @@ public override void WriteByte(byte value)
     {
         if (Length + 1 > maxCapacity)
         {
-            throw new NotSupportedException("Exceeding maximum memory stream size.");
+            throw new CapacityExceededException("Exceeding maximum memory stream size.");
         }
 
         base.WriteByte(value);

diff --git a/jose-jwt/compression/DeflateCompression.cs b/jose-jwt/compression/DeflateCompression.cs
@@ -1,4 +1,3 @@
-using System;
 using System.IO;
 using System.IO.Compression;
 
@@ -27,26 +26,19 @@ public byte[] Compress(byte[] plainText)
         }
 
         public byte[] Decompress(byte[] compressedText)
-        {
-            try
+        {
+            using (MemoryStream ms = new CappedMemoryStream(maxBufferSizeBytes))
             {
-                using (MemoryStream ms = new CappedMemoryStream(maxBufferSizeBytes))
+                using (MemoryStream compressedStream = new MemoryStream(compressedText))
                 {
-                    using (MemoryStream compressedStream = new MemoryStream(compressedText))
+                    using (DeflateStream deflater = new DeflateStream(compressedStream, CompressionMode.Decompress))
                     {
-                        using (DeflateStream deflater = new DeflateStream(compressedStream, CompressionMode.Decompress))
-                        {
-                            deflater.CopyTo(ms);
-                        }
-                    }
-
-                    return ms.ToArray();
-                }
-            }
-            catch (NotSupportedException e)
-            {
-                throw new JoseException("Unable to deflate compressed payload, most likely exceeded decompression buffer size.", e);
-            }
+                        deflater.CopyTo(ms);
+                    }
+                }
+
+                return ms.ToArray();
+            }
         }
     }
 }