xref: /freebsd/bin/timeout/timeout.c (revision d37eb51047221dc3322b34db1038ff3aa533883f)
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 
45 static sig_atomic_t sig_chld = 0;
46 static sig_atomic_t sig_term = 0;
47 static sig_atomic_t sig_alrm = 0;
48 static sig_atomic_t sig_ign = 0;
49 
50 static void
51 usage(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 
61 static double
62 parse_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 
99 static int
100 parse_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 
120 static void
121 sig_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 
145 static void
146 set_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 
159 int
160 main(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