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