// 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);
    }
}