nssm.cpp 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312
  1. #include "nssm.h"
  2. extern unsigned long tls_index;
  3. extern bool is_admin;
  4. extern imports_t imports;
  5. static TCHAR unquoted_imagepath[PATH_LENGTH];
  6. static TCHAR imagepath[PATH_LENGTH];
  7. static TCHAR imageargv0[PATH_LENGTH];
  8. void nssm_exit(int status) {
  9. free_imports();
  10. unsetup_utf8();
  11. exit(status);
  12. }
  13. /* Are two strings case-insensitively equivalent? */
  14. int str_equiv(const TCHAR *a, const TCHAR *b) {
  15. size_t len = _tcslen(a);
  16. if (_tcslen(b) != len) return 0;
  17. if (_tcsnicmp(a, b, len)) return 0;
  18. return 1;
  19. }
  20. /* Convert a string to a number. */
  21. int str_number(const TCHAR *string, unsigned long *number, TCHAR **bogus) {
  22. if (! string) return 1;
  23. *number = _tcstoul(string, bogus, 0);
  24. if (**bogus) return 2;
  25. return 0;
  26. }
  27. int str_number(const TCHAR *string, unsigned long *number) {
  28. TCHAR *bogus;
  29. return str_number(string, number, &bogus);
  30. }
  31. /* Does a char need to be escaped? */
  32. static bool needs_escape(const TCHAR c) {
  33. if (c == _T('"')) return true;
  34. if (c == _T('&')) return true;
  35. if (c == _T('%')) return true;
  36. if (c == _T('^')) return true;
  37. if (c == _T('<')) return true;
  38. if (c == _T('>')) return true;
  39. if (c == _T('|')) return true;
  40. return false;
  41. }
  42. /* Does a char need to be quoted? */
  43. static bool needs_quote(const TCHAR c) {
  44. if (c == _T(' ')) return true;
  45. if (c == _T('\t')) return true;
  46. if (c == _T('\n')) return true;
  47. if (c == _T('\v')) return true;
  48. if (c == _T('"')) return true;
  49. if (c == _T('*')) return true;
  50. return needs_escape(c);
  51. }
  52. /* https://blogs.msdn.microsoft.com/twistylittlepassagesallalike/2011/04/23/everyone-quotes-command-line-arguments-the-wrong-way/ */
  53. /* http://www.robvanderwoude.com/escapechars.php */
  54. int quote(const TCHAR *unquoted, TCHAR *buffer, size_t buflen) {
  55. size_t i, j, n;
  56. size_t len = _tcslen(unquoted);
  57. if (len > buflen - 1) return 1;
  58. bool escape = false;
  59. bool quotes = false;
  60. for (i = 0; i < len; i++) {
  61. if (needs_escape(unquoted[i])) {
  62. escape = quotes = true;
  63. break;
  64. }
  65. if (needs_quote(unquoted[i])) quotes = true;
  66. }
  67. if (! quotes) {
  68. memmove(buffer, unquoted, (len + 1) * sizeof(TCHAR));
  69. return 0;
  70. }
  71. /* "" */
  72. size_t quoted_len = 2;
  73. if (escape) quoted_len += 2;
  74. for (i = 0; ; i++) {
  75. n = 0;
  76. while (i != len && unquoted[i] == _T('\\')) {
  77. i++;
  78. n++;
  79. }
  80. if (i == len) {
  81. quoted_len += n * 2;
  82. break;
  83. }
  84. else if (unquoted[i] == _T('"')) quoted_len += n * 2 + 2;
  85. else quoted_len += n + 1;
  86. if (needs_escape(unquoted[i])) quoted_len += n;
  87. }
  88. if (quoted_len > buflen - 1) return 1;
  89. TCHAR *s = buffer;
  90. if (escape) *s++ = _T('^');
  91. *s++ = _T('"');
  92. for (i = 0; ; i++) {
  93. n = 0;
  94. while (i != len && unquoted[i] == _T('\\')) {
  95. i++;
  96. n++;
  97. }
  98. if (i == len) {
  99. for (j = 0; j < n * 2; j++) {
  100. if (escape) *s++ = _T('^');
  101. *s++ = _T('\\');
  102. }
  103. break;
  104. }
  105. else if (unquoted[i] == _T('"')) {
  106. for (j = 0; j < n * 2 + 1; j++) {
  107. if (escape) *s++ = _T('^');
  108. *s++ = _T('\\');
  109. }
  110. if (escape && needs_escape(unquoted[i])) *s++ = _T('^');
  111. *s++ = unquoted[i];
  112. }
  113. else {
  114. for (j = 0; j < n; j++) {
  115. if (escape) *s++ = _T('^');
  116. *s++ = _T('\\');
  117. }
  118. if (escape && needs_escape(unquoted[i])) *s++ = _T('^');
  119. *s++ = unquoted[i];
  120. }
  121. }
  122. if (escape) *s++ = _T('^');
  123. *s++ = _T('"');
  124. *s++ = _T('\0');
  125. return 0;
  126. }
  127. /* Remove basename of a path. */
  128. void strip_basename(TCHAR *buffer) {
  129. size_t len = _tcslen(buffer);
  130. size_t i;
  131. for (i = len; i && buffer[i] != _T('\\') && buffer[i] != _T('/'); i--);
  132. /* X:\ is OK. */
  133. if (i && buffer[i - 1] == _T(':')) i++;
  134. buffer[i] = _T('\0');
  135. }
  136. /* How to use me correctly */
  137. int usage(int ret) {
  138. if ((! GetConsoleWindow() || ! GetStdHandle(STD_OUTPUT_HANDLE)) && GetProcessWindowStation()) popup_message(0, MB_OK, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);
  139. else print_message(stderr, NSSM_MESSAGE_USAGE, NSSM_VERSION, NSSM_CONFIGURATION, NSSM_DATE);
  140. return(ret);
  141. }
  142. void check_admin() {
  143. is_admin = false;
  144. /* Lifted from MSDN examples */
  145. PSID AdministratorsGroup;
  146. SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
  147. if (! AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &AdministratorsGroup)) return;
  148. CheckTokenMembership(0, AdministratorsGroup, /*XXX*/(PBOOL) &is_admin);
  149. FreeSid(AdministratorsGroup);
  150. }
  151. static int elevate(int argc, TCHAR **argv, unsigned long message) {
  152. print_message(stderr, message);
  153. SHELLEXECUTEINFO sei;
  154. ZeroMemory(&sei, sizeof(sei));
  155. sei.cbSize = sizeof(sei);
  156. sei.lpVerb = _T("runas");
  157. sei.lpFile = (TCHAR *) nssm_imagepath();
  158. TCHAR *args = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, EXE_LENGTH * sizeof(TCHAR));
  159. if (! args) {
  160. print_message(stderr, NSSM_MESSAGE_OUT_OF_MEMORY, _T("GetCommandLine()"), _T("elevate()"));
  161. return 111;
  162. }
  163. /* Get command line, which includes the path to NSSM, and skip that part. */
  164. _sntprintf_s(args, EXE_LENGTH, _TRUNCATE, _T("%s"), GetCommandLine());
  165. size_t s = _tcslen(argv[0]) + 1;
  166. if (args[0] == _T('"')) s += 2;
  167. while (isspace(args[s])) s++;
  168. sei.lpParameters = args + s;
  169. sei.nShow = SW_SHOW;
  170. unsigned long exitcode = 0;
  171. if (! ShellExecuteEx(&sei)) exitcode = 100;
  172. HeapFree(GetProcessHeap(), 0, (void *) args);
  173. return exitcode;
  174. }
  175. int num_cpus() {
  176. DWORD_PTR i, affinity, system_affinity;
  177. if (! GetProcessAffinityMask(GetCurrentProcess(), &affinity, &system_affinity)) return 64;
  178. for (i = 0; system_affinity & (1LL << i); i++);
  179. return (int) i;
  180. }
  181. const TCHAR *nssm_unquoted_imagepath() {
  182. return unquoted_imagepath;
  183. }
  184. const TCHAR *nssm_imagepath() {
  185. return imagepath;
  186. }
  187. const TCHAR *nssm_exe() {
  188. return imageargv0;
  189. }
  190. int _tmain(int argc, TCHAR **argv) {
  191. if (check_console()) setup_utf8();
  192. /* Remember if we are admin */
  193. check_admin();
  194. /* Set up function pointers. */
  195. if (get_imports()) nssm_exit(111);
  196. /* Remember our path for later. */
  197. _sntprintf_s(imageargv0, _countof(imageargv0), _TRUNCATE, _T("%s"), argv[0]);
  198. PathQuoteSpaces(imageargv0);
  199. GetModuleFileName(0, unquoted_imagepath, _countof(unquoted_imagepath));
  200. GetModuleFileName(0, imagepath, _countof(imagepath));
  201. PathQuoteSpaces(imagepath);
  202. /* Elevate */
  203. if (argc > 1) {
  204. /*
  205. Valid commands are:
  206. start, stop, pause, continue, install, edit, get, set, reset, unset, remove
  207. */
  208. if (str_equiv(argv[1], _T("start"))) nssm_exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));
  209. if (str_equiv(argv[1], _T("stop"))) nssm_exit(control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2));
  210. if (str_equiv(argv[1], _T("restart"))) {
  211. int ret = control_service(SERVICE_CONTROL_STOP, argc - 2, argv + 2);
  212. if (ret) nssm_exit(ret);
  213. nssm_exit(control_service(NSSM_SERVICE_CONTROL_START, argc - 2, argv + 2));
  214. }
  215. if (str_equiv(argv[1], _T("pause"))) nssm_exit(control_service(SERVICE_CONTROL_PAUSE, argc - 2, argv + 2));
  216. if (str_equiv(argv[1], _T("continue"))) nssm_exit(control_service(SERVICE_CONTROL_CONTINUE, argc - 2, argv + 2));
  217. if (str_equiv(argv[1], _T("status"))) nssm_exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2));
  218. if (str_equiv(argv[1], _T("statuscode"))) nssm_exit(control_service(SERVICE_CONTROL_INTERROGATE, argc - 2, argv + 2, true));
  219. if (str_equiv(argv[1], _T("rotate"))) nssm_exit(control_service(NSSM_SERVICE_CONTROL_ROTATE, argc - 2, argv + 2));
  220. if (str_equiv(argv[1], _T("install"))) {
  221. if (! is_admin) nssm_exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_INSTALL));
  222. create_messages();
  223. nssm_exit(pre_install_service(argc - 2, argv + 2));
  224. }
  225. 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")) || str_equiv(argv[1], _T("dump"))) {
  226. int ret = pre_edit_service(argc - 1, argv + 1);
  227. if (ret == 3 && ! is_admin && argc == 3) nssm_exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_EDIT));
  228. /* There might be a password here. */
  229. for (int i = 0; i < argc; i++) SecureZeroMemory(argv[i], _tcslen(argv[i]) * sizeof(TCHAR));
  230. nssm_exit(ret);
  231. }
  232. if (str_equiv(argv[1], _T("list"))) nssm_exit(list_nssm_services(argc - 2, argv + 2));
  233. if (str_equiv(argv[1], _T("processes"))) nssm_exit(service_process_tree(argc - 2, argv + 2));
  234. if (str_equiv(argv[1], _T("remove"))) {
  235. if (! is_admin) nssm_exit(elevate(argc, argv, NSSM_MESSAGE_NOT_ADMINISTRATOR_CANNOT_REMOVE));
  236. nssm_exit(pre_remove_service(argc - 2, argv + 2));
  237. }
  238. }
  239. /* Thread local storage for error message buffer */
  240. tls_index = TlsAlloc();
  241. /* Register messages */
  242. if (is_admin) create_messages();
  243. /*
  244. Optimisation for Windows 2000:
  245. When we're run from the command line the StartServiceCtrlDispatcher() call
  246. will time out after a few seconds on Windows 2000. On newer versions the
  247. call returns instantly. Check for stdin first and only try to call the
  248. function if there's no input stream found. Although it's possible that
  249. we're running with input redirected it's much more likely that we're
  250. actually running as a service.
  251. This will save time when running with no arguments from a command prompt.
  252. */
  253. if (! GetStdHandle(STD_INPUT_HANDLE)) {
  254. /* Start service magic */
  255. SERVICE_TABLE_ENTRY table[] = { { NSSM, service_main }, { 0, 0 } };
  256. if (! StartServiceCtrlDispatcher(table)) {
  257. unsigned long error = GetLastError();
  258. /* User probably ran nssm with no argument */
  259. if (error == ERROR_FAILED_SERVICE_CONTROLLER_CONNECT) nssm_exit(usage(1));
  260. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_DISPATCHER_FAILED, error_string(error), 0);
  261. nssm_exit(100);
  262. }
  263. }
  264. else nssm_exit(usage(1));
  265. /* And nothing more to do */
  266. nssm_exit(0);
  267. }