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