gdipp_server/gdipp_server.cpp

   1 #include "stdafx.h"
   2 #include "gdipp_config/constant_hook.h"
   3 #include "gdipp_lib/helper.h"
   4 #include "gdipp_server/global.h"
   5 #include "gdipp_server/rpc_server.h"
   6
   7 namespace gdipp
   8 {
   9
  10 #define SVC_NAME L"gdipp_svc"
  11
  12 SERVICE_STATUS svc_status = {};
  13 SERVICE_STATUS_HANDLE h_svc_status = NULL;
  14
  15 HANDLE h_svc_events, h_wait_cleanup, h_rpc_thread;
  16
  17 std::map<ULONG, HANDLE> h_user_tokens;
  18 std::map<ULONG, PROCESS_INFORMATION> pi_hooks_32, pi_hooks_64;
  19
  20 BOOL hook_proc(HANDLE h_user_token, char *hook_env_str, const wchar_t *gdipp_hook_name, PROCESS_INFORMATION &pi)
  21 {
  22         wchar_t gdipp_hook_path[MAX_PATH];
  23
  24         if (!get_dir_file_path(NULL, gdipp_hook_name, gdipp_hook_path))
  25                 return FALSE;
  26
  27         STARTUPINFOW si = {sizeof(STARTUPINFO)};
  28
  29         return CreateProcessAsUserW(h_user_token, gdipp_hook_path, NULL, NULL, NULL, TRUE, NULL, hook_env_str, NULL, &si, &pi);
  30 }
  31
  32 BOOL start_hook(ULONG session_id)
  33 {
  34         // make the event handle inheritable
  35         SECURITY_ATTRIBUTES inheritable_sa = {sizeof(SECURITY_ATTRIBUTES), NULL, TRUE};
  36         BOOL b_ret, hook_success = TRUE;
  37         HANDLE h_user_token;
  38
  39         b_ret = WTSQueryUserToken(session_id, &h_user_token);
  40         if (!b_ret)
  41         {
  42                 hook_success = FALSE;
  43                 goto post_hook;
  44         }
  45
  46         // use the linked token if exists
  47         // needed in UAC-enabled scenarios and Run As Administrator
  48         TOKEN_LINKED_TOKEN linked_token;
  49         DWORD token_info_len;
  50         b_ret = GetTokenInformation(h_user_token, TokenLinkedToken, &linked_token, sizeof(TOKEN_LINKED_TOKEN), &token_info_len);
  51         if (b_ret)
  52         {
  53                 CloseHandle(h_user_token);
  54                 h_user_token = linked_token.LinkedToken;
  55         }
  56
  57         char hook_env_str[64];
  58         sprintf_s(hook_env_str, "gdipp_svc_proc_id=%d%c", GetCurrentProcessId(), 0);
  59
  60         if (!!config_instance.get_number(L"/gdipp/hook/include/proc_32_bit/text()", static_cast<int>(gdipp::hook_config::PROC_32_BIT)))
  61         {
  62                 const wchar_t *gdipp_hook_name_32 = L"gdipp_hook_32.exe";
  63                 PROCESS_INFORMATION pi;
  64
  65                 if (hook_proc(h_user_token, hook_env_str, gdipp_hook_name_32, pi))
  66                         pi_hooks_32[session_id] = pi;
  67                 else
  68                         hook_success = FALSE;
  69         }
  70
  71         if (!!config_instance.get_number(L"/gdipp/hook/include/proc_64_bit/text()", static_cast<int>(gdipp::hook_config::PROC_64_BIT)))
  72         {
  73                 const wchar_t *gdipp_hook_name_64 = L"gdipp_hook_64.exe";
  74                 PROCESS_INFORMATION pi;
  75
  76                 if (hook_proc(h_user_token, hook_env_str, gdipp_hook_name_64, pi))
  77                         pi_hooks_64[session_id] = pi;
  78                 else
  79                         hook_success = FALSE;
  80         }
  81
  82 post_hook:
  83         if (hook_success)
  84         {
  85                 h_user_tokens[session_id] = h_user_token;
  86         }
  87         else
  88         {
  89                 if (h_user_token)
  90                         CloseHandle(h_user_token);
  91         }
  92
  93         return b_ret;
  94 }
  95
  96 void stop_hook(ULONG session_id)
  97 {
  98         std::map<ULONG, PROCESS_INFORMATION>::const_iterator pi_iter_32, pi_iter_64;
  99         HANDLE h_hook_processes[2] = {};
 100
 101         pi_iter_32 = pi_hooks_32.find(session_id);
 102         if (pi_iter_32 != pi_hooks_32.end())
 103                 h_hook_processes[0] = pi_iter_32->second.hProcess;
 104
 105         pi_iter_64 = pi_hooks_64.find(session_id);
 106         if (pi_iter_64 != pi_hooks_64.end())
 107                 h_hook_processes[1] = pi_iter_64->second.hProcess;
 108
 109         // notify and wait hook subprocesses to exit
 110         WaitForMultipleObjects(2, h_hook_processes, TRUE, INFINITE);
 111
 112         // clean up
 113
 114         if (pi_iter_32 != pi_hooks_32.end())
 115         {
 116                 CloseHandle(pi_iter_32->second.hThread);
 117                 CloseHandle(pi_iter_32->second.hProcess);
 118                 pi_hooks_32.erase(pi_iter_32);
 119         }
 120
 121         if (pi_iter_64 != pi_hooks_64.end())
 122         {
 123                 CloseHandle(pi_iter_64->second.hThread);
 124                 CloseHandle(pi_iter_64->second.hProcess);
 125                 pi_hooks_64.erase(pi_iter_64);
 126         }
 127
 128         CloseHandle(h_user_tokens[session_id]);
 129         h_user_tokens.erase(session_id);
 130 }
 131
 132 void set_svc_status(DWORD dwCurrentState, DWORD dwWin32ExitCode, DWORD dwWaitHint)
 133 {
 134         static DWORD dwCheckPoint = 1;
 135
 136         // fill in the SERVICE_STATUS structure
 137         svc_status.dwCurrentState = dwCurrentState;
 138         svc_status.dwWin32ExitCode = dwWin32ExitCode;
 139         svc_status.dwWaitHint = dwWaitHint;
 140
 141         if (dwCurrentState == SERVICE_RUNNING)
 142                 svc_status.dwControlsAccepted = SERVICE_ACCEPT_STOP | SERVICE_ACCEPT_SESSIONCHANGE;
 143         else
 144                 svc_status.dwControlsAccepted = 0;
 145
 146         if (dwCurrentState == SERVICE_RUNNING || dwCurrentState == SERVICE_STOPPED)
 147                 svc_status.dwCheckPoint = 0;
 148         else
 149                 svc_status.dwCheckPoint = ++dwCheckPoint;
 150
 151         // report the status of the service to the SCM
 152         SetServiceStatus(h_svc_status, &svc_status);
 153 }
 154
 155 DWORD WINAPI svc_ctrl_handler(DWORD dwCtrl, DWORD dwEventType, LPVOID lpEventData, LPVOID lpContext)
 156 {
 157         BOOL b_ret;
 158
 159         // handle the requested control code
 160         switch (dwCtrl)
 161         {
 162         case SERVICE_CONTROL_STOP:
 163                 set_svc_status(SERVICE_STOP_PENDING, NO_ERROR, 0);
 164
 165                 b_ret = SetEvent(h_svc_events);
 166                 assert(b_ret);
 167
 168                 return NO_ERROR;
 169         case SERVICE_CONTROL_INTERROGATE:
 170                 return NO_ERROR;
 171         case SERVICE_CONTROL_SESSIONCHANGE:
 172                 if (dwEventType == WTS_SESSION_LOGON)
 173                 {
 174                         b_ret = start_hook(reinterpret_cast<WTSSESSION_NOTIFICATION *>(lpEventData)->dwSessionId);
 175
 176                         if (b_ret)
 177                                 return NO_ERROR;
 178                         else
 179                                 return GetLastError();
 180                 }
 181                 else if (dwEventType == WTS_SESSION_LOGOFF)
 182                 {
 183                         stop_hook(reinterpret_cast<WTSSESSION_NOTIFICATION *>(lpEventData)->dwSessionId);
 184
 185                         return NO_ERROR;
 186                 }
 187                 else
 188                         return ERROR_CALL_NOT_IMPLEMENTED;
 189         default:
 190                 return ERROR_CALL_NOT_IMPLEMENTED;
 191         }
 192 }
 193
 194 VOID CALLBACK exit_cleanup(PVOID lpParameter, BOOLEAN TimerOrWaitFired)
 195 {
 196         BOOL b_ret;
 197
 198         b_ret = UnregisterWait(h_wait_cleanup);
 199         assert(b_ret || GetLastError() == ERROR_IO_PENDING);
 200
 201         b_ret = stop_gdipp_rpc_server();
 202         if (b_ret)
 203         {
 204                 const DWORD wait_ret = WaitForSingleObject(h_rpc_thread, INFINITE);
 205                 assert(wait_ret == WAIT_OBJECT_0);
 206         }
 207
 208         set_svc_status(SERVICE_STOPPED, NO_ERROR, 0);
 209 }
 210
 211 void svc_init()
 212 {
 213         h_svc_events = CreateEvent(NULL, TRUE, FALSE, NULL);
 214         if (h_svc_events == NULL)
 215         {
 216                 set_svc_status(SERVICE_STOPPED, NO_ERROR, 0);
 217                 return;
 218         }
 219
 220         // clean up when event is set
 221         if (!RegisterWaitForSingleObject(&h_wait_cleanup, h_svc_events, exit_cleanup, NULL, INFINITE, WT_EXECUTEDEFAULT | WT_EXECUTEONLYONCE))
 222         {
 223                 set_svc_status(SERVICE_STOPPED, NO_ERROR, 0);
 224                 return;
 225         }
 226
 227         // report running status when initialization is complete
 228         set_svc_status(SERVICE_RUNNING, NO_ERROR, 0);
 229
 230         // initialize RPC for font service
 231         h_rpc_thread = CreateThread(NULL, 0, start_gdipp_rpc_server, NULL, 0, NULL);
 232         if (h_rpc_thread == NULL)
 233         {
 234                 set_svc_status(SERVICE_STOPPED, NO_ERROR, 0);
 235                 return;
 236         }
 237
 238         /*
 239         service process and its child processes run in session 0
 240         some functions in gdipp may require interactive session (session ID > 0)
 241         use CreateProcessAsUser to create process in the active user's session
 242         */
 243         const DWORD active_session_id = WTSGetActiveConsoleSessionId();
 244         if (active_session_id != 0xFFFFFFFF)
 245                 start_hook(active_session_id);
 246 }
 247
 248 VOID WINAPI svc_main(DWORD dwArgc, LPTSTR *lpszArgv)
 249 {
 250         // register the handler function for the service
 251         h_svc_status = RegisterServiceCtrlHandlerExW(SVC_NAME, svc_ctrl_handler, NULL);
 252         if (h_svc_status == NULL)
 253                 return;
 254
 255         // these SERVICE_STATUS members remain as set here
 256         svc_status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
 257         svc_status.dwWin32ExitCode = NO_ERROR;
 258
 259         // report initial status to the SCM
 260         set_svc_status(SERVICE_START_PENDING, NO_ERROR, 3000);
 261
 262         svc_init();
 263 }
 264
 265 }
 266
 267 // #define svc_debug
 268
 269 int APIENTRY wWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPWSTR lpCmdLine, int nCmdShow)
 270 {
 271 #ifdef svc_debug
 272         Sleep(5000);
 273 #endif  // svc_debug
 274
 275         SERVICE_TABLE_ENTRY dispatch_table[] =
 276         {
 277                 { SVC_NAME, gdipp::svc_main },
 278                 { NULL, NULL },
 279         };
 280
 281         if (!StartServiceCtrlDispatcherW(dispatch_table))
 282                 return EXIT_FAILURE;
 283
 284         return EXIT_SUCCESS;
 285 }