Pool `MemoryStream`, `BinaryReader`, `BinaryWriter`, and `SHA256` instances

- #23 part 2
- reduce `byte[]` and `char[]` allocations because all have internal buffers
 - fortunately, only `MemoryStream` has an unbounded buffer

src/Microsoft.AspNetCore.Antiforgery/AntiforgerySerializationContext.cs

@@ -0,0 +1,109 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.IO;
+using System.Security.Cryptography;
+using System.Text;
+
+namespace Microsoft.AspNetCore.Antiforgery
+{
+    public class AntiforgerySerializationContext
+    {
+        // Avoid allocating 256 bytes (the default) and using 18 (the AntiforgeryToken minimum). 64 bytes is enough for
+        // a short username or claim UID and some additional data. MemoryStream bumps capacity to 256 if exceeded.
+        private const int InitialStreamSize = 64;
+
+        // Don't let the MemoryStream grow beyond 1 MB.
+        private const int MaximumStreamSize = 0x100000;
+
+        private MemoryStream _memory;
+        private BinaryReader _reader;
+        private BinaryWriter _writer;
+        private SHA256 _sha256;
+
+        public MemoryStream Stream
+        {
+            get
+            {
+                if (_memory == null)
+                {
+                    _memory = new MemoryStream(InitialStreamSize);
+                }
+
+                return _memory;
+            }
+            private set
+            {
+                _memory = value;
+            }
+        }
+
+        public BinaryReader Reader
+        {
+            get
+            {
+                if (_reader == null)
+                {
+                    // Leave open to clean up correctly even if only one of the reader or writer has been created.
+                    _reader = new BinaryReader(Stream, Encoding.UTF8, leaveOpen: true);
+                }
+
+                return _reader;
+            }
+            private set
+            {
+                _reader = value;
+            }
+        }
+
+        public BinaryWriter Writer
+        {
+            get
+            {
+                if (_writer == null)
+                {
+                    // Leave open to clean up correctly even if only one of the reader or writer has been created.
+                    _writer = new BinaryWriter(Stream, Encoding.UTF8, leaveOpen: true);
+                }
+
+                return _writer;
+            }
+            private set
+            {
+                _writer = value;
+            }
+        }
+
+        public SHA256 Sha256
+        {
+            get
+            {
+                if (_sha256 == null)
+                {
+                    _sha256 = SHA256.Create();
+                }
+
+                return _sha256;
+            }
+            private set
+            {
+                _sha256 = value;
+            }
+        }
+
+        public void Reset()
+        {
+            if (Stream.Capacity > MaximumStreamSize)
+            {
+                Stream = null;
+                Reader = null;
+                Writer = null;
+            }
+            else
+            {
+                Stream.Position = 0L;
+                Stream.SetLength(0L);
+            }
+        }
+    }
+}

src/Microsoft.AspNetCore.Antiforgery/BinaryBlob.cs

@@ -72,8 +72,8 @@ namespace Microsoft.AspNetCore.Antiforgery
                 return false;
             }
 
-            Debug.Assert(this._data.Length == other._data.Length);
+            Debug.Assert(_data.Length == other._data.Length);
-            return AreByteArraysEqual(this._data, other._data);
+            return AreByteArraysEqual(_data, other._data);
         }
 
         public byte[] GetData()

src/Microsoft.AspNetCore.Antiforgery/DefaultAntiforgeryTokenSerializer.cs

@@ -5,42 +5,54 @@ using System;
 using System.IO;
 using Microsoft.AspNetCore.DataProtection;
 using Microsoft.AspNetCore.WebUtilities;
