From 8ad7af86c0ad130395019c540884ca0da12a7ed2 Mon Sep 17 00:00:00 2001
From: Cameron Gutman <aicommander@gmail.com>
Date: Thu, 29 Dec 2022 08:32:23 -0600
Subject: [PATCH] Graceful termination on shutdown, logoff, and service stop
 (#647)

---
 src/main.cpp          |  70 ++++++++++++++++++++++++
 tools/sunshinesvc.cpp | 124 +++++++++++++++++++++++++++++++++---------
 2 files changed, 168 insertions(+), 26 deletions(-)

diff --git a/src/main.cpp b/src/main.cpp
index 58271f83..2411640e 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -136,16 +136,86 @@ std::map<std::string_view, std::function<int(const char *name, int argc, char **
   { "version"sv, version::entry }
 };
 
+#ifdef _WIN32
+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
+    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());
+    return 0;
+  }
+  default:
+    return DefWindowProc(hwnd, uMsg, wParam, lParam);
+  }
+}
+#endif
+
 int main(int argc, char *argv[]) {
   util::TaskPool::task_id_t force_shutdown = nullptr;
 
   bool shutdown_by_interrupt = false;
 
+#ifdef _WIN32
+  // Wait as long as possible to terminate Sunshine.exe during logoff/shutdown
+  SetProcessShutdownParameters(0x100, SHUTDOWN_NORETRY);
+
+  // We must create a hidden window to receive shutdown notifications since we load gdi32.dll
+  std::thread window_thread([]() {
+    WNDCLASSA wnd_class {};
+    wnd_class.lpszClassName = "SunshineSessionMonitorClass";
+    wnd_class.lpfnWndProc   = SessionMonitorWindowProc;
+    if(!RegisterClassA(&wnd_class)) {
+      std::cout << "Failed to register session monitor window class"sv << std::endl;
+      return;
+    }
+
+    auto wnd = CreateWindowExA(
+      0,
+      wnd_class.lpszClassName,
+      "Sunshine Session Monitor Window",
+      0,
+      CW_USEDEFAULT,
+      CW_USEDEFAULT,
+      CW_USEDEFAULT,
+      CW_USEDEFAULT,
+      nullptr,
+      nullptr,
+      nullptr,
+      nullptr);
+    if(!wnd) {
+      std::cout << "Failed to create session monitor window"sv << std::endl;
+      return;
+    }
+
+    ShowWindow(wnd, SW_HIDE);
+
+    // Run the message loop for our window
+    MSG msg {};
+    while(GetMessage(&msg, nullptr, 0, 0) > 0) {
+      TranslateMessage(&msg);
+      DispatchMessage(&msg);
+    }
+  });
+  window_thread.detach();
+#endif
+
   auto exit_guard = util::fail_guard([&shutdown_by_interrupt, &force_shutdown]() {
     if(!shutdown_by_interrupt) {
       return;
     }
 
+#ifdef _WIN32
+    // If this is running from a service with no console window, don't wait for user input to exit
+    if(GetConsoleWindow() == NULL) {
+      return;
+    }
+#endif
+
     task_pool.cancel(force_shutdown);
 
     std::cout << "Sunshine exited: Press enter to continue"sv << std::endl;
diff --git a/tools/sunshinesvc.cpp b/tools/sunshinesvc.cpp
index c92c9ebb..9089be30 100644
--- a/tools/sunshinesvc.cpp
+++ b/tools/sunshinesvc.cpp
@@ -2,6 +2,8 @@
 #include <Windows.h>
 #include <wtsapi32.h>
 
+#include <string>
+
 // PROC_THREAD_ATTRIBUTE_JOB_LIST is currently missing from MinGW headers
 #ifndef PROC_THREAD_ATTRIBUTE_JOB_LIST
 #define PROC_THREAD_ATTRIBUTE_JOB_LIST ProcThreadAttributeValue(13, FALSE, TRUE, FALSE)
@@ -18,6 +20,8 @@ DWORD WINAPI HandlerEx(DWORD dwControl, DWORD dwEventType, LPVOID lpEventData, L
   case SERVICE_CONTROL_INTERROGATE:
     return NO_ERROR;
 
+  case SERVICE_CONTROL_PRESHUTDOWN:
+    // The system is shutting down
   case SERVICE_CONTROL_STOP:
     // Let SCM know we're stopping in up to 30 seconds
     service_status.dwCurrentState     = SERVICE_STOP_PENDING;
@@ -125,6 +129,49 @@ HANDLE OpenLogFileHandle() {
     NULL);
 }
 
+bool RunTerminationHelper(HANDLE console_token, DWORD pid) {
+  WCHAR module_path[MAX_PATH];
+  GetModuleFileNameW(NULL, module_path, _countof(module_path));
+  std::wstring command { module_path };
+
+  command += L" --terminate " + std::to_wstring(pid);
+
+  STARTUPINFOW startup_info = {};
+  startup_info.cb           = sizeof(startup_info);
+  startup_info.lpDesktop    = (LPWSTR)L"winsta0\\default";
+
+  // Execute ourselves as a detached process in the user session with the --terminate argument.
+  // This will allow us to attach to Sunshine's console and send it a Ctrl-C event.
+  PROCESS_INFORMATION process_info;
+  if(!CreateProcessAsUserW(console_token,
+       NULL,
+       (LPWSTR)command.c_str(),
+       NULL,
+       NULL,
+       FALSE,
+       CREATE_UNICODE_ENVIRONMENT | DETACHED_PROCESS,
+       NULL,
+       NULL,
+       &startup_info,
+       &process_info)) {
+    return false;
+  }
+
+  // Wait for the termination helper to complete
+  WaitForSingleObject(process_info.hProcess, INFINITE);
+
+  // Check the exit status of the helper process
+  DWORD exit_code;
+  GetExitCodeProcess(process_info.hProcess, &exit_code);
+
+  // Cleanup handles
+  CloseHandle(process_info.hProcess);
+  CloseHandle(process_info.hThread);
+
+  // If the helper process returned 0, it succeeded
+  return exit_code == 0;
+}
+
 VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
   service_status_handle = RegisterServiceCtrlHandlerEx(SERVICE_NAME, HandlerEx, NULL);
   if(service_status_handle == NULL) {
@@ -161,15 +208,6 @@ VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
     return;
   }
 
-  auto job_handle = CreateJobObjectForChildProcess();
-  if(job_handle == NULL) {
-    // Tell SCM we failed to start
-    service_status.dwWin32ExitCode = GetLastError();
-    service_status.dwCurrentState  = SERVICE_STOPPED;
-    SetServiceStatus(service_status_handle, &service_status);
-    return;
-  }
-
   // We can use a single STARTUPINFOEXW for all the processes that we launch
   STARTUPINFOEXW startup_info         = {};
   startup_info.StartupInfo.cb         = sizeof(startup_info);
@@ -198,17 +236,8 @@ VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
     NULL,
     NULL);
 
-  // Start Sunshine.exe inside our job object
-  UpdateProcThreadAttribute(startup_info.lpAttributeList,
-    0,
-    PROC_THREAD_ATTRIBUTE_JOB_LIST,
-    &job_handle,
-    sizeof(job_handle),
-    NULL,
-    NULL);
-
   // Tell SCM we're running (and stoppable now)
-  service_status.dwControlsAccepted = SERVICE_ACCEPT_STOP;
+  service_status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_PRESHUTDOWN;
   service_status.dwCurrentState     = SERVICE_RUNNING;
   SetServiceStatus(service_status_handle, &service_status);
 
@@ -219,6 +248,22 @@ VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
       continue;
     }
 
+    // Job objects cannot span sessions, so we must create one for each process
+    auto job_handle = CreateJobObjectForChildProcess();
+    if(job_handle == NULL) {
+      CloseHandle(console_token);
+      continue;
+    }
+
+    // Start Sunshine.exe inside our job object
+    UpdateProcThreadAttribute(startup_info.lpAttributeList,
+      0,
+      PROC_THREAD_ATTRIBUTE_JOB_LIST,
+      &job_handle,
+      sizeof(job_handle),
+      NULL,
+      NULL);
+
     PROCESS_INFORMATION process_info;
     if(!CreateProcessAsUserW(console_token,
          L"Sunshine.exe",
@@ -232,20 +277,21 @@ VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
          (LPSTARTUPINFOW)&startup_info,
          &process_info)) {
       CloseHandle(console_token);
+      CloseHandle(job_handle);
       continue;
     }
 
-    // Close handles that are no longer needed
-    CloseHandle(console_token);
-    CloseHandle(process_info.hThread);
-
     // Wait for either the stop event to be set or Sunshine.exe to terminate
     const HANDLE wait_objects[] = { stop_event, process_info.hProcess };
     switch(WaitForMultipleObjects(_countof(wait_objects), wait_objects, FALSE, INFINITE)) {
     case WAIT_OBJECT_0:
-      // The service is shutting down, so terminate Sunshine.exe.
-      // TODO: Send a graceful exit request and only terminate forcefully as a last resort.
-      TerminateProcess(process_info.hProcess, ERROR_PROCESS_ABORTED);
+      // The service is shutting down, so try to gracefully terminate Sunshine.exe.
+      // If it doesn't terminate in 20 seconds, we will forcefully terminate it.
+      if(!RunTerminationHelper(console_token, process_info.dwProcessId) ||
+         WaitForSingleObject(process_info.hProcess, 20000) != WAIT_OBJECT_0) {
+        // If it won't terminate gracefully, kill it now
+        TerminateProcess(process_info.hProcess, ERROR_PROCESS_ABORTED);
+      }
       break;
 
     case WAIT_OBJECT_0 + 1:
@@ -253,7 +299,10 @@ VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
       break;
     }
 
+    CloseHandle(process_info.hThread);
     CloseHandle(process_info.hProcess);
+    CloseHandle(console_token);
+    CloseHandle(job_handle);
   }
 
   // Let SCM know we've stopped
