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