+using Microsoft.Extensions.ObjectPool;
 
 namespace Microsoft.AspNetCore.Antiforgery
 {
     public class DefaultAntiforgeryTokenSerializer : IAntiforgeryTokenSerializer
     {
         private static readonly string Purpose = "Microsoft.AspNetCore.Antiforgery.AntiforgeryToken.v1";
-
-        private readonly IDataProtector _cryptoSystem;
         private const byte TokenVersion = 0x01;
 
-        public DefaultAntiforgeryTokenSerializer(IDataProtectionProvider provider)
+        private readonly IDataProtector _cryptoSystem;
+        private readonly ObjectPool<AntiforgerySerializationContext> _pool;
+
+        public DefaultAntiforgeryTokenSerializer(
+            IDataProtectionProvider provider,
+            ObjectPool<AntiforgerySerializationContext> pool)
         {
             if (provider == null)
             {
                 throw new ArgumentNullException(nameof(provider));
             }
 
+            if (pool == null)
+            {
+                throw new ArgumentNullException(nameof(pool));
+            }
+
             _cryptoSystem = provider.CreateProtector(Purpose);
+            _pool = pool;
         }
 
         public AntiforgeryToken Deserialize(string serializedToken)
         {
+            var serializationContext = _pool.Get();
+
             Exception innerException = null;
             try
             {
                 var tokenBytes = WebEncoders.Base64UrlDecode(serializedToken);
-                using (var stream = new MemoryStream(_cryptoSystem.Unprotect(tokenBytes)))
+                var unprotectedBytes = _cryptoSystem.Unprotect(tokenBytes);
+                var stream = serializationContext.Stream;
+                stream.Write(unprotectedBytes, 0, unprotectedBytes.Length);
+                stream.Position = 0L;
+
+                var reader = serializationContext.Reader;
+                var token = Deserialize(reader);
+                if (token != null)
                 {
-                    using (var reader = new BinaryReader(stream))
+                    return token;
-                    {
-                        var token = DeserializeImpl(reader);
-                        if (token != null)
-                        {
-                            return token;
-                        }
-                    }
                 }
             }
             catch (Exception ex)
@@ -48,6 +60,10 @@ namespace Microsoft.AspNetCore.Antiforgery
                 // swallow all exceptions - homogenize error if something went wrong
                 innerException = ex;
             }
+            finally
+            {
+                _pool.Return(serializationContext);
+            }
 
             // if we reached this point, something went wrong deserializing
             throw new InvalidOperationException(Resources.AntiforgeryToken_DeserializationFailed, innerException);
@@ -65,7 +81,7 @@ namespace Microsoft.AspNetCore.Antiforgery
          *   |    `- Username: UTF-8 string with 7-bit integer length prefix
          *   `- AdditionalData: UTF-8 string with 7-bit integer length prefix
          */
-        private static AntiforgeryToken DeserializeImpl(BinaryReader reader)
+        private static AntiforgeryToken Deserialize(BinaryReader reader)
         {
             // we can only consume tokens of the same serialized version that we generate
             var embeddedVersion = reader.ReadByte();
@@ -113,33 +129,38 @@ namespace Microsoft.AspNetCore.Antiforgery
                 throw new ArgumentNullException(nameof(token));
             }
 
-            using (var stream = new MemoryStream())
+            var serializationContext = _pool.Get();
+
+            try
             {
-                using (var writer = new BinaryWriter(stream))
+                var writer = serializationContext.Writer;
+                writer.Write(TokenVersion);
+                writer.Write(token.SecurityToken.GetData());
+                writer.Write(token.IsCookieToken);
+
+                if (!token.IsCookieToken)
                 {
-                    writer.Write(TokenVersion);
+                    if (token.ClaimUid != null)
-                    writer.Write(token.SecurityToken.GetData());
-                    writer.Write(token.IsCookieToken);
-
-                    if (!token.IsCookieToken)
                     {
-                        if (token.ClaimUid != null)
+                        writer.Write(true /* isClaimsBased */);
-                        {
+                        writer.Write(token.ClaimUid.GetData());
-                            writer.Write(true /* isClaimsBased */);
+                    }
-                            writer.Write(token.ClaimUid.GetData());
+                    else
-                        }
+                    {
-                        else
+                        writer.Write(false /* isClaimsBased */);
-                        {
+                        writer.Write(token.Username);
-                            writer.Write(false /* isClaimsBased */);
-                            writer.Write(token.Username);
-                        }
-
-                        writer.Write(token.AdditionalData);
                     }
 
-                    writer.Flush();
+                    writer.Write(token.AdditionalData);
-                    return WebEncoders.Base64UrlEncode(_cryptoSystem.Protect(stream.ToArray()));
                 }
+
+                writer.Flush();
+                var stream = serializationContext.Stream;
+                return WebEncoders.Base64UrlEncode(_cryptoSystem.Protect(stream.ToArray()));
+            }
+            finally
+            {
+                _pool.Return(serializationContext);
             }
         }
     }

src/Microsoft.AspNetCore.Antiforgery/DefaultClaimUidExtractor.cs

@@ -3,10 +3,9 @@
 
 using System;
 using System.Collections.Generic;
-using System.IO;
 using System.Linq;
 using System.Security.Claims;
