From 447e7d4f8e3efde1385d08e7e314ef14e6db0d21 Mon Sep 17 00:00:00 2001 From: Roman Konecny Date: Mon, 31 Mar 2025 15:59:59 +0200 Subject: [PATCH 1/4] Cleanup connection after managed webscoket dispose --- .../WebSockets/HttpWebSocket.Managed.cs | 70 ++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs index 7c925cb4e315a6..1660e414347263 100644 --- a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs +++ b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs @@ -1,6 +1,7 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. +using System.Threading; using System.Threading.Tasks; namespace System.Net.WebSockets @@ -56,7 +57,12 @@ internal static async Task AcceptWebSocketAsyncCor // Send websocket handshake headers await responseStream.WriteWebSocketHandshakeHeadersAsync().ConfigureAwait(false); - WebSocket webSocket = WebSocket.CreateFromStream(context.Connection.ConnectedStream, isServer: true, subProtocol, keepAliveInterval); + WebSocket rawWebSocket = WebSocket.CreateFromStream(context.Connection.ConnectedStream, isServer: true, subProtocol, keepAliveInterval); + + // Wrap the raw websocket so its Dispose() automatically closes the connection. + // This is important as common and recommended usage is to just Dispose the HttpListenerWebSocketContext.WebSocket + // and without it connection would not be cleaned up causing memory leaks and possible other issues. + HttpListenerWrappedWebSocket webSocket = new(rawWebSocket, context); HttpListenerWebSocketContext webSocketContext = new HttpListenerWebSocketContext( request.Url!, @@ -76,5 +82,67 @@ internal static async Task AcceptWebSocketAsyncCor } private const bool WebSocketsSupported = true; + + /// + /// Wraps the real WebSocket so that when Dispose() is called, + /// we also call context.Response.Close(), triggering UnregisterContext. + /// + private sealed class HttpListenerWrappedWebSocket : WebSocket + { + private readonly WebSocket _inner; + private readonly HttpListenerContext _context; + private bool _disposed; + + internal HttpListenerWrappedWebSocket(WebSocket inner, HttpListenerContext context) + { + _inner = inner; + _context = context; + } + + public override WebSocketCloseStatus? CloseStatus => _inner.CloseStatus; + public override string? CloseStatusDescription => _inner.CloseStatusDescription; + public override WebSocketState State => _inner.State; + public override string? SubProtocol => _inner.SubProtocol; + + public override void Abort() => _inner.Abort(); + + public override Task CloseAsync( + WebSocketCloseStatus closeStatus, + string? statusDescription, + CancellationToken cancellationToken) => + _inner.CloseAsync(closeStatus, statusDescription, cancellationToken); + + public override Task CloseOutputAsync( + WebSocketCloseStatus closeStatus, + string? statusDescription, + CancellationToken cancellationToken) => + _inner.CloseOutputAsync(closeStatus, statusDescription, cancellationToken); + + public override void Dispose() + { + if (_disposed) + return; + _disposed = true; + + // Dispose the underlying raw WebSocket + _inner.Dispose(); + + // Close the response stream, which will close and clean connection + // unregistering the context and freeing it from memory. + _context.Response.Close(); + } + + public override Task ReceiveAsync( + ArraySegment buffer, + CancellationToken cancellationToken) => + _inner.ReceiveAsync(buffer, cancellationToken); + + public override Task SendAsync( + ArraySegment buffer, + WebSocketMessageType messageType, + bool endOfMessage, + CancellationToken cancellationToken) => + _inner.SendAsync(buffer, messageType, endOfMessage, cancellationToken); + } } } From 200aee570cf0800dca64ec5bc13ca4f963f6bc23 Mon Sep 17 00:00:00 2001 From: Roman Konecny Date: Tue, 1 Apr 2025 14:10:32 +0200 Subject: [PATCH 2/4] Minor commenting changes --- .../System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs index 1660e414347263..981315d692d657 100644 --- a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs +++ b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs @@ -61,7 +61,7 @@ internal static async Task AcceptWebSocketAsyncCor // Wrap the raw websocket so its Dispose() automatically closes the connection. // This is important as common and recommended usage is to just Dispose the HttpListenerWebSocketContext.WebSocket - // and without it connection would not be cleaned up causing memory leaks and possible other issues. + // and without it connection would not be cleaned up causing memory leaks and other issues. HttpListenerWrappedWebSocket webSocket = new(rawWebSocket, context); HttpListenerWebSocketContext webSocketContext = new HttpListenerWebSocketContext( @@ -127,8 +127,9 @@ public override void Dispose() // Dispose the underlying raw WebSocket _inner.Dispose(); - // Close the response stream, which will close and clean connection - // unregistering the context and freeing it from memory. + // Ensure we remove the context from the HttpListener's tracking dictionary + // by forcibly calling Response.Close(). This closes the underlying connection + // and also calls UnregisterContext(context), preventing stale HttpListenerContext objects causing memory leak. _context.Response.Close(); } From 8ee17bf8d0f7c4d7549c67ff80c094f6618135d5 Mon Sep 17 00:00:00 2001 From: Roman Konecny Date: Wed, 9 Apr 2025 17:07:00 +0200 Subject: [PATCH 3/4] Clean up in finalizer to avoid memory leak if user is not disposing properly. --- .../WebSockets/HttpWebSocket.Managed.cs | 40 ++++++++++++++++--- 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs index 981315d692d657..8fcf6c6b974301 100644 --- a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs +++ b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs @@ -119,18 +119,46 @@ public override Task CloseOutputAsync( _inner.CloseOutputAsync(closeStatus, statusDescription, cancellationToken); public override void Dispose() + { + Dispose(true); + GC.SuppressFinalize(this); + } + + private void Dispose(bool disposing) { if (_disposed) return; _disposed = true; - // Dispose the underlying raw WebSocket - _inner.Dispose(); + if (disposing) + { + // Dispose the underlying raw WebSocket + _inner.Dispose(); + // Ensure we remove the context from the HttpListener's tracking dictionary + // by forcibly calling Response.Close(). This closes the underlying connection + // and also calls UnregisterContext(context), preventing stale HttpListenerContext objects causing memory leak. + _context.Response.Close(); + } + else + { + // Technically we shouldn't be doing the following work when disposing == false, + // as the following work relies on other finalizable objects. + // But given we have little choice: if someone drops the websocket without + // disposing of it we need to close the context to prevent memory leaks. + try + { + _context.Response.Close(); + } + catch + { + // explain why we ignore exceptions here + } + } + } - // Ensure we remove the context from the HttpListener's tracking dictionary - // by forcibly calling Response.Close(). This closes the underlying connection - // and also calls UnregisterContext(context), preventing stale HttpListenerContext objects causing memory leak. - _context.Response.Close(); + ~HttpListenerWrappedWebSocket() + { + Dispose(false); } public override Task ReceiveAsync( From c60128603316baa8b50f49699942e87aa57b46e1 Mon Sep 17 00:00:00 2001 From: Roman Konecny Date: Wed, 9 Apr 2025 23:38:15 +0200 Subject: [PATCH 4/4] Missing comment --- .../System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs index 8fcf6c6b974301..aa83f0eb3f8803 100644 --- a/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs +++ b/src/libraries/System.Net.HttpListener/src/System/Net/Managed/WebSockets/HttpWebSocket.Managed.cs @@ -151,7 +151,9 @@ private void Dispose(bool disposing) } catch { - // explain why we ignore exceptions here + // We are doing best effort here to handle the case where the user does not properly Dispose the WebSocket. + // If we fail to close the reponse, we are not going to throw an exception and possibly crash. + // We are just going to ignore it and let the GC do its best. } } }