From 9955890023e4f855231d1690555c7ffda7264cdf Mon Sep 17 00:00:00 2001 From: Cameron Gutman Date: Sun, 14 May 2023 19:27:57 -0500 Subject: [PATCH] Fix console session changes via fast user switching We need to respawn Sunshine.exe in the new console session. --- tools/sunshinesvc.cpp | 92 +++++++++++++++++++++++++++++-------------- 1 file changed, 62 insertions(+), 30 deletions(-) diff --git a/tools/sunshinesvc.cpp b/tools/sunshinesvc.cpp index b117ec91..d2829399 100644 --- a/tools/sunshinesvc.cpp +++ b/tools/sunshinesvc.cpp @@ -16,6 +16,7 @@ SERVICE_STATUS_HANDLE service_status_handle; SERVICE_STATUS service_status; HANDLE stop_event; +HANDLE session_change_event; #define SERVICE_NAME "SunshineSvc" @@ -25,6 +26,14 @@ HandlerEx(DWORD dwControl, DWORD dwEventType, LPVOID lpEventData, LPVOID lpConte case SERVICE_CONTROL_INTERROGATE: return NO_ERROR; + case SERVICE_CONTROL_SESSIONCHANGE: + // If a new session connects to the console, restart Sunshine + // to allow it to spawn inside the new console session. + if (dwEventType == WTS_CONSOLE_CONNECT) { + SetEvent(session_change_event); + } + return NO_ERROR; + case SERVICE_CONTROL_PRESHUTDOWN: // The system is shutting down case SERVICE_CONTROL_STOP: @@ -39,7 +48,7 @@ HandlerEx(DWORD dwControl, DWORD dwEventType, LPVOID lpEventData, LPVOID lpConte return NO_ERROR; default: - return NO_ERROR; + return ERROR_CALL_NOT_IMPLEMENTED; } } @@ -87,13 +96,7 @@ AllocateProcThreadAttributeList(DWORD attribute_count) { } HANDLE -DuplicateTokenForConsoleSession() { - auto console_session_id = WTSGetActiveConsoleSessionId(); - if (console_session_id == 0xFFFFFFFF) { - // No console session yet - return NULL; - } - +DuplicateTokenForSession(DWORD console_session_id) { HANDLE current_token; if (!OpenProcessToken(GetCurrentProcess(), TOKEN_DUPLICATE, ¤t_token)) { return NULL; @@ -201,6 +204,7 @@ ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) { service_status.dwCurrentState = SERVICE_START_PENDING; SetServiceStatus(service_status_handle, &service_status); + // Create a manual-reset stop event stop_event = CreateEventA(NULL, TRUE, FALSE, NULL); if (stop_event == NULL) { // Tell SCM we failed to start @@ -210,6 +214,16 @@ ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) { return; } + // Create an auto-reset session change event + session_change_event = CreateEventA(NULL, FALSE, FALSE, NULL); + if (session_change_event == NULL) { + // Tell SCM we failed to start + service_status.dwWin32ExitCode = GetLastError(); + service_status.dwCurrentState = SERVICE_STOPPED; + SetServiceStatus(service_status_handle, &service_status); + return; + } + auto log_file_handle = OpenLogFileHandle(); if (log_file_handle == INVALID_HANDLE_VALUE) { // Tell SCM we failed to start @@ -248,13 +262,19 @@ ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) { NULL); // Tell SCM we're running (and stoppable now) - service_status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_PRESHUTDOWN; + service_status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_PRESHUTDOWN | SERVICE_ACCEPT_SESSIONCHANGE; service_status.dwCurrentState = SERVICE_RUNNING; SetServiceStatus(service_status_handle, &service_status); // Loop every 3 seconds until the stop event is set or Sunshine.exe is running while (WaitForSingleObject(stop_event, 3000) != WAIT_OBJECT_0) { - auto console_token = DuplicateTokenForConsoleSession(); + auto console_session_id = WTSGetActiveConsoleSessionId(); + if (console_session_id == 0xFFFFFFFF) { + // No console session yet + continue; + } + + auto console_token = DuplicateTokenForSession(console_session_id); if (console_token == NULL) { continue; } @@ -292,30 +312,42 @@ ServiceMain(DWORD dwArgc, LPTSTR *lpszArgv) { continue; } - // 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 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; + bool still_running; + do { + // Wait for the stop event to be set, Sunshine.exe to terminate, or the console session to change + const HANDLE wait_objects[] = { stop_event, process_info.hProcess, session_change_event }; + switch (WaitForMultipleObjects(_countof(wait_objects), wait_objects, FALSE, INFINITE)) { + case WAIT_OBJECT_0 + 2: + if (WTSGetActiveConsoleSessionId() == console_session_id) { + // The active console session didn't actually change. Let Sunshine keep running. + still_running = true; + continue; + } + // Fall-through to terminate Sunshine.exe and start it again. + case WAIT_OBJECT_0: + // 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); + } + still_running = false; + break; - case WAIT_OBJECT_0 + 1: { - // Sunshine terminated itself. + 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); + 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); + } + still_running = false; + break; } - break; } - } + } while (still_running); CloseHandle(process_info.hThread); CloseHandle(process_info.hProcess);