Fix for #3252 - Issues with pooled buffer + unicode

- Dispose buffers when Flush throws inside the Dispose method
- Compute size of buffers correctly
- Throw earlier when handed invalid-sized buffers

src/Microsoft.AspNet.Mvc.Core/HttpResponseStreamWriter.cs

@@ -5,6 +5,7 @@
 using System.IO;
 using System.Text;
 using System.Threading.Tasks;
+using Microsoft.AspNet.Mvc.Core;
 using Microsoft.Extensions.MemoryPool;
 
 namespace Microsoft.AspNet.Mvc
@@ -58,6 +59,7 @@ public HttpResponseStreamWriter(Stream stream, Encoding encoding, int bufferSize
         public HttpResponseStreamWriter(
             Stream stream,
             Encoding encoding,
+            int bufferSize,
             LeasedArraySegment<byte> leasedByteBuffer,
             LeasedArraySegment<char> leasedCharBuffer)
         {
@@ -81,21 +83,27 @@ public HttpResponseStreamWriter(
                 throw new ArgumentNullException(nameof(leasedCharBuffer));
             }
 
+            var requiredLength = encoding.GetMaxByteCount(bufferSize);
+            if (requiredLength > leasedByteBuffer.Data.Count)
+            {
+                var message = Resources.FormatHttpResponseStreamWriter_InvalidBufferSize(
+                    requiredLength,
+                    bufferSize,
+                    encoding.EncodingName,
+                    typeof(Encoding).FullName,
+                    nameof(Encoding.GetMaxByteCount));
+                throw new ArgumentException(message, nameof(leasedByteBuffer));
+            }
+
             _stream = stream;
             Encoding = encoding;
+            _charBufferSize = bufferSize;
             _leasedByteBuffer = leasedByteBuffer;
             _leasedCharBuffer = leasedCharBuffer;
 
             _encoder = encoding.GetEncoder();
             _byteBuffer = leasedByteBuffer.Data;
             _charBuffer = leasedCharBuffer.Data;
-
-            // We need to compute the usable size of the char buffer based on the size of the byte buffer.
-            // Encoder.GetBytes assumes that the entirety of the byte[] passed in can be used, and that's not the
-            // case with ArraySegments.
-            _charBufferSize = Math.Min(
-                leasedCharBuffer.Data.Count,
-                encoding.GetMaxCharCount(leasedByteBuffer.Data.Count));
         }
 
         public override Encoding Encoding { get; }
@@ -215,16 +223,24 @@ public override Task FlushAsync()
         // sent in chunked encoding in case of Helios.
         protected override void Dispose(bool disposing)
         {
-            FlushInternal(flushStream: false, flushEncoder: true);
-
-            if (_leasedByteBuffer != null)
+            if (disposing)
             {
-                _leasedByteBuffer.Owner.Return(_leasedByteBuffer);
-            }
-
-            if (_leasedCharBuffer != null)
-            {
-                _leasedCharBuffer.Owner.Return(_leasedCharBuffer);
+                try
+                {
+                    FlushInternal(flushStream: false, flushEncoder: true);
+                }
+                finally
+                {
+                    if (_leasedByteBuffer != null)
+                    {
+                        _leasedByteBuffer.Owner.Return(_leasedByteBuffer);
+                    }
+
+                    if (_leasedCharBuffer != null)
+                    {
+                        _leasedCharBuffer.Owner.Return(_leasedCharBuffer);
+                    }
+                }
             }
         }
 

src/Microsoft.AspNet.Mvc.Core/Infrastructure/MemoryPoolHttpResponseStreamWriterFactory.cs

@@ -14,9 +14,9 @@ namespace Microsoft.AspNet.Mvc.Infrastructure
     public class MemoryPoolHttpResponseStreamWriterFactory : IHttpResponseStreamWriterFactory
     {
         /// <summary>
-        /// The default size of created buffers.
+        /// The default size of created char buffers.
         /// </summary>
-        public static readonly int DefaultBufferSize = 4 * 1024; // 4KB
+        public static readonly int DefaultBufferSize = 1024; // 1KB - results in a 4KB byte array for UTF8.
 
         private readonly IArraySegmentPool<byte> _bytePool;
         private readonly IArraySegmentPool<char> _charPool;
@@ -66,10 +66,14 @@ public TextWriter CreateWriter(Stream stream, Encoding encoding)
 
             try
             {
-                bytes = _bytePool.Lease(DefaultBufferSize);
                 chars = _charPool.Lease(DefaultBufferSize);
 
-                return new HttpResponseStreamWriter(stream, encoding, bytes, chars);
+                // We need to compute the minimum size of the byte buffer based on the size of the char buffer,
+                // so that we have enough room to encode the buffer in one shot.
+                var minimumSize = encoding.GetMaxByteCount(DefaultBufferSize);
+                bytes = _bytePool.Lease(minimumSize);
+
+                return new HttpResponseStreamWriter(stream, encoding, DefaultBufferSize, bytes, chars);
             }
             catch
             {

src/Microsoft.AspNet.Mvc.Core/Resources.resx

@@ -312,4 +312,7 @@
   <data name="ValueProviderResult_NoConverterExists" xml:space="preserve">
     <value>The parameter conversion from type '{0}' to type '{1}' failed because no type converter can convert between these types.</value>
   </data>
+  <data name="HttpResponseStreamWriter_InvalidBufferSize" xml:space="preserve">
+    <value>The byte buffer must have a length of at least '{0}' to be used with a char buffer of size '{1}' and encoding '{2}'. Use '{3}.{4}' to compute the correct size for the byte buffer.</value>
+  </data>
 </root>

test/Microsoft.AspNet.Mvc.Core.Test/HttpResponseStreamWriterTest.cs

@@ -6,6 +6,7 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNet.Testing;
 using Microsoft.Extensions.MemoryPool;
 using Xunit;
 
@@ -381,9 +382,9 @@ public void HttpResponseStreamWriter_UsingPooledBuffers()
                     try
                     {
                         bytes = bytePool.Lease(4096);
-                        chars = charPool.Lease(4096);
+                        chars = charPool.Lease(1024);
 
-                        writer = new HttpResponseStreamWriter(stream, encoding, bytes, chars);
+                        writer = new HttpResponseStreamWriter(stream, encoding, 1024, bytes, chars);
                     }
                     catch
                     {
@@ -412,36 +413,39 @@ public void HttpResponseStreamWriter_UsingPooledBuffers()
             Assert.Equal(expectedBytes, stream.ToArray());
         }
 
-        // This covers the case where we need to limit the usable region of the char buffer
-        // based on the size of the byte buffer. See comments in the constructor.
+        // This covers the error case where the byte buffer is too small. This is a safeguard, and shouldn't happen 
+        // if we're using the writer factory.
         [Fact]
         public void HttpResponseStreamWriter_UsingPooledBuffers_SmallByteBuffer()
         {
             // Arrange
             var encoding = Encoding.UTF8;
             var stream = new MemoryStream();
 
-            var charBufferSize = encoding.GetMaxCharCount(1024);
-
-            // This content is bigger than the byte buffer can hold, so it will need to be split
-            // into two separate encoding operations.
-            var content = new string('a', charBufferSize + 1);
-            var expectedBytes = encoding.GetBytes(content);
+            var message =
+                "The byte buffer must have a length of at least '12291' to be used with a char buffer of " +
+                "size '4096' and encoding 'Unicode (UTF-8)'. Use 'System.Text.Encoding.GetMaxByteCount' " +
+                "to compute the correct size for the byte buffer.";
 
             using (var bytePool = new DefaultArraySegmentPool<byte>())
             {
                 using (var charPool = new DefaultArraySegmentPool<char>())
                 {
                     LeasedArraySegment<byte> bytes = null;
                     LeasedArraySegment<char> chars = null;
-                    HttpResponseStreamWriter writer;
+                    HttpResponseStreamWriter writer = null;
 
                     try
                     {
                         bytes = bytePool.Lease(1024);
                         chars = charPool.Lease(4096);
 
-                        writer = new HttpResponseStreamWriter(stream, encoding, bytes, chars);
+                        // Act & Assert
+                        ExceptionAssert.ThrowsArgument(
+                            () => writer = new HttpResponseStreamWriter(stream, encoding, chars.Data.Count, bytes, chars),
+                            "byteBuffer",
+                            message);
+                        writer.Dispose();
                     }
                     catch
                     {
@@ -454,29 +458,9 @@ public void HttpResponseStreamWriter_UsingPooledBuffers_SmallByteBuffer()
                         {
                             chars.Owner.Return(chars);
                         }
-
-                        throw;
-                    }
-
-                    // Zero the byte buffer because we're going to examine it.
-                    Array.Clear(bytes.Data.Array, 0, bytes.Data.Array.Length);
-
-                    // Act
-                    using (writer)
-                    {
-                        writer.Write(content);
-                    }
-
-                    // Verify that we didn't buffer overflow 'our' region of the underlying array.
-                    if (bytes.Data.Array.Length > bytes.Data.Count)
-                    {
-                        Assert.Equal((byte)0, bytes.Data.Array[bytes.Data.Count]);
                     }
                 }
             }
-
-            // Assert
-            Assert.Equal(expectedBytes, stream.ToArray());
         }
 
         private class TestMemoryStream : MemoryStream

test/Microsoft.AspNet.Mvc.Core.Test/Infrastructure/MemoryPoolHttpResponseStreamWriterFactoryTest.cs

@@ -18,7 +18,7 @@ public void CreateWriter_BuffersReturned_OnException()
             // Arrange
             var bytePool = new Mock<IArraySegmentPool<byte>>(MockBehavior.Strict);
             bytePool
-                .Setup(p => p.Lease(MemoryPoolHttpResponseStreamWriterFactory.DefaultBufferSize))
+                .Setup(p => p.Lease(It.IsAny<int>()))
                 .Returns(new LeasedArraySegment<byte>(new ArraySegment<byte>(new byte[0]), bytePool.Object));
             bytePool
                 .Setup(p => p.Return(It.IsAny<LeasedArraySegment<byte>>()))
@@ -32,7 +32,7 @@ public void CreateWriter_BuffersReturned_OnException()
                 .Setup(p => p.Return(It.IsAny<LeasedArraySegment<char>>()))
                 .Verifiable();
 
-            var encoding = new Mock<Encoding>(MockBehavior.Strict);
+            var encoding = new Mock<Encoding>();
             encoding
                 .Setup(e => e.GetEncoder())
                 .Throws(new InvalidOperationException());