1/*	$OpenBSD: fork-exit.c,v 1.7 2021/07/06 11:50:34 bluhm Exp $	*/
2
3/*
4 * Copyright (c) 2021 Alexander Bluhm <bluhm@openbsd.org>
5 *
6 * Permission to use, copy, modify, and distribute this software for any
7 * purpose with or without fee is hereby granted, provided that the above
8 * copyright notice and this permission notice appear in all copies.
9 *
10 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
11 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
12 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
13 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
14 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
15 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
16 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
17 */
18
19#include <sys/mman.h>
20#include <sys/select.h>
21#include <sys/wait.h>
22
23#include <err.h>
24#include <errno.h>
25#include <fcntl.h>
26#include <limits.h>
27#include <pthread.h>
28#include <signal.h>
29#include <stdio.h>
30#include <stdlib.h>
31#include <string.h>
32#include <unistd.h>
33
34int execute = 0;
35int daemonize = 0;
36int heap = 0;
37int procs = 1;
38int stack = 0;
39int threads = 0;
40int timeout = 30;
41
42pthread_barrier_t thread_barrier;
43char timeoutstr[sizeof("-2147483647")];
44
45static void __dead
46usage(void)
47{
48	fprintf(stderr, "fork-exit [-ed] [-p procs] [-t threads] [-T timeout]\n"
49	    "    -e          child execs sleep(1), default call sleep(3)\n"
50	    "    -d          daemonize, use if already process group leader\n"
51	    "    -h heap     allocate pages of heap memory, default 0\n"
52	    "    -p procs    number of processes to fork, default 1\n"
53	    "    -s stack    allocate pages of stack memory, default 0\n"
54	    "    -t threads  number of threads to create, default 0\n"
55	    "    -T timeout  parent and children will exit, default 30 sec\n");
56	exit(2);
57}
58
59static void
60recurse_page(int depth)
61{
62	int p[4096 / sizeof(int)];
63
64	if (depth == 0)
65		return;
66	p[1] = 0x9abcdef0;
67	explicit_bzero(p, sizeof(int));
68	recurse_page(depth - 1);
69}
70
71static void
72alloc_stack(void)
73{
74	recurse_page(stack);
75}
76
77static void
78alloc_heap(void)
79{
80	int *p;
81	int i;
82
83	for(i = 0; i < heap; i++) {
84		p = mmap(0, 4096, PROT_WRITE, MAP_SHARED|MAP_ANON, -1, 0);
85		if (p == MAP_FAILED)
86			err(1, "mmap");
87		p[1] = 0x12345678;
88		explicit_bzero(p, sizeof(int));
89	}
90}
91
92static void * __dead
93run_thread(void *arg)
94{
95	int error;
96
97	if (heap)
98		alloc_heap();
99	if (stack)
100		alloc_stack();
101
102	error = pthread_barrier_wait(&thread_barrier);
103	if (error && error != PTHREAD_BARRIER_SERIAL_THREAD)
104		errc(1, error, "pthread_barrier_wait");
105
106	if (sleep(timeout) != 0)
107		err(1, "sleep %d", timeout);
108
109	/* should not happen */
110	_exit(0);
111}
112
113static void
114create_threads(void)
115{
116	pthread_attr_t tattr;
117	pthread_t *thrs;
118	int i, error;
119
120	error = pthread_attr_init(&tattr);
121	if (error)
122		errc(1, error, "pthread_attr_init");
123	if (stack) {
124		/* thread start and function call overhead needs a bit more */
125		error = pthread_attr_setstacksize(&tattr,
126		    (stack + 2) * (4096 + 64));
127		if (error)
128			errc(1, error, "pthread_attr_setstacksize");
129	}
130
131	error = pthread_barrier_init(&thread_barrier, NULL, threads + 1);
132	if (error)
133		errc(1, error, "pthread_barrier_init");
134
135	thrs = reallocarray(NULL, threads, sizeof(pthread_t));
136	if (thrs == NULL)
137		err(1, "thrs");
138
139	for (i = 0; i < threads; i++) {
140		error = pthread_create(&thrs[i], &tattr, run_thread, NULL);
141		if (error)
142			errc(1, error, "pthread_create");
143	}
144
145	error = pthread_barrier_wait(&thread_barrier);
146	if (error && error != PTHREAD_BARRIER_SERIAL_THREAD)
147		errc(1, error, "pthread_barrier_wait");
148
149	/* return to close child's pipe and sleep */
150}
151
152static void __dead
153exec_sleep(void)
154{
155	execl("/bin/sleep", "sleep", timeoutstr, NULL);
156	err(1, "exec sleep");
157}
158
159static void __dead
160run_child(int fd)
161{
162	/* close pipe to parent and sleep until killed */
163	if (execute) {
164		if (fcntl(fd, F_SETFD, FD_CLOEXEC))
165			err(1, "fcntl FD_CLOEXEC");
166		exec_sleep();
167	} else {
168		if (threads) {
169			create_threads();
170		} else {
171			if (heap)
172				alloc_heap();
173			if (stack)
174				alloc_stack();
175		}
176		if (close(fd) == -1)
177			err(1, "close child");
178		if (sleep(timeout) != 0)
179			err(1, "sleep %d", timeout);
180	}
181
182	/* should not happen */
183	_exit(0);
184}
185
186static void
187sigexit(int sig)
188{
189	int i, status;
190	pid_t pid;
191
192	/* all children must terminate in time */
193	alarm(timeout);
194
195	for (i = 0; i < procs; i++) {
196		pid = wait(&status);
197		if (pid == -1)
198			err(1, "wait");
199		if (!WIFSIGNALED(status))
200			errx(1, "child %d not killed", pid);
201		if(WTERMSIG(status) != SIGTERM)
202			errx(1, "child %d signal %d", pid, WTERMSIG(status));
203	}
204	exit(0);
205}
206
207int
208main(int argc, char *argv[])
209{
210	const char *errstr;
211	int ch, i, fdmax, fdlen, *rfds, waiting;
212	fd_set *fdset;
213	pid_t pgrp;
214	struct timeval tv;
215
216	while ((ch = getopt(argc, argv, "edh:p:s:T:t:")) != -1) {
217	switch (ch) {
218		case 'e':
219			execute = 1;
220			break;
221		case 'd':
222			daemonize = 1;
223			break;
224		case 'h':
225			heap = strtonum(optarg, 0, INT_MAX, &errstr);
226			if (errstr != NULL)
227				errx(1, "number of heap allocations is %s: %s",
228				    errstr, optarg);
229			break;
230		case 'p':
231			procs = strtonum(optarg, 0, INT_MAX / 4096, &errstr);
232			if (errstr != NULL)
233				errx(1, "number of procs is %s: %s", errstr,
234				    optarg);
235			break;
236		case 's':
237			stack = strtonum(optarg, 0,
238			    (INT_MAX / (4096 + 64)) - 2, &errstr);
239			if (errstr != NULL)
240				errx(1, "number of stack allocations is %s: %s",
241				    errstr, optarg);
242			break;
243		case 't':
244			threads = strtonum(optarg, 0, INT_MAX, &errstr);
245			if (errstr != NULL)
246				errx(1, "number of threads is %s: %s", errstr,
247				    optarg);
248			break;
249		case 'T':
250			timeout = strtonum(optarg, 0, INT_MAX, &errstr);
251			if (errstr != NULL)
252				errx(1, "timeout is %s: %s", errstr, optarg);
253			break;
254		default:
255			usage();
256		}
257	}
258	if (execute) {
259		int ret;
260
261		if (threads > 0)
262			errx(1, "execute sleep cannot be used with threads");
263
264		ret = snprintf(timeoutstr, sizeof(timeoutstr), "%d", timeout);
265		if (ret < 0 || (size_t)ret >= sizeof(timeoutstr))
266			err(1, "snprintf");
267	}
268
269	/* become process group leader */
270	if (daemonize) {
271		/* get rid of process leadership */
272		switch (fork()) {
273		case -1:
274			err(1, "fork parent");
275		case 0:
276			break;
277		default:
278			/* parent leaves orphan behind to do the work */
279			_exit(0);
280		}
281	}
282	pgrp = setsid();
283	if (pgrp == -1) {
284		if (!daemonize)
285			warnx("try -d to become process group leader");
286		err(1, "setsid");
287	}
288
289	/* create pipes to keep in contact with children */
290	rfds = reallocarray(NULL, procs, sizeof(int));
291	if (rfds == NULL)
292		err(1, "rfds");
293	fdmax = 0;
294
295	/* fork child processes and pass writing end of pipe */
296	for (i = 0; i < procs; i++) {
297		int pipefds[2], error;
298
299		if (pipe(pipefds) == -1)
300			err(1, "pipe");
301		if (fdmax < pipefds[0])
302			fdmax = pipefds[0];
303		rfds[i] = pipefds[0];
304
305		switch (fork()) {
306		case -1:
307			/* resource temporarily unavailable may happen */
308			error = errno;
309			/* reap children, but not parent */
310			signal(SIGTERM, SIG_IGN);
311			kill(-pgrp, SIGTERM);
312			errc(1, error, "fork child");
313		case 0:
314			/* child closes reading end, read is for the parent */
315			if (close(pipefds[0]) == -1)
316				err(1, "close read");
317			run_child(pipefds[1]);
318			/* cannot happen */
319			_exit(0);
320		default:
321			/* parent closes writing end, write is for the child */
322			if (close(pipefds[1]) == -1)
323				err(1, "close write");
324			break;
325		}
326	}
327
328	/* create select mask with all reading ends of child pipes */
329	fdlen = howmany(fdmax + 1, NFDBITS);
330	fdset = calloc(fdlen, sizeof(fd_mask));
331	if (fdset == NULL)
332		err(1, "fdset");
333	waiting = 0;
334	for (i = 0; i < procs; i++) {
335		FD_SET(rfds[i], fdset);
336		waiting = 1;
337	}
338
339	/* wait until all child processes are waiting */
340	while (waiting) {
341		tv.tv_sec = timeout;
342		tv.tv_usec = 0;
343		errno = ETIMEDOUT;
344		if (select(fdmax + 1, fdset, NULL, NULL, &tv) <= 0)
345			err(1, "select");
346
347		waiting = 0;
348		/* remove fd of children that closed their end  */
349		for (i = 0; i < procs; i++) {
350			if (rfds[i] >= 0) {
351				if (FD_ISSET(rfds[i], fdset)) {
352					if (close(rfds[i]) == -1)
353						err(1, "close parent");
354					FD_CLR(rfds[i], fdset);
355					rfds[i] = -1;
356				} else {
357					FD_SET(rfds[i], fdset);
358					waiting = 1;
359				}
360			}
361		}
362	}
363
364	/* kill all children simultaneously, parent exits in signal handler */
365	if (signal(SIGTERM, sigexit) == SIG_ERR)
366		err(1, "signal SIGTERM");
367	if (kill(-pgrp, SIGTERM) == -1)
368		err(1, "kill %d", -pgrp);
369
370	errx(1, "alive");
371}
372