Call HttpContext.Abort() when WebSocket.Abort() is called. (#48892)

* Call HttpContext.Abort() when WebSocket.Abort() is called.

Author: mitchdenny

src/Middleware/WebSockets/src/ServerWebSocket.cs

@@ -0,0 +1,80 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Net.WebSockets;
+using Microsoft.AspNetCore.Http;
+
+namespace Microsoft.AspNetCore.WebSockets;
+
+/// <summary>
+/// Used in ASP.NET Core to wrap a WebSocket with its associated HttpContext so that when the WebSocket is aborted
+/// the underlying HttpContext is aborted. All other methods are delegated to the underlying WebSocket.
+/// </summary>
+internal sealed class ServerWebSocket : WebSocket
+{
+    private readonly WebSocket _wrappedSocket;
+    private readonly HttpContext _context;
+
+    internal ServerWebSocket(WebSocket wrappedSocket, HttpContext context)
+    {
+        ArgumentNullException.ThrowIfNull(wrappedSocket);
+        ArgumentNullException.ThrowIfNull(context);
+
+        _wrappedSocket = wrappedSocket;
+        _context = context;
+    }
+
+    public override WebSocketCloseStatus? CloseStatus => _wrappedSocket.CloseStatus;
+
+    public override string? CloseStatusDescription => _wrappedSocket.CloseStatusDescription;
+
+    public override WebSocketState State => _wrappedSocket.State;
+
+    public override string? SubProtocol => _wrappedSocket.SubProtocol;
+
+    public override void Abort()
+    {
+        _wrappedSocket.Abort();
+        _context.Abort();
+    }
+
+    public override Task CloseAsync(WebSocketCloseStatus closeStatus, string? statusDescription, CancellationToken cancellationToken)
+    {
+        return _wrappedSocket.CloseAsync(closeStatus, statusDescription, cancellationToken);
+    }
+
+    public override Task CloseOutputAsync(WebSocketCloseStatus closeStatus, string? statusDescription, CancellationToken cancellationToken)
+    {
+        return _wrappedSocket.CloseOutputAsync(closeStatus, statusDescription, cancellationToken);
+    }
+
+    public override void Dispose()
+    {
+        _wrappedSocket.Dispose();
+    }
+
+    public override Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> buffer, CancellationToken cancellationToken)
+    {
+        return _wrappedSocket.ReceiveAsync(buffer, cancellationToken);
+    }
+
+    public override ValueTask<ValueWebSocketReceiveResult> ReceiveAsync(Memory<byte> buffer, CancellationToken cancellationToken)
+    {
+        return _wrappedSocket.ReceiveAsync(buffer, cancellationToken);
+    }
+
+    public override ValueTask SendAsync(ReadOnlyMemory<byte> buffer, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancellationToken)
+    {
+        return _wrappedSocket.SendAsync(buffer, messageType, endOfMessage, cancellationToken);
+    }
+
+    public override ValueTask SendAsync(ReadOnlyMemory<byte> buffer, WebSocketMessageType messageType, WebSocketMessageFlags messageFlags, CancellationToken cancellationToken)
+    {
+        return _wrappedSocket.SendAsync(buffer, messageType, messageFlags, cancellationToken);
+    }
+
+    public override Task SendAsync(ArraySegment<byte> buffer, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancellationToken)
+    {
+        return _wrappedSocket.SendAsync(buffer, messageType, endOfMessage, cancellationToken);
+    }
+}

src/Middleware/WebSockets/src/WebSocketMiddleware.cs

@@ -208,13 +208,15 @@ public async Task<WebSocket> AcceptAsync(WebSocketAcceptContext acceptContext)
             // Disable request timeout, if there is one, after the websocket has been accepted
             _context.Features.Get<IHttpRequestTimeoutFeature>()?.DisableTimeout();
 
-            return WebSocket.CreateFromStream(opaqueTransport, new WebSocketCreationOptions()
+            var wrappedSocket = WebSocket.CreateFromStream(opaqueTransport, new WebSocketCreationOptions()
             {
                 IsServer = true,
                 KeepAliveInterval = keepAliveInterval,
                 SubProtocol = subProtocol,
                 DangerousDeflateOptions = deflateOptions
             });
+
+            return new ServerWebSocket(wrappedSocket, _context);
         }
 
         public static bool CheckSupportedWebSocketRequest(string method, IHeaderDictionary requestHeaders)