From 5aee5247685b3e3f9157e229dfc2c7ba9248dcb3 Mon Sep 17 00:00:00 2001
From: Stephen Halter <halter73@gmail.com>
Date: Thu, 28 Jul 2016 16:15:12 -0700
Subject: [PATCH] Ensure Begin/End Read/Write methods are correctly implemented

- Add Begin/End Write methods to FrameResponseStream
- Fix state check of Begin/End Read methods in FrameRequestStream
---
 .../Internal/Http/FrameRequestStream.cs       |  8 +---
 .../Internal/Http/FrameResponseStream.cs      | 40 +++++++++++++++++++
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameRequestStream.cs b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameRequestStream.cs
index e9e295a593..09a4834ef8 100644
--- a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameRequestStream.cs
+++ b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameRequestStream.cs
@@ -76,8 +76,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
 #if NET451
         public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback callback, object state)
         {
-            ValidateState(default(CancellationToken));
-
             var task = ReadAsync(buffer, offset, count, default(CancellationToken), state);
             if (callback != null)
             {
@@ -93,11 +91,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
 
         private Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken, object state)
         {
-            ValidateState(cancellationToken);
-
             var tcs = new TaskCompletionSource<int>(state);
-            var task = _body.ReadAsync(new ArraySegment<byte>(buffer, offset, count), cancellationToken);
-            task.AsTask().ContinueWith((task2, state2) =>
+            var task = ReadAsync(buffer, offset, count, cancellationToken);
+            task.ContinueWith((task2, state2) =>
             {
                 var tcs2 = (TaskCompletionSource<int>)state2;
                 if (task2.IsCanceled)
diff --git a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameResponseStream.cs b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameResponseStream.cs
index 61013ac16c..5a3c9159d5 100644
--- a/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameResponseStream.cs
+++ b/src/Microsoft.AspNetCore.Server.Kestrel/Internal/Http/FrameResponseStream.cs
@@ -85,6 +85,46 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Http
             _frameControl.Write(new ArraySegment<byte>(buffer, offset, count));
         }
 
+#if NET451
+        public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback callback, object state)
+        {
+            var task = WriteAsync(buffer, offset, count, default(CancellationToken), state);
+            if (callback != null)
+            {
+                task.ContinueWith(t => callback.Invoke(t));
+            }
+            return task;
+        }
+
+        public override void EndWrite(IAsyncResult asyncResult)
+        {
+            ((Task<object>)asyncResult).GetAwaiter().GetResult();
+        }
+
+        private Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken, object state)
+        {
+            var tcs = new TaskCompletionSource<object>(state);
+            var task = WriteAsync(buffer, offset, count, cancellationToken);
+            task.ContinueWith((task2, state2) =>
+            {
+                var tcs2 = (TaskCompletionSource<object>)state2;
+                if (task2.IsCanceled)
+                {
+                    tcs2.SetCanceled();
+                }
+                else if (task2.IsFaulted)
+                {
+                    tcs2.SetException(task2.Exception);
+                }
+                else
+                {
+                    tcs2.SetResult(null);
+                }
+            }, tcs, cancellationToken);
+            return tcs.Task;
+        }
+#endif
+
         public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
             var task = ValidateState(cancellationToken);