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       create_messages();\r
 150       exit(pre_install_service(argc - 2, argv + 2));\r
 151     }\r
 152     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
 153       int ret = pre_edit_service(argc - 1, argv + 1);\r
 154       if (ret == 3 && ! is_admin && argc == 3) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT));\r
 155       /* There might be a password here. */\r
 156       for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));\r
 157       exit(ret);\r
 158     }\r
 159     if (str_equiv(argv[1], _T("list"))) exit(list_nssm_services());\r
 160     if (str_equiv(argv[1], _T("remove"))) {\r
 161       if (! is_admin) exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE));\r
 162       exit(pre_remove_service(argc - 2, argv + 2));\r
 163     }\r
 164   }\r
 165 \r
 166   /* Thread local storage for error message buffer */\r
 167   tls_index = TlsAlloc();\r
 168 \r
 169   /* Register messages */\r
 170   if (is_admin) create_messages();\r
 171 \r
 172   /*\r
 173     Optimisation for Windows 2000:\r
 174     When we're run from the command line the StartServiceCtrlDispatcher() call\r
 175     will time out after a few seconds on Windows 2000.  On newer versions the\r
 176     call returns instantly.  Check for stdin first and only try to call the\r
 177     function if there's no input stream found.  Although it's possible that\r
 178     we're running with input redirected it's much more likely that we're\r
 179     actually running as a service.\r
 180     This will save time when running with no arguments from a command prompt.\r
 181   */\r
 182   if (! GetStdHandle(STD_INPUT_HANDLE)) {\r
 183     /* Start service magic */\r
 184     SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };\r
 185     if (! StartServiceCtrlDispatcher(table)) {\r
 186       unsigned long error = GetLastError();\r
 187       /* User probably ran nssm with no argument */\r
 188       if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) exit(usage(1));\r
 189       log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);\r
 190       free_imports();\r
 191       exit(100);\r
 192     }\r
 193   }\r
 194   else exit(usage(1));\r
 195 \r
 196   /* And nothing more to do */\r
 197   exit(0);\r
 198 }\r