React to ANCM OnAsyncCompletion changes (#447)

samples/NativeIISSample/Startup.cs

@@ -1,7 +1,8 @@
+// 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;
-using System.Collections.Generic;
 using System.Linq;
-using System.Threading.Tasks;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;

src/Microsoft.AspNetCore.Server.IISIntegration/NativeMethods.cs

@@ -34,23 +34,27 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
         public delegate REQUEST_NOTIFICATION_STATUS PFN_REQUEST_HANDLER(IntPtr pHttpContext, IntPtr pvRequestContext);
         public delegate bool PFN_SHUTDOWN_HANDLER(IntPtr pvRequestContext);
-        public delegate REQUEST_NOTIFICATION_STATUS PFN_ASYNC_COMPLETION(IntPtr pHttpContext, IntPtr completionInfo, IntPtr pvCompletionContext);
+        public delegate REQUEST_NOTIFICATION_STATUS PFN_ASYNC_COMPLETION(IntPtr pvManagedHttpContext, int hr, int bytes);
+        public delegate REQUEST_NOTIFICATION_STATUS PFN_WEBSOCKET_ASYNC_COMPLETION(IntPtr pHttpContext, IntPtr completionInfo, IntPtr pvCompletionContext);
 
         // TODO make this all internal
         [DllImport(AspNetCoreModuleDll)]
         public static extern int http_post_completion(IntPtr pHttpContext, int cbBytes);
 
+        [DllImport(AspNetCoreModuleDll)]
+        public static extern int http_set_completion_status(IntPtr pHttpContext, REQUEST_NOTIFICATION_STATUS rquestNotificationStatus);
+
         [DllImport(AspNetCoreModuleDll)]
         public static extern void http_indicate_completion(IntPtr pHttpContext, REQUEST_NOTIFICATION_STATUS notificationStatus);
 
         [DllImport(AspNetCoreModuleDll)]
-        public static extern void register_callbacks(PFN_REQUEST_HANDLER request_callback, PFN_SHUTDOWN_HANDLER shutdown_callback, IntPtr pvRequestContext, IntPtr pvShutdownContext);
+        public static extern void register_callbacks(PFN_REQUEST_HANDLER request_callback, PFN_SHUTDOWN_HANDLER shutdown_callback, PFN_ASYNC_COMPLETION managed_context_handler, IntPtr pvRequestContext, IntPtr pvShutdownContext);
 
         [DllImport(AspNetCoreModuleDll)]
-        internal unsafe static extern int http_write_response_bytes(IntPtr pHttpContext, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, PFN_ASYNC_COMPLETION pfnCompletionCallback, IntPtr pvCompletionContext, out bool fCompletionExpected);
+        internal unsafe static extern int http_write_response_bytes(IntPtr pHttpContext, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
-        public unsafe static extern int http_flush_response_bytes(IntPtr pHttpContext, PFN_ASYNC_COMPLETION pfnCompletionCallback, IntPtr pvCompletionContext, out bool fCompletionExpected);
+        public unsafe static extern int http_flush_response_bytes(IntPtr pHttpContext, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
         internal unsafe static extern HttpApiTypes.HTTP_REQUEST_V2* http_get_raw_request(IntPtr pHttpContext);
@@ -62,11 +66,14 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
         public unsafe static extern void http_set_response_status_code(IntPtr pHttpContext, ushort statusCode, byte* pszReason);
 
         [DllImport(AspNetCoreModuleDll)]
-        public unsafe static extern int http_read_request_bytes(IntPtr pHttpContext, byte* pvBuffer, int cbBuffer, PFN_ASYNC_COMPLETION pfnCompletionCallback, IntPtr pvCompletionContext, out int dwBytesReceived, out bool fCompletionExpected);
+        public unsafe static extern int http_read_request_bytes(IntPtr pHttpContext, byte* pvBuffer, int cbBuffer, out int dwBytesReceived, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
         public unsafe static extern bool http_get_completion_info(IntPtr pCompletionInfo, out int cbBytes, out int hr);
 
+        [DllImport(AspNetCoreModuleDll)]
+        public unsafe static extern bool http_set_managed_context(IntPtr pHttpContext, IntPtr pvManagedContext);
+
         [DllImport(AspNetCoreModuleDll)]
         [return: MarshalAs(UnmanagedType.BStr)]
         public unsafe static extern string http_get_application_full_path();
@@ -74,6 +81,18 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
         [DllImport(AspNetCoreModuleDll)]
         public unsafe static extern bool http_shutdown();
 
+        [DllImport(AspNetCoreModuleDll)]
+        public unsafe static extern int http_websockets_read_bytes(IntPtr pHttpContext, byte* pvBuffer, int cbBuffer, PFN_WEBSOCKET_ASYNC_COMPLETION pfnCompletionCallback, IntPtr pvCompletionContext, out int dwBytesReceived, out bool fCompletionExpected);
+
+        [DllImport(AspNetCoreModuleDll)]
+        internal unsafe static extern int http_websockets_write_bytes(IntPtr pHttpContext, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, PFN_WEBSOCKET_ASYNC_COMPLETION pfnCompletionCallback, IntPtr pvCompletionContext, out bool fCompletionExpected);
+
+        [DllImport(AspNetCoreModuleDll)]
+        public unsafe static extern int http_enable_websockets(IntPtr pHttpContext);
+
+        [DllImport(AspNetCoreModuleDll)]
+        public unsafe static extern int http_cancel_io(IntPtr pHttpContext);
+
         [DllImport("kernel32.dll")]
         public static extern IntPtr GetModuleHandle(string lpModuleName);
 

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISAwaitable.cs

@@ -21,35 +21,24 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
         private int _cbBytes;
 
-        public static readonly NativeMethods.PFN_ASYNC_COMPLETION ReadCallback = (IntPtr pHttpContext, IntPtr pCompletionInfo, IntPtr pvCompletionContext) =>
+        public static readonly NativeMethods.PFN_WEBSOCKET_ASYNC_COMPLETION ReadCallback = (IntPtr pHttpContext, IntPtr pCompletionInfo, IntPtr pvCompletionContext) =>
         {
-            var context = (HttpProtocol)GCHandle.FromIntPtr(pvCompletionContext).Target;
+            var context = (IISHttpContext)GCHandle.FromIntPtr(pvCompletionContext).Target;
 
             NativeMethods.http_get_completion_info(pCompletionInfo, out int cbBytes, out int hr);
 
-            context.CompleteRead(hr, cbBytes);
+            context.CompleteReadWebSockets(hr, cbBytes);
 
             return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_PENDING;
         };
 
-        public static readonly NativeMethods.PFN_ASYNC_COMPLETION WriteCallback = (IntPtr pHttpContext, IntPtr pCompletionInfo, IntPtr pvCompletionContext) =>
+        public static readonly NativeMethods.PFN_WEBSOCKET_ASYNC_COMPLETION WriteCallback = (IntPtr pHttpContext, IntPtr pCompletionInfo, IntPtr pvCompletionContext) =>
         {
-            var context = (HttpProtocol)GCHandle.FromIntPtr(pvCompletionContext).Target;
+            var context = (IISHttpContext)GCHandle.FromIntPtr(pvCompletionContext).Target;
 
             NativeMethods.http_get_completion_info(pCompletionInfo, out int cbBytes, out int hr);
 
-            context.CompleteWrite(hr, cbBytes);
+            context.CompleteWriteWebSockets(hr, cbBytes);
-
-            return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_PENDING;
-        };
-
-        public static readonly NativeMethods.PFN_ASYNC_COMPLETION FlushCallback = (IntPtr pHttpContext, IntPtr pCompletionInfo, IntPtr pvCompletionContext) =>
-        {
-            var context = (HttpProtocol)GCHandle.FromIntPtr(pvCompletionContext).Target;
-
-            NativeMethods.http_get_completion_info(pCompletionInfo, out int cbBytes, out int hr);
-
-            context.CompleteFlush(hr, cbBytes);
 
             return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_PENDING;
         };

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISHttpContext.FeatureCollection.cs

@@ -16,13 +16,13 @@ using Microsoft.Extensions.Primitives;
 
 namespace Microsoft.AspNetCore.Server.IISIntegration
 {
-    internal partial class HttpProtocol : IFeatureCollection,
+    internal partial class IISHttpContext : IFeatureCollection,
-                                          IHttpRequestFeature,
+                                            IHttpRequestFeature,
-                                          IHttpResponseFeature,
+                                            IHttpResponseFeature,
-                                          IHttpUpgradeFeature,
+                                            IHttpUpgradeFeature,
-                                          IHttpConnectionFeature,
+                                            IHttpConnectionFeature,
-                                          IHttpRequestLifetimeFeature,
+                                            IHttpRequestLifetimeFeature,
-                                          IHttpRequestIdentifierFeature
+                                            IHttpRequestIdentifierFeature
     {
         // NOTE: When feature interfaces are added to or removed from this HttpProtocol implementation,
         // then the list of `implementedFeatures` in the generated code project MUST also be updated.
@@ -266,11 +266,14 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
                 throw new InvalidOperationException("CoreStrings.UpgradeCannotBeCalledMultipleTimes");
             }
 
-            _wasUpgraded = true;
-
             StatusCode = StatusCodes.Status101SwitchingProtocols;
             ReasonPhrase = ReasonPhrases.GetReasonPhrase(StatusCodes.Status101SwitchingProtocols);
             await UpgradeAsync();
+            NativeMethods.http_enable_websockets(_pHttpContext);
+
+            _wasUpgraded = true;
+            _readWebSocketsOperation = new IISAwaitable();
+            _writeWebSocketsOperation = new IISAwaitable();
 
             return new DuplexStream(RequestBody, ResponseBody);
         }

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISHttpContext.Features.cs

@@ -8,7 +8,7 @@ using System.Threading.Tasks;
 
 namespace Microsoft.AspNetCore.Server.IISIntegration
 {
-    internal partial class HttpProtocol
+    internal partial class IISHttpContext
     {
         private static readonly Type IHttpRequestFeatureType = typeof(global::Microsoft.AspNetCore.Http.Features.IHttpRequestFeature);
         private static readonly Type IHttpResponseFeatureType = typeof(global::Microsoft.AspNetCore.Http.Features.IHttpResponseFeature);

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISHttpContext.cs

@@ -20,18 +20,20 @@ using Microsoft.Extensions.Primitives;
 
 namespace Microsoft.AspNetCore.Server.IISIntegration
 {
-    internal abstract partial class HttpProtocol : NativeRequestContext, IDisposable
+    internal abstract partial class IISHttpContext : NativeRequestContext, IDisposable
     {
         private const int MinAllocBufferSize = 2048;
 
         private static bool UpgradeAvailable = (Environment.OSVersion.Version >= new Version(6, 2));
 
         protected readonly IntPtr _pHttpContext;
+
         private bool _wasUpgraded;
         private int _statusCode;
         private string _reasonPhrase;
         private readonly object _onStartingSync = new object();
         private readonly object _onCompletedSync = new object();
+
         protected Stack<KeyValuePair<Func<object, Task>, object>> _onStarting;
         protected Stack<KeyValuePair<Func<object, Task>, object>> _onCompleted;
 
@@ -42,9 +44,10 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
         private GCHandle _thisHandle;
         private BufferHandle _inputHandle;
-        private IISAwaitable _readOperation = new IISAwaitable();
+        private IISAwaitable _operation = new IISAwaitable();
-        private IISAwaitable _writeOperation = new IISAwaitable();
+
-        private IISAwaitable _flushOperation = new IISAwaitable();
+        private IISAwaitable _readWebSocketsOperation;
+        private IISAwaitable _writeWebSocketsOperation;
 
         private TaskCompletionSource<object> _upgradeTcs;
 
@@ -53,14 +56,18 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
         protected int _requestAborted;
 
-        internal unsafe HttpProtocol(PipeFactory pipeFactory, IntPtr pHttpContext)
+        private CurrentOperationType _currentOperationType;
-            : base((HttpApiTypes.HTTP_REQUEST*) NativeMethods.http_get_raw_request(pHttpContext))
+        private Task _currentOperation = Task.CompletedTask;
+
+        internal unsafe IISHttpContext(PipeFactory pipeFactory, IntPtr pHttpContext)
+            : base((HttpApiTypes.HTTP_REQUEST*)NativeMethods.http_get_raw_request(pHttpContext))
         {
             _thisHandle = GCHandle.Alloc(this);
 
             _pipeFactory = pipeFactory;
             _pHttpContext = pHttpContext;
 
+            NativeMethods.http_set_managed_context(pHttpContext, (IntPtr)_thisHandle);
             unsafe
             {
                 Method = GetVerb();
@@ -193,14 +200,15 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
             unsafe
             {
                 var hr = 0;
-                hr = NativeMethods.http_flush_response_bytes(_pHttpContext, IISAwaitable.FlushCallback, (IntPtr)_thisHandle, out var fCompletionExpected);
+                hr = NativeMethods.http_flush_response_bytes(_pHttpContext, out var fCompletionExpected);
-
                 if (!fCompletionExpected)
                 {
-                    CompleteFlush(hr, 0);
+                    FreePinnedHeaders(_pinnedHeaders);
+                    _pinnedHeaders = null;
+                    _operation.Complete(hr, 0);
                 }
+                return _operation;
             }
-            return _flushOperation;
         }
 
         public async Task FlushAsync(CancellationToken cancellationToken = default(CancellationToken))
@@ -219,6 +227,35 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
             }
         }
 
+        public async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            StartReadingRequestBody();
+
+            while (true)
+            {
+                var result = await Input.Reader.ReadAsync();
+                var readableBuffer = result.Buffer;
+                try
+                {
+                    if (!readableBuffer.IsEmpty)
+                    {
+                        var actual = Math.Min(readableBuffer.Length, count);
+                        readableBuffer = readableBuffer.Slice(0, actual);
+                        readableBuffer.CopyTo(buffer);
+                        return (int)actual;
+                    }
+                    else if (result.IsCompleted)
+                    {
+                        return 0;
+                    }
+                }
+                finally
+                {
+                    Input.Reader.Advance(readableBuffer.End, readableBuffer.End);
+                }
+            }
+        }
+
         public Task WriteAsync(ArraySegment<byte> data, CancellationToken cancellationToken = default(CancellationToken))
         {
             if (!HasResponseStarted)
@@ -530,7 +567,20 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
                     try
                     {
-                        int read = await ReadAsync(wb.Buffer.Length);
+                        int read = 0;
+                        if (_wasUpgraded)
+                        {
+                            read = await ReadWebSocketsAsync(wb.Buffer.Length);
+                        }
+                        else
+                        {
+                            _currentOperation = _currentOperation.ContinueWith(async (t) =>
+                            {
+                                _currentOperationType = CurrentOperationType.Read;
+                                read = await ReadAsync(wb.Buffer.Length);
+                            }).Unwrap();
+                            await _currentOperation;
+                        }
 
                         if (read == 0)
                         {
@@ -597,7 +647,19 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
                     if (!buffer.IsEmpty)
                     {
-                        await WriteAsync(buffer);
+                        if (_wasUpgraded)
+                        {
+                            await WriteAsync(buffer);
+                        }
+                        else
+                        {
+                            _currentOperation = _currentOperation.ContinueWith(async (t) =>
+                            {
+                                _currentOperationType = CurrentOperationType.Write;
+                                await WriteAsync(buffer);
+                            }).Unwrap();
+                            await _currentOperation;
+                        }
                     }
                     else if (result.IsCompleted)
                     {
@@ -605,7 +667,12 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
                     }
                     else
                     {
-                        await DoFlushAsync();
+                        _currentOperation = _currentOperation.ContinueWith(async (t) =>
+                        {
+                            _currentOperationType = CurrentOperationType.Flush;
+                            await DoFlushAsync();
+                        }).Unwrap();
+                        await _currentOperation;
                     }
 
                     _upgradeTcs?.TrySetResult(null);
@@ -647,8 +714,14 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
                     chunk.DataChunkType = HttpApiTypes.HTTP_DATA_CHUNK_TYPE.HttpDataChunkFromMemory;
                     chunk.fromMemory.pBuffer = (IntPtr)pBuffer;
                     chunk.fromMemory.BufferLength = (uint)buffer.Length;
-
+                    if (_wasUpgraded)
-                    hr = NativeMethods.http_write_response_bytes(_pHttpContext, pDataChunks, nChunks, IISAwaitable.WriteCallback, (IntPtr)_thisHandle, out fCompletionExpected);
+                    {
+                        hr = NativeMethods.http_websockets_write_bytes(_pHttpContext, pDataChunks, nChunks, IISAwaitable.WriteCallback, (IntPtr)_thisHandle, out fCompletionExpected);
+                    }
+                    else
+                    {
+                        hr = NativeMethods.http_write_response_bytes(_pHttpContext, pDataChunks, nChunks, out fCompletionExpected);
+                    }
                 }
             }
             else
@@ -674,9 +747,14 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
                     currentChunk++;
                 }
-
+                if (_wasUpgraded)
-                hr = NativeMethods.http_write_response_bytes(_pHttpContext, pDataChunks, nChunks, IISAwaitable.WriteCallback, (IntPtr)_thisHandle, out fCompletionExpected);
+                {
-
+                    hr = NativeMethods.http_websockets_write_bytes(_pHttpContext, pDataChunks, nChunks, IISAwaitable.WriteCallback, (IntPtr)_thisHandle, out fCompletionExpected);
+                }
+                else
+                {
+                    hr = NativeMethods.http_write_response_bytes(_pHttpContext, pDataChunks, nChunks, out fCompletionExpected);
+                }
                 // Free the handles
                 foreach (var handle in handles)
                 {
@@ -684,31 +762,57 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
                 }
             }
 
-            if (!fCompletionExpected)
+            if (_wasUpgraded)
             {
-                CompleteWrite(hr, cbBytes: 0);
+                if (!fCompletionExpected)
+                {
+                    CompleteWriteWebSockets(hr, 0);
+                }
+                return _writeWebSocketsOperation;
+            }
+            else
+            {
+                if (!fCompletionExpected)
+                {
+                    _operation.Complete(hr, 0);
+                }
+                return _operation;
             }
-
-            return _writeOperation;
         }
 
         private unsafe IISAwaitable ReadAsync(int length)
         {
             var hr = NativeMethods.http_read_request_bytes(
-                                _pHttpContext,
+                            _pHttpContext,
-                                (byte*)_inputHandle.PinnedPointer,
+                            (byte*)_inputHandle.PinnedPointer,
-                                length,
+                            length,
-                                IISAwaitable.ReadCallback,
+                            out var dwReceivedBytes,
-                                (IntPtr)_thisHandle,
+                            out bool fCompletionExpected);
-                                out var dwReceivedBytes,
-                                out var fCompletionExpected);
-
             if (!fCompletionExpected)
             {
-                CompleteRead(hr, dwReceivedBytes);
+                _operation.Complete(hr, dwReceivedBytes);
             }
+            return _operation;
+        }
 
-            return _readOperation;
+        private unsafe IISAwaitable ReadWebSocketsAsync(int length)
+        {
+            var hr = 0;
+            int dwReceivedBytes;
+            bool fCompletionExpected;
+            hr = NativeMethods.http_websockets_read_bytes(
+                                      _pHttpContext,
+                                      (byte*)_inputHandle.PinnedPointer,
+                                      length,
+                                      IISAwaitable.ReadCallback,
+                                      (IntPtr)_thisHandle,
+                                      out dwReceivedBytes,
+                                      out fCompletionExpected);
+            if (!fCompletionExpected)
+            {
+                CompleteReadWebSockets(hr, dwReceivedBytes);
+            }
+            return _readWebSocketsOperation;
         }
 
         public abstract Task ProcessRequestAsync();
@@ -806,13 +910,17 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
             }
         }
 
-        public void PostCompletion()
+        public void PostCompletion(NativeMethods.REQUEST_NOTIFICATION_STATUS requestNotificationStatus)
         {
-            Debug.Assert(!_readOperation.HasContinuation, "Pending read async operation!");
+            Debug.Assert(!_operation.HasContinuation, "Pending async operation!");
-            Debug.Assert(!_writeOperation.HasContinuation, "Pending write async operation!");
 
-            var hr = NativeMethods.http_post_completion(_pHttpContext, 0);
+            var hr = NativeMethods.http_set_completion_status(_pHttpContext, requestNotificationStatus);
+            if (hr != NativeMethods.S_OK)
+            {
+                throw Marshal.GetExceptionForHR(hr);
+            }
 
+            hr = NativeMethods.http_post_completion(_pHttpContext, 0);
             if (hr != NativeMethods.S_OK)
             {
                 throw Marshal.GetExceptionForHR(hr);
@@ -824,22 +932,30 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
             NativeMethods.http_indicate_completion(_pHttpContext, notificationStatus);
         }
 
-        internal void CompleteWrite(int hr, int cbBytes)
+        internal void OnAsyncCompletion(int hr, int cbBytes)
         {
-            _writeOperation.Complete(hr, cbBytes);
+            switch (_currentOperationType)
+            {
+                case CurrentOperationType.Read:
+                case CurrentOperationType.Write:
+                    _operation.Complete(hr, cbBytes);
+                    break;
+                case CurrentOperationType.Flush:
+                    FreePinnedHeaders(_pinnedHeaders);
+                    _pinnedHeaders = null;
+                    _operation.Complete(hr, cbBytes);
+                    break;
+            }
         }
 
-        internal void CompleteRead(int hr, int cbBytes)
+        internal void CompleteWriteWebSockets(int hr, int cbBytes)
         {
-            _readOperation.Complete(hr, cbBytes);
+            _writeWebSocketsOperation.Complete(hr, cbBytes);
         }
 
-        internal void CompleteFlush(int hr, int cbBytes)
+        internal void CompleteReadWebSockets(int hr, int cbBytes)
         {
-            FreePinnedHeaders(_pinnedHeaders);
+            _readWebSocketsOperation.Complete(hr, cbBytes);
-            _pinnedHeaders = null;
-
-            _flushOperation.Complete(hr, cbBytes);
         }
 
         private bool disposedValue = false; // To detect redundant calls
@@ -870,5 +986,12 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
             throw new InvalidOperationException("Response already started");
         }
 
+        private enum CurrentOperationType
+        {
+            None,
+            Read,
+            Write,
+            Flush
+        }
     }
 }

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISHttpContextOfT.cs

