From 6f50ec66a9f6731741e7e168e84851ca878c1304 Mon Sep 17 00:00:00 2001
From: flashwave <me@flash.moe>
Date: Tue, 30 Aug 2022 18:29:11 +0000
Subject: [PATCH] Added /shutdown and /restart commands for server maintenance.

---
 SharpChat/ChatUserSession.cs | 23 ++++++++++++++---------
 SharpChat/Program.cs         |  2 +-
 SharpChat/SockChatServer.cs  | 36 ++++++++++++++++++++++++++++++++++--
 3 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/SharpChat/ChatUserSession.cs b/SharpChat/ChatUserSession.cs
index 69c6658..93ba5a0 100644
--- a/SharpChat/ChatUserSession.cs
+++ b/SharpChat/ChatUserSession.cs
@@ -20,6 +20,8 @@ namespace SharpChat {
         public DateTimeOffset LastPing { get; set; } = DateTimeOffset.MinValue;
         public ChatUser User { get; set; }
 
+        private static int CloseCode { get; set; } = 1000;
+
         public string TargetName => @"@log";
 
 
@@ -27,8 +29,8 @@ namespace SharpChat {
 
         public IPAddress RemoteAddress {
             get {
-                if (_RemoteAddress == null) {
-                    if ((Connection.ConnectionInfo.ClientIpAddress == @"127.0.0.1" || Connection.ConnectionInfo.ClientIpAddress == @"::1")
+                if(_RemoteAddress == null) {
+                    if((Connection.ConnectionInfo.ClientIpAddress == @"127.0.0.1" || Connection.ConnectionInfo.ClientIpAddress == @"::1")
                         && Connection.ConnectionInfo.Headers.ContainsKey(@"X-Real-IP"))
                         _RemoteAddress = IPAddress.Parse(Connection.ConnectionInfo.Headers[@"X-Real-IP"]);
                     else
@@ -52,14 +54,14 @@ namespace SharpChat {
         }
 
         public void Send(IServerPacket packet) {
-            if (!Connection.IsAvailable)
+            if(!Connection.IsAvailable)
                 return;
 
             IEnumerable<string> data = packet.Pack();
 
-            if (data != null)
-                foreach (string line in data)
-                    if (!string.IsNullOrWhiteSpace(line))
+            if(data != null)
+                foreach(string line in data)
+                    if(!string.IsNullOrWhiteSpace(line))
                         Connection.Send(line);
         }
 
@@ -69,6 +71,9 @@ namespace SharpChat {
         public bool HasTimedOut
             => DateTimeOffset.Now - LastPing > SessionTimeOut;
 
+        public void PrepareForRestart()
+            => CloseCode = 1012;
+
         public void Dispose()
             => Dispose(true);
 
@@ -76,13 +81,13 @@ namespace SharpChat {
             => Dispose(false);
 
         private void Dispose(bool disposing) {
-            if (IsDisposed)
+            if(IsDisposed)
                 return;
 
             IsDisposed = true;
-            Connection.Close();
+            Connection.Close(CloseCode);
 
-            if (disposing)
+            if(disposing)
                 GC.SuppressFinalize(this);
         }
     }
diff --git a/SharpChat/Program.cs b/SharpChat/Program.cs
index e460bee..b5f250b 100644
--- a/SharpChat/Program.cs
+++ b/SharpChat/Program.cs
@@ -22,7 +22,7 @@ namespace SharpChat {
             Database.ReadConfig();
 
             using ManualResetEvent mre = new ManualResetEvent(false);
-            using SockChatServer scs = new SockChatServer(PORT);
+            using SockChatServer scs = new SockChatServer(mre, PORT);
             Console.CancelKeyPress += (s, e) => { e.Cancel = true; mre.Set(); };
             mre.WaitOne();
         }
diff --git a/SharpChat/SockChatServer.cs b/SharpChat/SockChatServer.cs
index 60593b4..8398898 100644
--- a/SharpChat/SockChatServer.cs
+++ b/SharpChat/SockChatServer.cs
@@ -11,6 +11,7 @@ using System.Net;
 using System.Net.Http;
 using System.Runtime.CompilerServices;
 using System.Text;
+using System.Threading;
 
 namespace SharpChat {
     public class SockChatServer : IDisposable {
@@ -58,9 +59,14 @@ namespace SharpChat {
             HttpClient.DefaultRequestHeaders.UserAgent.ParseAdd(@"SharpChat");
         }
 
-        public SockChatServer(ushort port) {
+        private ManualResetEvent Shutdown { get; }
+        private bool IsShuttingDown = false;
+
+        public SockChatServer(ManualResetEvent mre, ushort port) {
             Logger.Write("Starting Sock Chat server...");
 
+            Shutdown = mre ?? throw new ArgumentNullException(nameof(mre));
+
             Context = new ChatContext(this);
 
             Context.Channels.Add(new ChatChannel(@"Lounge"));
@@ -75,6 +81,11 @@ namespace SharpChat {
             Server = new SharpChatWebSocketServer($@"ws://0.0.0.0:{port}");
 
             Server.Start(sock => {
+                if(IsShuttingDown || IsDisposed) {
+                    sock.Close(1013);
+                    return;
+                }
+
                 sock.OnOpen = () => OnOpen(sock);
                 sock.OnClose = () => OnClose(sock);
                 sock.OnError = err => OnError(sock, err);
@@ -821,6 +832,25 @@ namespace SharpChat {
                         user.Send(new LegacyCommandResponse(LCR.IP_ADDRESS, false, ipUser.Username, ip));
                     break;
 
+                case @"shutdown":
+                case @"restart":
+                    if(user.UserId != 1) {
+                        user.Send(new LegacyCommandResponse(LCR.COMMAND_NOT_ALLOWED, true, $@"/{commandName}"));
+                        break;
+                    }
+
+                    if(IsShuttingDown)
+                        break;
+                    IsShuttingDown = true;
+
+                    if(commandName == @"restart")
+                        lock(SessionsLock)
+                            Sessions.ForEach(s => s.PrepareForRestart());
+
+                    Context.Update();
+                    Shutdown.Set();
+                    break;
+
                 default:
                     user.Send(new LegacyCommandResponse(LCR.COMMAND_NOT_FOUND, true, commandName));
                     break;
@@ -840,7 +870,9 @@ namespace SharpChat {
                 return;
             IsDisposed = true;
 
-            Sessions?.Clear();
+            lock(SessionsLock)
+                Sessions.ForEach(s => s.Dispose());
+
             Server?.Dispose();
             Context?.Dispose();
             HttpClient?.Dispose();