Wait for an arbitrary process.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

311 lines
7.2 KiB

  1. /*
  2. * wait4pid - Wait for an arbitrary process
  3. * Copyright (C) 2013 Damien Goutte-Gattat
  4. *
  5. * This program is free software: you can redistribute it and/or modify
  6. * it under the terms of the GNU General Public License as published by
  7. * the Free Software Foundation, either version 3 of the License, or
  8. * (at your option) any later version.
  9. *
  10. * This program is distributed in the hope that it will be useful,
  11. * but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. * GNU General Public License for more details.
  14. *
  15. * You should have received a copy of the GNU General Public License
  16. * along with this program. If not, see <http://www.gnu.org/licenses/>.
  17. */
  18. #ifdef HAVE_CONFIG_H
  19. #include <config.h>
  20. #endif
  21. #if defined HAVE_LINUX_NETLINK_H /* Netlink */
  22. #include <string.h>
  23. #include <unistd.h>
  24. #include <signal.h>
  25. #include <sys/socket.h>
  26. #include <sys/wait.h> /* for WEXITSTATUS */
  27. #include <linux/netlink.h>
  28. #include <linux/connector.h>
  29. #include <linux/cn_proc.h>
  30. /* Netlink message to (un)register a proc events listener. */
  31. struct __attribute__ ((__packed__)) cn_proc_op {
  32. struct cn_msg msg;
  33. enum proc_cn_mcast_op op;
  34. };
  35. /* Netlink message sent on a proc event. */
  36. struct __attribute__ ((__packed__)) cn_proc_msg {
  37. struct cn_msg msg;
  38. struct proc_event evt;
  39. };
  40. int sock = 0;
  41. int listener_set = 0;
  42. static ssize_t
  43. do_send(int sock, const void *buf, size_t len, int flags)
  44. {
  45. int ret;
  46. do {
  47. ret = send(sock, buf, len, flags);
  48. } while ( ret == -1 && errno == EINTR );
  49. return ret;
  50. }
  51. static int
  52. do_recv(int sock, void *buf, size_t len, int flags)
  53. {
  54. int ret;
  55. do {
  56. ret = recv(sock, buf, len, flags);
  57. } while ( ret == -1 && errno == EINTR );
  58. return ret;
  59. }
  60. static int
  61. set_proc_listen(int sock, int enable)
  62. {
  63. struct nlmsghdr *hdr;
  64. struct cn_proc_op *payload;
  65. char buf[NLMSG_SPACE(sizeof(struct cn_proc_op))];
  66. memset(buf, 0, sizeof(buf));
  67. hdr = (struct nlmsghdr *)buf;
  68. hdr->nlmsg_type = NLMSG_DONE;
  69. hdr->nlmsg_flags = NLM_F_REQUEST;
  70. hdr->nlmsg_len = NLMSG_LENGTH(sizeof(*payload));
  71. hdr->nlmsg_pid = getpid();
  72. payload = (struct cn_proc_op *) NLMSG_DATA(hdr);
  73. payload->msg.id.idx = CN_IDX_PROC;
  74. payload->msg.id.val = CN_VAL_PROC;
  75. payload->msg.len = sizeof(payload->op);
  76. payload->op = enable ? PROC_CN_MCAST_LISTEN : PROC_CN_MCAST_IGNORE;
  77. return do_send(sock, hdr, NLMSG_SPACE(sizeof(*payload)), 0);
  78. }
  79. static void
  80. wait4pid_close(void)
  81. {
  82. if ( listener_set ) {
  83. set_proc_listen(sock, 0);
  84. listener_set = 0;
  85. }
  86. if ( sock ) {
  87. close(sock);
  88. sock = 0;
  89. }
  90. }
  91. static int
  92. wait4pid_init(void)
  93. {
  94. struct sockaddr_nl sa;
  95. if ( (sock = socket(PF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR)) == -1 )
  96. return -1;
  97. atexit(wait4pid_close);
  98. memset(&sa, 0, sizeof(sa));
  99. sa.nl_family = AF_NETLINK;
  100. sa.nl_groups = CN_IDX_PROC;
  101. sa.nl_pid = getpid();
  102. if ( bind(sock, (struct sockaddr *)&sa, sizeof(sa)) == -1 )
  103. return -1;
  104. if ( setgid(getgid()) == -1 || setuid(getuid()) == -1 )
  105. return -1;
  106. if ( set_proc_listen(sock, 1) == -1 )
  107. return -1;
  108. return 0;
  109. }
  110. #ifdef HAVE_LINUX_FILTER_H
  111. #include <linux/filter.h>
  112. #include <arpa/inet.h>
  113. #include <stddef.h>
  114. static int
  115. set_filter(int sock, pid_t *pids, size_t len)
  116. {
  117. int i;
  118. struct sock_fprog flt;
  119. struct sock_filter filter[256] = {
  120. /* Load event type. */
  121. BPF_STMT(BPF_LD+BPF_W+BPF_ABS,
  122. NLMSG_LENGTH(0) + offsetof(struct cn_proc_msg, evt.what)),
  123. /* If this is a PROC_EVENT_EXIT, continue to the next step;
  124. * otherwise go to the "block" step at the end of the filter. */
  125. BPF_JUMP(BPF_JMP+BPF_JEQ+BPF_K, htonl(PROC_EVENT_EXIT), 0, 1 + len),
  126. /* Load the process identifier. */
  127. BPF_STMT(BPF_LD+BPF_W+BPF_ABS,
  128. NLMSG_LENGTH(0) +
  129. offsetof(struct cn_proc_msg, evt.event_data.exit.process_pid))
  130. };
  131. if ( len > 251 ) {
  132. errno = EINVAL;
  133. return -1;
  134. }
  135. flt.len = 5 + len;
  136. flt.filter = filter;
  137. for ( i = 0; i < len; i++ ) {
  138. /* If the loaded PID matches pids[i], go to the "accept" step
  139. * at the end of the filter; otherwise continue to the next step. */
  140. filter[i+3].code = BPF_JMP+BPF_JEQ+BPF_K;
  141. filter[i+3].k = htonl(pids[i]);
  142. filter[i+3].jt = len - i;
  143. filter[i+3].jf = 0;
  144. }
  145. /* No match, block the packet. */
  146. filter[len+3].code = BPF_RET+BPF_K;
  147. filter[len+3].k = 0;
  148. /* Accept the packet. */
  149. filter[len+4].code = BPF_RET+BPF_K;
  150. filter[len+4].k = (uint)-1;
  151. return setsockopt(sock, SOL_SOCKET, SO_ATTACH_FILTER, &flt, sizeof(flt));
  152. }
  153. #endif
  154. int
  155. wait4all(pid_t *pids, size_t len)
  156. {
  157. int ec, nproc, i;
  158. struct cn_proc_msg *payload;
  159. char buf[NLMSG_SPACE(sizeof(struct cn_proc_msg))];
  160. if ( ! listener_set && wait4pid_init() == -1 )
  161. return -1;
  162. for ( i = 0; i < len; i++ ) {
  163. if ( kill(pids[i], 0) == -1 && errno == ESRCH )
  164. return -1;
  165. }
  166. #ifdef HAVE_LINUX_FILTER_H
  167. if ( set_filter(sock, pids, len) == -1 )
  168. return -1;
  169. #endif
  170. nproc = len;
  171. while ( nproc > 0 ) {
  172. if ( do_recv(sock, buf, sizeof(buf), 0) <= 0 )
  173. return -1;
  174. payload = (struct cn_proc_msg *) NLMSG_DATA(buf);
  175. if ( payload->evt.what != PROC_EVENT_EXIT )
  176. continue;
  177. #ifdef HAVE_LINUX_FILTER_H
  178. nproc -= 1;
  179. if ( payload->evt.event_data.exit.process_pid == pids[len-1] )
  180. ec = WEXITSTATUS(payload->evt.event_data.exit.exit_code);
  181. #else
  182. for ( i = 0; i < len; i++ ) {
  183. if ( pids[i] != -1
  184. && pids[i] == payload->evt.event_data.exit.process_pid ) {
  185. pids[i] = -1;
  186. nproc -= 1;
  187. if ( i == len - 1 )
  188. ec = WEXITSTATUS(payload->evt.event_data.exit.exit_code);
  189. }
  190. }
  191. #endif
  192. }
  193. return ec;
  194. }
  195. #elif defined HAVE_SYS_EVENT_H /* BSD Kevent */
  196. #include <sys/event.h>
  197. #include <xmem.h>
  198. int
  199. wait4all(pid_t *pids, size_t len)
  200. {
  201. int fd, i, nproc, ec;
  202. struct kevent *evts, evt;
  203. if ( (fd = kqueue()) == -1 )
  204. return -1;
  205. evts = xmalloc(sizeof(struct kevent) * len);
  206. for ( i = 0; i < len; i++ ) {
  207. EV_SET(evts[i], pids[i], EVFILT_PROC, EV_ENABLE, NOTE_EXIT, 0, 0):
  208. }
  209. i = kevent(fd, evts, len, NULL, 0, 0);
  210. free(evts);
  211. if ( i == -1 )
  212. return -1;
  213. nproc = len;
  214. ec = 0;
  215. while ( nproc > 0 ) {
  216. if ( kevent(fd, NULL, 0, &evt, 1, 0) <= 0 )
  217. return -1;
  218. if ( evt->ident == pids[len - 1] )
  219. ec = evt->data;
  220. nproc -= 1;
  221. }
  222. return ec;
  223. }
  224. #else /* Fallback */
  225. #include <time.h>
  226. #include <signal.h>
  227. int
  228. wait4all(pid_t *pids, size_t len)
  229. {
  230. int nproc, i;
  231. struct timespec ts = { 1, 0 };
  232. nproc = len;
  233. while ( nproc > 0 ) {
  234. for ( i = 0; i < len; i++ ) {
  235. if ( pids[i] == -1 )
  236. continue;
  237. if ( kill(pids[i], 0) == -1 && errno == ESRCH ) {
  238. pids[i] = -1;
  239. nproc -= 1;
  240. }
  241. }
  242. nanosleep(&ts, NULL);
  243. }
  244. return 0;
  245. }
  246. #endif