A middleware to manage Web Socket in ASP.NET Core

Introduction

ASP.NET Core SignalR is a useful library that simplifies the management of real-time in web applications. In this case I rather use WebSockets because I wanted more flexibility and compatibility with any WebSocket client. 
In Microsoft's documentation I found a great working example of WebSockets. It remains to manage the connections to be able to broadcast messages from one connection  to the others, feature that come out of the box with SignalR. Expecting this logic to be really complex, I wanted to remove it from the Startup class itself.

Background

To read about WebSockets support in ASP.NET Core, you can check here.
If you want to read about middleware and how to write it in ASP.NET Core, read this link.

Using the code

First, you have to add the Microsoft.AspNetCore.WebSockets package to your project.

Now, you can create an extension method and class to manage the WebSockets:

public static class WebSocketExtensions
{
    public static IApplicationBuilder UseCustomWebSocketManager(this IApplicationBuilder app)
    {
       return app.UseMiddleware<CustomWebSocketManager>();
    }
}

public class CustomWebSocketManager
{
    private readonly RequestDelegate _next;

    public CustomWebSocketManager(RequestDelegate next)
    {
       _next = next;
    }

    public async Task Invoke(HttpContext context, ICustomWebSocketFactory wsFactory, ICustomWebSocketMessageHandler wsmHandler)
    {
        if (context.Request.Path == "/ws")
        {
            if (context.WebSockets.IsWebSocketRequest)
            {
                string username = context.Request.Query["u"];
                if (!string.IsNullOrEmpty(username))
                {
                    WebSocket webSocket = await context.WebSockets.AcceptWebSocketAsync();
                    CustomWebSocket userWebSocket = new CustomWebSocket()
                    {
                       WebSocket = webSocket,
                       Username = username
                    };
                    wsFactory.Add(userWebSocket);
                    await wsmHandler.SendInitialMessages(userWebSocket);
                    await Listen(context, userWebSocket, wsFactory, wsmHandler);
                }
            }
            else
            {
                 context.Response.StatusCode = 400;
            }
        }
        await _next(context);
    }

    private async Task Listen(HttpContext context, CustomWebSocket userWebSocket, ICustomWebSocketFactory wsFactory, ICustomWebSocketMessageHandler wsmHandler)
    {
        WebSocket webSocket = userWebSocket.WebSocket;
        var buffer = new byte[1024 * 4];
        WebSocketReceiveResult result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
        while (!result.CloseStatus.HasValue)
        {
             await wsmHandler.HandleMessage(result, buffer, userWebSocket, wsFactory);
             buffer = new byte[1024 * 4];
             result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
        } 
        wsFactory.Remove(userWebSocket.Username);
        await webSocket.CloseAsync(result.CloseStatus.Value, result.CloseStatusDescription, CancellationToken.None);
    }
}

In this case, WebSockets requests contains always "/ws" in the url. The query string contains a parameter u for the username that associate the WebSocket to the login user.

CustomWebSocket is a class containing a WebSocket and the username:

public class CustomWebSocket
{
   public WebSocket WebSocket { get; set; }
   public string Username { get; set; }
}

I also create a custom WebSocket message:

class CustomWebSocketMessage
{
   public string Text { get; set; }
   public DateTime MessagDateTime { get; set; }
   public string Username { get; set; }
   public WSMessageType Type { get; set; }
}

where the Type is an enumeration of the different type of message you may have.

In the Startup class you have to register the following services:

services.AddSingleton<ICustomWebSocketFactory, CustomWebSocketFactory>();
services.AddSingleton<ICustomWebSocketMessageHandler, CustomWebSocketMessageHandler>();

where CustomWebSocketFactory take care of gathering together the list of connected WebSockets: 

public interface ICustomWebSocketFactory
{
   void Add(CustomWebSocket uws);
   void Remove(string username);
   List<CustomWebSocket> All();
   List<CustomWebSocket> Others(CustomWebSocket client);
   CustomWebSocket Client(string username);
}

public class CustomWebSocketFactory : ICustomWebSocketFactory
{
   List<CustomWebSocket> List;

