From 34b93ad55df68b4f1204220f45341f04265fc9bb Mon Sep 17 00:00:00 2001
From: cs8425 <cs8425@gmail.com>
Date: Tue, 28 Apr 2026 21:09:14 +0800
Subject: [PATCH] fix infinite loop when pressed ctrl+c

---
 Common/Util/IConsole.cs         | 36 ++++++++++++++++++++++++---------
 MikuSB/Program/LoaderManager.cs |  4 ++--
 MikuSB/Program/MikuSB.cs        | 28 ++++++++++++++++++-------
 3 files changed, 49 insertions(+), 19 deletions(-)

diff --git a/Common/Util/IConsole.cs b/Common/Util/IConsole.cs
index ac0c0e9..395561e 100644
--- a/Common/Util/IConsole.cs
+++ b/Common/Util/IConsole.cs
@@ -53,7 +53,7 @@ public class IConsole
         Console.WriteLine();
         Input = [];
         CursorIndex = 0;
-        if (InputHistory.Count >= HistoryMaxCount) 
+        if (InputHistory.Count >= HistoryMaxCount)
             InputHistory.RemoveAt(0);
         InputHistory.Add(input);
         HistoryIndex = InputHistory.Count;
@@ -80,7 +80,7 @@ public class IConsole
     public static void HandleUpArrow()
     {
         if (InputHistory.Count == 0) return;
-        
+
         if (HistoryIndex > 0)
         {
             HistoryIndex--;
@@ -94,7 +94,7 @@ public class IConsole
     public static void HandleDownArrow()
     {
         if (HistoryIndex >= InputHistory.Count) return;
-        
+
         HistoryIndex++;
         if (HistoryIndex >= InputHistory.Count)
         {
@@ -102,7 +102,7 @@ public class IConsole
             Input = [];
             CursorIndex = 0;
         }
-        else 
+        else
         {
             var history = InputHistory[HistoryIndex];
             Input = [.. history];
@@ -114,7 +114,7 @@ public class IConsole
     public static void HandleLeftArrow()
     {
         if (CursorIndex <= 0) return;
-        
+
         var (left, _) = Console.GetCursorPosition();
         CursorIndex--;
         Console.SetCursorPosition(left - GetWidth(Input[CursorIndex].ToString()), Console.CursorTop);
@@ -123,7 +123,7 @@ public class IConsole
     public static void HandleRightArrow()
     {
         if (CursorIndex >= Input.Count) return;
-        
+
         var (left, _) = Console.GetCursorPosition();
         CursorIndex++;
         Console.SetCursorPosition(left + GetWidth(Input[CursorIndex - 1].ToString()), Console.CursorTop);
@@ -148,13 +148,29 @@ public class IConsole
 
     #endregion
 
-    public static string ListenConsole()
+    public static async Task ListenConsole(CancellationToken exitToken)
     {
-        while (true)
+        while (!exitToken.IsCancellationRequested)
         {
             ConsoleKeyInfo keyInfo;
-            try { keyInfo = Console.ReadKey(true); }
-            catch (InvalidOperationException) { continue; }
+            try
+            {
+                if (!Console.KeyAvailable)
+                {
+                    await Task.Delay(10, exitToken);
+                    continue;
+                }
+                keyInfo = Console.ReadKey(true);
+            }
+            catch (OperationCanceledException)
+            {
+                break;
+            }
+            catch (InvalidOperationException)
+            {
+                await Task.Delay(50, exitToken);
+                continue;
+            }
 
             switch (keyInfo.Key)
             {
diff --git a/MikuSB/Program/LoaderManager.cs b/MikuSB/Program/LoaderManager.cs
index 5d60ce0..04bfc29 100644
--- a/MikuSB/Program/LoaderManager.cs
+++ b/MikuSB/Program/LoaderManager.cs
@@ -160,7 +160,7 @@ public class LoaderManager : MikuSB
         }
     }
 
-    public static void InitCommand()
+    public static async Task InitCommand(CancellationToken exitToken)
     {
         // Register the command handlers
         try
@@ -178,6 +178,6 @@ public class LoaderManager : MikuSB
         IConsole.OnConsoleExcuteCommand += CommandExecutor.ConsoleExcuteCommand;
         CommandExecutor.OnRunCommand += (sender, e) => { CommandManager.HandleCommand(e, sender); };
 
-        IConsole.ListenConsole();
+        await IConsole.ListenConsole(exitToken);
     }
 }
\ No newline at end of file
diff --git a/MikuSB/Program/MikuSB.cs b/MikuSB/Program/MikuSB.cs
index d20e124..edd0778 100644
--- a/MikuSB/Program/MikuSB.cs
+++ b/MikuSB/Program/MikuSB.cs
@@ -18,6 +18,10 @@ public class MikuSB
     public static readonly Listener Listener = new();
     public static readonly CommandManager CommandManager = new();
 
+    // for exit signal
+    private static readonly CancellationTokenSource _cts = new();
+    private static int _exitCode = 0;
+
     public static async Task Main()
     {
         var time = DateTime.Now;
@@ -50,44 +54,54 @@ public class MikuSB
         ResourceManager.IsLoaded = true;
 
         HandbookGenerator.GenerateAll();
-        LoaderManager.InitCommand();
+        var consoleTask = Task.Run(() => LoaderManager.InitCommand(_cts.Token), _cts.Token);
 
         var elapsed = DateTime.Now - time;
         Logger.Info(I18NManager.Translate("Server.ServerInfo.ServerStarted",
             Math.Round(elapsed.TotalSeconds, 2).ToString(CultureInfo.InvariantCulture)));
+
+        await consoleTask;
+
+        ProcessExit(Volatile.Read(ref _exitCode));
     }
 
-    # region Exit
+    #region Exit
 
     private static void RegisterExitEvent()
     {
         AppDomain.CurrentDomain.ProcessExit += (_, _) =>
         {
             Logger.Info(I18NManager.Translate("Server.ServerInfo.Shutdown"));
-            ProcessExit();
+            RequestShutdown(0);
         };
         AppDomain.CurrentDomain.UnhandledException += (obj, arg) =>
         {
             Logger.Error(I18NManager.Translate("Server.ServerInfo.UnhandledException", obj.GetType().Name),
                 (Exception)arg.ExceptionObject);
             Logger.Info(I18NManager.Translate("Server.ServerInfo.Shutdown"));
-            ProcessExit();
-            Environment.Exit(1);
+            RequestShutdown(1);
         };
 
         Console.CancelKeyPress += (_, eventArgs) =>
         {
             Logger.Info(I18NManager.Translate("Server.ServerInfo.CancelKeyPressed"));
             eventArgs.Cancel = true;
-            Environment.Exit(0);
+            RequestShutdown(0);
         };
     }
 
-    private static void ProcessExit()
+    private static void RequestShutdown(int exitCode)
+    {
+        Interlocked.Exchange(ref _exitCode, exitCode);
+        _cts.Cancel();
+    }
+
+    private static void ProcessExit(int exitCode)
     {
         SocketListener.Connections.Values.ToList().ForEach(x => x.Stop(true));
         DatabaseHelper.SaveThread?.Interrupt();
         DatabaseHelper.SaveDatabase();
+        Environment.Exit(exitCode);
     }
 
     # endregion