1/*-
2 * Copyright (c) 2014 Baptiste Daroussin <bapt@FreeBSD.org>
3 * Copyright (c) 2014 Vsevolod Stakhov <vsevolod@FreeBSD.org>
4 * All rights reserved.
5 *
6 * Redistribution and use in source and binary forms, with or without
7 * modification, are permitted provided that the following conditions
8 * are met:
9 * 1. Redistributions of source code must retain the above copyright
10 *    notice, this list of conditions and the following disclaimer
11 *    in this position and unchanged.
12 * 2. Redistributions in binary form must reproduce the above copyright
13 *    notice, this list of conditions and the following disclaimer in the
14 *    documentation and/or other materials provided with the distribution.
15 *
16 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR(S) ``AS IS'' AND ANY EXPRESS OR
17 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
18 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
19 * IN NO EVENT SHALL THE AUTHOR(S) BE LIABLE FOR ANY DIRECT, INDIRECT,
20 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
21 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
22 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
23 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
25 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26 */
27
28#include <sys/cdefs.h>
29#include <sys/procctl.h>
30#include <sys/time.h>
31#include <sys/wait.h>
32
33#include <err.h>
34#include <errno.h>
35#include <getopt.h>
36#include <signal.h>
37#include <stdbool.h>
38#include <stdio.h>
39#include <stdlib.h>
40#include <string.h>
41#include <unistd.h>
42
43#define EXIT_TIMEOUT 124
44
45static sig_atomic_t sig_chld = 0;
46static sig_atomic_t sig_term = 0;
47static sig_atomic_t sig_alrm = 0;
48static sig_atomic_t sig_ign = 0;
49
50static void
51usage(void)
52{
53
54	fprintf(stderr, "Usage: %s [--signal sig | -s sig] [--preserve-status]"
55	    " [--kill-after time | -k time] [--foreground] <duration> <command>"
56	    " <arg ...>\n", getprogname());
57
58	exit(EXIT_FAILURE);
59}
60
61static double
62parse_duration(const char *duration)
63{
64	double ret;
65	char *end;
66
67	ret = strtod(duration, &end);
68	if (ret == 0 && end == duration)
69		errx(125, "invalid duration");
70
71	if (end == NULL || *end == '\0')
72		return (ret);
73
74	if (end != NULL && *(end + 1) != '\0')
75		errx(125, "invalid duration");
76
77	switch (*end) {
78	case 's':
79		break;
80	case 'm':
81		ret *= 60;
82		break;
83	case 'h':
84		ret *= 60 * 60;
85		break;
86	case 'd':
87		ret *= 60 * 60 * 24;
88		break;
89	default:
90		errx(125, "invalid duration");
91	}
92
93	if (ret < 0 || ret >= 100000000UL)
94		errx(125, "invalid duration");
95
96	return (ret);
97}
98
99static int
100parse_signal(const char *str)
101{
102	int sig, i;
103	const char *errstr;
104
105	sig = strtonum(str, 1, sys_nsig - 1, &errstr);
106
107	if (errstr == NULL)
108		return (sig);
109	if (strncasecmp(str, "SIG", 3) == 0)
110		str += 3;
111
112	for (i = 1; i < sys_nsig; i++) {
113		if (strcasecmp(str, sys_signame[i]) == 0)
114			return (i);
115	}
116
117	errx(125, "invalid signal");
118}
119
120static void
121sig_handler(int signo)
122{
123	if (sig_ign != 0 && signo == sig_ign) {
124		sig_ign = 0;
125		return;
126	}
127
128	switch (signo) {
129	case 0:
130	case SIGINT:
131	case SIGHUP:
132	case SIGQUIT:
133	case SIGTERM:
134		sig_term = signo;
135		break;
136	case SIGCHLD:
137		sig_chld = 1;
138		break;
139	case SIGALRM:
140		sig_alrm = 1;
141		break;
142	}
143}
144
145static void
146set_interval(double iv)
147{
148	struct itimerval tim;
149
150	memset(&tim, 0, sizeof(tim));
151	tim.it_value.tv_sec = (time_t)iv;
152	iv -= (time_t)iv;
153	tim.it_value.tv_usec = (suseconds_t)(iv * 1000000UL);
154
155	if (setitimer(ITIMER_REAL, &tim, NULL) == -1)
156		err(EXIT_FAILURE, "setitimer()");
157}
158
159int
160main(int argc, char **argv)
161{
162	int ch;
163	unsigned long i;
164	int foreground, preserve;
165	int error, pstat, status;
166	int killsig = SIGTERM;
167	pid_t pid, cpid;
168	double first_kill;
169	double second_kill;
170	bool timedout = false;
171	bool do_second_kill = false;
172	bool child_done = false;
173	struct sigaction signals;
174	struct procctl_reaper_status info;
175	struct procctl_reaper_kill killemall;
176	int signums[] = {
177		-1,
178		SIGTERM,
179		SIGINT,
180		SIGHUP,
181		SIGCHLD,
182		SIGALRM,
183		SIGQUIT,
184	};
185
186	foreground = preserve = 0;
187	second_kill = 0;
188
189	const struct option longopts[] = {
190		{ "preserve-status", no_argument,       &preserve,    1 },
191		{ "foreground",      no_argument,       &foreground,  1 },
192		{ "kill-after",      required_argument, NULL,        'k'},
193		{ "signal",          required_argument, NULL,        's'},
194		{ "help",            no_argument,       NULL,        'h'},
195		{ NULL,              0,                 NULL,         0 }
196	};
197
198	while ((ch = getopt_long(argc, argv, "+k:s:h", longopts, NULL)) != -1) {
199		switch (ch) {
200			case 'k':
201				do_second_kill = true;
202				second_kill = parse_duration(optarg);
203				break;
204			case 's':
205				killsig = parse_signal(optarg);
206				break;
207			case 0:
208				break;
209			case 'h':
210			default:
211				usage();
212		}
213	}
214
215	argc -= optind;
216	argv += optind;
217
218	if (argc < 2)
219		usage();
220
221	first_kill = parse_duration(argv[0]);
222	argc--;
223	argv++;
224
225	if (!foreground) {
226		/* Acquire a reaper */
227		if (procctl(P_PID, getpid(), PROC_REAP_ACQUIRE, NULL) == -1)
228			err(EXIT_FAILURE, "Fail to acquire the reaper");
229	}
230
231	memset(&signals, 0, sizeof(signals));
232	sigemptyset(&signals.sa_mask);
233
234	if (killsig != SIGKILL && killsig != SIGSTOP)
235		signums[0] = killsig;
236
237	for (i = 0; i < sizeof(signums) / sizeof(signums[0]); i++)
238		sigaddset(&signals.sa_mask, signums[i]);
239
240	signals.sa_handler = sig_handler;
241	signals.sa_flags = SA_RESTART;
242
243	for (i = 0; i < sizeof(signums) / sizeof(signums[0]); i++)
244		if (signums[i] != -1 && signums[i] != 0 &&
245		    sigaction(signums[i], &signals, NULL) == -1)
246			err(EXIT_FAILURE, "sigaction()");
247
248	signal(SIGTTIN, SIG_IGN);
249	signal(SIGTTOU, SIG_IGN);
250
251	pid = fork();
252	if (pid == -1)
253		err(EXIT_FAILURE, "fork()");
254	else if (pid == 0) {
255		/* child process */
256		signal(SIGTTIN, SIG_DFL);
257		signal(SIGTTOU, SIG_DFL);
258
259		error = execvp(argv[0], argv);
260		if (error == -1) {
261			if (errno == ENOENT)
262				err(127, "exec(%s)", argv[0]);
263			else
264				err(126, "exec(%s)", argv[0]);
265		}
266	}
267
268	if (sigprocmask(SIG_BLOCK, &signals.sa_mask, NULL) == -1)
269		err(EXIT_FAILURE, "sigprocmask()");
270
271	/* parent continues here */
272	set_interval(first_kill);
273
274	for (;;) {
275		sigemptyset(&signals.sa_mask);
276		sigsuspend(&signals.sa_mask);
277
278		if (sig_chld) {
279			sig_chld = 0;
280
281			while ((cpid = waitpid(-1, &status, WNOHANG)) != 0) {
282				if (cpid < 0) {
283					if (errno == EINTR)
284						continue;
285					else
286						break;
287				} else if (cpid == pid) {
288					pstat = status;
289					child_done = true;
290				}
291			}
292			if (child_done) {
293				if (foreground) {
294					break;
295				} else {
296					procctl(P_PID, getpid(),
297					    PROC_REAP_STATUS, &info);
298					if (info.rs_children == 0)
299						break;
300				}
301			}
302		} else if (sig_alrm) {
303			sig_alrm = 0;
304
305			timedout = true;
306			if (!foreground) {
307				killemall.rk_sig = killsig;
308				killemall.rk_flags = 0;
309				procctl(P_PID, getpid(), PROC_REAP_KILL,
310				    &killemall);
311			} else
312				kill(pid, killsig);
313
314			if (do_second_kill) {
315				set_interval(second_kill);
316				second_kill = 0;
317				sig_ign = killsig;
318				killsig = SIGKILL;
319			} else
320				break;
321
322		} else if (sig_term) {
323			if (!foreground) {
324				killemall.rk_sig = sig_term;
325				killemall.rk_flags = 0;
326				procctl(P_PID, getpid(), PROC_REAP_KILL,
327				    &killemall);
328			} else
329				kill(pid, sig_term);
330
331			if (do_second_kill) {
332				set_interval(second_kill);
333				second_kill = 0;
334				sig_ign = killsig;
335				killsig = SIGKILL;
336			} else
337				break;
338		}
339	}
340
341	while (!child_done && wait(&pstat) == -1) {
342		if (errno != EINTR)
343			err(EXIT_FAILURE, "waitpid()");
344	}
345
346	if (!foreground)
347		procctl(P_PID, getpid(), PROC_REAP_RELEASE, NULL);
348
349	if (WEXITSTATUS(pstat))
350		pstat = WEXITSTATUS(pstat);
351	else if (WIFSIGNALED(pstat))
352		pstat = 128 + WTERMSIG(pstat);
353
354	if (timedout && !preserve)
355		pstat = EXIT_TIMEOUT;
356
357	return (pstat);
358}
359