-using System.Security.Cryptography;
+using Microsoft.Extensions.ObjectPool;
 
 namespace Microsoft.AspNetCore.Antiforgery
 {
@@ -15,6 +14,13 @@ namespace Microsoft.AspNetCore.Antiforgery
     /// </summary>
     public class DefaultClaimUidExtractor : IClaimUidExtractor
     {
+        private readonly ObjectPool<AntiforgerySerializationContext> _pool;
+
+        public DefaultClaimUidExtractor(ObjectPool<AntiforgerySerializationContext> pool)
+        {
+            _pool = pool;
+        }
+
         /// <inheritdoc />
         public string ExtractClaimUid(ClaimsIdentity claimsIdentity)
         {
@@ -25,16 +31,15 @@ namespace Microsoft.AspNetCore.Antiforgery
             }
 
             var uniqueIdentifierParameters = GetUniqueIdentifierParameters(claimsIdentity);
-            var claimUidBytes = ComputeSHA256(uniqueIdentifierParameters);
+            var claimUidBytes = ComputeSha256(uniqueIdentifierParameters);
             return Convert.ToBase64String(claimUidBytes);
         }
 
         // Internal for testing
         internal static IEnumerable<string> GetUniqueIdentifierParameters(ClaimsIdentity claimsIdentity)
         {
-            var nameIdentifierClaim = claimsIdentity.FindFirst(claim =>
+            var nameIdentifierClaim = claimsIdentity.FindFirst(
-                                                            String.Equals(ClaimTypes.NameIdentifier,
+                claim => string.Equals(ClaimTypes.NameIdentifier, claim.Type, StringComparison.Ordinal));
-                                                                        claim.Type, StringComparison.Ordinal));
             if (nameIdentifierClaim != null && !string.IsNullOrEmpty(nameIdentifierClaim.Value))
             {
                 return new string[]
@@ -47,7 +52,8 @@ namespace Microsoft.AspNetCore.Antiforgery
             // We do not understand this ClaimsIdentity, fallback on serializing the entire claims Identity.
             var claims = claimsIdentity.Claims.ToList();
             claims.Sort((a, b) => string.Compare(a.Type, b.Type, StringComparison.Ordinal));
-            var identifierParameters = new List<string>();
+
+            var identifierParameters = new List<string>(claims.Count * 2);
             foreach (var claim in claims)
             {
                 identifierParameters.Add(claim.Type);
@@ -57,25 +63,29 @@ namespace Microsoft.AspNetCore.Antiforgery
             return identifierParameters;
         }
 
-        private static byte[] ComputeSHA256(IEnumerable<string> parameters)
+        private byte[] ComputeSha256(IEnumerable<string> parameters)
         {
-            using (var stream = new MemoryStream())
+            var serializationContext = _pool.Get();
+
+            try
             {
-                using (var writer = new BinaryWriter(stream))
+                var writer = serializationContext.Writer;
+                foreach (string parameter in parameters)
                 {
-                    foreach (string parameter in parameters)
+                    writer.Write(parameter); // also writes the length as a prefix; unambiguous
-                    {
-                        writer.Write(parameter); // also writes the length as a prefix; unambiguous
-                    }
-
-                    writer.Flush();
-
-                    using (var sha256 = SHA256.Create())
-                    {
-                        var bytes = sha256.ComputeHash(stream.ToArray(), 0, checked((int)stream.Length));
-                        return bytes;
-                    }
                 }
+
+                writer.Flush();
+
+                var sha256 = serializationContext.Sha256;
+                var stream = serializationContext.Stream;
+                var bytes = sha256.ComputeHash(stream.ToArray(), 0, checked((int)stream.Length));
+
+                return bytes;
+            }
+            finally
+            {
+                _pool.Return(serializationContext);
             }
         }
     }

src/Microsoft.AspNetCore.Antiforgery/Internal/AntiforgerySerializationContextPooledObjectPolicy.cs

@@ -0,0 +1,23 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using Microsoft.Extensions.ObjectPool;
+
+namespace Microsoft.AspNetCore.Antiforgery.Internal
+{
+    public class AntiforgerySerializationContextPooledObjectPolicy
+        : IPooledObjectPolicy<AntiforgerySerializationContext>
+    {
+        public AntiforgerySerializationContext Create()
+        {
+            return new AntiforgerySerializationContext();
+        }
+
+        public bool Return(AntiforgerySerializationContext obj)
+        {
+            obj.Reset();
+
+            return true;
+        }
+    }
+}

src/Microsoft.AspNetCore.Antiforgery/ServiceCollectionExtensions.cs

@@ -3,7 +3,9 @@
 
 using System;
 using Microsoft.AspNetCore.Antiforgery;
+using Microsoft.AspNetCore.Antiforgery.Internal;
 using Microsoft.Extensions.DependencyInjection.Extensions;
+using Microsoft.Extensions.ObjectPool;
 using Microsoft.Extensions.Options;
 
 namespace Microsoft.Extensions.DependencyInjection
@@ -30,6 +32,15 @@ namespace Microsoft.Extensions.DependencyInjection
             services.TryAddSingleton<IClaimUidExtractor, DefaultClaimUidExtractor>();
             services.TryAddScoped<IAntiforgeryContextAccessor, DefaultAntiforgeryContextAccessor>();
             services.TryAddSingleton<IAntiforgeryAdditionalDataProvider, DefaultAntiforgeryAdditionalDataProvider>();
+
+            services.TryAddSingleton<ObjectPoolProvider>(new DefaultObjectPoolProvider());
+            services.TryAddSingleton<ObjectPool<AntiforgerySerializationContext>>(serviceProvider =>
+            {
+                var provider = serviceProvider.GetRequiredService<ObjectPoolProvider>();
+                var policy = new AntiforgerySerializationContextPooledObjectPolicy();
+                return provider.Create(policy);
+            });
+
             return services;
         }
 

src/Microsoft.AspNetCore.Antiforgery/project.json

@@ -13,7 +13,8 @@
     "Microsoft.AspNetCore.DataProtection": "1.0.0-*",
     "Microsoft.AspNetCore.Html.Abstractions": "1.0.0-*",
     "Microsoft.AspNetCore.Http.Abstractions": "1.0.0-*",
-    "Microsoft.AspNetCore.WebUtilities": "1.0.0-*"
+    "Microsoft.AspNetCore.WebUtilities": "1.0.0-*",
+    "Microsoft.Extensions.ObjectPool": "1.0.0-*"
   },
   "frameworks": {
     "dotnet5.4": { },

test/Microsoft.AspNetCore.Antiforgery.Test/DefaultAntiforgeryTokenSerializerTest.cs

@@ -4,7 +4,9 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using Microsoft.AspNetCore.Antiforgery.Internal;
 using Microsoft.AspNetCore.DataProtection;
+using Microsoft.Extensions.ObjectPool;
 using Moq;
 using Xunit;
 
@@ -15,6 +17,8 @@ namespace Microsoft.AspNetCore.Antiforgery
         private static readonly Mock<IDataProtectionProvider> _dataProtector = GetDataProtector();
         private static readonly BinaryBlob _claimUid = new BinaryBlob(256, new byte[] { 0x6F, 0x16, 0x48, 0xE9, 0x72, 0x49, 0xAA, 0x58, 0x75, 0x40, 0x36, 0xA6, 0x7E, 0x24, 0x8C, 0xF0, 0x44, 0xF0, 0x7E, 0xCF, 0xB0, 0xED, 0x38, 0x75, 0x56, 0xCE, 0x02, 0x9A, 0x4F, 0x9A, 0x40, 0xE0 });
         private static readonly BinaryBlob _securityToken = new BinaryBlob(128, new byte[] { 0x70, 0x5E, 0xED, 0xCC, 0x7D, 0x42, 0xF1, 0xD6, 0xB3, 0xB9, 0x8A, 0x59, 0x36, 0x25, 0xBB, 0x4C });
+        private static readonly ObjectPool<AntiforgerySerializationContext> _pool =
+            new DefaultObjectPoolProvider().Create(new AntiforgerySerializationContextPooledObjectPolicy());
         private const byte _salt = 0x05;
 
         [Theory]
@@ -45,7 +49,7 @@ namespace Microsoft.AspNetCore.Antiforgery
         public void Deserialize_BadToken_Throws(string serializedToken)
         {
             // Arrange
-            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object);
+            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object, _pool);
 
             // Act & assert
             var ex = Assert.Throws<InvalidOperationException>(() => testSerializer.Deserialize(serializedToken));
