nssm.cpp

   1 #include "nssm.h"\r
   2 \r
   3 extern unsigned long tls_index;\r
   4 extern bool is_admin;\r
   5 extern imports_t imports;\r
   6 \r
   7 static TCHAR unquoted_imagepath[PATH_LENGTH];\r
   8 static TCHAR imagepath[PATH_LENGTH];\r
   9 static TCHAR imageargv0[PATH_LENGTH];\r
  10 \r
  11 /* Are two strings case-insensitively equivalent? */\r
  12 int str_equiv(const TCHAR *a, const TCHAR *b) {\r
  13   size_t len = _tcslen(a);\r
  14   if (_tcslen(b) != len) return 0;\r
  15   if (_tcsnicmp(a, b, len)) return 0;\r
  16   return 1;\r
  17 }\r
  18 \r
  19 /* Convert a string to a number. */\r
  20 int str_number(const TCHAR *string, unsigned long *number, TCHAR **bogus) {\r
  21   if (! string) return 1;\r
  22 \r
  23   *number = _tcstoul(string, bogus, 0);\r
  24   if (**bogus) return 2;\r
  25 \r
  26   return 0;\r
  27 }\r
  28 \r
  29 int str_number(const TCHAR *string, unsigned long *number) {\r
  30   TCHAR *bogus;\r
  31   return str_number(string, number, &bogus);\r
  32 }\r
  33 \r
  34 /* Does a char need to be escaped? */\r
  35 static bool needs_escape(const TCHAR c) {\r
  36   if (c == _T('"')) return true;\r
  37   if (c == _T('&')) return true;\r
  38   if (c == _T('%')) return true;\r
  39   if (c == _T('^')) return true;\r
  40   if (c == _T('<')) return true;\r
  41   if (c == _T('>')) return true;\r
  42   if (c == _T('|')) return true;\r
  43   return false;\r
  44 }\r
  45 \r
  46 /* Does a char need to be quoted? */\r
  47 static bool needs_quote(const TCHAR c) {\r
  48   if (c == _T(' ')) return true;\r
  49   if (c == _T('\t')) return true;\r
  50   if (c == _T('\n')) return true;\r
  51   if (c == _T('\v')) return true;\r
  52   if (c == _T('"')) return true;\r
  53   if (c == _T('*')) return true;\r
  54   return needs_escape(c);\r
  55 }\r
  56 \r
  57 /* https://blogs.msdn.microsoft.com/twistylittlepassagesallalike/2011/04/23/everyone-quotes-command-line-arguments-the-wrong-way/ */\r
  58 /* http://www.robvanderwoude.com/escapechars.php */\r
  59 int quote(const TCHAR *unquoted, TCHAR *buffer, size_t buflen) {\r
  60   size_t i, j, n;\r
  61   size_t len = _tcslen(unquoted);\r
  62   if (len > buflen - 1) return 1;\r
  63 \r
  64   bool escape = false;\r
  65   bool quotes = false;\r
  66 \r
  67   for (i = 0; i < len; i++) {\r
  68     if (needs_escape(unquoted[i])) {\r
  69       escape = quotes = true;\r
  70       break;\r
  71     }\r
  72     if (needs_quote(unquoted[i])) quotes = true;\r
  73   }\r
  74   if (! quotes) {\r
  75     memmove(buffer, unquoted, (len + 1) * sizeof(TCHAR));\r
  76     return 0;\r
  77   }\r
  78 \r
  79   /* "" */\r
  80   size_t quoted_len = 2;\r
  81   if (escape) quoted_len += 2;\r
  82   for (i = 0; ; i++) {\r
  83     n = 0;\r
  84 \r
  85     while (i != len && unquoted[i] == _T('\\')) {\r
  86       i++;\r
  87       n++;\r
  88     }\r
  89 \r
  90     if (i == len) {\r
  91       quoted_len += n * 2;\r
  92       break;\r
  93     }\r
  94     else if (unquoted[i] == _T('"')) quoted_len += n * 2 + 2;\r
  95     else quoted_len += n + 1;\r
  96     if (needs_escape(unquoted[i])) quoted_len += n;\r
  97   }\r
  98   if (quoted_len > buflen - 1) return 1;\r
  99 \r
 100   TCHAR *s = buffer;\r
 101   if (escape) *s++ = _T('^');\r
 102   *s++ = _T('"');\r
 103 \r
 104   for (i = 0; ; i++) {\r
 105     n = 0;\r
 106 \r
 107     while (i != len && unquoted[i] == _T('\\')) {\r
 108       i++;\r
 109       n++;\r
 110     }\r
 111 \r
 112     if (i == len) {\r
 113       for (j = 0; j < n * 2; j++) {\r
 114         if (escape) *s++ = _T('^');\r
 115         *s++ = _T('\\');\r
 116       }\r
 117       break;\r
 118     }\r
 119     else if (unquoted[i] == _T('"')) {\r
 120       for (j = 0; j < n * 2 + 1; j++) {\r
 121         if (escape) *s++ = _T('^');\r
 122         *s++ = _T('\\');\r
 123       }\r
 124       if (escape && needs_escape(unquoted[i])) *s++ = _T('^');\r
 125       *s++ = unquoted[i];\r
 126     }\r
 127     else {\r
 128       for (j = 0; j < n; j++) {\r
 129         if (escape) *s++ = _T('^');\r
 130         *s++ = _T('\\');\r
 131       }\r
 132       if (escape && needs_escape(unquoted[i])) *s++ = _T('^');\r
 133       *s++ = unquoted[i];\r
 134     }\r
 135   }\r
 136   if (escape) *s++ = _T('^');\r
 137   *s++ = _T('"');\r
 138   *s++ = _T('\0');\r
 139 \r
 140   return 0;\r
 141 }\r
 142 \r
 143 /* Remove basename of a path. */\r
 144 void strip_basename(TCHAR *buffer) {\r
 145   size_t len = _tcslen(buffer);\r
 146   size_t i;\r
 147   for (i = len; i && buffer[i] != _T('\\') && buffer[i] != _T('/'); i--);\r
 148   /* X:\ is OK. */\r
 149   if (i && buffer[i - 1] == _T(':')) i++;\r
 150   buffer[i] = _T('\0');\r
 151 }\r
 152 \r
 153 /* How to use me correctly */\r
 154 int usage(int ret) {\r
 155   if (GetConsoleWindow()) print_message(stderr, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);\r
 156   else popup_message(0, MB_OK, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);\r
 157   return(ret);\r
 158 }\r
 159 \r
 160 void check_admin() {\r
 161   is_admin = false;\r
 162 \r
 163   /* Lifted from MSDN examples */\r
 164   PSID AdministratorsGroup;\r
 165   SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;\r
 166   if (! AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &AdministratorsGroup)) return;\r
 167   CheckTokenMembership(0, AdministratorsGroup, /*XXX*/(PBOOL) &is_admin);\r
 168   FreeSid(AdministratorsGroup);\r
 169 }\r
 170 \r
 171 static int elevate(int argc, TCHAR **argv, unsigned long message) {\r
 172   print_message(stderr, message);\r
 173 \r
 174   SHELLEXECUTEINFO sei;\r
 175   ZeroMemory(&sei, sizeof(sei));\r
 176   sei.cbSize = sizeof(sei);\r
 177   sei.lpVerb = _T("runas");\r
 178   sei.lpFile = (TCHAR *) nssm_imagepath();\r
 179 \r
 180   TCHAR *args = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, EXE_LENGTH * sizeof(TCHAR));\r
 181   if (! args) {\r
 182     print_message(stderr, NSSM_MESSAGE_OUT_OF_MEMORY, _T("GetCommandLine()"), _T("elevate()"));\r
 183     return 111;\r
 184   }\r
 185 \r
 186   /* Get command line, which includes the path to NSSM, and skip that part. */\r
 187   _sntprintf_s(args, EXE_LENGTH, _TRUNCATE, _T("%s"), GetCommandLine());\r
 188   size_t s = _tcslen(argv[0]) + 1;\r
 189   if (args[0] == _T('"')) s += 2;\r
 190   while (isspace(args[s])) s++;\r
 191 \r
 192   sei.lpParameters = args + s;\r
 193   sei.nShow = SW_SHOW;\r
 194 \r
 195   unsigned long exitcode = 0;\r
 196   if (! ShellExecuteEx(&sei)) exitcode = 100;\r
 197 \r
 198   HeapFree(GetProcessHeap(), 0, (void *) args);\r
 199   return exitcode;\r
 200 }\r
 201 \r
 202 int num_cpus() {\r
 203   DWORD_PTR i, affinity, system_affinity;\r
 204   if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;\r
 205   for (i = 0; system_affinity & (1LL << i); i++);\r
 206   return (int) i;\r
 207 }\r
 208 \r
 209 const TCHAR *nssm_unquoted_imagepath() {\r
 210   return unquoted_imagepath;\r
 211 }\r
 212 \r
 213 const TCHAR *nssm_imagepath() {\r
 214   return imagepath;\r
 215 }\r
 216 \r
 217 const TCHAR *nssm_exe() {\r
 218   return imageargv0;\r
 219 }\r
 220 \r
 221 int _tmain(int argc, TCHAR **argv) {\r
 222   check_console();\r
 223 \r
 224 #ifdef UNICODE\r
 225   /*\r
 226     Ensure we write in UTF-16 mode, so that non-ASCII characters don't get\r
 227     mangled.  If we were compiled in ANSI mode it won't work.\r
 228    */\r
 229   _setmode(_fileno(stdout), _O_U16TEXT);\r
 230   _setmode(_fileno(stderr), _O_U16TEXT);\r
 231 #endif\r
 232 \r
 233   /* Remember if we are admin */\r
 234   check_admin();\r
 235 \r
 236   /* Set up function pointers. */\r
 237   if (get_imports()) exit(111);\r
 238 \r
 239   /* Remember our path for later. */\r
 240   _sntprintf_s(imageargv0, _countof(imageargv0), _TRUNCATE, _T("%s"), argv[0]);\r
 241   PathQuoteSpaces(imageargv0);\r
 242   GetModuleFileName(0, unquoted_imagepath, _countof(unquoted_imagepath));\r
 243   GetModuleFileName(0, imagepath, _countof(imagepath));\r
 244   PathQuoteSpaces(imagepath);\r
 245 \r
 246   /* Elevate */\r
 247   if (argc > 1) {\r
 248     /*\r
 249       Valid commands are:\r
 250       start, stop, pause, continue, install, edit, get, set, reset, unset, remove\r
 251     */\r
 252     if (str_equiv(argv[1], _T("start"))) exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 253     if (str_equiv(argv[1], _T("stop"))) exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));\r
 254     if (str_equiv(argv[1], _T("restart"))) {\r
 255       int ret = control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2);\r
 256       if (ret) exit(ret);\r
 257       exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 258     }\r
 259     if (str_equiv(argv[1], _T("pause"))) exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));\r
 260     if (str_equiv(argv[1], _T("continue"))) exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));\r
 261     if (str_equiv(argv[1], _T("status"))) exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));\r
 262     if (str_equiv(argv[1], _T("rotate"))) exit(control_service(NSSM_SERVICE_CONTROL_ROTATE, argc - 2, argv + 2));\r
 263     if (str_equiv(argv[1], _T("install"))) {\r
 264       if (! is_admin) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL));\r
 265       create_messages();\r
 266       exit(pre_install_service(argc - 2, argv + 2));\r
 267     }\r
 268     if (str_equiv(argv[1], _T("edit")) || str_equiv(argv[1], _T("get")) || str_equiv(argv[1], _T("set")) || str_equiv(argv[1], _T("reset")) || str_equiv(argv[1], _T("unset")) || str_equiv(argv[1], _T("dump"))) {\r
 269       int ret = pre_edit_service(argc - 1, argv + 1);\r
 270       if (ret == 3 && ! is_admin && argc == 3) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT));\r
 271       /* There might be a password here. */\r
 272       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 273       exit(ret);\r
 274     }\r
 275     if (str_equiv(argv[1], _T("list"))) exit(list_nssm_services());\r
 276     if (str_equiv(argv[1], _T("processes"))) exit(service_process_tree(argc - 2, argv + 2));\r
 277     if (str_equiv(argv[1], _T("remove"))) {\r
 278       if (! is_admin) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE));\r
 279       exit(pre_remove_service(argc - 2, argv + 2));\r
 280     }\r
 281   }\r
 282 \r
 283   /* Thread local storage for error message buffer */\r
 284   tls_index = TlsAlloc();\r
 285 \r
 286   /* Register messages */\r
 287   if (is_admin) create_messages();\r
 288 \r
 289   /*\r
 290     Optimisation for Windows 2000:\r
 291     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 292     will time out after a few seconds on Windows 2000.  On newer versions the\r
 293     call returns instantly.  Check for stdin first and only try to call the\r
 294     function if there's no input stream found.  Although it's possible that\r
 295     we're running with input redirected it's much more likely that we're\r
 296     actually running as a service.\r
 297     This will save time when running with no arguments from a command prompt.\r
 298   */\r
 299   if (! GetStdHandle(STD_INPUT_HANDLE)) {\r
 300     /* Start service magic */\r
 301     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 302     if (! StartServiceCtrlDispatcher(table)) {\r
 303       unsigned long error = GetLastError();\r
 304       /* User probably ran nssm with no argument */\r
 305       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 306       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 307       free_imports();\r
 308       exit(100);\r
 309     }\r
 310   }\r
 311   else exit(usage(1));\r
 312 \r
 313   /* And nothing more to do */\r
 314   exit(0);\r
 315 }\r