Constant stackalloc can be optimized by the JIT. (#569)

* Constant `stackalloc` can be optimized by the JIT. dotnet/runtime#54186 (comment) * Fix ESXXX KID computation cause by incorrect canonicalized JWK size ('crv' was computed as base64-encoded) * `EllipticalCurve` is now a class instead of a readonly struct. This reduce the cost of struct copy.
uruk-project · Aug 3, 2021 · 8f7d759 · 8f7d759
1 parent 6db64a2
commit 8f7d759
Show file tree

Hide file tree

Showing 51 changed files with 678 additions and 340 deletions.
diff --git a/perf/Sandbox/StructCopyBenchmark.cs b/perf/Sandbox/StructCopyBenchmark.cs
@@ -0,0 +1,69 @@
+using System.Security.Cryptography;
+using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Diagnosers;
+
+namespace JsonWebToken.Performance
+{
+    [MemoryDiagnoser]
+    [DisassemblyDiagnoser]
+    public class StructCopyBenchmark
+    {
+        [Params(true, false)]
+        public bool includePrivateParameters;
+
+        private ECParameters _parameters = new ECParameters { Curve = ECCurve.NamedCurves.nistP256, D = new byte[32], Q = new ECPoint { X = new byte[32], Y = new byte[32] } };
+
+        [Benchmark(Baseline = true)]
+        public ECParameters ExportParameters_Old()
+        {
+            var parameters = new ECParameters
+            {
+                Q = _parameters.Q,
+                Curve = _parameters.Curve
+            };
+            if (includePrivateParameters)
+            {
+                parameters.D = _parameters.D;
+            }
+
+            return parameters;
+        }
+
+        [Benchmark]
+        public ECParameters ExportParameters_Old2()
+        {
+            var parameters = new ECParameters
+            {
+                Q = _parameters.Q,
+                Curve = _parameters.Curve
+            };
+                parameters.D = includePrivateParameters ? _parameters.D : null;
+
+            return parameters;
+        }
+        [Benchmark]
+        public ECParameters ExportParameters_Brute()
+        {
+            return _parameters;
+        }
+        [Benchmark]
+        public ECParameters ExportParameters_Brute2()
+        {
+            var p = _parameters;
+            return p;
+        }
+
+        [Benchmark]
+        public ECParameters ExportParameters_New()
+        {
+            var parameters = _parameters;
+
+            if (!includePrivateParameters)
+            {
+                parameters.D = null;
+            }
+
+            return parameters;
+        }
+    }
+}
diff --git a/src/JsonWebToken/Base64.cs b/src/JsonWebToken/Base64.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Buffers;
+using System.Diagnostics;
 using System.Runtime.CompilerServices;
 
 namespace JsonWebToken
@@ -46,7 +47,7 @@ public static int Decode(ReadOnlySpan<char> base64, Span<byte> data)
             byte[]? arrayToReturn = null;
             var buffer = base64.Length > Constants.MaxStackallocBytes
                 ? (arrayToReturn = ArrayPool<byte>.Shared.Rent(base64.Length))
-                : stackalloc byte[base64.Length];
+                : stackalloc byte[Constants.MaxStackallocBytes];
             try
             {
                 int length = Utf8.GetBytes(base64, buffer);
@@ -83,7 +84,7 @@ public static OperationStatus Decode(ReadOnlySpan<byte> base64, Span<byte> data,
                 byte[]? utf8ArrayToReturn = null;
                 Span<byte> utf8Data = base64.Length > Constants.MaxStackallocBytes
                     ? (utf8ArrayToReturn = ArrayPool<byte>.Shared.Rent(base64.Length))
-                    : stackalloc byte[base64.Length];
+                    : stackalloc byte[Constants.MaxStackallocBytes];
                 try
                 {
                     int length = 0;
@@ -118,7 +119,7 @@ public static OperationStatus Decode(ReadOnlySpan<byte> base64, Span<byte> data,
             }
         }
 
-        private static bool IsWhiteSpace(byte c) 
+        private static bool IsWhiteSpace(byte c)
             => c == ' ' || (c >= '\t' && c <= '\r');
 
         private static ReadOnlySpan<byte> WhiteSpace => new byte[] { (byte)' ', (byte)'\t', (byte)'\r', (byte)'\n', (byte)'\v', (byte)'\f' };
@@ -173,7 +174,7 @@ public static byte[] Encode(ReadOnlySpan<char> data)
                 int length = Utf8.GetMaxByteCount(data.Length);
                 var utf8Data = length > Constants.MaxStackallocBytes
                     ? (utf8ArrayToReturn = ArrayPool<byte>.Shared.Rent(length))
-                    : stackalloc byte[length];
+                    : stackalloc byte[Constants.MaxStackallocBytes];
 
                 int written = Utf8.GetBytes(data, utf8Data);
                 return Encode(utf8Data.Slice(0, written));

diff --git a/src/JsonWebToken/Base64Url.cs b/src/JsonWebToken/Base64Url.cs
@@ -26,7 +26,7 @@ public static byte[] Decode(string data)
             {
                 Span<byte> tmp = length > Constants.MaxStackallocBytes
                     ? (utf8ArrayToReturn = ArrayPool<byte>.Shared.Rent(length))
-                    : stackalloc byte[length];
+                    : stackalloc byte[Constants.MaxStackallocBytes];
                 int written = Utf8.GetBytes(data, tmp);
                 return Decode(tmp.Slice(0, written));
             }
@@ -70,7 +70,7 @@ public static int Decode(ReadOnlySpan<char> base64Url, Span<byte> data)
             byte[]? arrayToReturn = null;
             var buffer = base64Url.Length > Constants.MaxStackallocBytes
                 ? (arrayToReturn = ArrayPool<byte>.Shared.Rent(base64Url.Length))
-                : stackalloc byte[base64Url.Length];
+                : stackalloc byte[Constants.MaxStackallocBytes];
             try
             {
                 int length = Utf8.GetBytes(base64Url, buffer);
@@ -163,7 +163,7 @@ public static byte[] Encode(ReadOnlySpan<char> data)
                 int length = Utf8.GetMaxByteCount(data.Length);
                 var utf8Data = length > Constants.MaxStackallocBytes
                     ? (utf8ArrayToReturn = ArrayPool<byte>.Shared.Rent(length))
-                    : stackalloc byte[length];
+                    : stackalloc byte[Constants.MaxStackallocBytes];
 
                 int written = Utf8.GetBytes(data, utf8Data);
                 return Encode(utf8Data.Slice(0, written));

diff --git a/src/JsonWebToken/Cryptography/AesCbcHmacDecryptor.cs b/src/JsonWebToken/Cryptography/AesCbcHmacDecryptor.cs
@@ -74,11 +74,12 @@ private bool VerifyAuthenticationTag(ReadOnlySpan<byte> key, ReadOnlySpan<byte>
         {
             byte[]? byteArrayToReturnToPool = null;
             int macLength = associatedData.Length + iv.Length + ciphertext.Length + sizeof(long);
-            Span<byte> macBytes = macLength <= Constants.MaxStackallocBytes
-                                    ? stackalloc byte[macLength]
-                                    : (byteArrayToReturnToPool = ArrayPool<byte>.Shared.Rent(macLength)).AsSpan(0, macLength);
+            Span<byte> macBytes = macLength > Constants.MaxStackallocBytes
+                                    ? (byteArrayToReturnToPool = ArrayPool<byte>.Shared.Rent(macLength))
+                                    : stackalloc byte[Constants.MaxStackallocBytes];
             try
             {
+                macBytes = macBytes.Slice(0, macLength);
                 associatedData.CopyTo(macBytes);
                 var bytes = macBytes.Slice(associatedData.Length);
                 iv.CopyTo(bytes);
@@ -88,9 +89,11 @@ private bool VerifyAuthenticationTag(ReadOnlySpan<byte> key, ReadOnlySpan<byte>
                 BinaryPrimitives.WriteInt64BigEndian(bytes, associatedData.Length << 3);
 
                 Sha2 hashAlgorithm = _encryptionAlgorithm.SignatureAlgorithm.Sha;
-                Span<byte> hmacKey = stackalloc byte[hashAlgorithm.BlockSize * 2];
+                Span<byte> hmacKey = stackalloc byte[Sha2.BlockSizeStackallocThreshold * 2]
+                    .Slice(0, hashAlgorithm.BlockSize * 2); 
                 Hmac hmac = new Hmac(hashAlgorithm, key, hmacKey);
-                Span<byte> hash = stackalloc byte[authenticationTag.Length * 2];
+                Span<byte> hash = stackalloc byte[AuthenticatedEncryptor.TagSizeStackallocThreshold]
+                    .Slice(0, authenticationTag.Length * 2);
                 hmac.ComputeHash(macBytes, hash);
                 CryptographicOperations.ZeroMemory(hmacKey);
 

diff --git a/src/JsonWebToken/Cryptography/AesCbcHmacEncryptor.cs b/src/JsonWebToken/Cryptography/AesCbcHmacEncryptor.cs
@@ -89,9 +89,9 @@ private void ComputeAuthenticationTag(ReadOnlySpan<byte> key, ReadOnlySpan<byte>
             try
             {
                 int macLength = associatedData.Length + iv.Length + ciphertext.Length + sizeof(long);
-                Span<byte> macBytes = macLength <= Constants.MaxStackallocBytes
-                    ? stackalloc byte[macLength]
-                    : (arrayToReturnToPool = ArrayPool<byte>.Shared.Rent(macLength)).AsSpan(0, macLength);
+                Span<byte> macBytes = macLength > Constants.MaxStackallocBytes
+                    ? (arrayToReturnToPool = ArrayPool<byte>.Shared.Rent(macLength))
+                    : stackalloc byte[Constants.MaxStackallocBytes];
 
                 associatedData.CopyTo(macBytes);
                 var bytes = macBytes.Slice(associatedData.Length);
@@ -102,9 +102,9 @@ private void ComputeAuthenticationTag(ReadOnlySpan<byte> key, ReadOnlySpan<byte>
                 BinaryPrimitives.WriteInt64BigEndian(bytes, associatedData.Length << 3);
 
                 Sha2 hashAlgorithm = _encryptionAlgorithm.SignatureAlgorithm.Sha;
-                Span<byte> hmacKey = stackalloc byte[hashAlgorithm.BlockSize * 2];
-                Hmac hmac = new Hmac(hashAlgorithm, key, hmacKey);
-                hmac.ComputeHash(macBytes, authenticationTag);
+                Span<byte> hmacKey = stackalloc byte[Sha2.BlockSizeStackallocThreshold * 2];
+                Hmac hmac = new Hmac(hashAlgorithm, key, hmacKey.Slice(0, hashAlgorithm.BlockSize * 2));
+                hmac.ComputeHash(macBytes.Slice(0, macLength), authenticationTag);
                 authenticationTagBytesWritten = authenticationTag.Length >> 1;
             }
             finally

diff --git a/src/JsonWebToken/Cryptography/AesGcmKeyUnwrapper.cs b/src/JsonWebToken/Cryptography/AesGcmKeyUnwrapper.cs
@@ -10,6 +10,9 @@ namespace JsonWebToken.Cryptography
 {
     internal sealed class AesGcmKeyUnwrapper : KeyUnwrapper
     {
+        private const int IVSize = 12;
+        private const int TagSize = 16;
+
         private readonly SymmetricJwk _key;
 
         public AesGcmKeyUnwrapper(SymmetricJwk key, EncryptionAlgorithm encryptionAlgorithm, KeyManagementAlgorithm algorithm)
@@ -38,9 +41,9 @@ public override bool TryUnwrapKey(ReadOnlySpan<byte> keyBytes, Span<byte> destin
             }
 
             var rawIV = encodedIV.GetRawValue();
-            Span<byte> nonce = stackalloc byte[Base64Url.GetArraySizeRequiredToDecode(rawIV.Length)];
+            Span<byte> nonce = stackalloc byte[IVSize];
             var rawTag = encodedTag.GetRawValue();
-            Span<byte> tag = stackalloc byte[Base64Url.GetArraySizeRequiredToDecode(rawTag.Length)];
+            Span<byte> tag = stackalloc byte[TagSize];
             try
             {
                 Base64Url.Decode(rawIV.Span, nonce);

diff --git a/src/JsonWebToken/Cryptography/AesKeyUnwrapper.cs b/src/JsonWebToken/Cryptography/AesKeyUnwrapper.cs
@@ -12,6 +12,7 @@ namespace JsonWebToken.Cryptography
     internal sealed class AesKeyUnwrapper : KeyUnwrapper
     {
         private const int BlockSizeInBytes = 8;
+        private const int KeyWrappedSizeThreshold = 64;
 
         // The default initialization vector from RFC3394
         private const ulong _defaultIV = 0XA6A6A6A6A6A6A6A6;
@@ -85,7 +86,7 @@ public override bool TryUnwrapKey(ReadOnlySpan<byte> key, Span<byte> destination
             int n = key.Length - BlockSizeInBytes;
 
             // The set of input blocks
-            Span<byte> r = stackalloc byte[n];
+            Span<byte> r = stackalloc byte[KeyWrappedSizeThreshold].Slice(0, n);
             r.Clear();
             ref byte rRef = ref MemoryMarshal.GetReference(r);
             Unsafe.CopyBlockUnaligned(ref rRef, ref Unsafe.AddByteOffset(ref input, (IntPtr)8), (uint)n);

diff --git a/src/JsonWebToken/Cryptography/AesKeyWrapper.cs b/src/JsonWebToken/Cryptography/AesKeyWrapper.cs
@@ -12,6 +12,7 @@ namespace JsonWebToken.Cryptography
     internal sealed class AesKeyWrapper : KeyWrapper
     {
         private const int BlockSizeInBytes = 8;
+        private const int KeyWrappedSizeThreshold = 64;
 
         // The default initialization vector from RFC3394
         private const ulong _defaultIV = 0XA6A6A6A6A6A6A6A6;
@@ -84,7 +85,7 @@ public override SymmetricJwk WrapKey(Jwk? staticKey, JwtHeader header, Span<byte
 
             ulong a = _defaultIV;
             ref byte input = ref MemoryMarshal.GetReference(inputBuffer);
-            Span<byte> r = stackalloc byte[n];
+            Span<byte> r = stackalloc byte[KeyWrappedSizeThreshold].Slice(0, n);
             ref byte rRef = ref MemoryMarshal.GetReference(r);
             Unsafe.CopyBlockUnaligned(ref rRef, ref input, (uint)n);
             TryWrapKey(ref a, n, ref rRef);

diff --git a/src/JsonWebToken/Cryptography/AuthenticatedEncryptor.cs b/src/JsonWebToken/Cryptography/AuthenticatedEncryptor.cs
@@ -30,5 +30,9 @@ public abstract class AuthenticatedEncryptor
 
         /// <summary>Gets the required size of the base64-URL nonce.</summary>
         public abstract int GetBase64NonceSize();
+
+        internal const int TagSizeStackallocThreshold = 64;
+
+        internal const int NonceSizeStackallocThreshold = 16;
     }
 }
diff --git a/src/JsonWebToken/Cryptography/EcdsaSignatureVerifier.cs b/src/JsonWebToken/Cryptography/EcdsaSignatureVerifier.cs
@@ -50,7 +50,7 @@ public override bool Verify(ReadOnlySpan<byte> data, ReadOnlySpan<byte> signatur
 
             var ecdsa = _ecdsaPool.Get();
 #if SUPPORT_SPAN_CRYPTO
-            Span<byte> hash = stackalloc byte[_sha.HashSize];
+            Span<byte> hash = stackalloc byte[Sha2.BlockSizeStackallocThreshold].Slice(0, _sha.HashSize);
             _sha.ComputeHash(data, hash);
             return ecdsa.VerifyHash(hash, signature);
 #else

diff --git a/src/JsonWebToken/Cryptography/EcdsaSigner.cs b/src/JsonWebToken/Cryptography/EcdsaSigner.cs
@@ -3,6 +3,7 @@
 
 #if SUPPORT_ELLIPTIC_CURVE_SIGNATURE
 using System;
+using System.Buffers;
 using System.Diagnostics;
 using System.Security.Cryptography;
 
@@ -45,9 +46,10 @@ public override bool TrySign(ReadOnlySpan<byte> data, Span<byte> destination, ou
 
             var ecdsa = _ecdsaPool.Get();
 #if SUPPORT_SPAN_CRYPTO
-            Span<byte> hash = stackalloc byte[_sha.HashSize];
-            _sha.ComputeHash(data, hash);
-            return ecdsa.TrySignHash(hash, destination, out bytesWritten);
+                Span<byte> hash = stackalloc byte[Sha2.HashSizeStackallocThreshold];
+                hash = hash.Slice(0, _sha.HashSize);
+                _sha.ComputeHash(data, hash);
+                return ecdsa.TrySignHash(hash, destination, out bytesWritten);
 #else
             byte[] hash = new byte[_sha.HashSize];
             _sha.ComputeHash(data, hash);
@@ -72,6 +74,6 @@ protected override void Dispose(bool disposing)
             }
         }
 
-    }   
+    }
 }
 #endif