nitroctl: use notification sockets

1 files changed, 113 insertions, 81 deletions

diff --git a/nitroctl.c b/nitroctl.c
index 498ed6f..cce867a 100644
--- a/nitroctl.c
+++ b/nitroctl.c
@@ -2,6 +2,10 @@
 #include <sys/un.h>
 
 #include <ctype.h>
+#include <errno.h>
+#include <libgen.h>
+#include <limits.h>
+#include <signal.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -9,7 +13,9 @@
 #include <unistd.h>
 
 int connfd;
-char template[256] = "/tmp/nitroctl.XXXXXX";
+const char *sockpath;
+char notifypath[PATH_MAX];
+volatile sig_atomic_t timed_out;
 
 enum process_state {
 	PROC_DOWN = 1,
@@ -24,28 +30,109 @@ enum process_state {
 };
 
 void
-cleanup()
+noop()
 {
-	unlink(template);
-	*strrchr(template, '/') = 0;
-	rmdir(template);
+	timed_out = 1;
 }
 
-enum process_state
-check(char *name)
+void
+cleanup_notify()
 {
-	dprintf(connfd, "?%s", name);
+	unlink(notifypath);
+}
 
-	ssize_t rd;
-	char buf[64];
-	rd = read(connfd, buf, sizeof buf);
-	if (rd < 0) {
-		perror("read");
+void
+notifysock(const char *service)
+{
+	static char default_sock[] = "/run/nitro/nitro.sock";
+	char *path = strdup(sockpath);
+	if (!path || !*path)
+		path = default_sock;
+
+	snprintf(notifypath, sizeof notifypath,
+	    "%s/notify/%s,%ld", dirname(path), service, (long)getpid());
+
+	struct sockaddr_un my_addr = { 0 };
+	my_addr.sun_family = AF_UNIX;
+	strncpy(my_addr.sun_path, notifypath, sizeof my_addr.sun_path - 1);
+again:
+	if (bind(connfd, (struct sockaddr *)&my_addr, sizeof my_addr) < 0) {
+		if (errno == EADDRINUSE) {
+			// ok to delete this, only we can have the current pid.
+			if (unlink(notifypath) == 0)
+				goto again;
+		}
+		perror("bind");
+		exit(111);
+	}
+	atexit(cleanup_notify);
+
+	struct sockaddr_un addr = { 0 };
+	addr.sun_family = AF_UNIX;
+	strncpy(addr.sun_path, sockpath, sizeof addr.sun_path - 1);
+	if (connect(connfd, (struct sockaddr *)&addr, sizeof addr) < 0) {
+		perror("connect");
 		exit(111);
 	}
-	if (buf[0] >= 'A' && buf[0] <= 'Z')
-		return buf[0] - 64;
-	return -1;
+}
+
+int
+send_and_wait(char cmd, const char *service)
+{
+	notifysock(service);
+
+	dprintf(connfd, "%c%s", cmd, service);
+
+	struct sigaction sa = {
+		.sa_handler = noop,
+		.sa_flags = 0,
+	};
+	sigaction(SIGALRM, &sa, 0);
+	alarm(5);
+
+	while(!timed_out) {
+		ssize_t rd;
+		char buf[64];
+
+		rd = read(connfd, buf, sizeof buf); /* block */
+		if (rd < 0) {
+			if (errno == EINTR)
+				continue;
+			perror("read");
+			exit(111);
+		}
+
+		buf[rd] = 0;
+		printf("got %s\n", buf);
+
+		int state = 0;
+		if (buf[0] >= 'A' && buf[0] <= 'Z')
+			state = buf[0] - 64;
+
+		if (buf[0] == 'e') {
+			return 111;
+		}
+
+		switch (cmd) {
+		case 'u':
+		case 'r':
+			if (state == PROC_UP || state == PROC_ONESHOT)
+				return 0;
+			if (state == PROC_FATAL) {
+				fprintf(stderr,
+				    "nitroctl: failed to %sstart '%s'\n", service, cmd == 'u' ? "" : "re");
+				return 1;
+			}
+			break;
+		case 'd':
+			if (state == PROC_DOWN || state == PROC_FATAL)
+				return 0;
+			break;
+		}
+	}
+
+	fprintf(stderr, "nitroctl: action timed out\n");
+	return 2;
 }
 
 int
@@ -68,6 +155,7 @@ main(int argc, char *argv[])
 	    strcmp(argv[1], "2") != 0 &&
 	    strcmp(argv[1], "check") != 0 &&
 	    strcmp(argv[1], "start") != 0 &&
+	    strcmp(argv[1], "restart") != 0 &&
 	    strcmp(argv[1], "stop") != 0 &&
 	    strcmp(argv[1], "Reboot") != 0 &&
 	    strcmp(argv[1], "Shutdown") != 0)) {
@@ -76,47 +164,29 @@ main(int argc, char *argv[])
 	}
 
 	static const char default_sock[] = "/run/nitro/nitro.sock";
