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 \r
  10 /* Are two strings case-insensitively equivalent? */\r
  11 int str_equiv(const TCHAR *a, const TCHAR *b) {\r
  12   size_t len = _tcslen(a);\r
  13   if (_tcslen(b) != len) return 0;\r
  14   if (_tcsnicmp(a, b, len)) return 0;\r
  15   return 1;\r
  16 }\r
  17 \r
  18 /* Convert a string to a number. */\r
  19 int str_number(const TCHAR *string, unsigned long *number, TCHAR **bogus) {\r
  20   if (! string) return 1;\r
  21 \r
  22   *number = _tcstoul(string, bogus, 0);\r
  23   if (**bogus) return 2;\r
  24 \r
  25   return 0;\r
  26 }\r
  27 \r
  28 int str_number(const TCHAR *string, unsigned long *number) {\r
  29   TCHAR *bogus;\r
  30   return str_number(string, number, &bogus);\r
  31 }\r
  32 \r
  33 /* Remove basename of a path. */\r
  34 void strip_basename(TCHAR *buffer) {\r
  35   size_t len = _tcslen(buffer);\r
  36   size_t i;\r
  37   for (i = len; i && buffer[i] != _T('\\') && buffer[i] != _T('/'); i--);\r
  38   /* X:\ is OK. */\r
  39   if (i && buffer[i - 1] == _T(':')) i++;\r
  40   buffer[i] = _T('\0');\r
  41 }\r
  42 \r
  43 /* How to use me correctly */\r
  44 int usage(int ret) {\r
  45   if (GetConsoleWindow()) print_message(stderr, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);\r
  46   else popup_message(0, MB_OK, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);\r
  47   return(ret);\r
  48 }\r
  49 \r
  50 void check_admin() {\r
  51   is_admin = false;\r
  52 \r
  53   /* Lifted from MSDN examples */\r
  54   PSID AdministratorsGroup;\r
  55   SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;\r
  56   if (! AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &AdministratorsGroup)) return;\r
  57   CheckTokenMembership(0, AdministratorsGroup, /*XXX*/(PBOOL) &is_admin);\r
  58   FreeSid(AdministratorsGroup);\r
  59 }\r
  60 \r
  61 static int elevate(int argc, TCHAR **argv, unsigned long message) {\r
  62   print_message(stderr, message);\r
  63 \r
  64   SHELLEXECUTEINFO sei;\r
  65   ZeroMemory(&sei, sizeof(sei));\r
  66   sei.cbSize = sizeof(sei);\r
  67   sei.lpVerb = _T("runas");\r
  68   sei.lpFile = (TCHAR *) nssm_imagepath();\r
  69 \r
  70   TCHAR *args = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, EXE_LENGTH * sizeof(TCHAR));\r
  71   if (! args) {\r
  72     print_message(stderr, NSSM_MESSAGE_OUT_OF_MEMORY, _T("GetCommandLine()"), _T("elevate()"));\r
  73     return 111;\r
  74   }\r
  75 \r
  76   /* Get command line, which includes the path to NSSM, and skip that part. */\r
  77   _sntprintf_s(args, EXE_LENGTH, _TRUNCATE, _T("%s"), GetCommandLine());\r
  78   size_t s = _tcslen(argv[0]) + 1;\r
  79   if (args[0] == _T('"')) s += 2;\r
  80   while (isspace(args[s])) s++;\r
  81 \r
  82   sei.lpParameters = args + s;\r
  83   sei.nShow = SW_SHOW;\r
  84 \r
  85   unsigned long exitcode = 0;\r
  86   if (! ShellExecuteEx(&sei)) exitcode = 100;\r
  87 \r
  88   HeapFree(GetProcessHeap(), 0, (void *) args);\r
  89   return exitcode;\r
  90 }\r
  91 \r
  92 int num_cpus() {\r
  93   DWORD_PTR i, affinity, system_affinity;\r
  94   if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;\r
  95   for (i = 0; system_affinity & (1LL << i); i++);\r
  96   return (int) i;\r
  97 }\r
  98 \r
  99 const TCHAR *nssm_unquoted_imagepath() {\r
 100   return unquoted_imagepath;\r
 101 }\r
 102 \r
 103 const TCHAR *nssm_imagepath() {\r
 104   return imagepath;\r
 105 }\r
 106 \r
 107 int _tmain(int argc, TCHAR **argv) {\r
 108   check_console();\r
 109 \r
 110 #ifdef UNICODE\r
 111   /*\r
 112     Ensure we write in UTF-16 mode, so that non-ASCII characters don't get\r
 113     mangled.  If we were compiled in ANSI mode it won't work.\r
 114    */\r
 115   _setmode(_fileno(stdout), _O_U16TEXT);\r
 116   _setmode(_fileno(stderr), _O_U16TEXT);\r
 117 #endif\r
 118 \r
 119   /* Remember if we are admin */\r
 120   check_admin();\r
 121 \r
 122   /* Set up function pointers. */\r
 123   if (get_imports()) exit(111);\r
 124 \r
 125   /* Remember our path for later. */\r
 126   GetModuleFileName(0, unquoted_imagepath, _countof(unquoted_imagepath));\r
 127   GetModuleFileName(0, imagepath, _countof(imagepath));\r
 128   PathQuoteSpaces(imagepath);\r
 129 \r
 130   /* Elevate */\r
 131   if (argc > 1) {\r
 132     /*\r
 133       Valid commands are:\r
 134       start, stop, pause, continue, install, edit, get, set, reset, unset, remove\r
 135     */\r
 136     if (str_equiv(argv[1], _T("start"))) exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 137     if (str_equiv(argv[1], _T("stop"))) exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));\r
 138     if (str_equiv(argv[1], _T("restart"))) {\r
 139       int ret = control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2);\r
 140       if (ret) exit(ret);\r
 141       exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 142     }\r
 143     if (str_equiv(argv[1], _T("pause"))) exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));\r
 144     if (str_equiv(argv[1], _T("continue"))) exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));\r
 145     if (str_equiv(argv[1], _T("status"))) exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));\r
 146     if (str_equiv(argv[1], _T("rotate"))) exit(control_service(NSSM_SERVICE_CONTROL_ROTATE, argc - 2, argv + 2));\r
 147     if (str_equiv(argv[1], _T("install"))) {\r
 148       if (! is_admin) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL));\r
 149       exit(pre_install_service(argc - 2, argv + 2));\r
 150     }\r
 151     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"))) {\r
 152       int ret = pre_edit_service(argc - 1, argv + 1);\r
 153       if (ret == 3 && ! is_admin && argc == 3) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT));\r
 154       /* There might be a password here. */\r
 155       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 156       exit(ret);\r
 157     }\r
 158     if (str_equiv(argv[1], _T("remove"))) {\r
 159       if (! is_admin) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE));\r
 160       exit(pre_remove_service(argc - 2, argv + 2));\r
 161     }\r
 162   }\r
 163 \r
 164   /* Thread local storage for error message buffer */\r
 165   tls_index = TlsAlloc();\r
 166 \r
 167   /* Register messages */\r
 168   if (is_admin) create_messages();\r
 169 \r
 170   /*\r
 171     Optimisation for Windows 2000:\r
 172     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 173     will time out after a few seconds on Windows 2000.  On newer versions the\r
 174     call returns instantly.  Check for stdin first and only try to call the\r
 175     function if there's no input stream found.  Although it's possible that\r
 176     we're running with input redirected it's much more likely that we're\r
 177     actually running as a service.\r
 178     This will save time when running with no arguments from a command prompt.\r
 179   */\r
 180   if (! GetStdHandle(STD_INPUT_HANDLE)) {\r
 181     /* Start service magic */\r
 182     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 183     if (! StartServiceCtrlDispatcher(table)) {\r
 184       unsigned long error = GetLastError();\r
 185       /* User probably ran nssm with no argument */\r
 186       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 187       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 188       free_imports();\r
 189       exit(100);\r
 190     }\r
 191   }\r
 192   else exit(usage(1));\r
 193 \r
 194   /* And nothing more to do */\r
 195   exit(0);\r
 196 }\r