1 files changed, 392 insertions, 0 deletions
diff --git a/src/lib-test/test-subprocess.c b/src/lib-test/test-subprocess.c
new file mode 100644
index 0000000..eeaae80
--- /dev/null
+++ b/src/lib-test/test-subprocess.c
@@ -0,0 +1,392 @@
+/* Copyright (c) 2020 Dovecot authors, see the included COPYING file */
+
+#include "lib.h"
+#include "lib-signals.h"
+#include "hostpid.h"
+#include "array.h"
+#include "ioloop.h"
+#include "sleep.h"
+#include "test-common.h"
+#include "test-subprocess.h"
+
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+struct test_subprocess {
+	pid_t pid;
+};
+
+volatile sig_atomic_t test_subprocess_is_child = 0;
+static bool test_subprocess_lib_init = FALSE;
+static volatile bool test_subprocess_notification_signal_received[SIGUSR1 + 1];
+static struct event *test_subprocess_event = NULL;
+static ARRAY(struct test_subprocess *) test_subprocesses = ARRAY_INIT;
+static void (*test_subprocess_cleanup_callback)(void) = NULL;
+
+static void
+test_subprocess_signal(const siginfo_t *si ATTR_UNUSED,
+		       void *context ATTR_UNUSED)
+{
+	io_loop_stop(current_ioloop);
+}
+
+static void test_subprocess_free_all(void)
+{
+	struct test_subprocess *subp;
+
+	array_foreach_elem(&test_subprocesses, subp)
+		i_free(subp);
+	array_free(&test_subprocesses);
+}
+
+static void ATTR_NORETURN
+test_subprocess_child(int (*func)(void *context), void *context,
+		      bool continue_test)
+{
+	int ret;
+
+	if (!continue_test)
+		test_forked_end();
+
+	hostpid_init();
+
+	lib_signals_deinit();
+	lib_signals_init();
+
+	lib_signals_set_handler(SIGTERM,
+		LIBSIG_FLAG_DELAYED | LIBSIG_FLAG_IOLOOP_AUTOMOVE,
+		test_subprocess_signal, NULL);
+	lib_signals_set_handler(SIGQUIT,
+		LIBSIG_FLAG_DELAYED | LIBSIG_FLAG_IOLOOP_AUTOMOVE,
+		test_subprocess_signal, NULL);
+	lib_signals_set_handler(SIGINT,
+		LIBSIG_FLAG_DELAYED | LIBSIG_FLAG_IOLOOP_AUTOMOVE,
+		test_subprocess_signal, NULL);
+
+	ret = func(context);
+
+	/* Prevent race condition */
+	lib_signals_clear_handlers_and_ignore(SIGTERM);
+
+	event_unref(&test_subprocess_event);
+	lib_signals_deinit();
+
+	if (!continue_test) {
+		lib_deinit();
+		lib_exit(ret);
+	}
+	test_exit((test_has_failed() ? 1 : 0));
+}
+
+#undef test_subprocess_fork
+void test_subprocess_fork(int (*func)(void *context), void *context,
+			  bool continue_test)
+{
+	struct test_subprocess *subprocess;
+
+	subprocess = i_new(struct test_subprocess, 1);
+
+	lib_signals_ioloop_detach();
+
+	/* avoid races: fork the child process with test_subprocess_is_child
+	   set to 1 in case it immediately receives a signal. */
+	test_subprocess_is_child = 1;
+	if ((subprocess->pid = fork()) == (pid_t)-1)
+		i_fatal("test: sub-process: fork() failed: %m");
+	if (subprocess->pid == 0) {
+		/* cannot include this in the list to avoid accidental
+		 * kill of PID 0, so just free it here explicitly. */
+		i_free(subprocess);
+		test_subprocess_free_all();
+
+		test_subprocess_child(func, context, continue_test);
+		i_unreached();
+	}
+	test_subprocess_is_child = 0;
+
+	array_push_back(&test_subprocesses, &subprocess);
+	lib_signals_ioloop_attach();
+}
+
+static void test_subprocess_verify_exit_status(int status)
+{
+	test_out_quiet("sub-process ended properly",
+		       WIFEXITED(status) && WEXITSTATUS(status) == 0);
+	if (WIFEXITED(status)) {
+		if (WEXITSTATUS(status) != 0) {
+			e_warning(test_subprocess_event,
+				  "Sub-process exited with status %d",
+				  WEXITSTATUS(status));
+		}
+	} else if (WIFSIGNALED(status)) {
+		e_warning(test_subprocess_event,
+			  "Sub-process forcibly terminated with signal %d",
+		          WTERMSIG(status));
+	} else if (WIFSTOPPED(status)) {
+		e_warning(test_subprocess_event,
+			  "Sub-process stopped with signal %d",
+			  WSTOPSIG(status));
+	} else {
+		e_warning(test_subprocess_event,
+			  "Sub-process terminated abnormally with status %d",
+			  status);
+	}
+}
+
+static void test_subprocess_kill_forced(struct test_subprocess *subp)
+{
+	i_assert(subp->pid > 0);
+	(void)kill(subp->pid, SIGKILL);
+	(void)waitpid(subp->pid, NULL, 0);
+}
+
+void test_subprocess_kill_all(unsigned int timeout_secs)
+{
+	struct test_subprocess **subps;
+	unsigned int subps_count, subps_left, i;
+
+	subps = array_get_modifiable(&test_subprocesses, &subps_count);
+
+	/* Request children to terminate gently */
+	for (i = 0; i < subps_count; i++) {
+		if (subps[i] == NULL || subps[i]->pid == (pid_t)-1)
+			continue;
+
+		e_debug(test_subprocess_event,
+			"Terminating sub-process [%u]", i);
+		if (kill(subps[i]->pid, SIGTERM) < 0) {
+			e_error(test_subprocess_event,
+				"Failed to kill sub-process [%u] with SIGTERM: "
+				"%m", i);
+		}
+	}
+
+	/* Wait for children */
+	subps_left = subps_count;
+	while (subps_left > 0) {
+		int status;
+		pid_t wret = (pid_t)-1;
+
+		alarm(timeout_secs);
+		wret = waitpid(-1, &status, 0);
+		alarm(0);
+
+		test_assert(wret > 0);
+		if (wret < 0 && errno == EINTR)
+			e_warning(test_subprocess_event,
+				  "Wait for sub-processes timed out");
+		if (wret > 0)
+			test_subprocess_verify_exit_status(status);
+
+		if (wret == 0)
+			break;
+		if (wret < 0) {
+			if (errno == ECHILD)
+				continue;
+			e_warning(test_subprocess_event,
+				  "Wait for sub-processes failed: %m");
+			break;
+		}
+		for (i = 0; i < subps_count; i++) {
+			if (subps[i] == NULL || subps[i]->pid != wret)
+				continue;
+			e_debug(test_subprocess_event,
+				"Terminated sub-process [%u]", i);
+			i_free(subps[i]);
+			subps_left--;
+		}
+	}
+
+	/* Kill disobedient ones with fire */
+	for (i = 0; i < subps_count; i++) {
+		if (subps[i] == NULL || subps[i]->pid == (pid_t)-1)
+			continue;
+		e_warning(test_subprocess_event,
+			  "Forcibly killed sub-process [%u]", i);
+		test_subprocess_kill_forced(subps[i]);
+		i_assert(subps_left > 0);
+		i_free(subps[i]);
+		subps_left--;
+	}
+	i_assert(subps_left == 0);
+
+	array_clear(&test_subprocesses);
+}
+
+static void test_subprocess_kill_all_forced(void)
+{
+	struct test_subprocess **subps;
+	unsigned int subps_count, i;
+
+	if (!array_is_created(&test_subprocesses))
+		return;
+
+	/* This is also called from signal handler context, so no debug logging
+	   here.
+	 */
+
+	subps = array_get_modifiable(&test_subprocesses, &subps_count);
+
+	if (subps_count == 0)
+		return;
+
+	for (i = 0; i < subps_count; i++) {
+		if (subps[i] == NULL || subps[i]->pid == (pid_t)-1)
+			continue;
+		test_subprocess_kill_forced(subps[i]);
+		subps[i]->pid = (pid_t)-1;
+	}
+}
+
+/*
+ * Main
+ */
+
+volatile sig_atomic_t terminating = 0;
+
+static void test_subprocess_cleanup(void)
+{
+	if (test_subprocess_is_child != 0) {
+		/* Child processes must not execute the cleanups */
+		return;
+	}
+
+	/* We get here when the test ended normally, badly failed, crashed,
+	   terminated, or executed exit() unexpectedly. The cleanups performed
+	   here are important and must be executed at all times. */
+
+	/* While any unfreed memory will be handled by the system, lingering
+	   child processes will not be handled so well. So, we need to make sure
+	   here that we don't leave any pesky child processes alive. */
+	test_subprocess_kill_all_forced();
+
+	/* Perform any additional important cleanup specific to the test. */
+	if (test_subprocess_cleanup_callback != NULL)
+		test_subprocess_cleanup_callback();
+}
+
+static void
+test_subprocess_alarm(const siginfo_t *si ATTR_UNUSED,
+		      void *context ATTR_UNUSED)
+{
+	/* We use alarm() to implement a simple timeout on waitpid(), which will
+	   exit with EINTR when SIGALRM is received. This handler overrides the
+	   default SIGALRM handler, so that the process is not killed and no
+	   messages are printed to terminal.
+	 */
+}
+
+static void
+test_subprocess_terminate(const siginfo_t *si, void *context ATTR_UNUSED)
+{
+	int signo = si->si_signo;
+
+	if (terminating != 0)
+		raise(signo);
+	terminating = 1;
+
+	/* Perform important cleanups */
+	test_subprocess_cleanup();
+
+	(void)signal(signo, SIG_DFL);
+	if (signo == SIGTERM)
+		_exit(0);
+	else
+		raise(signo);
+}
+
+static void test_atexit(void)
+{
+	/* NOTICE: This is also called by children, so be careful. */
+
+	/* Perform important cleanups */
+	test_subprocess_cleanup();
+}
+
+void test_subprocess_set_cleanup_callback(void (*callback)(void))
+{
+	test_subprocess_cleanup_callback = callback;
+}
+
+void test_subprocess_notify_signal_send(int signo, pid_t pid)
+{
+	if (kill(pid, signo) < 0)
+		i_fatal("kill(%ld, SIGHUP) failed: %m", (long)pid);
+}
+
+void test_subprocess_notify_signal_send_parent(int signo)
+{
+	test_subprocess_notify_signal_send(signo, getppid());
+}
+
+void test_subprocess_notify_signal_reset(int signo)
+{
+	i_assert(signo >= 0 &&
+		 (unsigned int)signo < N_ELEMENTS(test_subprocess_notification_signal_received));
+	test_subprocess_notification_signal_received[signo] = FALSE;
+}
+
+void test_subprocess_notify_signal_wait(int signo, unsigned int timeout_msecs)
+{
+	unsigned int i, count = timeout_msecs / 10;
+
+	for (i = 0; i < count; i++) {
+		if (test_subprocess_notification_signal_received[signo])
+			return;
+		i_sleep_msecs(10);
+	}
+	i_fatal("Didn't receive wait notification signal from server");
+}
+
+static void
+test_subprocess_notification_signal(const siginfo_t *si,
+				    void *context ATTR_UNUSED)
+{
+	int signo = si->si_signo;
+
+	i_assert(signo >= 0 &&
+		 (unsigned int)signo < N_ELEMENTS(test_subprocess_notification_signal_received));
+	test_subprocess_notification_signal_received[signo] = TRUE;
+}
+
+void test_subprocesses_init(bool debug)
+{
+	if (!lib_is_initialized()) {
+		lib_init();
+		test_subprocess_lib_init = TRUE;
+	}
+	lib_signals_init();
+
+	atexit(test_atexit);
+	lib_signals_ignore(SIGPIPE, TRUE);
+	lib_signals_set_handler(SIGALRM, 0, test_subprocess_alarm, NULL);
+	lib_signals_set_handler(SIGTERM, 0, test_subprocess_terminate, NULL);
+	lib_signals_set_handler(SIGQUIT, 0, test_subprocess_terminate, NULL);
+	lib_signals_set_handler(SIGINT, 0, test_subprocess_terminate, NULL);
+	lib_signals_set_handler(SIGSEGV, 0, test_subprocess_terminate, NULL);
+	lib_signals_set_handler(SIGABRT, 0, test_subprocess_terminate, NULL);
+	lib_signals_set_handler(SIGHUP, LIBSIG_FLAG_RESTART,
+				test_subprocess_notification_signal, NULL);
+	lib_signals_set_handler(SIGUSR1, LIBSIG_FLAG_RESTART,
+				test_subprocess_notification_signal, NULL);
+
+	i_array_init(&test_subprocesses, 8);
+
+	test_subprocess_event = event_create(NULL);
+	event_set_forced_debug(test_subprocess_event, debug);
+	event_set_append_log_prefix(test_subprocess_event, "test: ");
+}
+
+void test_subprocesses_deinit(void)
+{
+	test_subprocess_cleanup();
+	test_subprocess_free_all();
+	array_free(&test_subprocesses);
+
+	event_unref(&test_subprocess_event);
+	lib_signals_deinit();
+
+	if (test_subprocess_lib_init)
+		lib_deinit();
+}