/*
* Copyright (c) 2024 Vaughn Nugent
*
* Library: VNLib
* Package: VNLib.Plugins.Essentials
* File: WebSocketSession.cs
*
* WebSocketSession.cs is part of VNLib.Plugins.Essentials which is part of the larger
* VNLib collection of libraries and utilities.
*
* VNLib.Plugins.Essentials is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as
* published by the Free Software Foundation, either version 3 of the
* License, or (at your option) any later version.
*
* VNLib.Plugins.Essentials is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see https://www.gnu.org/licenses/.
*/
using System;
using System.IO;
using System.Threading;
using System.Net.WebSockets;
using System.Threading.Tasks;
using VNLib.Net.Http;
namespace VNLib.Plugins.Essentials
{
///
/// A callback method to invoke when an HTTP service successfully transfers protocols to
/// the WebSocket protocol and the socket is ready to be used
///
/// The open websocket session instance
///
/// A that will be awaited by the HTTP layer. When the task completes, the transport
/// will be closed and the session disposed
///
public delegate Task WebSocketAcceptedCallback(WebSocketSession session);
///
/// A callback method to invoke when an HTTP service successfully transfers protocols to
/// the WebSocket protocol and the socket is ready to be used
///
/// The type of the user state object
/// The open websocket session instance
///
/// A that will be awaited by the HTTP layer. When the task completes, the transport
/// will be closed and the session disposed
///
public delegate Task WebSocketAcceptedCallback(WebSocketSession session);
///
/// Represents a wrapper to manage the lifetime of the captured
/// connection context and the underlying transport. This session is managed by the parent
/// that it was created on.
///
public class WebSocketSession(WebSocketAcceptedCallback callback) : AlternateProtocolBase
{
internal WebSocket? WsHandle;
internal readonly WebSocketAcceptedCallback AcceptedCallback = callback;
///
/// A cancellation token that can be monitored to reflect the state
/// of the webscocket
///
public CancellationToken Token => CancelSource.Token;
///
/// Id assigned to this instance on creation
///
public required string SocketID { get; init; }
///
/// Negotiated sub-protocol
///
public string? SubProtocol { get; internal init; }
///
/// The websocket keep-alive interval
///
internal TimeSpan KeepAlive { get; init; }
///
/// Initialzes the created websocket with the specified protocol
///
/// Transport stream to use for the websocket
/// The accept callback function specified during object initialization
protected override async Task RunAsync(Stream transport)
{
try
{
WebSocketCreationOptions ce = new()
{
IsServer = true,
KeepAliveInterval = KeepAlive,
SubProtocol = SubProtocol,
};
//Create a new websocket from the context stream
WsHandle = WebSocket.CreateFromStream(transport, ce);
//Register token to abort the websocket so the managed ws uses the non-fallback send/recv method
using CancellationTokenRegistration abortReg = Token.Register(WsHandle.Abort);
//Return the callback function to explcitly invoke it
await AcceptedCallback(this);
}
finally
{
WsHandle?.Dispose();
}
}
///
/// Asynchronously receives data from the Websocket and copies the data to the specified buffer
///
/// The buffer to store read data
/// A task that resolves a which contains the status of the operation
///
public Task ReceiveAsync(ArraySegment buffer) => WsHandle!.ReceiveAsync(buffer, CancellationToken.None);
///
/// Asynchronously receives data from the Websocket and copies the data to the specified buffer
///
/// The buffer to store read data
///
///
public ValueTask ReceiveAsync(Memory buffer) => WsHandle!.ReceiveAsync(buffer, CancellationToken.None);
///
/// Asynchronously sends the specified buffer to the client of the specified type
///
/// The buffer containing data to send
/// The message/data type of the packet to send
/// A value that indicates this message is the final message of the transaction
///
///
public Task SendAsync(ArraySegment buffer, WebSocketMessageType type, bool endOfMessage) => WsHandle!.SendAsync(buffer, type, endOfMessage, CancellationToken.None);
///
/// Asynchronously sends the specified buffer to the client of the specified type
///
/// The buffer containing data to send
/// The message/data type of the packet to send
/// A value that indicates this message is the final message of the transaction
///
///
public ValueTask SendAsync(ReadOnlyMemory buffer, WebSocketMessageType type, bool endOfMessage)
{
//Begin receive operation only with the internal token
return SendAsync(buffer, type, endOfMessage ? WebSocketMessageFlags.EndOfMessage : WebSocketMessageFlags.None);
}
///
/// Asynchronously sends the specified buffer to the client of the specified type
///
/// The buffer containing data to send
/// The message/data type of the packet to send
/// Websocket message flags
///
///
public ValueTask SendAsync(ReadOnlyMemory buffer, WebSocketMessageType type, WebSocketMessageFlags flags)
=> WsHandle!.SendAsync(buffer, type, flags, CancellationToken.None);
///
/// Properly closes a currently connected websocket
///
/// Set the close status
/// Set the close reason
///
public Task CloseSocketAsync(WebSocketCloseStatus status, string reason) => WsHandle!.CloseAsync(status, reason, CancellationToken.None);
///
///
///
///
///
///
///
public Task CloseSocketOutputAsync(WebSocketCloseStatus status, string reason, CancellationToken cancellation = default)
{
if (WsHandle!.State == WebSocketState.Open || WsHandle.State == WebSocketState.CloseSent)
{
return WsHandle.CloseOutputAsync(status, reason, cancellation);
}
return Task.CompletedTask;
}
}
///
///
///
/// The user-state type
public sealed class WebSocketSession : WebSocketSession
{
#nullable disable
///
/// A user-defined state object passed during socket accept handshake
///
public T UserState { get; internal init; }
#nullable enable
internal WebSocketSession(WebSocketAcceptedCallback callback)
: base((ses) => callback((ses as WebSocketSession)!))
{
UserState = default;
}
}
}