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