From 31e70a08b979ae935052e905030a94b73a76f01f Mon Sep 17 00:00:00 2001
From: Sanae <atmatm6000@outlook.com>
Date: Mon, 21 Feb 2022 22:05:13 -0600
Subject: [PATCH] Add proper exit handling

---
 Server/Program.cs | 30 ++++++++++++++++++++++++------
 Server/Server.cs  | 42 ++++++++++++++++++++++++++++--------------
 2 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/Server/Program.cs b/Server/Program.cs
index d9c9f0f..70a965a 100644
--- a/Server/Program.cs
+++ b/Server/Program.cs
@@ -7,7 +7,10 @@ using Timer = System.Timers.Timer;
 
 Server.Server server = new Server.Server();
 HashSet<int> shineBag = new HashSet<int>();
-int shineTx = 0; // used for logging
+// int shineTx = 0; // used for logging
+CancellationTokenSource cts = new CancellationTokenSource();
+Task listenTask = server.Listen(cts.Token);
+Logger consoleLogger = new Logger("Console");
 
 server.ClientJoined += async (c, type) => {
     c.Metadata["shineSync"] = new ConcurrentBag<int>();
@@ -33,8 +36,8 @@ timer.Start();
 bool flipEnabled = Settings.Instance.Flip.EnabledOnStart;
 
 float MarioSize(bool is2d) => is2d ? 180 : 160;
-
 server.PacketHandler = (c, p) => {
+    Console.WriteLine($"{c.Id} {p}");
     switch (p) {
         case CostumePacket:
             ClientSyncShineBag(c);
@@ -141,13 +144,28 @@ CommandHandler.RegisterCommand("savesettings", _ => {
     return "Saved settings.json";
 });
 
+Console.CancelKeyPress += (_, e) => {
+    e.Cancel = true;
+    consoleLogger.Info("Received Ctrl+C");
+    cts.Cancel();
+};
+
+CommandHandler.RegisterCommand("exit", _ => {
+    cts.Cancel();
+    return "Shutting down clients";
+});
+
+CommandHandler.RegisterCommand("quit", _ => {
+    cts.Cancel();
+    return "Shutting down clients";
+});
+
 Task.Run(() => {
-    Logger logger = new Logger("Console");
-    logger.Info("Run help command for valid commands.");
+    consoleLogger.Info("Run help command for valid commands.");
     while (true) {
         string? text = Console.ReadLine();
-        if (text != null) logger.Info(CommandHandler.GetResult(text));
+        if (text != null) consoleLogger.Info(CommandHandler.GetResult(text));
     }
 });
 
-await server.Listen();
\ No newline at end of file
+await listenTask;
\ No newline at end of file
diff --git a/Server/Server.cs b/Server/Server.cs
index f6c49d8..dd690b1 100644
--- a/Server/Server.cs
+++ b/Server/Server.cs
@@ -15,7 +15,7 @@ public class Server {
     public Func<Client, IPacket, bool>? PacketHandler = null!;
     public event Action<Client, ConnectPacket> ClientJoined = null!;
 
-    public async Task Listen() {
+    public async Task Listen(CancellationToken? token = null) {
         Socket serverSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
         serverSocket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
         serverSocket.Bind(new IPEndPoint(IPAddress.Parse(Settings.Instance.Server.Address), Settings.Instance.Server.Port));
@@ -23,24 +23,38 @@ public class Server {
 
         Logger.Info($"Listening on {serverSocket.LocalEndPoint}");
 
-        while (true) {
-            Socket socket = await serverSocket.AcceptAsync();
-            socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.NoDelay, true);
+        try {
+            while (true) {
+                Socket socket = token.HasValue ? await serverSocket.AcceptAsync(token.Value) : await serverSocket.AcceptAsync();
+                socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.NoDelay, true);
 
-            Logger.Warn($"Accepted connection for client {socket.RemoteEndPoint}");
+                Logger.Warn($"Accepted connection for client {socket.RemoteEndPoint}");
 
-            try {
-                if (Clients.Count > Constants.MaxClients) {
-                    Logger.Warn("Turned away client due to max clients");
-                    await socket.DisconnectAsync(false);
-                    continue;
+                try {
+                    if (Clients.Count > Constants.MaxClients) {
+                        Logger.Warn("Turned away client due to max clients");
+                        await socket.DisconnectAsync(false);
+                        continue;
+                    }
+
+                    Task.Run(() => HandleSocket(socket));
+                } catch {
+                    // super ignore this
                 }
-
-                Task.Run(() => HandleSocket(socket));
-            } catch {
-                // super ignore this
             }
+        } catch (OperationCanceledException) {
+            // ignore the exception, it's just for closing the server
         }
+        Logger.Info("Server closing");
+
+        try {
+            serverSocket.Shutdown(SocketShutdown.Both);
+        } catch (Exception) {
+            // ignore
+        } finally {
+            serverSocket.Close();
+        }
+        Logger.Info("Server closed");
     }
 
     public static void FillPacket<T>(PacketHeader header, T packet, Memory<byte> memory) where T : struct, IPacket {