@@ -56,7 +60,7 @@ namespace Microsoft.AspNetCore.Antiforgery
         public void Serialize_FieldToken_WithClaimUid_TokenRoundTripSuccessful()
         {
             // Arrange
-            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object);
+            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object, _pool);
 
             //"01" // Version
             //+ "705EEDCC7D42F1D6B3B98A593625BB4C" // SecurityToken
@@ -86,7 +90,7 @@ namespace Microsoft.AspNetCore.Antiforgery
         public void Serialize_FieldToken_WithUsername_TokenRoundTripSuccessful()
         {
             // Arrange
-            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object);
+            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object, _pool);
 
             //"01" // Version
             //+ "705EEDCC7D42F1D6B3B98A593625BB4C" // SecurityToken
@@ -117,7 +121,7 @@ namespace Microsoft.AspNetCore.Antiforgery
         public void Serialize_CookieToken_TokenRoundTripSuccessful()
         {
             // Arrange
-            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object);
+            var testSerializer = new DefaultAntiforgeryTokenSerializer(_dataProtector.Object, _pool);
 
             //"01" // Version
             //+ "705EEDCC7D42F1D6B3B98A593625BB4C" // SecurityToken

test/Microsoft.AspNetCore.Antiforgery.Test/DefaultAntiforgeryTokenStoreTest.cs

