General cleanup in IIS (#24754)

Author: jkotalik

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

@@ -271,7 +271,7 @@ string IServerVariablesFeature.this[string variableName]
                 // 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 @@ string IServerVariablesFeature.this[string variableName]
                 // 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 @@ async Task<Stream> IHttpUpgradeFeature.UpgradeAsync()
             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 @@ unsafe X509Certificate2 ITlsConnectionFeature.ClientCertificate
         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 @@ IHeaderDictionary IHttpResponseTrailersFeature.Trailers
         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 @@ void IHttpResetFeature.Reset(int errorCode)
 
         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 @@ public void Abort(Exception reason)
         {
             _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 @@ internal abstract partial class IISHttpContext : NativeRequestContext, IThreadPo
         private const int PauseWriterThreshold = 65536;
         private const int ResumeWriterTheshold = PauseWriterThreshold / 2;
 
-        protected readonly IntPtr _pInProcessHandler;
+        protected readonly NativeSafeHandle _requestNativeHandle;
 
         private readonly IISServerOptions _options;
 
@@ -75,15 +75,15 @@ internal abstract partial class IISHttpContext : NativeRequestContext, IThreadPo
 
         internal unsafe IISHttpContext(
             MemoryPool<byte> memoryPool,
-            IntPtr pInProcessHandler,
+            NativeSafeHandle pInProcessHandler,
             IISServerOptions options,
             IISHttpServer server,
             ILogger logger,
             bool useLatin1)
             : 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 @@ protected void InitializeContext()
 
                 ResetFeatureCollection();
 
-                if (!_server.IsWebSocketAvailable(_pInProcessHandler))
+                if (!_server.IsWebSocketAvailable(_requestNativeHandle))
                 {
                     _currentIHttpUpgradeFeature = null;
                 }
@@ -198,7 +198,7 @@ protected void InitializeContext()
                 _bodyOutput = new OutputProducer(pipe);
             }
 
-            NativeMethods.HttpSetManagedContext(_pInProcessHandler, (IntPtr)_thisHandle);
+            NativeMethods.HttpSetManagedContext(_requestNativeHandle, (IntPtr)_thisHandle);
         }
 
         private string GetOriginalPath()
@@ -324,7 +324,7 @@ private void EnsureIOInitialized()
             // 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 @@ public unsafe void SetResponseHeaders()
             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 @@ public unsafe void SetResponseHeaders()
                             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 @@ public unsafe void SetResponseTrailers()
                         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 @@ protected void ReportApplicationError(Exception ex)
 
         public void PostCompletion(NativeMethods.REQUEST_NOTIFICATION_STATUS requestNotificationStatus)
         {
-            NativeMethods.HttpSetCompletionStatus(_pInProcessHandler, requestNotificationStatus);
-            NativeMethods.HttpPostCompletion(_pInProcessHandler, 0);
+            NativeMethods.HttpSetCompletionStatus(_requestNativeHandle, requestNotificationStatus);
+            NativeMethods.HttpPostCompletion(_requestNativeHandle, 0);
         }
 
         internal void OnAsyncCompletion(int hr, int bytes)
@@ -628,7 +628,7 @@ private void ThrowResponseAlreadyStartedException(string name)
 
         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 @@ private async Task HandleRequest()
                 // 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 @@ internal class IISHttpContextOfT<TContext> : IISHttpContext
     {
         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 @@ public override async Task<bool> ProcessRequestAsync()
                     // 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 @@ internal class IISHttpServer : IServer
         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 @@ public Task StopAsync(CancellationToken cancellationToken)
 
         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 @@ private static NativeMethods.REQUEST_NOTIFICATION_STATUS HandleRequest(IntPtr pI
                 // 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 @@ private static bool HandleShutdown(IntPtr pvRequestContext)
             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 @@ private static void OnDisconnect(IntPtr pvManagedHttpContext)
             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 @@ private static NativeMethods.REQUEST_NOTIFICATION_STATUS OnAsyncCompletion(IntPt
             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 @@ private static void OnRequestsDrained(IntPtr serverContext)
             {
                 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 @@ public IISContextFactory(MemoryPool<byte> memoryPool, IHttpApplication<T> applic
                 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 @@ public IISHttpContext CreateHttpContext(IntPtr pInProcessHandler)
     // 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 @@ public void RegisterCallbacks(
         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,15 +11,15 @@ internal class AsyncFlushOperation : AsyncIOOperation
         {
             private readonly AsyncIOEngine _engine;
 
-            private IntPtr _requestHandler;
+            private NativeSafeHandle _requestHandler;
             private bool _moreData;
 
             public AsyncFlushOperation(AsyncIOEngine engine)
             {
                 _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 @@ internal class AsyncReadOperation : AsyncIOOperation
 
             private MemoryHandle _inputHandle;
 
-            private IntPtr _requestHandler;
+            private NativeSafeHandle _requestHandler;
 
             private Memory<byte> _memory;
 
@@ -23,7 +23,7 @@ public AsyncReadOperation(AsyncIOEngine engine)
                 _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 @@ public AsyncWriteOperation(AsyncIOEngine engine)
                 _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);