@@ -261,12 +310,35 @@ VOID WINAPI ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) {
   SetServiceStatus(service_status_handle, &service_status);
 }
 
+// This will run in a child process in the user session
+int DoGracefulTermination(DWORD pid) {
+  // Attach to Sunshine's console
+  if(!AttachConsole(pid)) {
+    return GetLastError();
+  }
+
+  // Disable our own Ctrl-C handling
+  SetConsoleCtrlHandler(NULL, TRUE);
+
+  // Send a Ctrl-C event to Sunshine
+  if(!GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0)) {
+    return GetLastError();
+  }
+
+  return 0;
+}
+
 int main(int argc, char *argv[]) {
   static const SERVICE_TABLE_ENTRY service_table[] = {
     { (LPSTR)SERVICE_NAME, ServiceMain },
     { NULL, NULL }
   };
 
+  // Check if this is a reinvocation of ourselves to send Ctrl-C to Sunshine.exe
+  if(argc == 3 && strcmp(argv[1], "--terminate") == 0) {
+    return DoGracefulTermination(atol(argv[2]));
+  }
+
   // By default, services have their current directory set to %SYSTEMROOT%\System32.
   // We want to use the directory where Sunshine.exe is located instead of system32.
   // This requires stripping off 2 path components: the file name and the last folder