@@ -4,10 +4,12 @@
 using System;
 using System.Collections.Generic;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Antiforgery.Internal;
 using Microsoft.AspNetCore.DataProtection;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Internal;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.ObjectPool;
 using Microsoft.Extensions.Primitives;
 using Moq;
 using Xunit;
@@ -16,6 +18,8 @@ namespace Microsoft.AspNetCore.Antiforgery
 {
     public class DefaultAntiforgeryTokenStoreTest
     {
+        private static readonly ObjectPool<AntiforgerySerializationContext> _pool =
+            new DefaultObjectPoolProvider().Create(new AntiforgerySerializationContextPooledObjectPolicy());
         private readonly string _cookieName = "cookie-name";
 
         [Fact]
@@ -209,7 +213,7 @@ namespace Microsoft.AspNetCore.Antiforgery
 
             var tokenStore = new DefaultAntiforgeryTokenStore(
                 optionsAccessor: new TestOptionsManager(options),
-                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider()));
+                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider(), _pool));
 
             // Act
             var exception = await Assert.ThrowsAsync<AntiforgeryValidationException>(
@@ -244,7 +248,7 @@ namespace Microsoft.AspNetCore.Antiforgery
 
             var tokenStore = new DefaultAntiforgeryTokenStore(
                 optionsAccessor: new TestOptionsManager(options),
-                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider()));
+                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider(), _pool));
 
             // Act
             var tokens = await tokenStore.GetRequestTokensAsync(httpContext);
@@ -279,7 +283,7 @@ namespace Microsoft.AspNetCore.Antiforgery
 
             var tokenStore = new DefaultAntiforgeryTokenStore(
                 optionsAccessor: new TestOptionsManager(options),
-                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider()));
+                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider(), _pool));
 
             // Act
             var tokens = await tokenStore.GetRequestTokensAsync(httpContext);
@@ -312,7 +316,7 @@ namespace Microsoft.AspNetCore.Antiforgery
 
             var tokenStore = new DefaultAntiforgeryTokenStore(
                 optionsAccessor: new TestOptionsManager(options),
-                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider()));
+                tokenSerializer: new DefaultAntiforgeryTokenSerializer(new EphemeralDataProtectionProvider(), _pool));
 
             // Act
             var exception = await Assert.ThrowsAsync<AntiforgeryValidationException>(

test/Microsoft.AspNetCore.Antiforgery.Test/DefaultClaimUidExtractorTest.cs

@@ -4,6 +4,8 @@
 using System;
 using System.Linq;
 using System.Security.Claims;
+using Microsoft.AspNetCore.Antiforgery.Internal;
+using Microsoft.Extensions.ObjectPool;
 using Moq;
 using Xunit;
 
@@ -11,11 +13,14 @@ namespace Microsoft.AspNetCore.Antiforgery
 {
     public class DefaultClaimUidExtractorTest
     {
+        private static readonly ObjectPool<AntiforgerySerializationContext> _pool =
+            new DefaultObjectPoolProvider().Create(new AntiforgerySerializationContextPooledObjectPolicy());
+
         [Fact]
         public void ExtractClaimUid_NullIdentity()
         {
             // Arrange
-            IClaimUidExtractor extractor = new DefaultClaimUidExtractor();
+            var extractor = new DefaultClaimUidExtractor(_pool);
 
             // Act
             var claimUid = extractor.ExtractClaimUid(null);
@@ -28,7 +33,7 @@ namespace Microsoft.AspNetCore.Antiforgery
         public void ExtractClaimUid_Unauthenticated()
         {
             // Arrange
-            IClaimUidExtractor extractor = new DefaultClaimUidExtractor();
+            var extractor = new DefaultClaimUidExtractor(_pool);
 
             var mockIdentity = new Mock<ClaimsIdentity>();
             mockIdentity.Setup(o => o.IsAuthenticated)
@@ -49,7 +54,7 @@ namespace Microsoft.AspNetCore.Antiforgery
             mockIdentity.Setup(o => o.IsAuthenticated)
                         .Returns(true);
 
-            IClaimUidExtractor extractor = new DefaultClaimUidExtractor();
+            var extractor = new DefaultClaimUidExtractor(_pool);
 
             // Act
             var claimUid = extractor.ExtractClaimUid(mockIdentity.Object);