General cleanup in IIS (#24754)

src/Servers/IIS/IIS/src/Core/IISHttpContext.FeatureCollection.cs

@@ -271,7 +271,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                 // Synchronize access to native methods that might run in parallel with IO loops
                 lock (_contextLock)
                 {
-                    return NativeMethods.HttpTryGetServerVariable(_pInProcessHandler, variableName, out var value) ? value : null;
+                    return NativeMethods.HttpTryGetServerVariable(_requestNativeHandle, variableName, out var value) ? value : null;
                 }
             }
             set
@@ -284,7 +284,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                 // Synchronize access to native methods that might run in parallel with IO loops
                 lock (_contextLock)
                 {
-                    NativeMethods.HttpSetServerVariable(_pInProcessHandler, variableName, value);
+                    NativeMethods.HttpSetServerVariable(_requestNativeHandle, variableName, value);
                 }
             }
         }
@@ -345,7 +345,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             HasStartedConsumingRequestBody = false;
 
             // Upgrade async will cause the stream processing to go into duplex mode
-            AsyncIO = new WebSocketsAsyncIOEngine(_contextLock, _pInProcessHandler);
+            AsyncIO = new WebSocketsAsyncIOEngine(_contextLock, _requestNativeHandle);
 
             await InitializeResponse(flushHeaders: true);
 
@@ -419,7 +419,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
         internal IHttpResponseTrailersFeature GetResponseTrailersFeature()
         {
             // Check version is above 2.
-            if (HttpVersion >= System.Net.HttpVersion.Version20 && NativeMethods.HttpSupportTrailer(_pInProcessHandler))
+            if (HttpVersion >= System.Net.HttpVersion.Version20 && NativeMethods.HttpSupportTrailer(_requestNativeHandle))
             {
                 return this;
             }
@@ -436,7 +436,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
         internal IHttpResetFeature GetResetFeature()
         {
             // Check version is above 2.
-            if (HttpVersion >= System.Net.HttpVersion.Version20 && NativeMethods.HttpSupportTrailer(_pInProcessHandler))
+            if (HttpVersion >= System.Net.HttpVersion.Version20 && NativeMethods.HttpSupportTrailer(_requestNativeHandle))
             {
                 return this;
             }
@@ -457,12 +457,12 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
         internal unsafe void SetResetCode(int errorCode)
         {
-            NativeMethods.HttpResetStream(_pInProcessHandler, (ulong)errorCode);
+            NativeMethods.HttpResetStream(_requestNativeHandle, (ulong)errorCode);
         }
 
         void IHttpResponseBodyFeature.DisableBuffering()
         {
-            NativeMethods.HttpDisableBuffering(_pInProcessHandler);
+            NativeMethods.HttpDisableBuffering(_requestNativeHandle);
             DisableCompression();
         }
 

src/Servers/IIS/IIS/src/Core/IISHttpContext.IO.cs

@@ -272,7 +272,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
         {
             _bodyOutput.Abort(reason);
             _streams.Abort(reason);
-            NativeMethods.HttpCloseConnection(_pInProcessHandler);
+            NativeMethods.HttpCloseConnection(_requestNativeHandle);
 
             AbortIO(clientDisconnect: false);
         }

src/Servers/IIS/IIS/src/Core/IISHttpContext.cs

@@ -34,7 +34,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
         private const int PauseWriterThreshold = 65536;
         private const int ResumeWriterTheshold = PauseWriterThreshold / 2;
 
-        protected readonly IntPtr _pInProcessHandler;
+        protected readonly NativeSafeHandle _requestNativeHandle;
 
         private readonly IISServerOptions _options;
 
@@ -75,7 +75,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
         internal unsafe IISHttpContext(
             MemoryPool<byte> memoryPool,
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             IISServerOptions options,
             IISHttpServer server,
             ILogger logger,
@@ -83,7 +83,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             : base((HttpApiTypes.HTTP_REQUEST*)NativeMethods.HttpGetRawRequest(pInProcessHandler), useLatin1: useLatin1)
         {
             _memoryPool = memoryPool;
-            _pInProcessHandler = pInProcessHandler;
+            _requestNativeHandle = pInProcessHandler;
             _options = options;
             _server = server;
             _logger = logger;
@@ -179,7 +179,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
                 ResetFeatureCollection();
 
-                if (!_server.IsWebSocketAvailable(_pInProcessHandler))
+                if (!_server.IsWebSocketAvailable(_requestNativeHandle))
                 {
                     _currentIHttpUpgradeFeature = null;
                 }
@@ -198,7 +198,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                 _bodyOutput = new OutputProducer(pipe);
             }
 
-            NativeMethods.HttpSetManagedContext(_pInProcessHandler, (IntPtr)_thisHandle);
+            NativeMethods.HttpSetManagedContext(_requestNativeHandle, (IntPtr)_thisHandle);
         }
 
         private string GetOriginalPath()
@@ -324,7 +324,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             // If at this point request was not upgraded just start a normal IO engine
             if (AsyncIO == null)
             {
-                AsyncIO = new AsyncIOEngine(_contextLock, _pInProcessHandler);
+                AsyncIO = new AsyncIOEngine(_contextLock, _requestNativeHandle);
             }
         }
 
@@ -383,7 +383,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             var reasonPhrase = string.IsNullOrEmpty(ReasonPhrase) ? ReasonPhrases.GetReasonPhrase(StatusCode) : ReasonPhrase;
 
             // This copies data into the underlying buffer
-            NativeMethods.HttpSetResponseStatusCode(_pInProcessHandler, (ushort)StatusCode, reasonPhrase);
+            NativeMethods.HttpSetResponseStatusCode(_requestNativeHandle, (ushort)StatusCode, reasonPhrase);
 
             HttpResponseHeaders.IsReadOnly = true;
             foreach (var headerPair in HttpResponseHeaders)
@@ -412,12 +412,12 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                             var headerNameBytes = Encoding.UTF8.GetBytes(headerPair.Key);
                             fixed (byte* pHeaderName = headerNameBytes)
                             {
-                                NativeMethods.HttpResponseSetUnknownHeader(_pInProcessHandler, pHeaderName, pHeaderValue, (ushort)headerValueBytes.Length, fReplace: isFirst);
+                                NativeMethods.HttpResponseSetUnknownHeader(_requestNativeHandle, pHeaderName, pHeaderValue, (ushort)headerValueBytes.Length, fReplace: isFirst);
                             }
                         }
                         else
                         {
-                            NativeMethods.HttpResponseSetKnownHeader(_pInProcessHandler, knownHeaderIndex, pHeaderValue, (ushort)headerValueBytes.Length, fReplace: isFirst);
+                            NativeMethods.HttpResponseSetKnownHeader(_requestNativeHandle, knownHeaderIndex, pHeaderValue, (ushort)headerValueBytes.Length, fReplace: isFirst);
                         }
                     }
                 }