@@ -9,7 +9,7 @@ using System.IO.Pipelines;
 
 namespace Microsoft.AspNetCore.Server.IISIntegration
 {
-    internal class IISHttpContextOfT<TContext> : HttpProtocol
+    internal class IISHttpContextOfT<TContext> : IISHttpContext
     {
         private readonly IHttpApplication<TContext> _application;
 

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISHttpRequestBody.cs

@@ -13,9 +13,9 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 {
     internal class IISHttpRequestBody : Stream
     {
-        private readonly HttpProtocol _httpContext;
+        private readonly IISHttpContext _httpContext;
 
-        public IISHttpRequestBody(HttpProtocol httpContext)
+        public IISHttpRequestBody(IISHttpContext httpContext)
         {
             _httpContext = httpContext;
         }
@@ -40,33 +40,9 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
             return ReadAsync(buffer, offset, count, CancellationToken.None).GetAwaiter().GetResult();
         }
 
-        public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        public override unsafe Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
         {
-            _httpContext.StartReadingRequestBody();
+            return _httpContext.ReadAsync(buffer, offset, count, cancellationToken);
-
-            while (true)
-            {
-                var result = await _httpContext.Input.Reader.ReadAsync();
-                var readableBuffer = result.Buffer;
-                try
-                {
-                    if (!readableBuffer.IsEmpty)
-                    {
-                        var actual = Math.Min(readableBuffer.Length, count);
-                        readableBuffer = readableBuffer.Slice(0, actual);
-                        readableBuffer.CopyTo(buffer);
-                        return (int)actual;
-                    }
-                    else if (result.IsCompleted)
-                    {
-                        return 0;
-                    }
-                }
-                finally
-                {
-                    _httpContext.Input.Reader.Advance(readableBuffer.End, readableBuffer.End);
-                }
-            }
         }
 
         public override long Seek(long offset, SeekOrigin origin)

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISHttpResponseBody.cs

@@ -10,9 +10,9 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 {
     internal class IISHttpResponseBody : Stream
     {
-        private readonly HttpProtocol _httpContext;
+        private readonly IISHttpContext _httpContext;
 
-        public IISHttpResponseBody(HttpProtocol httpContext)
+        public IISHttpResponseBody(IISHttpContext httpContext)
         {
             _httpContext = httpContext;
         }

src/Microsoft.AspNetCore.Server.IISIntegration/Server/IISHttpServer.cs

@@ -17,6 +17,7 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
     {
         private static NativeMethods.PFN_REQUEST_HANDLER _requestHandler = HandleRequest;
         private static NativeMethods.PFN_SHUTDOWN_HANDLER _shutdownHandler = HandleShutdown;
+        private static NativeMethods.PFN_ASYNC_COMPLETION _onAsyncCompletion = OnAsyncCompletion;
 
         private IISContextFactory _iisContextFactory;
 
@@ -38,7 +39,7 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
 
             // Start the server by registering the callback
             // TODO the context may change here for shutdown.
-            NativeMethods.register_callbacks(_requestHandler, _shutdownHandler, (IntPtr)_httpServerHandle, (IntPtr)_httpServerHandle);
+            NativeMethods.register_callbacks(_requestHandler, _shutdownHandler, _onAsyncCompletion, (IntPtr)_httpServerHandle, (IntPtr)_httpServerHandle);
 
             return Task.CompletedTask;
         }
@@ -78,7 +79,7 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
                 return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_CONTINUE;
             }
 
-            task.ContinueWith((t, state) => CompleteRequest((HttpProtocol)state), context);
+            task.ContinueWith((t, state) => CompleteRequest((IISHttpContext)state), context);
 
             return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_PENDING;
         }
@@ -90,10 +91,17 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
             return true;
         }
 
-        private static void CompleteRequest(HttpProtocol context)
+        private static NativeMethods.REQUEST_NOTIFICATION_STATUS OnAsyncCompletion(IntPtr pvManagedHttpContext, int hr, int bytes)
+        {
+            var context = (IISHttpContext)GCHandle.FromIntPtr(pvManagedHttpContext).Target;
+            context.OnAsyncCompletion(hr, bytes);
+            return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_PENDING;
+        }
+
+        private static void CompleteRequest(IISHttpContext context)
         {
             // Post completion after completing the request to resume the state machine
-            context.PostCompletion();
+            context.PostCompletion(NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_CONTINUE);
 
             // Dispose the context
             context.Dispose();
@@ -110,7 +118,7 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
                 _pipeFactory = pipeFactory;
             }
 
-            public HttpProtocol CreateHttpContext(IntPtr pHttpContext)
+            public IISHttpContext CreateHttpContext(IntPtr pHttpContext)
             {
                 return new IISHttpContextOfT<T>(_pipeFactory, _application, pHttpContext);
             }
@@ -120,6 +128,6 @@ namespace Microsoft.AspNetCore.Server.IISIntegration
     // Over engineering to avoid allocations...
     internal interface IISContextFactory
     {
-        HttpProtocol CreateHttpContext(IntPtr pHttpContext);
+        IISHttpContext CreateHttpContext(IntPtr pHttpContext);
     }
 }