From f51876893e731422b70e75751ea9dc5749075670 Mon Sep 17 00:00:00 2001
From: Cameron Gutman <aicommander@gmail.com>
Date: Sat, 15 Apr 2023 17:51:32 -0500
Subject: [PATCH] Stop the service if the user quits via the tray icon

---
 src/main.cpp                  | 37 +++++++++++++++++++++++++++--------
 src/main.h                    |  6 ++++++
 src/platform/windows/misc.cpp |  7 ++++---
 src/system_tray.cpp           | 10 +++++++++-
 tools/sunshinesvc.cpp         |  9 ++++++++-
 5 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/src/main.cpp b/src/main.cpp
index 4c1ae2ab..0733778f 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -107,6 +107,31 @@ namespace version {
   }
 }  // namespace version
 
+namespace lifetime {
+  static std::atomic_int desired_exit_code;
+
+  /**
+   * @brief Terminates Sunshine gracefully with the provided exit code
+   * @param exit_code The exit code to return from main()
+   * @param async Specifies whether our termination will be non-blocking
+   */
+  void
+  exit_sunshine(int exit_code, bool async) {
+    // Store the exit code of the first exit_sunshine() call
+    int zero = 0;
+    desired_exit_code.compare_exchange_strong(zero, exit_code);
+
+    // Raise SIGINT to start termination
+    std::raise(SIGINT);
+
+    // Termination will happen asynchronously, but the caller may
+    // have wanted synchronous behavior.
+    while (!async) {
+      std::this_thread::sleep_for(1s);
+    }
+  }
+}  // namespace lifetime
+
 /**
  * @brief Flush the log.
  *
@@ -159,13 +184,9 @@ LRESULT CALLBACK
 SessionMonitorWindowProc(HWND hwnd, UINT uMsg, WPARAM wParam, LPARAM lParam) {
   switch (uMsg) {
     case WM_ENDSESSION: {
-      // Raise a SIGINT to trigger our cleanup logic and terminate ourselves
+      // Terminate ourselves with a blocking exit call
       std::cout << "Received WM_ENDSESSION"sv << std::endl;
-      std::raise(SIGINT);
-
-      // The signal handling is asynchronous, so we will wait here to be terminated.
-      // If for some reason we don't terminate in a few seconds, Windows will kill us.
-      SuspendThread(GetCurrentThread());
+      lifetime::exit_sunshine(0, false);
       return 0;
     }
     default:
@@ -376,7 +397,7 @@ main(int argc, char *argv[]) {
 
   // FIXME: Temporary workaround: Simple-Web_server needs to be updated or replaced
   if (shutdown_event->peek()) {
-    return 0;
+    return lifetime::desired_exit_code;
   }
 
   std::thread httpThread { nvhttp::start };
@@ -395,7 +416,7 @@ main(int argc, char *argv[]) {
   system_tray::end_tray();
 #endif
 
-  return 0;
+  return lifetime::desired_exit_code;
 }
 
 /**
diff --git a/src/main.h b/src/main.h
index aff1cfdd..05c0f15f 100644
--- a/src/main.h
+++ b/src/main.h
@@ -79,4 +79,10 @@ namespace mail {
 #undef MAIL
 
 }  // namespace mail
+
+namespace lifetime {
+  void
+  exit_sunshine(int exit_code, bool async);
+}  // namespace lifetime
+
 #endif  // SUNSHINE_MAIN_H
diff --git a/src/platform/windows/misc.cpp b/src/platform/windows/misc.cpp
index 45b25e3d..5335cae2 100644
--- a/src/platform/windows/misc.cpp
+++ b/src/platform/windows/misc.cpp
@@ -727,9 +727,10 @@ namespace platf {
 
   bool
   restart() {
-    // Raise SIGINT to trigger the graceful exit logic. The service will
-    // restart us in a few seconds.
-    std::raise(SIGINT);
+    // Gracefully exit. The service will restart us in a few seconds.
+    // We use an async exit call here because we can't block the
+    // HTTP thread or we'll hang shutdown.
+    lifetime::exit_sunshine(0, true);
     return true;
   }
 
diff --git a/src/system_tray.cpp b/src/system_tray.cpp
index b42acc19..b25d6d38 100644
--- a/src/system_tray.cpp
+++ b/src/system_tray.cpp
@@ -131,7 +131,15 @@ namespace system_tray {
   tray_quit_cb(struct tray_menu *item) {
     BOOST_LOG(info) << "Quiting from system tray"sv;
 
-    std::raise(SIGINT);
+  #ifdef _WIN32
+    // If we're running in a service, return a special status to
+    // tell it to terminate too, otherwise it will just respawn us.
+    if (GetConsoleWindow() == NULL) {
+      lifetime::exit_sunshine(ERROR_SHUTDOWN_IN_PROGRESS, false);
+    }
+  #endif
+
+    lifetime::exit_sunshine(0, false);
   }
 
   // Tray menu
diff --git a/tools/sunshinesvc.cpp b/tools/sunshinesvc.cpp
index 19a7b9dc..e02fc8b9 100644
--- a/tools/sunshinesvc.cpp
+++ b/tools/sunshinesvc.cpp
@@ -301,9 +301,16 @@ ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
         }
         break;
 
-      case WAIT_OBJECT_0 + 1:
+      case WAIT_OBJECT_0 + 1: {
         // Sunshine terminated itself.
+
+        DWORD exit_code;
+        if (GetExitCodeProcess(process_info.hProcess, &exit_code) && exit_code == ERROR_SHUTDOWN_IN_PROGRESS) {
+          // Sunshine is asking for us to shut down, so gracefully stop ourselves.
+          SetEvent(stop_event);
+        }
         break;
+      }
     }
 
     CloseHandle(process_info.hThread);