Skip to content

Instantly share code, notes, and snippets.

@camnewnham
Last active August 20, 2024 06:28
Show Gist options
  • Select an option

  • Save camnewnham/261312127a9d2bc653b60d029f9652db to your computer and use it in GitHub Desktop.

Select an option

Save camnewnham/261312127a9d2bc653b60d029f9652db to your computer and use it in GitHub Desktop.
Simple WebSocket server and client extensions.
using System;
using System.IO;
using System.Net;
using System.Net.WebSockets;
using System.Threading;
using System.Threading.Tasks;
/// <summary>
/// Extension methods to make working with websockets easier.
/// </summary>
public static class WebSockets
{
/// <summary>
/// Creates a websocket server that runs until the task is cancelled.
/// </summary>
/// <param name="url">The http URL URL to listen on. Typically http://localhost:8080/ </param>
/// <param name="onConnection">A method to call when a connection is established</param>
/// <param name="token">A cancellation token to close the server.</param>
/// <returns>A task that completes when the server closes.</returns>
public static async Task Listen(string url, Action<WebSocket> onConnection, CancellationToken token = default)
{
HttpListener httpListener = new HttpListener();
httpListener.Prefixes.Add(url);
httpListener.Start();
while (!token.IsCancellationRequested)
{
Task<HttpListenerContext> getContextTask = httpListener.GetContextAsync();
CancellationTokenSource localCts = CancellationTokenSource.CreateLinkedTokenSource(token);
Task completed = await Task.WhenAny(getContextTask, Task.Delay(Timeout.Infinite, localCts.Token));
localCts.Dispose();
if (completed != getContextTask)
{
break;
}
HttpListenerContext context = getContextTask.Result;
if (context.Request.IsWebSocketRequest)
{
Task<HttpListenerWebSocketContext> wsAcceptTask = context.AcceptWebSocketAsync(null);
CancellationTokenSource localCts2 = CancellationTokenSource.CreateLinkedTokenSource(token);
_ = Task.WhenAny(wsAcceptTask, Task.Delay(Timeout.Infinite, localCts2.Token)).ContinueWith(t =>
{
if (t.Result == wsAcceptTask)
{
onConnection?.Invoke(wsAcceptTask.Result.WebSocket);
}
});
}
}
httpListener?.Close();
}
/// <summary>
/// Begins receiving messages from the websocket
/// </summary>
/// <param name="ws">The websocket</param>
/// <param name="onText">A method to call when a text message is received</param>
/// <param name="onBinary">A method to call when a binary message is received</param>
/// <param name="token">A cancellation token</param>
/// <returns>A task that completes when the websocket is closed.</returns>
public static async Task Receive(this WebSocket ws, Action<string> onText = null, Action<MemoryStream> onBinary = null, CancellationToken token = default)
{
ArraySegment<byte> buffer = new ArraySegment<byte>(new byte[1024]);
MemoryStream memoryStream = new MemoryStream();
try
{
while (ws.State == WebSocketState.Open)
{
WebSocketReceiveResult result = await ws.ReceiveAsync(buffer, token);
if (result.MessageType == WebSocketMessageType.Close)
{
return;
}
if (memoryStream.Position == 0 && result.EndOfMessage)
{
if (result.MessageType == WebSocketMessageType.Text)
{
if (onText != null)
{
string data = System.Text.Encoding.UTF8.GetString(buffer.Array, 0, result.Count);
onText?.Invoke(data);
}
}
else
{
onBinary?.Invoke(new MemoryStream(buffer.Array, 0, result.Count));
}
}
else
{
await memoryStream.WriteAsync(buffer.Array, 0, result.Count); // Append stream
if (result.EndOfMessage)
{
if (result.MessageType == WebSocketMessageType.Text)
{
if (onText != null)
{
using (StreamReader reader = new StreamReader(memoryStream)) // End stream
{
memoryStream.Seek(0, SeekOrigin.Begin);
onText?.Invoke(reader.ReadToEnd());
}
memoryStream = new MemoryStream();
}
}
else
{
onBinary?.Invoke(memoryStream);
memoryStream.SetLength(0);
}
}
}
}
}
finally
{
memoryStream.Dispose();
}
}
/// <summary>
/// Sends a message on the websocket as text
/// </summary>
/// <param name="ws">The websocket</param>
/// <param name="message">The message to send</param>
/// <param name="token">A token to cancel sending</param>
/// <returns>A task that completes when sending is complete.</returns>
public static async Task SendAsync(this WebSocket ws, string message, CancellationToken token = default)
{
await ws.SendAsync(new ArraySegment<byte>(System.Text.Encoding.UTF8.GetBytes(message)), WebSocketMessageType.Text, true, token);
}
/// <summary>
/// Sends a message on the websocket as binary
/// </summary>
/// <param name="ws">The websocket</param>
/// <param name="data">The message to send</param>
/// <param name="token">A token to cancel sending</param>
/// <returns>A task that completes when sending is complete.</returns>
public static async Task SendAsync(this WebSocket ws, byte[] data, CancellationToken token = default)
{
await SendAsync(ws, new ArraySegment<byte>(data), token);
}
/// <summary>
/// Sends a message on the websocket as binary
/// </summary>
/// <param name="ws">The websocket</param>
/// <param name="data">The message to send</param>
/// <param name="token">A token to cancel sending</param>
/// <returns>A task that completes when sending is complete.</returns>
public static async Task SendAsync(this WebSocket ws, ArraySegment<byte> data, CancellationToken token = default)
{
await ws.SendAsync(data, WebSocketMessageType.Binary, true, token);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment