xref: /freebsd/bin/timeout/timeout.c (revision b077aed33b7b6aefca7b17ddb250cf521f938613)
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 __FBSDID("$FreeBSD$");
30 
31 #include <sys/procctl.h>
32 #include <sys/time.h>
33 #include <sys/wait.h>
34 
35 #include <err.h>
36 #include <errno.h>
37 #include <getopt.h>
38 #include <signal.h>
39 #include <stdbool.h>
40 #include <stdio.h>
41 #include <stdlib.h>
42 #include <string.h>
43 #include <unistd.h>
44 
45 #define EXIT_TIMEOUT 124
46 
47 static sig_atomic_t sig_chld = 0;
48 static sig_atomic_t sig_term = 0;
49 static sig_atomic_t sig_alrm = 0;
50 static sig_atomic_t sig_ign = 0;
51 
52 static void
53 usage(void)
54 {
55 
56 	fprintf(stderr, "Usage: %s [--signal sig | -s sig] [--preserve-status]"
57 	    " [--kill-after time | -k time] [--foreground] <duration> <command>"
58 	    " <arg ...>\n", getprogname());
59 
60 	exit(EXIT_FAILURE);
61 }
62 
63 static double
64 parse_duration(const char *duration)
65 {
66 	double ret;
67 	char *end;
68 
69 	ret = strtod(duration, &end);
70 	if (ret == 0 && end == duration)
71 		errx(125, "invalid duration");
72 
73 	if (end == NULL || *end == '\0')
74 		return (ret);
75 
76 	if (end != NULL && *(end + 1) != '\0')
77 		errx(125, "invalid duration");
78 
79 	switch (*end) {
80 	case 's':
81 		break;
82 	case 'm':
83 		ret *= 60;
84 		break;
85 	case 'h':
86 		ret *= 60 * 60;
87 		break;
88 	case 'd':
89 		ret *= 60 * 60 * 24;
90 		break;
91 	default:
92 		errx(125, "invalid duration");
93 	}
94 
95 	if (ret < 0 || ret >= 100000000UL)
96 		errx(125, "invalid duration");
97 
98 	return (ret);
99 }
100 
101 static int
102 parse_signal(const char *str)
103 {
104 	int sig, i;
105 	const char *errstr;
106 
107 	sig = strtonum(str, 1, sys_nsig - 1, &errstr);
108 
109 	if (errstr == NULL)
110 		return (sig);
111 	if (strncasecmp(str, "SIG", 3) == 0)
112 		str += 3;
113 
114 	for (i = 1; i < sys_nsig; i++) {
115 		if (strcasecmp(str, sys_signame[i]) == 0)
116 			return (i);
117 	}
118 
119 	errx(125, "invalid signal");
120 }
121 
122 static void
123 sig_handler(int signo)
124 {
125 	if (sig_ign != 0 && signo == sig_ign) {
126 		sig_ign = 0;
127 		return;
128 	}
129 
130 	switch (signo) {
131 	case 0:
132 	case SIGINT:
133 	case SIGHUP:
134 	case SIGQUIT:
135 	case SIGTERM:
136 		sig_term = signo;
137 		break;
138 	case SIGCHLD:
139 		sig_chld = 1;
140 		break;
141 	case SIGALRM:
142 		sig_alrm = 1;
143 		break;
144 	}
145 }
146 
147 static void
148 set_interval(double iv)
149 {
150 	struct itimerval tim;
151 
152 	memset(&tim, 0, sizeof(tim));
153 	tim.it_value.tv_sec = (time_t)iv;
154 	iv -= (time_t)iv;
155 	tim.it_value.tv_usec = (suseconds_t)(iv * 1000000UL);
156 
157 	if (setitimer(ITIMER_REAL, &tim, NULL) == -1)
158 		err(EXIT_FAILURE, "setitimer()");
159 }
160 
161 int
162 main(int argc, char **argv)
163 {
164 	int ch;
165 	unsigned long i;
166 	int foreground, preserve;
167 	int error, pstat, status;
168 	int killsig = SIGTERM;
169 	pid_t pid, cpid;
170 	double first_kill;
171 	double second_kill;
172 	bool timedout = false;
173 	bool do_second_kill = false;
174 	bool child_done = false;
175 	struct sigaction signals;
176 	struct procctl_reaper_status info;
177 	struct procctl_reaper_kill killemall;
178 	int signums[] = {
179 		-1,
180 		SIGTERM,
181 		SIGINT,
182 		SIGHUP,
183 		SIGCHLD,
184 		SIGALRM,
185 		SIGQUIT,
186 	};
187 
188 	foreground = preserve = 0;
189 	second_kill = 0;
190 
191 	const struct option longopts[] = {
192 		{ "preserve-status", no_argument,       &preserve,    1 },
193 		{ "foreground",      no_argument,       &foreground,  1 },
194 		{ "kill-after",      required_argument, NULL,        'k'},
195 		{ "signal",          required_argument, NULL,        's'},
196 		{ "help",            no_argument,       NULL,        'h'},
197 		{ NULL,              0,                 NULL,         0 }
198 	};
199 
200 	while ((ch = getopt_long(argc, argv, "+k:s:h", longopts, NULL)) != -1) {
201 		switch (ch) {
202 			case 'k':
203 				do_second_kill = true;
204 				second_kill = parse_duration(optarg);
205 				break;
206 			case 's':
207 				killsig = parse_signal(optarg);
208 				break;
209 			case 0:
210 				break;
211 			case 'h':
212 			default:
213 				usage();
214 				break;
215 		}
216 	}
217 
218 	argc -= optind;
219 	argv += optind;
220 
221 	if (argc < 2)
222 		usage();
223 
224 	first_kill = parse_duration(argv[0]);
225 	argc--;
226 	argv++;
227 
228 	if (!foreground) {
229 		/* Acquire a reaper */
230 		if (procctl(P_PID, getpid(), PROC_REAP_ACQUIRE, NULL) == -1)
231 			err(EXIT_FAILURE, "Fail to acquire the reaper");
232 	}
233 
234 	memset(&signals, 0, sizeof(signals));
235 	sigemptyset(&signals.sa_mask);
236 
237 	if (killsig != SIGKILL && killsig != SIGSTOP)
238 		signums[0] = killsig;
239 
240 	for (i = 0; i < sizeof(signums) / sizeof(signums[0]); i++)
241 		sigaddset(&signals.sa_mask, signums[i]);
242 
243 	signals.sa_handler = sig_handler;
244 	signals.sa_flags = SA_RESTART;
245 
246 	for (i = 0; i < sizeof(signums) / sizeof(signums[0]); i++)
247 		if (signums[i] != -1 && signums[i] != 0 &&
248 		    sigaction(signums[i], &signals, NULL) == -1)
249 			err(EXIT_FAILURE, "sigaction()");
250 
251 	signal(SIGTTIN, SIG_IGN);
252 	signal(SIGTTOU, SIG_IGN);
253 
254 	pid = fork();
255 	if (pid == -1)
256 		err(EXIT_FAILURE, "fork()");
257 	else if (pid == 0) {
258 		/* child process */
259 		signal(SIGTTIN, SIG_DFL);
260 		signal(SIGTTOU, SIG_DFL);
261 
262 		error = execvp(argv[0], argv);
263 		if (error == -1) {
264 			if (errno == ENOENT)
265 				err(127, "exec(%s)", argv[0]);
266 			else
267 				err(126, "exec(%s)", argv[0]);
268 		}
269 	}
270 
271 	if (sigprocmask(SIG_BLOCK, &signals.sa_mask, NULL) == -1)
272 		err(EXIT_FAILURE, "sigprocmask()");
273 
274 	/* parent continues here */
275 	set_interval(first_kill);
276 
277 	for (;;) {
278 		sigemptyset(&signals.sa_mask);
279 		sigsuspend(&signals.sa_mask);
280 
281 		if (sig_chld) {
282 			sig_chld = 0;
283 
284 			while ((cpid = waitpid(-1, &status, WNOHANG)) != 0) {
285 				if (cpid < 0) {
286 					if (errno == EINTR)
287 						continue;
288 					else
289 						break;
290 				} else if (cpid == pid) {
291 					pstat = status;
292 					child_done = true;
293 				}
294 			}
295 			if (child_done) {
296 				if (foreground) {
297 					break;
298 				} else {
299 					procctl(P_PID, getpid(),
300 					    PROC_REAP_STATUS, &info);
301 					if (info.rs_children == 0)
302 						break;
303 				}
304 			}
305 		} else if (sig_alrm) {
306 			sig_alrm = 0;
307 
308 			timedout = true;
309 			if (!foreground) {
310 				killemall.rk_sig = killsig;
311 				killemall.rk_flags = 0;
312 				procctl(P_PID, getpid(), PROC_REAP_KILL,
313 				    &killemall);
314 			} else
315 				kill(pid, killsig);
316 
317 			if (do_second_kill) {
318 				set_interval(second_kill);
319 				second_kill = 0;
320 				sig_ign = killsig;
321 				killsig = SIGKILL;
322 			} else
323 				break;
324 
325 		} else if (sig_term) {
326 			if (!foreground) {
327 				killemall.rk_sig = sig_term;
328 				killemall.rk_flags = 0;
329 				procctl(P_PID, getpid(), PROC_REAP_KILL,
330 				    &killemall);
331 			} else
332 				kill(pid, sig_term);
333 
334 			if (do_second_kill) {
335 				set_interval(second_kill);
336 				second_kill = 0;
337 				sig_ign = killsig;
338 				killsig = SIGKILL;
339 			} else
340 				break;
341 		}
342 	}
343 
344 	while (!child_done && wait(&pstat) == -1) {
345 		if (errno != EINTR)
346 			err(EXIT_FAILURE, "waitpid()");
347 	}
348 
349 	if (!foreground)
350 		procctl(P_PID, getpid(), PROC_REAP_RELEASE, NULL);
351 
352 	if (WEXITSTATUS(pstat))
353 		pstat = WEXITSTATUS(pstat);
354 	else if (WIFSIGNALED(pstat))
355 		pstat = 128 + WTERMSIG(pstat);
356 
357 	if (timedout && !preserve)
358 		pstat = EXIT_TIMEOUT;
359 
360 	return (pstat);
361 }
362