Implement a cleaner shutdown procedure

README.md

@@ -28,8 +28,11 @@ may not be the same on the remote system.
 
 ## Installing
 
-Build with meson[0]. A recent version of wayland is required. Man pages are
-generated with scdoc[1].
+Build with meson[0]. Requirements:
+
+* wayland (>= 1.15, to support absolute paths in WAYLAND_DISPLAY)
+* scdoc (to generate a man page)
+* ssh (OpenSSH >= 6.7, for Unix domain socket forwarding)
 
 [0] [https://mesonbuild.com/](https://mesonbuild.com/)
 [1] [https://git.sr.ht/~sircmpwn/scdoc](https://git.sr.ht/~sircmpwn/scdoc)

client.c

@@ -74,7 +74,7 @@ static int run_client_child(int chanfd, const char *socket_path)
 	const int maxmsg = 4096;
 	char *buffer = calloc(1, maxmsg + 1);
 	struct pollfd *pfds = NULL;
-	while (1) {
+	while (!shutdown_flag) {
 		int npoll = 2 + count_npipes(&fdtransmap);
 		free(pfds);
 		// todo: resizing logic
@@ -85,7 +85,7 @@ static int run_client_child(int chanfd, const char *socket_path)
 		pfds[1].events = POLL_IN;
 		fill_with_pipes(&fdtransmap, pfds + 2);
 
-		int r = poll(pfds, (nfds_t)npoll, 700);
+		int r = poll(pfds, (nfds_t)npoll, -1);
 		if (r == -1) {
 			wp_log(WP_ERROR, "poll failed, stopping\n");
 			break;
@@ -99,7 +99,7 @@ static int run_client_child(int chanfd, const char *socket_path)
 			wp_log(WP_DEBUG, "Channel read begun\n");
 			ssize_t nbytes = read_size_then_buf(chanfd, &tmpbuf);
 			if (nbytes == 0) {
-				wp_log(WP_ERROR,
+				wp_log(WP_DEBUG,
 						"Channel read connection closed\n");
 				break;
 			}
@@ -245,30 +245,31 @@ int run_client(const char *socket_path, bool oneshot, pid_t eol_pid)
 	cs.fd = channelsock;
 	cs.events = POLL_IN;
 	cs.revents = 0;
-	while (1) {
+	while (!shutdown_flag) {
 		// TODO: figure out a safe, non-polling solution
-		int r = poll(&cs, 1, 1000);
-		if (r == -1) {
-			if (errno == EINTR) {
-				continue;
-			}
-			retcode = EXIT_FAILURE;
-			break;
-		}
 		if (eol_pid) {
-			int stat;
-			int wp = waitpid(eol_pid, &stat, WNOHANG);
+			int wp = waitpid(eol_pid, NULL, WNOHANG);
 			if (wp > 0) {
-				wp_log(WP_ERROR, "Child (ssh) died early\n");
+				wp_log(WP_DEBUG, "Child (ssh) died, exiting\n");
 				eol_pid = 0; // < recycled
-				retcode = EXIT_FAILURE;
+				retcode = EXIT_SUCCESS;
 				break;
 			}
 		}
+
 		// scan stack for children, and clean them up!
 		wait_on_children(&children, WNOHANG);
 
-		if (r == 0) {
+		int r = poll(&cs, 1, -1);
+		if (r == -1) {
+			if (errno == EINTR) {
+				// If SIGCHLD, we will check the child.
+				// If SIGINT, the loop ends
+				continue;
+			}
+			retcode = EXIT_FAILURE;
+			break;
+		} else if (r == 0) {
 			// Nothing to read
 			continue;
 		}
@@ -311,6 +312,9 @@ int run_client(const char *socket_path, bool oneshot, pid_t eol_pid)
 					retcode = EXIT_FAILURE;
 					break;
 				} else {
+					// Remove connection from this process
+					close(chanclient);
+
 					struct kstack *kd = calloc(1,
 							sizeof(struct kstack));
 					kd->pid = npid;
@@ -324,11 +328,18 @@ int run_client(const char *socket_path, bool oneshot, pid_t eol_pid)
 
 	close(channelsock);
 	unlink(socket_path);
+	int cleanup_type = shutdown_flag ? WNOHANG : 0;
 	if (eol_pid) {
 		// Don't return until the child process completes
-		int status;
-		waitpid(eol_pid, &status, 0);
+		waitpid(eol_pid, NULL, cleanup_type);
+	}
+	wait_on_children(&children, cleanup_type);
+	// Free stack, in case we suddenly shutdown and fail to clean up
+	// children
+	while (children) {
+		struct kstack *nxt = children->nxt;
+		free(children);
+		children = nxt;
 	}
-	wait_on_children(&children, 0);
 	return retcode;
 }

server.c

@@ -50,7 +50,7 @@ static int run_server_child(int chanfd, int appfd)
 	struct fd_translation_map fdtransmap = {
 			.local_sign = -1, .list = NULL, .max_local_id = 1};
 	struct pollfd *pfds = NULL;
-	while (1) {
+	while (!shutdown_flag) {
 		int npoll = 2 + count_npipes(&fdtransmap);
 		free(pfds);
 		// todo: resizing logic
@@ -61,10 +61,18 @@ static int run_server_child(int chanfd, int appfd)
 		pfds[1].events = POLL_IN;
 		fill_with_pipes(&fdtransmap, pfds + 2);
 
-		int r = poll(pfds, (nfds_t)npoll, 700);
+		int r = poll(pfds, (nfds_t)npoll, -1);
 		if (r == -1) {
-			wp_log(WP_ERROR, "poll failed, stopping\n");
-			break;
+			if (errno == EINTR) {
+				wp_log(WP_ERROR,
+						"poll interrupted: shutdown=%c\n",
+						shutdown_flag ? 'Y' : 'n');
+			} else {
+				wp_log(WP_ERROR,
+						"poll failed due to, stopping: %s\n",
+						strerror(errno));
+				break;
+			}
 		}
 
 		mark_pipe_object_statuses(&fdtransmap, npoll - 2, pfds + 2);
@@ -295,6 +303,7 @@ int run_server(const char *socket_path, bool oneshot, char *const app_argv[])
 		} else {
 			retval = EXIT_FAILURE;
 		}
+		close(server_link);
 	} else {
 		struct kstack *children = NULL;
 
@@ -304,24 +313,34 @@ int run_server(const char *socket_path, bool oneshot, char *const app_argv[])
 		pf.fd = wdisplay_socket;
 		pf.events = POLL_IN;
 		pf.revents = 0;
-		while (1) {
-			int r = poll(&pf, 1, 1000);
-			if (r == -1) {
-				if (errno == EINTR) {
-					continue;
-				}
-				retval = EXIT_FAILURE;
+		while (!shutdown_flag) {
+			int wp = waitpid(pid, NULL, WNOHANG);
+			if (wp > 0) {
+				wp_log(WP_DEBUG,
+						"Child program has died, exiting\n");
+				retval = EXIT_SUCCESS;
 				break;
-			}
-			int stat;
-			if (waitpid(pid, &stat, WNOHANG) > 0) {
-				wp_log(WP_ERROR, "Child program died early\n");
+			} else if (wp == -1) {
+				wp_log(WP_ERROR, "Failed in waitpid: %s\n",
+						strerror(errno));
 				retval = EXIT_FAILURE;
 				break;
 			}
 			// scan stack for children, and clean them up!
 			wait_on_children(&children, WNOHANG);
-			if (r == 0) {
+
+			int r = poll(&pf, 1, -1);
+			if (r == -1) {
+				if (errno == EINTR) {
+					// If SIGCHLD, we will check the child.
+					// If SIGINT, the loop ends
+					continue;
+				}
+				fprintf(stderr, "Poll failed: %s\n",
+						strerror(errno));
+				retval = EXIT_FAILURE;
+				break;
+			} else if (r == 0) {
 				continue;
 			}
 
@@ -359,6 +378,10 @@ int run_server(const char *socket_path, bool oneshot, char *const app_argv[])
 					retval = EXIT_FAILURE;
 					break;
 				} else {
+					// This process no longer needs the
+					// application connection
+					close(appfd);
+
 					struct kstack *kd = calloc(1,
 							sizeof(struct kstack));
 					kd->pid = npid;
@@ -368,15 +391,25 @@ int run_server(const char *socket_path, bool oneshot, char *const app_argv[])
 				continue;
 			}
 		}
+		close(wdisplay_socket);
 		// Wait for child processes to exit
 		wp_log(WP_DEBUG, "Waiting for child handlers\n");
-		wait_on_children(&children, 0);
+		wait_on_children(&children, shutdown_flag ? WNOHANG : 0);
+		// Free stack, in case we suddenly shutdown and fail to clean up
+		// children
+		while (children) {
+			struct kstack *nxt = children->nxt;
+			free(children);
+			children = nxt;
+		}
 	}
 
+	if (!oneshot) {
+		unlink(displaypath);
+	}
 	// todo: scope manipulation, to ensure all cleanups are done
 	wp_log(WP_DEBUG, "Waiting for child process\n");
-	int status;
-	waitpid(pid, &status, 0);
+	waitpid(pid, NULL, shutdown_flag ? WNOHANG : 0);
 	wp_log(WP_DEBUG, "Program ended\n");
 	return retval;
 }

util.c

@@ -44,8 +44,26 @@
 #include <time.h>
 #include <unistd.h>
 
+bool shutdown_flag = false;
+char waypipe_log_mode = '?';
 log_cat_t waypipe_loglevel = WP_ERROR;
 
+void handle_sigint(int sig)
+{
+	(void)sig;
+	char buf[20];
+	int pid = getpid();
+	sprintf(buf, "SIGINT(%d)\n", pid);
+	write(STDOUT_FILENO, buf, strlen(buf));
+	if (!shutdown_flag) {
+		shutdown_flag = true;
+	} else {
+		const char msg[] = "Second SIGINT, aborting.\n";
+		write(STDERR_FILENO, msg, sizeof(msg));
+		abort();
+	}
+}
+
 int set_fnctl_flag(int fd, int the_flag)
 {
 	int flags = fcntl(fd, F_GETFL, 0);
@@ -254,8 +272,9 @@ static int translate_fd(struct fd_translation_map *map, int fd)
 		} else {
 			wp_log(WP_ERROR,
 					"double ended pipe head encountered; not implemented\n");
-			shadow->type = FDC_PIPE_IR; // <typical of data offers,
-						    // creator sends
+			shadow->type = FDC_PIPE_IR;
+			// TODO:  implement with socketpair, and use
+			// dir-specific close logic
 		}
 
 		// Make this end of the pipe nonblocking, so that we can include
@@ -766,9 +785,8 @@ void wait_on_children(struct kstack **children, int options)
 	struct kstack *cur = *children;
 	struct kstack **prv = children;
 	while (cur) {
-		int stat;
-		if (waitpid(cur->pid, &stat, options) > 0) {
-			wp_log(WP_ERROR, "Child handler %d has died\n",
+		if (waitpid(cur->pid, NULL, options) > 0) {
+			wp_log(WP_DEBUG, "Child handler %d has died\n",
 					cur->pid);
 			struct kstack *nxt = cur->nxt;
 			free(cur);

util.h

@@ -30,6 +30,11 @@
 #include <stdint.h>
 #include <sys/types.h>
 
+// On SIGINT, this is set to true. The main program should then cleanup ASAP
+extern bool shutdown_flag;
+
+void handle_sigint(int sig);
+
 /** Set the given flag with fcntl. Silently return -1 on failure. */
 int set_fnctl_flag(int fd, int the_flag);
 /** Create a nonblocking AF_UNIX/SOCK_STREAM socket, and listen with
@@ -43,6 +48,7 @@ ssize_t iovec_write(int conn, const char *buf, size_t buflen, const int *fds,
 
 typedef enum { WP_DEBUG = 1, WP_ERROR = 2 } log_cat_t;
 
+extern char waypipe_log_mode;
 extern log_cat_t waypipe_loglevel;
 
 // mutates a static local, hence can only be called singlethreaded
@@ -54,7 +60,8 @@ const char *static_timestamp(void);
 // no trailing ;, user must supply
 #define wp_log(level, fmt, ...)                                                \
 	if ((level) >= waypipe_loglevel)                                       \
-	fprintf(stderr, "%s [%s:%3d] " fmt, static_timestamp(),                \
+	fprintf(stderr, "%c%d:%s [%s:%3d] " fmt, waypipe_log_mode, getpid(),   \
+			static_timestamp(),                                    \
 			((const char *)__FILE__) + WAYPIPE_SRC_DIR_LENGTH,     \
 			__LINE__, ##__VA_ARGS__)
 

waypipe.c

@@ -23,9 +23,12 @@
  * SOFTWARE.
  */
 
+#define _XOPEN_SOURCE 700
+
 #include "util.h"
 
 #include <getopt.h>
+#include <signal.h>
 #include <stdarg.h>
 #include <stdbool.h>
 #include <stdint.h>
@@ -134,6 +137,7 @@ static int locate_openssh_cmd_hostname(int argc, char *const *argv)
 	return dstidx;
 }
 
+void handle_noop(int sig) { (void)sig; }
 int main(int argc, char **argv)
 {
 	bool help = false;
@@ -229,6 +233,25 @@ int main(int argc, char **argv)
 		argc--;
 	}
 	waypipe_loglevel = debug ? WP_DEBUG : WP_ERROR;
+	waypipe_log_mode = is_client ? 'C' : 'S';
+
+	// Setup signals
+	struct sigaction ia; // SIGINT: abort operations, and set a flag
+	ia.sa_handler = handle_sigint;
+	sigemptyset(&ia.sa_mask);
+	ia.sa_flags = 0;
+	struct sigaction ca; // SIGCHLD: restart operations, but EINTR on poll
+	ca.sa_handler = handle_noop;
+	sigemptyset(&ca.sa_mask);
+	ca.sa_flags = SA_RESTART | SA_NOCLDSTOP;
+	if (sigaction(SIGINT, &ia, NULL) == -1) {
+		wp_log(WP_ERROR, "Failed to set signal action for SIGINT\n");
+		return EXIT_FAILURE;
+	}
+	if (sigaction(SIGCHLD, &ca, NULL) == -1) {
+		wp_log(WP_ERROR, "Failed to set signal action for SIGCHLD\n");
+		return EXIT_FAILURE;
+	}
 
 	if (is_client) {
 		if (setup_ssh) {