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 *) HeapAlloc(GetProcessHeap(), 0, PATH_LENGTH);\r
  69   if (! sei.lpFile) {\r
  70     print_message(stderr, NSSM_MESSAGE_OUT_OF_MEMORY, _T("GetModuleFileName()"), _T("elevate()"));\r
  71     return 111;\r
  72   }\r
  73   GetModuleFileName(0, (TCHAR *) sei.lpFile, PATH_LENGTH);\r
  74 \r
  75   TCHAR *args = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, EXE_LENGTH * sizeof(TCHAR));\r
  76   if (! args) {\r
  77     HeapFree(GetProcessHeap(), 0, (void *) sei.lpFile);\r
  78     print_message(stderr, NSSM_MESSAGE_OUT_OF_MEMORY, _T("GetCommandLine()"), _T("elevate()"));\r
  79     return 111;\r
  80   }\r
  81 \r
  82   /* Get command line, which includes the path to NSSM, and skip that part. */\r
  83   _sntprintf_s(args, EXE_LENGTH, _TRUNCATE, _T("%s"), GetCommandLine());\r
  84   size_t s = _tcslen(argv[0]) + 1;\r
  85   if (args[0] == _T('"')) s += 2;\r
  86   while (isspace(args[s])) s++;\r
  87 \r
  88   sei.lpParameters = args + s;\r
  89   sei.nShow = SW_SHOW;\r
  90 \r
  91   unsigned long exitcode = 0;\r
  92   if (! ShellExecuteEx(&sei)) exitcode = 100;\r
  93 \r
  94   HeapFree(GetProcessHeap(), 0, (void *) sei.lpFile);\r
  95   HeapFree(GetProcessHeap(), 0, (void *) args);\r
  96   return exitcode;\r
  97 }\r
  98 \r
  99 int num_cpus() {\r
 100   DWORD_PTR i, affinity, system_affinity;\r
 101   if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;\r
 102   for (i = 0; system_affinity & (1LL << i); i++);\r
 103   return (int) i;\r
 104 }\r
 105 \r
 106 const TCHAR *nssm_unquoted_imagepath() {\r
 107   return unquoted_imagepath;\r
 108 }\r
 109 \r
 110 const TCHAR *nssm_imagepath() {\r
 111   return imagepath;\r
 112 }\r
 113 \r
 114 int _tmain(int argc, TCHAR **argv) {\r
 115   check_console();\r
 116 \r
 117 #ifdef UNICODE\r
 118   /*\r
 119     Ensure we write in UTF-16 mode, so that non-ASCII characters don't get\r
 120     mangled.  If we were compiled in ANSI mode it won't work.\r
 121    */\r
 122   _setmode(_fileno(stdout), _O_U16TEXT);\r
 123   _setmode(_fileno(stderr), _O_U16TEXT);\r
 124 #endif\r
 125 \r
 126   /* Remember if we are admin */\r
 127   check_admin();\r
 128 \r
 129   /* Set up function pointers. */\r
 130   if (get_imports()) exit(111);\r
 131 \r
 132   /* Remember our path for later. */\r
 133   GetModuleFileName(0, unquoted_imagepath, _countof(unquoted_imagepath));\r
 134   GetModuleFileName(0, imagepath, _countof(imagepath));\r
 135   PathQuoteSpaces(imagepath);\r
 136 \r
 137   /* Elevate */\r
 138   if (argc > 1) {\r
 139     /*\r
 140       Valid commands are:\r
 141       start, stop, pause, continue, install, edit, get, set, reset, unset, remove\r
 142     */\r
 143     if (str_equiv(argv[1], _T("start"))) exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 144     if (str_equiv(argv[1], _T("stop"))) exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));\r
 145     if (str_equiv(argv[1], _T("restart"))) {\r
 146       int ret = control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2);\r
 147       if (ret) exit(ret);\r
 148       exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 149     }\r
 150     if (str_equiv(argv[1], _T("pause"))) exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));\r
 151     if (str_equiv(argv[1], _T("continue"))) exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));\r
 152     if (str_equiv(argv[1], _T("status"))) exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));\r
 153     if (str_equiv(argv[1], _T("rotate"))) exit(control_service(NSSM_SERVICE_CONTROL_ROTATE, argc - 2, argv + 2));\r
 154     if (str_equiv(argv[1], _T("install"))) {\r
 155       if (! is_admin) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL));\r
 156       exit(pre_install_service(argc - 2, argv + 2));\r
 157     }\r
 158     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
 159       int ret = pre_edit_service(argc - 1, argv + 1);\r
 160       if (ret == 3 && ! is_admin && argc == 3) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT));\r
 161       /* There might be a password here. */\r
 162       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 163       exit(ret);\r
 164     }\r
 165     if (str_equiv(argv[1], _T("remove"))) {\r
 166       if (! is_admin) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE));\r
 167       exit(pre_remove_service(argc - 2, argv + 2));\r
 168     }\r
 169   }\r
 170 \r
 171   /* Thread local storage for error message buffer */\r
 172   tls_index = TlsAlloc();\r
 173 \r
 174   /* Register messages */\r
 175   if (is_admin) create_messages();\r
 176 \r
 177   /*\r
 178     Optimisation for Windows 2000:\r
 179     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 180     will time out after a few seconds on Windows 2000.  On newer versions the\r
 181     call returns instantly.  Check for stdin first and only try to call the\r
 182     function if there's no input stream found.  Although it's possible that\r
 183     we're running with input redirected it's much more likely that we're\r
 184     actually running as a service.\r
 185     This will save time when running with no arguments from a command prompt.\r
 186   */\r
 187   if (! GetStdHandle(STD_INPUT_HANDLE)) {\r
 188     /* Start service magic */\r
 189     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 190     if (! StartServiceCtrlDispatcher(table)) {\r
 191       unsigned long error = GetLastError();\r
 192       /* User probably ran nssm with no argument */\r
 193       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 194       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 195       free_imports();\r
 196       exit(100);\r
 197     }\r
 198   }\r
 199   else exit(usage(1));\r
 200 \r
 201   /* And nothing more to do */\r
 202   exit(0);\r
 203 }\r