Simplify the smtpctl code: add a try_connect() function to test if

the server is running, and flush()/next_message() helpers to make
imsg request/response loops easier.

ok gilles@

1 files changed, 183 insertions, 184 deletions

diff --git a/usr.sbin/smtpd/smtpctl.c b/usr.sbin/smtpd/smtpctl.c
index 757024da491..dca5f2941dc 100644
--- a/usr.sbin/smtpd/smtpctl.c
+++ b/usr.sbin/smtpd/smtpctl.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: smtpctl.c,v 1.95 2012/11/12 14:58:53 eric Exp $	*/
+/*	$OpenBSD: smtpctl.c,v 1.96 2012/11/13 09:35:18 eric Exp $	*/
 
 /*
  * Copyright (c) 2006 Pierre-Yves Ritschard <pyr@openbsd.org>
@@ -58,6 +58,10 @@ static void show_envelope(const char *);
 static void show_message(const char *);
 static void show_monitor(struct stat_digest *);
 
+static int try_connect(void);
+static void flush(void);
+static void next_message(struct imsg *);
+
 int proctype;
 struct imsgbuf	*ibuf;
 
@@ -96,104 +100,131 @@ setup_env(struct smtpd *smtpd)
 		errx(1, "invalid directory permissions");
 }
 
+static int
+try_connect(void)
+{
+	struct sockaddr_un	sun;
+	int			ctl_sock, saved_errno;
+
+	/* connect to smtpd control socket */
+	if ((ctl_sock = socket(AF_UNIX, SOCK_STREAM, 0)) == -1)
+		err(1, "socket");
+
+	bzero(&sun, sizeof(sun));
+	sun.sun_family = AF_UNIX;
+	strlcpy(sun.sun_path, SMTPD_SOCKET, sizeof(sun.sun_path));
+	if (connect(ctl_sock, (struct sockaddr *)&sun, sizeof(sun)) == -1) {
+		saved_errno = errno;
+		close(ctl_sock);
+		errno = saved_errno;
+		return (0);
+	}
+
+	if ((ibuf = calloc(1, sizeof(struct imsgbuf))) == NULL)
+		err(1, "calloc");
+	imsg_init(ibuf, ctl_sock);
+
+	return (1);
+}
+
+static void
+flush(void)
+{
+	if (imsg_flush(ibuf) == -1)
+		err(1, "write error");
+}
+
+static void
+next_message(struct imsg *imsg)
+{
+	ssize_t	n;
+
+	while(1) {
+		if ((n = imsg_get(ibuf, imsg)) == -1)
+			errx(1, "imsg_get error");
+		if (n)
+			return;
+
+		if ((n = imsg_read(ibuf)) == -1)
+			errx(1, "imsg_read error");
+		if (n == 0)
+			errx(1, "pipe closed");
+	}
+}
+
 int
 main(int argc, char *argv[])
 {
-	struct sockaddr_un	sun;
 	struct parse_result	*res = NULL;
 	struct imsg		imsg;
 	struct smtpd		smtpd;
-	int			ctl_sock;
+	uint64_t		ulval;
+	char			name[MAX_LINE_SIZE];
 	int			done = 0;
-	int			n, verbose = 0;
+	int			verbose = 0;
+	int			action = -1;
 
 	/* parse options */
-	if (strcmp(__progname, "sendmail") == 0 || strcmp(__progname, "send-mail") == 0)
+	if (strcmp(__progname, "sendmail") == 0 ||
+	    strcmp(__progname, "send-mail") == 0) {
 		sendmail = 1;
-	else if (strcmp(__progname, "mailq") == 0) {
-		if (geteuid())
-			errx(1, "need root privileges");
-		setup_env(&smtpd);
-		show_queue(0);
-		return 0;
-	} else if (strcmp(__progname, "smtpctl") == 0) {
-
-		/* check for root privileges */
-		if (geteuid())
-			errx(1, "need root privileges");
+		if (try_connect())
+			return (enqueue(argc, argv));
+		return (enqueue_offline(argc, argv));
+	}
 
-		setup_env(&smtpd);
+	if (geteuid())
+		errx(1, "need root privileges");
+	setup_env(&smtpd);
 
+	if (strcmp(__progname, "mailq") == 0)
+		action = SHOW_QUEUE;
+	else if (strcmp(__progname, "smtpctl") == 0) {
 		if ((res = parse(argc - 1, argv + 1)) == NULL)
 			exit(1);
-
-		/* handle "disconnected" commands */
-		switch (res->action) {
-		case SHOW_QUEUE:
-			show_queue(0);
-			break;
-		case SHOW_ENVELOPE:
-			show_envelope(res->data);
-			break;
-		case SHOW_MESSAGE:
-			show_message(res->data);
-			break;
-		default:
-			goto connected;
-		}
-		return 0;
+		action = res->action;
 	} else
 		errx(1, "unsupported mode");
 
-connected:
-	/* connect to smtpd control socket */
-	if ((ctl_sock = socket(AF_UNIX, SOCK_STREAM, 0)) == -1)
-		err(1, "socket");
-
-	bzero(&sun, sizeof(sun));
-	sun.sun_family = AF_UNIX;
-	strlcpy(sun.sun_path, SMTPD_SOCKET, sizeof(sun.sun_path));
-	if (connect(ctl_sock, (struct sockaddr *)&sun, sizeof(sun)) == -1) {
-		if (sendmail)
-			return enqueue_offline(argc, argv);
-		err(1, "connect: %s", SMTPD_SOCKET);
+	/* test for not connected actions */
+	switch (action) {
+	case SHOW_QUEUE:
+		show_queue(0);
+		return (0);
+	case SHOW_ENVELOPE:
+		show_envelope(res->data);
+		return (0);
+	case SHOW_MESSAGE:
+		show_message(res->data);
+		return (0);
+	default:
+		break;
 	}
 
-	if ((ibuf = calloc(1, sizeof(struct imsgbuf))) == NULL)
-		err(1, NULL);
-	imsg_init(ibuf, ctl_sock);
-
-	if (sendmail)
-		return enqueue(argc, argv);
+	if (!try_connect())
+		errx(1, "smtpd doesn't seem to be running");
 
 	/* process user request */
-	switch (res->action) {
+	switch (action) {
 	case NONE:
 		usage();
 		/* not reached */
 
 	case SCHEDULE:
-	case REMOVE: {
-		uint64_t ulval;
-
 		ulval = strtoevpid(res->data);
-		if (res->action == SCHEDULE)
-			imsg_compose(ibuf, IMSG_SCHEDULER_SCHEDULE, 0, 0, -1, &ulval,
-			    sizeof(ulval));
-		if (res->action == REMOVE)
-			imsg_compose(ibuf, IMSG_SCHEDULER_REMOVE, 0, 0, -1, &ulval,
-			    sizeof(ulval));
+		imsg_compose(ibuf, IMSG_SCHEDULER_SCHEDULE, 0, 0, -1, &ulval,
+		    sizeof(ulval));
 		break;
-	}
-
-	case SCHEDULE_ALL: {
-		uint64_t ulval = 0;
-
+	case REMOVE:
+		ulval = strtoevpid(res->data);
+		imsg_compose(ibuf, IMSG_SCHEDULER_REMOVE, 0, 0, -1, &ulval,
+		    sizeof(ulval));
+		break;
+	case SCHEDULE_ALL:
+		ulval = 0;
 		imsg_compose(ibuf, IMSG_SCHEDULER_SCHEDULE, 0, 0, -1, &ulval,
 		    sizeof(ulval));
 		break;
-	}
-
 	case SHUTDOWN:
 		imsg_compose(ibuf, IMSG_CTL_SHUTDOWN, 0, 0, -1, NULL, 0);
 		break;
@@ -218,19 +249,23 @@ connected:
 	case SHOW_STATS:
 		imsg_compose(ibuf, IMSG_STATS, 0, 0, -1, NULL, 0);
 		break;
-	case UPDATE_MAP: {
-		char	name[MAX_LINE_SIZE];
-
+	case UPDATE_MAP:
 		if (strlcpy(name, res->data, sizeof name) >= sizeof name)
 			errx(1, "map name too long.");
 		imsg_compose(ibuf, IMSG_LKA_UPDATE_MAP, 0, 0, -1,
 		    name, strlen(name) + 1);
 		done = 1;
 		break;
-	}
 	case MONITOR:
-    again:
-		imsg_compose(ibuf, IMSG_DIGEST, 0, 0, -1, NULL, 0);
+		done = 1;
+		while (1) {
+			imsg_compose(ibuf, IMSG_DIGEST, 0, 0, -1, NULL, 0);
+			flush();
+			next_message(&imsg);
+			show_monitor(imsg.data);
+			imsg_free(&imsg);
+			sleep(1);
+		}
 		break;
 	case LOG_VERBOSE:
 		verbose = 1;
@@ -242,62 +277,45 @@ connected:
 		done = 1;
 		break;
 	default:
-		errx(1, "unknown request (%d)", res->action);
+		errx(1, "unknown request (%d)", action);
 	}
 
-	while (ibuf->w.queued)
-		if (msgbuf_write(&ibuf->w) < 0)
-			err(1, "write error");
-
 	while (!done) {
-		if ((n = imsg_read(ibuf)) == -1)
-			errx(1, "imsg_read error");
-		if (n == 0)
-			errx(1, "pipe closed");
 
-		while (!done) {
-			if ((n = imsg_get(ibuf, &imsg)) == -1)
-				errx(1, "imsg_get error");
-			if (n == 0)
-				break;
-			switch(res->action) {
-			case REMOVE:
-			case SCHEDULE:
-			case SCHEDULE_ALL:
-			case SHUTDOWN:
-			case PAUSE_MDA:
-			case PAUSE_MTA:
-			case PAUSE_SMTP:
-			case RESUME_MDA:
-			case RESUME_MTA:
-			case RESUME_SMTP:
-			case LOG_VERBOSE:
-			case LOG_BRIEF:
-				done = show_command_output(&imsg);
-				break;
-			case SHOW_STATS:
-				show_stats_output();
-				done = 1;
-				break;
-			case NONE:
-				break;
-			case UPDATE_MAP:
-				break;
-			case MONITOR:
-				show_monitor(imsg.data);
-				imsg_free(&imsg);
-				sleep(1);
-				goto again;
-				break;
-			default:
-				err(1, "unexpected reply (%d)", res->action);
-			}
-			/* insert imsg replies switch here */
+		flush();
+		next_message(&imsg);
+
+		switch(action) {
+		case REMOVE:
+		case SCHEDULE:
+		case SCHEDULE_ALL:
+		case SHUTDOWN:
+		case PAUSE_MDA:
+		case PAUSE_MTA:
+		case PAUSE_SMTP:
+		case RESUME_MDA:
+		case RESUME_MTA:
+		case RESUME_SMTP:
+		case LOG_VERBOSE:
+		case LOG_BRIEF:
+			done = show_command_output(&imsg);
+			break;
+		case SHOW_STATS:
+			show_stats_output();
+			done = 1;
+			break;
+		case NONE:
+			break;
+		case UPDATE_MAP:
+			break;
+		case MONITOR:
 
-			imsg_free(&imsg);
+		default:
+			err(1, "unexpected reply (%d)", action);
 		}
+
+		imsg_free(&imsg);
 	}
-	close(ctl_sock);
 	free(ibuf);
 
 	return (0);
@@ -324,76 +342,57 @@ show_stats_output(void)
 {
 	struct stat_kv	kv, *kvp;
 	struct imsg	imsg;
-	int		n;
 	time_t		duration;
 
 	bzero(&kv, sizeof kv);
 
-again:
-	imsg_compose(ibuf, IMSG_STATS_GET, 0, 0, -1, &kv, sizeof kv);
-	
-	while (ibuf->w.queued)
-		if (msgbuf_write(&ibuf->w) < 0)
-			err(1, "write error");
+	while(1) {
+		imsg_compose(ibuf, IMSG_STATS_GET, 0, 0, -1, &kv, sizeof kv);
+		flush();
+		next_message(&imsg);
+		if (imsg.hdr.type != IMSG_STATS_GET)
+			errx(1, "invalid imsg type");
 
-	do {
-		if ((n = imsg_read(ibuf)) == -1)
-			errx(1, "imsg_read error");
-		if (n == 0)
-			errx(1, "pipe closed");
+		kvp = imsg.data;
+		if (kvp->iter == NULL) {
+			imsg_free(&imsg);
+			return;
+		}
 
-		do {
-			if ((n = imsg_get(ibuf, &imsg)) == -1)
-				errx(1, "imsg_get error");
-			if (n == 0)
+		if (strcmp(kvp->key, "uptime") == 0) {
+			duration = time(NULL) - kvp->val.u.counter;
+			printf("uptime=%zd\n", (size_t)duration); 
+			printf("uptime.human=%s\n",
+			    duration_to_text(duration));
+		}
+		else {
+			switch (kvp->val.type) {
+			case STAT_COUNTER:
+				printf("%s=%zd\n",
+				    kvp->key, kvp->val.u.counter);
+				break;
+			case STAT_TIMESTAMP:
+				printf("%s=%" PRId64 "\n",
+				    kvp->key, (int64_t)kvp->val.u.timestamp);
+				break;
+			case STAT_TIMEVAL:
+				printf("%s=%zd.%zd\n",
+				    kvp->key, kvp->val.u.tv.tv_sec,
+				    kvp->val.u.tv.tv_usec);
+				break;
+			case STAT_TIMESPEC:
+				printf("%s=%li.%06li\n",
+				    kvp->key,
+				    kvp->val.u.ts.tv_sec * 1000000 +
+				    kvp->val.u.ts.tv_nsec / 1000000,
+				    kvp->val.u.ts.tv_nsec % 1000000);
 				break;
-			if (imsg.hdr.type != IMSG_STATS_GET)
-				errx(1, "invalid imsg type");
-
-			kvp = imsg.data;
-			if (kvp->iter == NULL) {
-				imsg_free(&imsg);
-				return;
-			}
-
-			if (strcmp(kvp->key, "uptime") == 0) {
-				duration = time(NULL) - kvp->val.u.counter;
-				printf("uptime=%zd\n", (size_t)duration); 
-				printf("uptime.human=%s\n",
-				    duration_to_text(duration));
-			} else {
-				switch (kvp->val.type) {
-				case STAT_COUNTER:
-					printf("%s=%zd\n",
-					    kvp->key, kvp->val.u.counter);
-					break;
-				case STAT_TIMESTAMP:
-					printf("%s=%" PRId64 "\n",
-					    kvp->key, (int64_t)kvp->val.u.timestamp);
-					break;
-				case STAT_TIMEVAL:
-					printf("%s=%zd.%zd\n",
-					    kvp->key, kvp->val.u.tv.tv_sec,
-					    kvp->val.u.tv.tv_usec);
-					break;
-				case STAT_TIMESPEC:
-					printf("%s=%li.%06li\n",
-					    kvp->key,
-					    kvp->val.u.ts.tv_sec * 1000000 +
-					    kvp->val.u.ts.tv_nsec / 1000000,
-					    kvp->val.u.ts.tv_nsec % 1000000);
-					break;
-				}
 			}
+		}
 
-			kv = *kvp;
-
-			imsg_free(&imsg);
-			goto again;
-
-		} while (n != 0);
-
-	} while (n != 0);
+		kv = *kvp;
+		imsg_free(&imsg);
+	}
 }
 
 static void