damaged-thunder/server/Controllers/WebsocketController.cs

using System;
using System.Collections.Generic;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using gameapi.Commands;
using gameapi.Library;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.Hosting;

namespace gameapi.Controllers
{
	public class WebsocketController : IMiddleware
	{
		private static CancellationTokenRegistration AppShutdownHandler;

		private WebsocketCommandFactory websocketCommandFactory;

		public WebsocketController(IHostApplicationLifetime hostLifetime)
		{
			if (AppShutdownHandler.Token.Equals(CancellationToken.None))
			{
				AppShutdownHandler = hostLifetime.ApplicationStopping.Register(ApplicationShutdownHandler);
			}
			websocketCommandFactory = new WebsocketCommandFactory();
		}

		public async Task InvokeAsync(HttpContext context, RequestDelegate next)
		{
			if (context.WebSockets.IsWebSocketRequest)
			{
				WebSocket webSocket = await context.WebSockets.AcceptWebSocketAsync();
				GameState.GetInstance().AddPlayer(webSocket);
				await SocketConnectionLoop(context, webSocket);
			}
			else
			{
				await next(context);
			}
		}

		private async Task SocketConnectionLoop(HttpContext context, WebSocket webSocket)
		{
			//https://mcguirev10.com/2019/08/18/minimal-full-feature-kestrel-websocket-server.html
			//https://github.com/MV10/WebSocketExample/tree/master/KestrelWebSocketServer
			try
			{
				var buffer = new byte[1024 * 4];
				var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
				while (!result.CloseStatus.HasValue)
				{
					switch (result.MessageType)
					{
						case WebSocketMessageType.Text:
							var message = Encoding.UTF8.GetString(buffer, 0, result.Count);
							await websocketCommandFactory.Create(message).Run(webSocket);
							break;
						case WebSocketMessageType.Close:
							Console.WriteLine("Disconnecting");
							await webSocket.CloseAsync(result.CloseStatus.Value, result.CloseStatusDescription, CancellationToken.None);
							webSocket.Dispose();
							return;
					}

					result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
				}
			}
			catch (WebSocketException e)
			{
				Console.WriteLine("Something went wrong with this websocket." + e.Message + " " + e.StackTrace);
				if (webSocket.State == WebSocketState.Open)
				{
					await webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "server error", CancellationToken.None);
				}

			}
		}

		public static async void ApplicationShutdownHandler()
		{
			List<Task> tasksToComplete = new List<Task>();
			foreach (var player in GameState.GetInstance().GetPlayers())
			{
				var timeout = new CancellationTokenSource();
				var websocket = player.GetConnection();
				var message = "server shutdown";
				tasksToComplete.Add(websocket.SendAsync(new ArraySegment<byte>(Encoding.UTF8.GetBytes(message), 0, message.Length), WebSocketMessageType.Text, true, CancellationToken.None));
				tasksToComplete.Add(websocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "server shutdown", timeout.Token));
			}

			try
			{
				await Task.WhenAll(tasksToComplete);
			}
			catch (OperationCanceledException ex)
			{
				Console.WriteLine("Application Shutdown error: " + ex.Message);
			}
		}
	}
}