@@ -451,7 +451,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                         var headerValueBytes = Encoding.UTF8.GetBytes(headerValue);
                         fixed (byte* pHeaderValue = headerValueBytes)
                         {
-                            NativeMethods.HttpResponseSetTrailer(_pInProcessHandler, pHeaderName, pHeaderValue, (ushort)headerValueBytes.Length, replace: isFirst);
+                            NativeMethods.HttpResponseSetTrailer(_requestNativeHandle, pHeaderName, pHeaderValue, (ushort)headerValueBytes.Length, replace: isFirst);
                         }
 
                         isFirst = false;
@@ -576,8 +576,8 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
         public void PostCompletion(NativeMethods.REQUEST_NOTIFICATION_STATUS requestNotificationStatus)
         {
-            NativeMethods.HttpSetCompletionStatus(_pInProcessHandler, requestNotificationStatus);
+            NativeMethods.HttpSetCompletionStatus(_requestNativeHandle, requestNotificationStatus);
-            NativeMethods.HttpPostCompletion(_pInProcessHandler, 0);
+            NativeMethods.HttpPostCompletion(_requestNativeHandle, 0);
         }
 
         internal void OnAsyncCompletion(int hr, int bytes)
@@ -628,7 +628,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
         private WindowsPrincipal GetWindowsPrincipal()
         {
-            NativeMethods.HttpGetAuthenticationInformation(_pInProcessHandler, out var authenticationType, out var token);
+            NativeMethods.HttpGetAuthenticationInformation(_requestNativeHandle, out var authenticationType, out var token);
 
             if (token != IntPtr.Zero && authenticationType != null)
             {
@@ -664,6 +664,17 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                 // Post completion after completing the request to resume the state machine
                 PostCompletion(ConvertRequestCompletionResults(successfulRequest));
 
+                // After disposing a safe handle, Dispose() will not block waiting for the pinvokes to finish.
+                // Instead Safehandle will call ReleaseHandle on the pinvoke thread when the pinvokes complete
+                // and the reference count goes to zero.
+
+                // What this means is we need to wait until ReleaseHandle is called to finish disposal.
+                // This is to make sure it is safe to return back to native.
+                // The handle implements IValueTaskSource
+                _requestNativeHandle.Dispose();
+
+                await new ValueTask<object>(_requestNativeHandle, _requestNativeHandle.Version);
+
                 // Dispose the context
                 Dispose();
             }

src/Servers/IIS/IIS/src/Core/IISHttpContextOfT.cs

@@ -18,7 +18,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
     {
         private readonly IHttpApplication<TContext> _application;
 
-        public IISHttpContextOfT(MemoryPool<byte> memoryPool, IHttpApplication<TContext> application, IntPtr pInProcessHandler, IISServerOptions options, IISHttpServer server, ILogger logger, bool useLatin1)
+        public IISHttpContextOfT(MemoryPool<byte> memoryPool, IHttpApplication<TContext> application, NativeSafeHandle pInProcessHandler, IISServerOptions options, IISHttpServer server, ILogger logger, bool useLatin1)
             : base(memoryPool, pInProcessHandler, options, server, logger, useLatin1)
         {
             _application = application;
@@ -64,7 +64,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                     // Dispose
                 }
 
-                if (!success && HasResponseStarted && NativeMethods.HttpSupportTrailer(_pInProcessHandler))
+                if (!success && HasResponseStarted && NativeMethods.HttpSupportTrailer(_requestNativeHandle))
                 {
                     // HTTP/2 INTERNAL_ERROR = 0x2 https://tools.ietf.org/html/rfc7540#section-7
                     // Otherwise the default is Cancel = 0x8.

src/Servers/IIS/IIS/src/Core/IISHttpServer.cs

@@ -40,11 +40,12 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
         private readonly TaskCompletionSource<object> _shutdownSignal = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
         private bool? _websocketAvailable;
         private CancellationTokenRegistration _cancellationTokenRegistration;
+        private bool _disposed;
 
         public IFeatureCollection Features { get; } = new FeatureCollection();
 
         // TODO: Remove pInProcessHandler argument
-        public bool IsWebSocketAvailable(IntPtr pInProcessHandler)
+        public bool IsWebSocketAvailable(NativeSafeHandle pInProcessHandler)
         {
             // Check if the Http upgrade feature is available in IIS.
             // To check this, we can look at the server variable WEBSOCKET_VERSION
@@ -114,6 +115,12 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 
         public void Dispose()
         {
+            if (_disposed)
+            {
+                return;
+            }
+            _disposed = true;
+
             // Block any more calls into managed from native as we are unloading.
             _nativeApplication.StopCallsIntoManaged();
             _shutdownSignal.TrySetResult(null);
@@ -135,7 +142,14 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                 // Unwrap the server so we can create an http context and process the request
                 server = (IISHttpServer)GCHandle.FromIntPtr(pvRequestContext).Target;
 
-                var context = server._iisContextFactory.CreateHttpContext(pInProcessHandler);
+                // server can be null if ungraceful shutdown.
+                if (server == null)
+                {
+                    return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_FINISH_REQUEST;
+                }
+
+                var safehandle = new NativeSafeHandle(pInProcessHandler);
+                var context = server._iisContextFactory.CreateHttpContext(safehandle);
 
                 ThreadPool.UnsafeQueueUserWorkItem(context, preferLocal: false);
 
@@ -155,6 +169,14 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             try
             {
                 server = (IISHttpServer)GCHandle.FromIntPtr(pvRequestContext).Target;
+
+                // server can be null if ungraceful shutdown.
+                if (server ==  null)
+                {
+                    // return value isn't checked.
+                    return true;
+                }
+
                 server._applicationLifetime.StopApplication();
             }
             catch (Exception ex)
@@ -170,7 +192,14 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             try
             {
                 context = (IISHttpContext)GCHandle.FromIntPtr(pvManagedHttpContext).Target;
-                context?.AbortIO(clientDisconnect: true);
+
+                // Context can be null if ungraceful shutdown.
+                if (context == null)
+                {
+                    return;
+                }
+
+                context.AbortIO(clientDisconnect: true);
             }
             catch (Exception ex)
             {
@@ -184,7 +213,14 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             try
             {
                 context = (IISHttpContext)GCHandle.FromIntPtr(pvManagedHttpContext).Target;
-                context?.OnAsyncCompletion(hr, bytes);
+
+                // Context can be null if ungraceful shutdown.
+                if (context == null)
+                {
+                    return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_FINISH_REQUEST;
+                }
+
+                context.OnAsyncCompletion(hr, bytes);
                 return NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_PENDING;
             }
             catch (Exception ex)
@@ -202,6 +238,12 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
             {
                 server = (IISHttpServer)GCHandle.FromIntPtr(serverContext).Target;
 
+                // server can be null if ungraceful shutdown.
+                if (server == null)
+                {
+                    return;
+                }
+
                 server._nativeApplication.StopCallsIntoManaged();
                 server._shutdownSignal.TrySetResult(null);
                 server._cancellationTokenRegistration.Dispose();
@@ -233,7 +275,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
                 AppContext.TryGetSwitch(Latin1Suppport, out _useLatin1);
             }
 
-            public IISHttpContext CreateHttpContext(IntPtr pInProcessHandler)
+            public IISHttpContext CreateHttpContext(NativeSafeHandle pInProcessHandler)
             {
                 return new IISHttpContextOfT<T>(_memoryPool, _application, pInProcessHandler, _options, _server, _logger, _useLatin1);
             }
@@ -243,6 +285,6 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
     // Over engineering to avoid allocations...
     internal interface IISContextFactory
     {
-        IISHttpContext CreateHttpContext(IntPtr pInProcessHandler);
+        IISHttpContext CreateHttpContext(NativeSafeHandle pInProcessHandler);
     }
 }

src/Servers/IIS/IIS/src/Core/IISNativeApplication.cs

@@ -8,9 +8,9 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
 {
     internal class IISNativeApplication
     {
-        private readonly IntPtr _nativeApplication;
+        private readonly NativeSafeHandle _nativeApplication;
 
-        public IISNativeApplication(IntPtr nativeApplication)
+        public IISNativeApplication(NativeSafeHandle nativeApplication)
         {
             _nativeApplication = nativeApplication;
         }
@@ -48,6 +48,9 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
         public void Dispose()
         {
             GC.SuppressFinalize(this);
+
+            // Don't need to await here because pinvokes should never been called after disposing the safe handle.
+            _nativeApplication.Dispose();
         }
 
         ~IISNativeApplication()

src/Servers/IIS/IIS/src/Core/IO/AsyncIOEngine.Flush.cs

@@ -11,7 +11,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
         {
             private readonly AsyncIOEngine _engine;
 
-            private IntPtr _requestHandler;
+            private NativeSafeHandle _requestHandler;
             private bool _moreData;
 
             public AsyncFlushOperation(AsyncIOEngine engine)
@@ -19,7 +19,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
                 _engine = engine;
             }
 
-            public void Initialize(IntPtr requestHandler, bool moreData)
+            public void Initialize(NativeSafeHandle requestHandler, bool moreData)
             {
                 _requestHandler = requestHandler;
                 _moreData = moreData;

src/Servers/IIS/IIS/src/Core/IO/AsyncIOEngine.Read.cs

@@ -14,7 +14,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
 
             private MemoryHandle _inputHandle;
 
-            private IntPtr _requestHandler;
+            private NativeSafeHandle _requestHandler;
 
             private Memory<byte> _memory;
 
@@ -23,7 +23,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
                 _engine = engine;
             }
 
-            public void Initialize(IntPtr requestHandler, Memory<byte> memory)
+            public void Initialize(NativeSafeHandle requestHandler, Memory<byte> memory)
             {
                 _requestHandler = requestHandler;
                 _memory = memory;

src/Servers/IIS/IIS/src/Core/IO/AsyncIOEngine.Write.cs

@@ -17,7 +17,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
                 _engine = engine;
             }
 
-            protected override unsafe int WriteChunks(IntPtr requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks,
+            protected override unsafe int WriteChunks(NativeSafeHandle requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks,
                 out bool completionExpected)
             {
                 return NativeMethods.HttpWriteResponseBytes(requestHandler, dataChunks, chunkCount, out completionExpected);

src/Servers/IIS/IIS/src/Core/IO/AsyncIOEngine.cs

@@ -13,7 +13,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
     internal partial class AsyncIOEngine : IAsyncIOEngine
     {
         private readonly object _contextSync;
-        private readonly IntPtr _handler;
+        private readonly NativeSafeHandle _handler;
 
         private bool _stopped;
 
@@ -24,7 +24,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
         private AsyncWriteOperation _cachedAsyncWriteOperation;
         private AsyncFlushOperation _cachedAsyncFlushOperation;
 
-        public AsyncIOEngine(object contextSync, IntPtr handler)
+        public AsyncIOEngine(object contextSync, NativeSafeHandle handler)
         {
             _contextSync = contextSync;
             _handler = handler;

src/Servers/IIS/IIS/src/Core/IO/AsyncWriteOperationBase.cs

@@ -11,11 +11,11 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
     {
         private const int HttpDataChunkStackLimit = 128; // 16 bytes per HTTP_DATA_CHUNK
 
-        private IntPtr _requestHandler;
+        private NativeSafeHandle _requestHandler;
         private ReadOnlySequence<byte> _buffer;
         private MemoryHandle[] _handles;
 
-        public void Initialize(IntPtr requestHandler, ReadOnlySequence<byte> buffer)
+        public void Initialize(NativeSafeHandle requestHandler, ReadOnlySequence<byte> buffer)
         {
             _requestHandler = requestHandler;
             _buffer = buffer;
@@ -114,6 +114,6 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             return WriteChunks(_requestHandler, nChunks, pDataChunks, out fCompletionExpected);
         }
 
-        protected abstract unsafe int WriteChunks(IntPtr requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks, out bool completionExpected);
+        protected abstract unsafe int WriteChunks(NativeSafeHandle requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks, out bool completionExpected);
     }
 }

src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Initialize.cs

@@ -11,14 +11,14 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
         {
             private readonly WebSocketsAsyncIOEngine _engine;
 
-            private IntPtr _requestHandler;
+            private NativeSafeHandle _requestHandler;
 
             public AsyncInitializeOperation(WebSocketsAsyncIOEngine engine)
             {
                 _engine = engine;
             }
 
-            public void Initialize(IntPtr requestHandler)
+            public void Initialize(NativeSafeHandle requestHandler)
             {
                 _requestHandler = requestHandler;
             }

src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Read.cs

@@ -27,7 +27,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
             private readonly WebSocketsAsyncIOEngine _engine;
             private GCHandle _thisHandle;
             private MemoryHandle _inputHandle;
-            private IntPtr _requestHandler;
+            private NativeSafeHandle _requestHandler;
             private Memory<byte> _memory;
 
             public WebSocketReadOperation(WebSocketsAsyncIOEngine engine)
@@ -52,7 +52,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
                 return !completionExpected;
             }
 
-            public void Initialize(IntPtr requestHandler, Memory<byte> memory)
+            public void Initialize(NativeSafeHandle requestHandler, Memory<byte> memory)
             {
                 _requestHandler = requestHandler;
                 _memory = memory;

src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.Write.cs

@@ -32,7 +32,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
                 _engine = engine;
             }
 
-            protected override unsafe int WriteChunks(IntPtr requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks, out bool completionExpected)
+            protected override unsafe int WriteChunks(NativeSafeHandle requestHandler, int chunkCount, HttpApiTypes.HTTP_DATA_CHUNK* dataChunks, out bool completionExpected)
             {
                 _thisHandle = GCHandle.Alloc(this);
                 return NativeMethods.HttpWebsocketsWriteBytes(requestHandler, dataChunks, chunkCount, WriteCallback, (IntPtr)_thisHandle, out completionExpected);

src/Servers/IIS/IIS/src/Core/IO/WebSocketsAsyncIOEngine.cs

@@ -12,7 +12,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
     {
         private readonly object _contextLock;
 
-        private readonly IntPtr _handler;
+        private readonly NativeSafeHandle _handler;
 
         private bool _isInitialized;
 
@@ -24,7 +24,7 @@ namespace Microsoft.AspNetCore.Server.IIS.Core.IO
 
         private AsyncInitializeOperation _cachedAsyncInitializeOperation;
 
-        public WebSocketsAsyncIOEngine(object contextLock, IntPtr handler)
+        public WebSocketsAsyncIOEngine(object contextLock, NativeSafeHandle handler)
         {
             _contextLock = contextLock;
             _handler = handler;

src/Servers/IIS/IIS/src/Core/NativeSafeHandle.cs

@@ -0,0 +1,50 @@
+// 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.Runtime.InteropServices;
+using System.Threading.Tasks.Sources;
+
+namespace Microsoft.AspNetCore.Server.IIS.Core
+{
+    internal class NativeSafeHandle : SafeHandle, IValueTaskSource<object>
+    {
+        private ManualResetValueTaskSourceCore<object> _core; // mutable struct; do not make this readonly
+
+        public override bool IsInvalid => handle == IntPtr.Zero;
+        public short Version => _core.Version;
+
+        public NativeSafeHandle(IntPtr handle) : base(IntPtr.Zero, ownsHandle: true)
+        {
+            this.handle = handle;
+        }
+
+        protected override bool ReleaseHandle()
+        {
+            handle = IntPtr.Zero;
+
+            // Complete the ManualResetValueTaskSourceCore
+            if (_core.GetStatus(_core.Version) == ValueTaskSourceStatus.Pending)
+            {
+                _core.SetResult(null);
+            }
+
+            return true;
+        }
+
+        public object GetResult(short token)
+        {
+            return _core.GetResult(token);
+        }
+
+        public ValueTaskSourceStatus GetStatus(short token)
+        {
+            return _core.GetStatus(token);
+        }
+
+        public void OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags)
+        {
+            _core.OnCompleted(continuation, state, token, flags);
+        }
+    }
+}

src/Servers/IIS/IIS/src/NativeMethods.cs

@@ -47,16 +47,16 @@ namespace Microsoft.AspNetCore.Server.IIS
         public delegate void PFN_REQUESTS_DRAINED_HANDLER(IntPtr pvRequestContext);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_post_completion(IntPtr pInProcessHandler, int cbBytes);
+        private static extern int http_post_completion(NativeSafeHandle pInProcessHandler, int cbBytes);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_set_completion_status(IntPtr pInProcessHandler, REQUEST_NOTIFICATION_STATUS rquestNotificationStatus);
+        private static extern int http_set_completion_status(NativeSafeHandle pInProcessHandler, REQUEST_NOTIFICATION_STATUS rquestNotificationStatus);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern void http_indicate_completion(IntPtr pInProcessHandler, REQUEST_NOTIFICATION_STATUS notificationStatus);
+        private static extern void http_indicate_completion(NativeSafeHandle pInProcessHandler, REQUEST_NOTIFICATION_STATUS notificationStatus);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int register_callbacks(IntPtr pInProcessApplication,
+        private static extern int register_callbacks(NativeSafeHandle pInProcessApplication,
             PFN_REQUEST_HANDLER requestCallback,
             PFN_SHUTDOWN_HANDLER shutdownCallback,
             PFN_DISCONNECT_HANDLER disconnectCallback,
@@ -66,53 +66,53 @@ namespace Microsoft.AspNetCore.Server.IIS
             IntPtr pvShutdownContext);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe int http_write_response_bytes(IntPtr pInProcessHandler, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, out bool fCompletionExpected);
+        private static extern unsafe int http_write_response_bytes(NativeSafeHandle pInProcessHandler, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_flush_response_bytes(IntPtr pInProcessHandler, bool fMoreData, out bool fCompletionExpected);
+        private static extern int http_flush_response_bytes(NativeSafeHandle pInProcessHandler, bool fMoreData, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe HttpApiTypes.HTTP_REQUEST_V2* http_get_raw_request(IntPtr pInProcessHandler);
+        private static extern unsafe HttpApiTypes.HTTP_REQUEST_V2* http_get_raw_request(NativeSafeHandle pInProcessHandler);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_stop_calls_into_managed(IntPtr pInProcessApplication);
+        private static extern int http_stop_calls_into_managed(NativeSafeHandle pInProcessApplication);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_stop_incoming_requests(IntPtr pInProcessApplication);
+        private static extern int http_stop_incoming_requests(NativeSafeHandle pInProcessApplication);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_disable_buffering(IntPtr pInProcessApplication);
+        private static extern int http_disable_buffering(NativeSafeHandle pInProcessHandler);
 
         [DllImport(AspNetCoreModuleDll, CharSet = CharSet.Ansi)]
-        private static extern int http_set_response_status_code(IntPtr pInProcessHandler, ushort statusCode, string pszReason);
+        private static extern int http_set_response_status_code(NativeSafeHandle pInProcessHandler, ushort statusCode, string pszReason);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe int http_read_request_bytes(IntPtr pInProcessHandler, byte* pvBuffer, int cbBuffer, out int dwBytesReceived, out bool fCompletionExpected);
+        private static extern unsafe int http_read_request_bytes(NativeSafeHandle pInProcessHandler, byte* pvBuffer, int cbBuffer, out int dwBytesReceived, out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern void http_get_completion_info(IntPtr pCompletionInfo, out int cbBytes, out int hr);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_set_managed_context(IntPtr pInProcessHandler, IntPtr pvManagedContext);
+        private static extern int http_set_managed_context(NativeSafeHandle pInProcessHandler, IntPtr pvManagedContext);
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern int http_get_application_properties(ref IISConfigurationData iiConfigData);
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern int http_get_server_variable(
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             [MarshalAs(UnmanagedType.LPStr)] string variableName,
             [MarshalAs(UnmanagedType.BStr)] out string value);
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern int http_set_server_variable(
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             [MarshalAs(UnmanagedType.LPStr)] string variableName,
             [MarshalAs(UnmanagedType.LPWStr)] string value);
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern unsafe int http_websockets_read_bytes(
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             byte* pvBuffer,
             int cbBuffer,
             PFN_WEBSOCKET_ASYNC_COMPLETION pfnCompletionCallback,
@@ -122,7 +122,7 @@ namespace Microsoft.AspNetCore.Server.IIS
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern unsafe int http_websockets_write_bytes(
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks,
             int nChunks,
             PFN_WEBSOCKET_ASYNC_COMPLETION pfnCompletionCallback,
@@ -130,45 +130,45 @@ namespace Microsoft.AspNetCore.Server.IIS
             out bool fCompletionExpected);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_enable_websockets(IntPtr pInProcessHandler);
+        private static extern int http_enable_websockets(NativeSafeHandle pInProcessHandler);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_cancel_io(IntPtr pInProcessHandler);
+        private static extern int http_cancel_io(NativeSafeHandle pInProcessHandler);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_close_connection(IntPtr pInProcessHandler);
+        private static extern int http_close_connection(NativeSafeHandle pInProcessHandler);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe int http_response_set_unknown_header(IntPtr pInProcessHandler, byte* pszHeaderName, byte* pszHeaderValue, ushort usHeaderValueLength, bool fReplace);
+        private static extern unsafe int http_response_set_unknown_header(NativeSafeHandle pInProcessHandler, byte* pszHeaderName, byte* pszHeaderValue, ushort usHeaderValueLength, bool fReplace);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe int http_has_response4(IntPtr pInProcessHandler, out bool isResponse4);
+        private static extern unsafe int http_has_response4(NativeSafeHandle pInProcessHandler, out bool isResponse4);
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe int http_response_set_trailer(IntPtr pInProcessHandler, byte* pszHeaderName, byte* pszHeaderValue, ushort usHeaderValueLength, bool replace);
+        private static extern unsafe int http_response_set_trailer(NativeSafeHandle pInProcessHandler, byte* pszHeaderName, byte* pszHeaderValue, ushort usHeaderValueLength, bool replace);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe int http_reset_stream(IntPtr pInProcessHandler, ulong errorCode);
+        private static extern unsafe int http_reset_stream(NativeSafeHandle pInProcessHandler, ulong errorCode);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern unsafe int http_response_set_known_header(IntPtr pInProcessHandler, int headerId, byte* pHeaderValue, ushort length, bool fReplace);
+        private static extern unsafe int http_response_set_known_header(NativeSafeHandle pInProcessHandler, int headerId, byte* pHeaderValue, ushort length, bool fReplace);
 
         [DllImport(AspNetCoreModuleDll)]
-        private static extern int http_get_authentication_information(IntPtr pInProcessHandler, [MarshalAs(UnmanagedType.BStr)] out string authType, out IntPtr token);
+        private static extern int http_get_authentication_information(NativeSafeHandle pInProcessHandler, [MarshalAs(UnmanagedType.BStr)] out string authType, out IntPtr token);
 
         [DllImport(AspNetCoreModuleDll)]
         private static extern unsafe int http_set_startup_error_page_content(byte* content, int contentLength);
 
-        public static void HttpPostCompletion(IntPtr pInProcessHandler, int cbBytes)
+        public static void HttpPostCompletion(NativeSafeHandle pInProcessHandler, int cbBytes)
         {
             Validate(http_post_completion(pInProcessHandler, cbBytes));
         }
 
-        public static void HttpSetCompletionStatus(IntPtr pInProcessHandler, REQUEST_NOTIFICATION_STATUS rquestNotificationStatus)
+        public static void HttpSetCompletionStatus(NativeSafeHandle pInProcessHandler, REQUEST_NOTIFICATION_STATUS rquestNotificationStatus)
         {
             Validate(http_set_completion_status(pInProcessHandler, rquestNotificationStatus));
         }
 
-        public static void HttpRegisterCallbacks(IntPtr pInProcessApplication,
+        public static void HttpRegisterCallbacks(NativeSafeHandle pInProcessApplication,
             PFN_REQUEST_HANDLER requestCallback,
             PFN_SHUTDOWN_HANDLER shutdownCallback,
             PFN_DISCONNECT_HANDLER disconnectCallback,
@@ -180,42 +180,42 @@ namespace Microsoft.AspNetCore.Server.IIS
             Validate(register_callbacks(pInProcessApplication, requestCallback, shutdownCallback, disconnectCallback, asyncCallback, requestsDrainedHandler, pvRequestContext, pvShutdownContext));
         }
 
-        internal static unsafe int HttpWriteResponseBytes(IntPtr pInProcessHandler, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, out bool fCompletionExpected)
+        internal static unsafe int HttpWriteResponseBytes(NativeSafeHandle pInProcessHandler, HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks, int nChunks, out bool fCompletionExpected)
         {
             return http_write_response_bytes(pInProcessHandler, pDataChunks, nChunks, out fCompletionExpected);
         }
 
-        public static int HttpFlushResponseBytes(IntPtr pInProcessHandler, bool fMoreData, out bool fCompletionExpected)
+        public static int HttpFlushResponseBytes(NativeSafeHandle pInProcessHandler, bool fMoreData, out bool fCompletionExpected)
         {
             return http_flush_response_bytes(pInProcessHandler, fMoreData, out fCompletionExpected);
         }
 
-        internal static unsafe HttpApiTypes.HTTP_REQUEST_V2* HttpGetRawRequest(IntPtr pInProcessHandler)
+        internal static unsafe HttpApiTypes.HTTP_REQUEST_V2* HttpGetRawRequest(NativeSafeHandle pInProcessHandler)
         {
             return http_get_raw_request(pInProcessHandler);
         }
 
-        public static void HttpStopCallsIntoManaged(IntPtr pInProcessApplication)
+        public static void HttpStopCallsIntoManaged(NativeSafeHandle pInProcessApplication)
         {
             Validate(http_stop_calls_into_managed(pInProcessApplication));
         }
 
-        public static void HttpStopIncomingRequests(IntPtr pInProcessApplication)
+        public static void HttpStopIncomingRequests(NativeSafeHandle pInProcessApplication)
         {
             Validate(http_stop_incoming_requests(pInProcessApplication));
         }
 
-        public static void HttpDisableBuffering(IntPtr pInProcessApplication)
+        public static void HttpDisableBuffering(NativeSafeHandle pInProcessHandler)
         {
-            Validate(http_disable_buffering(pInProcessApplication));
+            Validate(http_disable_buffering(pInProcessHandler));
         }
 
-        public static void HttpSetResponseStatusCode(IntPtr pInProcessHandler, ushort statusCode, string pszReason)
+        public static void HttpSetResponseStatusCode(NativeSafeHandle pInProcessHandler, ushort statusCode, string pszReason)
         {
             Validate(http_set_response_status_code(pInProcessHandler, statusCode, pszReason));
         }
 
-        public static unsafe int HttpReadRequestBytes(IntPtr pInProcessHandler, byte* pvBuffer, int cbBuffer, out int dwBytesReceived, out bool fCompletionExpected)
+        public static unsafe int HttpReadRequestBytes(NativeSafeHandle pInProcessHandler, byte* pvBuffer, int cbBuffer, out int dwBytesReceived, out bool fCompletionExpected)
         {
             return http_read_request_bytes(pInProcessHandler, pvBuffer, cbBuffer, out dwBytesReceived, out fCompletionExpected);
         }
@@ -225,7 +225,7 @@ namespace Microsoft.AspNetCore.Server.IIS
             http_get_completion_info(pCompletionInfo, out cbBytes, out hr);
         }
 
-        public static void HttpSetManagedContext(IntPtr pInProcessHandler, IntPtr pvManagedContext)
+        public static void HttpSetManagedContext(NativeSafeHandle pInProcessHandler, IntPtr pvManagedContext)
         {
             Validate(http_set_managed_context(pInProcessHandler, pvManagedContext));
         }
@@ -237,18 +237,18 @@ namespace Microsoft.AspNetCore.Server.IIS
             return iisConfigurationData;
         }
 
-        public static bool HttpTryGetServerVariable(IntPtr pInProcessHandler, string variableName, out string value)
+        public static bool HttpTryGetServerVariable(NativeSafeHandle pInProcessHandler, string variableName, out string value)
         {
             return http_get_server_variable(pInProcessHandler, variableName, out value) == 0;
         }
 
-        public static void HttpSetServerVariable(IntPtr pInProcessHandler, string variableName, string value)
+        public static void HttpSetServerVariable(NativeSafeHandle pInProcessHandler, string variableName, string value)
         {
             Validate(http_set_server_variable(pInProcessHandler, variableName, value));
         }
 
         public static unsafe int HttpWebsocketsReadBytes(
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             byte* pvBuffer,
             int cbBuffer,
             PFN_WEBSOCKET_ASYNC_COMPLETION pfnCompletionCallback,
@@ -259,7 +259,7 @@ namespace Microsoft.AspNetCore.Server.IIS
         }
 
         internal static unsafe int HttpWebsocketsWriteBytes(
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             HttpApiTypes.HTTP_DATA_CHUNK* pDataChunks,
             int nChunks,
             PFN_WEBSOCKET_ASYNC_COMPLETION pfnCompletionCallback,
@@ -269,12 +269,12 @@ namespace Microsoft.AspNetCore.Server.IIS
             return http_websockets_write_bytes(pInProcessHandler, pDataChunks, nChunks, pfnCompletionCallback, pvCompletionContext, out fCompletionExpected);
         }
 
-        public static void HttpEnableWebsockets(IntPtr pInProcessHandler)
+        public static void HttpEnableWebsockets(NativeSafeHandle pInProcessHandler)
         {
             Validate(http_enable_websockets(pInProcessHandler));
         }
 
-        public static bool HttpTryCancelIO(IntPtr pInProcessHandler)
+        public static bool HttpTryCancelIO(NativeSafeHandle pInProcessHandler)
         {
             var hr = http_cancel_io(pInProcessHandler);
             // ERROR_NOT_FOUND is expected if async operation finished
@@ -288,22 +288,22 @@ namespace Microsoft.AspNetCore.Server.IIS
             return true;
         }
 
-        public static void HttpCloseConnection(IntPtr pInProcessHandler)
+        public static void HttpCloseConnection(NativeSafeHandle pInProcessHandler)
         {
             Validate(http_close_connection(pInProcessHandler));
         }
 
-        public static unsafe void HttpResponseSetUnknownHeader(IntPtr pInProcessHandler, byte* pszHeaderName, byte* pszHeaderValue, ushort usHeaderValueLength, bool fReplace)
+        public static unsafe void HttpResponseSetUnknownHeader(NativeSafeHandle pInProcessHandler, byte* pszHeaderName, byte* pszHeaderValue, ushort usHeaderValueLength, bool fReplace)
         {
             Validate(http_response_set_unknown_header(pInProcessHandler, pszHeaderName, pszHeaderValue, usHeaderValueLength, fReplace));
         }
 
-        public static unsafe void HttpResponseSetKnownHeader(IntPtr pInProcessHandler, int headerId, byte* pHeaderValue, ushort length, bool fReplace)
+        public static unsafe void HttpResponseSetKnownHeader(NativeSafeHandle pInProcessHandler, int headerId, byte* pHeaderValue, ushort length, bool fReplace)
         {
             Validate(http_response_set_known_header(pInProcessHandler, headerId, pHeaderValue, length, fReplace));
         }
 
-        public static void HttpGetAuthenticationInformation(IntPtr pInProcessHandler, out string authType, out IntPtr token)
+        public static void HttpGetAuthenticationInformation(NativeSafeHandle pInProcessHandler, out string authType, out IntPtr token)
         {
             Validate(http_get_authentication_information(pInProcessHandler, out authType, out token));
         }
@@ -316,17 +316,17 @@ namespace Microsoft.AspNetCore.Server.IIS
             }
         }
 
-        internal static unsafe void HttpResponseSetTrailer(IntPtr pInProcessHandler, byte* pHeaderName, byte* pHeaderValue, ushort length, bool replace)
+        internal static unsafe void HttpResponseSetTrailer(NativeSafeHandle pInProcessHandler, byte* pHeaderName, byte* pHeaderValue, ushort length, bool replace)
         {
             Validate(http_response_set_trailer(pInProcessHandler, pHeaderName, pHeaderValue, length, false));
         }
 
-        internal static unsafe void HttpResetStream(IntPtr pInProcessHandler, ulong errorCode)
+        internal static unsafe void HttpResetStream(NativeSafeHandle pInProcessHandler, ulong errorCode)
         {
             Validate(http_reset_stream(pInProcessHandler, errorCode));
         }
 
-        internal static unsafe bool HttpSupportTrailer(IntPtr pInProcessHandler)
+        internal static unsafe bool HttpSupportTrailer(NativeSafeHandle pInProcessHandler)
         {
             bool supportsTrailers;
             Validate(http_has_response4(pInProcessHandler, out supportsTrailers));

src/Servers/IIS/IIS/src/WebHostBuilderIISExtensions.cs

@@ -36,7 +36,7 @@ namespace Microsoft.AspNetCore.Hosting
                 hostBuilder.UseContentRoot(contentRoot);
                 return hostBuilder.ConfigureServices(
                     services => {
-                        services.AddSingleton(new IISNativeApplication(iisConfigData.pNativeApplication));
+                        services.AddSingleton(new IISNativeApplication(new NativeSafeHandle(iisConfigData.pNativeApplication)));
                         services.AddSingleton<IServer, IISHttpServer>();
                         services.AddSingleton<IStartupFilter>(new IISServerSetupFilter(iisConfigData.pwzVirtualApplicationPath));
                         services.AddAuthenticationCore();