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 /* See if we were launched from a console window. */\r
  59 static void check_console() {\r
  60   /* If we're running in a service context there will be no console window. */\r
  61   HWND console = GetConsoleWindow();\r
  62   if (! console) return;\r
  63 \r
  64   unsigned long pid;\r
  65   if (! GetWindowThreadProcessId(console, &pid)) return;\r
  66 \r
  67   /*\r
  68     If the process associated with the console window handle is the same as\r
  69     this process, we were not launched from an existing console.  The user\r
  70     probably double-clicked our executable.\r
  71   */\r
  72   if (GetCurrentProcessId() != pid) return;\r
  73 \r
  74   /* We close our new console so that subsequent messages appear in a popup. */\r
  75   FreeConsole();\r
  76 }\r
  77 \r
  78 int num_cpus() {\r
  79   DWORD_PTR i, affinity, system_affinity;\r
  80   if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;\r
  81   for (i = 0; system_affinity & (1LL << i); i++);\r
  82   return (int) i;\r
  83 }\r
  84 \r
  85 int _tmain(int argc, TCHAR **argv) {\r
  86   check_console();\r
  87 \r
  88 #ifdef UNICODE\r
  89   /*\r
  90     Ensure we write in UTF-16 mode, so that non-ASCII characters don't get\r
  91     mangled.  If we were compiled in ANSI mode it won't work.\r
  92    */\r
  93   _setmode(_fileno(stdout), _O_U16TEXT);\r
  94   _setmode(_fileno(stderr), _O_U16TEXT);\r
  95 #endif\r
  96 \r
  97   /* Remember if we are admin */\r
  98   check_admin();\r
  99 \r
 100   /* Elevate */\r
 101   if (argc > 1) {\r
 102     /*\r
 103       Valid commands are:\r
 104       start, stop, pause, continue, install, edit, get, set, reset, unset, remove\r
 105     */\r
 106     if (str_equiv(argv[1], _T("start"))) exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 107     if (str_equiv(argv[1], _T("stop"))) exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));\r
 108     if (str_equiv(argv[1], _T("restart"))) {\r
 109       int ret = control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2);\r
 110       if (ret) exit(ret);\r
 111       exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));\r
 112     }\r
 113     if (str_equiv(argv[1], _T("pause"))) exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));\r
 114     if (str_equiv(argv[1], _T("continue"))) exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));\r
 115     if (str_equiv(argv[1], _T("status"))) exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));\r
 116     if (str_equiv(argv[1], _T("rotate"))) exit(control_service(NSSM_SERVICE_CONTROL_ROTATE, argc - 2, argv + 2));\r
 117     if (str_equiv(argv[1], _T("install"))) {\r
 118       if (! is_admin) {\r
 119         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL);\r
 120         exit(100);\r
 121       }\r
 122       exit(pre_install_service(argc - 2, argv + 2));\r
 123     }\r
 124     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
 125       if (! is_admin) {\r
 126         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT);\r
 127         exit(100);\r
 128       }\r
 129       int ret = pre_edit_service(argc - 1, argv + 1);\r
 130       /* There might be a password here. */\r
 131       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 132       exit(ret);\r
 133     }\r
 134     if (str_equiv(argv[1], _T("remove"))) {\r
 135       if (! is_admin) {\r
 136         print_message(stderr, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE);\r
 137         exit(100);\r
 138       }\r
 139       exit(pre_remove_service(argc - 2, argv + 2));\r
 140     }\r
 141   }\r
 142 \r
 143   /* Thread local storage for error message buffer */\r
 144   tls_index = TlsAlloc();\r
 145 \r
 146   /* Register messages */\r
 147   if (is_admin) create_messages();\r
 148 \r
 149   /*\r
 150     Optimisation for Windows 2000:\r
 151     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 152     will time out after a few seconds on Windows 2000.  On newer versions the\r
 153     call returns instantly.  Check for stdin first and only try to call the\r
 154     function if there's no input stream found.  Although it's possible that\r
 155     we're running with input redirected it's much more likely that we're\r
 156     actually running as a service.\r
 157     This will save time when running with no arguments from a command prompt.\r
 158   */\r
 159   if (_fileno(stdin) < 0) {\r
 160     /* Set up function pointers. */\r
 161     if (get_imports()) exit(111);\r
 162 \r
 163     /* Start service magic */\r
 164     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 165     if (! StartServiceCtrlDispatcher(table)) {\r
 166       unsigned long error = GetLastError();\r
 167       /* User probably ran nssm with no argument */\r
 168       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 169       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 170       free_imports();\r
 171       exit(100);\r
 172     }\r
 173   }\r
 174   else exit(usage(1));\r
 175 \r
 176   /* And nothing more to do */\r
 177   exit(0);\r
 178 }\r