   public CustomWebSocketFactory()
   {
      List = new List<CustomWebSocket>();
   }

   public void Add(CustomWebSocket uws)
   {
      List.Add(uws);
   }

   //when disconnect
   public void Remove(string username) 
   {
      List.Remove(Client(username));
   }

   public List<CustomWebSocket> All()
   {
      return List;
   }
   
   public List<CustomWebSocket> Others(CustomWebSocket client)
   {
      return List.Where(c => c.Username != client.Username).ToList();
   }
 
   public CustomWebSocket Client(string username)
   {
      return List.First(c=>c.Username == username);
   }
}

and CustomWebSocketMessageHandler contains the logic regarding messages (i.e. if any message needs to be sent upon connection and how to react to the incoming messages)

public interface ICustomWebSocketMessageHandler
{
   Task SendInitialMessages(CustomWebSocket userWebSocket);
   Task HandleMessage(WebSocketReceiveResult result, byte[] buffer, CustomWebSocket userWebSocket, ICustomWebSocketFactory wsFactory);
   Task BroadcastOthers(byte[] buffer, CustomWebSocket userWebSocket, ICustomWebSocketFactory wsFactory);
   Task BroadcastAll(byte[] buffer, CustomWebSocket userWebSocket, ICustomWebSocketFactory wsFactory);
}

public class CustomWebSocketMessageHandler : ICustomWebSocketMessageHandler
{
   public async Task SendInitialMessages(CustomWebSocket userWebSocket)
   {
      WebSocket webSocket = userWebSocket.WebSocket;
      var msg = new CustomWebSocketMessage
      {
         MessagDateTime = DateTime.Now,
         Type = WSMessageType.anyType,
         Text = anyText,
         Username = "system"
      };

      string serialisedMessage = JsonConvert.SerializeObject(msg);
      byte[] bytes = Encoding.ASCII.GetBytes(serialisedMessage);
      await webSocket.SendAsync(new ArraySegment<byte>(bytes, 0, bytes.Length), WebSocketMessageType.Text, true, CancellationToken.None);
   }

   public async Task HandleMessage(WebSocketReceiveResult result, byte[] buffer, CustomWebSocket userWebSocket, ICustomWebSocketFactory wsFactory)
   {
      string msg = Encoding.ASCII.GetString(buffer);
      try
      {
         var message = JsonConvert.DeserializeObject<CustomWebSocketMessage>(msg);
         if (message.Type == WSMessageType.anyType)
         {
            await BroadcastOthers(buffer, userWebSocket, wsFactory);
         }
      }
      catch (Exception e)
      {
         await userWebSocket.WebSocket.SendAsync(new ArraySegment<byte>(buffer, 0, result.Count), result.MessageType, result.EndOfMessage, CancellationToken.None);
      }
   }

   public async Task BroadcastOthers(byte[] buffer, CustomWebSocket userWebSocket, ICustomWebSocketFactory wsFactory)
   {
      var others = wsFactory.Others(userWebSocket);
      foreach (var uws in others)
      {
         await uws.WebSocket.SendAsync(new ArraySegment<byte>(buffer, 0, buffer.Length), WebSocketMessageType.Text, true, CancellationToken.None);
      }
   }

   public async Task BroadcastAll(byte[] buffer, CustomWebSocket userWebSocket, ICustomWebSocketFactory wsFactory)
   {
      var all = wsFactory.All();
      foreach (var uws in all)
      {
         await uws.WebSocket.SendAsync(new ArraySegment<byte>(buffer, 0, buffer.Length), WebSocketMessageType.Text, true, CancellationToken.None);
      }
   }
}

Finally, in the Startup class in the Configure method add the following:

var webSocketOptions = new WebSocketOptions()
{
    KeepAliveInterval = TimeSpan.FromSeconds(120),
    ReceiveBufferSize = 4 * 1024
};

app.UseWebSockets(webSocketOptions);
app.UseCustomWebSocketManager();

So, in this way the Starup class remains clean and the logics to manage the WebSockets can grows giving you the flexibility to organize it as you prefer.

Points of Interest

History

Keep a running update of any changes or improvements you've made here.