-	const char *path = getenv("NITRO_SOCK");
-	if (!path || !*path)
-		path = default_sock;
+	sockpath = getenv("NITRO_SOCK");
+	if (!sockpath || !*sockpath)
+		sockpath = default_sock;
 
-	struct sockaddr_un addr = { 0 };
-	addr.sun_family = AF_UNIX;
-	strncpy(addr.sun_path, path, sizeof addr.sun_path - 1);
 	connfd = socket(AF_UNIX, SOCK_DGRAM | SOCK_CLOEXEC, 0);
 	if (connfd < 0) {
 		perror("socket");
 		exit(111);
 	}
 
-	struct sockaddr_un my_addr = { 0 };
-	my_addr.sun_family = AF_UNIX;
-	if (!mkdtemp(template)) {
-		perror("mkdtemp");
-		exit(111);
-	}
-	strcat(template, "/sock");
-	strncpy(my_addr.sun_path, template, sizeof my_addr.sun_path - 1);
-	if (bind(connfd, (struct sockaddr *)&my_addr, sizeof my_addr) < 0) {
-		perror("bind");
-		exit(111);
-	}
-	atexit(cleanup);
-
-	if (connect(connfd, (struct sockaddr *)&addr, sizeof addr) < 0) {
-		perror("connect");
-		exit(111);
-	}
-
 	char cmd = argv[1][0];
 
 	if (strcmp(argv[1], "start") == 0 && argv[2])
-		cmd = 'u';
+		return send_and_wait('u', argv[2]);
 	else if (strcmp(argv[1], "stop") == 0 && argv[2])
-		cmd = 'd';
+		return send_and_wait('d', argv[2]);
+	else if (strcmp(argv[1], "restart") == 0 && argv[2])
+		return send_and_wait('r', argv[2]);
 	else if (strcmp(argv[1], "check") == 0 && argv[2])
 		cmd = '?';
 
+	notifysock("");
+
 	dprintf(connfd, "%c%s", cmd, argv[2] ? argv[2] : "");
 
 	int status = 1;
@@ -132,43 +202,5 @@ main(int argc, char *argv[])
 		status = 0;
 	write(1, buf, rd);
 
-	if (strcmp(argv[1], "start") == 0 && argv[2]) {
-		int checks = 0;
-		while (1) {
-			switch (check(argv[2])) {
-			case PROC_DOWN:
-			case PROC_FATAL:
-				fprintf(stderr, "start failed\n");
-				return 1;
-			case PROC_UP:
-			case PROC_ONESHOT:
-				return 0;
-			default:
-				if (checks > 10) {
-					fprintf(stderr, "start timed out\n");
-					return 1;
-				}
-				nanosleep(&(struct timespec){0, 250000000}, 0);
-				checks++;
-			}
-		}
-	} else if (strcmp(argv[1], "stop") == 0 && argv[2]) {
-		int checks = 0;
-		while (1) {
-			switch (check(argv[2])) {
-			case PROC_DOWN:
-			case PROC_FATAL: /* ? */
-				return 0;
-			default:
-				if (checks > 10) {
-					fprintf(stderr, "stop timed out\n");
-					return 1;
-				}
-				nanosleep(&(struct timespec){0, 250000000}, 0);
-				checks++;
-			}
-		}
-	}
-
 	return status;
 }