Sunshine/tools/sunshinesvc.cpp

#define WIN32_LEAN_AND_MEAN
#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)
#endif

SERVICE_STATUS_HANDLE service_status_handle;
SERVICE_STATUS service_status;
HANDLE stop_event;

#define SERVICE_NAME "SunshineSvc"

DWORD WINAPI HandlerEx(DWORD dwControl, DWORD dwEventType, LPVOID lpEventData, LPVOID lpContext) {
  switch(dwControl) {
  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;
    service_status.dwControlsAccepted = 0;
    service_status.dwWaitHint         = 30 * 1000;
    SetServiceStatus(service_status_handle, &service_status);

    // Trigger ServiceMain() to start cleanup
    SetEvent(stop_event);
    return NO_ERROR;

  default:
    return NO_ERROR;
  }
}

HANDLE CreateJobObjectForChildProcess() {
  HANDLE job_handle = CreateJobObjectW(NULL, NULL);
  if(!job_handle) {
    return NULL;
  }

  JOBOBJECT_EXTENDED_LIMIT_INFORMATION job_limit_info = {};

  // Kill Sunshine.exe when the final job object handle is closed (which will happen if we terminate unexpectedly).
  // This ensures we don't leave an orphaned Sunshine.exe running with an inherited handle to our log file.
  job_limit_info.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;

  // Allow Sunshine.exe to use CREATE_BREAKAWAY_FROM_JOB when spawning processes to ensure they can to live beyond
  // the lifetime of SunshineSvc.exe. This avoids unexpected user data loss if we crash or are killed.
  job_limit_info.BasicLimitInformation.LimitFlags |= JOB_OBJECT_LIMIT_BREAKAWAY_OK;

  if(!SetInformationJobObject(job_handle, JobObjectExtendedLimitInformation, &job_limit_info, sizeof(job_limit_info))) {
    CloseHandle(job_handle);
    return NULL;
  }

  return job_handle;
}

LPPROC_THREAD_ATTRIBUTE_LIST AllocateProcThreadAttributeList(DWORD attribute_count) {
  SIZE_T size;
  InitializeProcThreadAttributeList(NULL, attribute_count, 0, &size);

  auto list = (LPPROC_THREAD_ATTRIBUTE_LIST)HeapAlloc(GetProcessHeap(), 0, size);
  if(list == NULL) {
    return NULL;
  }

  if(!InitializeProcThreadAttributeList(list, attribute_count, 0, &size)) {
    HeapFree(GetProcessHeap(), 0, list);
    return NULL;
  }

  return list;
}

HANDLE DuplicateTokenForConsoleSession() {
  auto console_session_id = WTSGetActiveConsoleSessionId();
  if(console_session_id == 0xFFFFFFFF) {
    // No console session yet
    return NULL;
  }

  HANDLE current_token;
  if(!OpenProcessToken(GetCurrentProcess(), TOKEN_DUPLICATE, &current_token)) {
    return NULL;
  }

  // Duplicate our own LocalSystem token
  HANDLE new_token;
  if(!DuplicateTokenEx(current_token, TOKEN_ALL_ACCESS, NULL, SecurityImpersonation, TokenPrimary, &new_token)) {
    CloseHandle(current_token);
    return NULL;
  }

  CloseHandle(current_token);

  // Change the duplicated token to the console session ID
  if(!SetTokenInformation(new_token, TokenSessionId, &console_session_id, sizeof(console_session_id))) {
    CloseHandle(new_token);
    return NULL;
  }

  return new_token;
}

HANDLE OpenLogFileHandle() {
  WCHAR log_file_name[MAX_PATH];

  // Create sunshine.log in the Temp folder (usually %SYSTEMROOT%\Temp)
  GetTempPathW(_countof(log_file_name), log_file_name);
  wcscat_s(log_file_name, L"sunshine.log");

  // The file handle must be inheritable for our child process to use it
  SECURITY_ATTRIBUTES security_attributes = { sizeof(security_attributes), NULL, TRUE };

  // Overwrite the old sunshine.log
  return CreateFileW(log_file_name,
    GENERIC_WRITE,
    FILE_SHARE_READ,
    &security_attributes,
    CREATE_ALWAYS,
    0,
    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) {
    // Nothing we can really do here but terminate ourselves
    ExitProcess(GetLastError());
    return;
  }

  // Tell SCM we're starting
  service_status.dwServiceType             = SERVICE_WIN32_OWN_PROCESS;
  service_status.dwServiceSpecificExitCode = 0;
  service_status.dwWin32ExitCode           = NO_ERROR;
  service_status.dwWaitHint                = 0;
  service_status.dwControlsAccepted        = 0;
  service_status.dwCheckPoint              = 0;
  service_status.dwCurrentState            = SERVICE_START_PENDING;
  SetServiceStatus(service_status_handle, &service_status);

  stop_event = CreateEventA(NULL, TRUE, FALSE, NULL);
  if(stop_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
    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);
  startup_info.StartupInfo.lpDesktop  = (LPWSTR)L"winsta0\\default";
  startup_info.StartupInfo.dwFlags    = STARTF_USESTDHANDLES;
  startup_info.StartupInfo.hStdInput  = NULL;
  startup_info.StartupInfo.hStdOutput = log_file_handle;
  startup_info.StartupInfo.hStdError  = log_file_handle;

  // Allocate an attribute list with space for 2 entries
  startup_info.lpAttributeList = AllocateProcThreadAttributeList(2);
  if(startup_info.lpAttributeList == NULL) {
    // Tell SCM we failed to start
    service_status.dwWin32ExitCode = GetLastError();
    service_status.dwCurrentState  = SERVICE_STOPPED;
    SetServiceStatus(service_status_handle, &service_status);
    return;
  }

  // Only allow Sunshine.exe to inherit the log file handle, not all inheritable handles
  UpdateProcThreadAttribute(startup_info.lpAttributeList,
    0,
    PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
    &log_file_handle,
    sizeof(log_file_handle),
    NULL,
    NULL);

  // Tell SCM we're running (and stoppable now)
  service_status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_PRESHUTDOWN;
  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();
    if(console_token == NULL) {
      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",
         NULL,
         NULL,
         NULL,
         TRUE,
         CREATE_UNICODE_ENVIRONMENT | CREATE_NO_WINDOW | EXTENDED_STARTUPINFO_PRESENT,
         NULL,
         NULL,
         (LPSTARTUPINFOW)&startup_info,
         &process_info)) {
      CloseHandle(console_token);
      CloseHandle(job_handle);
      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;

    case WAIT_OBJECT_0 + 1:
      // Sunshine terminated itself.
      break;
    }

    CloseHandle(process_info.hThread);
    CloseHandle(process_info.hProcess);
    CloseHandle(console_token);
    CloseHandle(job_handle);
  }

  // Let SCM know we've stopped
  service_status.dwCurrentState = SERVICE_STOPPED;
  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
  WCHAR module_path[MAX_PATH];
  GetModuleFileNameW(NULL, module_path, _countof(module_path));
  for(auto i = 0; i < 2; i++) {
    auto last_sep = wcsrchr(module_path, '\\');
    if(last_sep) {
      *last_sep = 0;
    }
  }
  SetCurrentDirectoryW(module_path);

  // Trigger our ServiceMain()
  return StartServiceCtrlDispatcher(service_table);
}