Make setproctitle(NULL) restore all of the original arguments

(if it's able).

2 files changed, 45 insertions, 30 deletions

diff --git a/lib/libutil/setproctitle.3 b/lib/libutil/setproctitle.3
index 2d85b29..10eae65 100644
--- a/lib/libutil/setproctitle.3
+++ b/lib/libutil/setproctitle.3
@@ -54,23 +54,12 @@ argument.
 .Pp
 If
 .Va fmt
-is NULL, the process title is reset to simply the name of the executable.
-.\" The following requests should be uncommented and used where appropriate.
-.\" This next request is for sections 2 and 3 function return values only.
-.\" .Sh RETURN VALUES
-.\" This next request is for sections 1, 6, 7 & 8 only
-.\" .Sh ENVIRONMENT
-.\" .Sh FILES
+is NULL, the process title is restored.
 .Sh EXAMPLES
 To set the title on a daemon to indicate its activity:
 .Bd -literal -offset indent
 setproctitle("talking to %s", inet_ntoa(addr));
 .Ed
-.\" This next request is for sections 1, 6, 7 & 8 only
-.\"     (command return values (to shell) and fprintf/stderr type diagnostics)
-.\" .Sh DIAGNOSTICS
-.\" The next request is for sections 2 and 3 error and signal handling only.
-.\" .Sh ERRORS
 .Sh SEE ALSO
 .Xr ps 1 ,
 .Xr w 1 ,
@@ -110,4 +99,3 @@ stole the idea from the
 .Sy "Sendmail 8.7.3"
 source code by
 .An Eric Allman Aq eric@sendmail.org .
-.\" .Sh BUGS
diff --git a/lib/libutil/setproctitle.c b/lib/libutil/setproctitle.c
index 36c21db..37b9fab 100644
--- a/lib/libutil/setproctitle.c
+++ b/lib/libutil/setproctitle.c
@@ -59,8 +59,6 @@ struct old_ps_strings {
 #define SPT_BUFSIZE 2048	/* from other parts of sendmail */
 extern char * __progname;	/* is this defined in a .h anywhere? */
 
-static struct ps_strings *ps_strings;
-
 void
 #if defined(__STDC__)
 setproctitle(const char *fmt, ...)
@@ -70,8 +68,14 @@ setproctitle(fmt, va_alist)
 	va_dcl
 #endif
 {
+	static struct ps_strings *ps_strings;
 	static char buf[SPT_BUFSIZE];
-	static char *ps_argv[2];
+	static char obuf[SPT_BUFSIZE];
+	static char **oargv, *kbuf;
+	static int oargc = -1;
+	static char *nargv[2] = { buf, NULL };
+	char **nargvp;
+	int nargc;
 	va_list ap;
 	size_t len;
 	unsigned long ul_ps_strings;
@@ -83,11 +87,11 @@ setproctitle(fmt, va_alist)
 	va_start(ap);
 #endif
 
-	buf[sizeof(buf) - 1] = '\0';
 	if (fmt) {
+		buf[sizeof(buf) - 1] = '\0';
 
 		/* print program name heading for grep */
-		(void) snprintf(buf, sizeof(buf) - 1, "%s: ", __progname);
+		(void) snprintf(buf, sizeof(buf), "%s: ", __progname);
 
 		/*
 		 * can't use return from sprintf, as that is the count of how
@@ -97,11 +101,19 @@ setproctitle(fmt, va_alist)
 		len = strlen(buf);
 
 		/* print the argument string */
-		(void) vsnprintf(buf + len, sizeof(buf) - 1 - len, fmt, ap);
-	} else {
-		/* Idea from NetBSD - reset the title on fmt == NULL */
-		strncpy(buf, __progname, sizeof(buf) - 1);
-	}
+		(void) vsnprintf(buf + len, sizeof(buf) - len, fmt, ap);
+
+		nargvp = nargv;
+		nargc = 1;
+		kbuf = buf;
+	} else if (*obuf != '\0') {
+  		/* Idea from NetBSD - reset the title on fmt == NULL */
+		nargvp = oargv;
+		nargc = oargc;
+		kbuf = obuf;
+	} else
+		/* Nothing to restore */
+		return;
 
 	va_end(ap);
 
@@ -110,7 +122,7 @@ setproctitle(fmt, va_alist)
 	oid[1] = KERN_PROC;
 	oid[2] = KERN_PROC_ARGS;
 	oid[3] = getpid();
-	sysctl(oid, 4, 0, 0, buf, strlen(buf) + 1);
+	sysctl(oid, 4, 0, 0, kbuf, strlen(kbuf) + 1);
 
 	if (ps_strings == NULL) {
 		len = sizeof(ul_ps_strings);
@@ -123,13 +135,28 @@ setproctitle(fmt, va_alist)
 	/* PS_STRINGS points to zeroed memory on a style #2 kernel */
 	if (ps_strings->ps_argvstr) {
 		/* style #3 */
-		ps_argv[0] = buf;
-		ps_argv[1] = NULL;
-		ps_strings->ps_nargvstr = 1;
-		ps_strings->ps_argvstr = ps_argv;
+		if (oargc == -1) {
+			/* Record our original args */
+			oargc = ps_strings->ps_nargvstr;
+			oargv = ps_strings->ps_argvstr;
+			for (nargc = len = 0; nargc < oargc; nargc++) {
+				snprintf(obuf + len, sizeof(obuf) - len, "%s%s",
+				    len ? " " : "", oargv[nargc]);
+				if (len)
+					len++;
+				len += strlen(oargv[nargc]);
+				if (len >= sizeof(obuf))
+					break;
+			}
+		}
+		ps_strings->ps_nargvstr = nargc;
+		ps_strings->ps_argvstr = nargvp;
 	} else {
-		/* style #2 */
+		/* style #2 - we can only restore our first arg :-( */
+		if (*obuf == '\0')
+			strncpy(obuf, OLD_PS_STRINGS->old_ps_argvstr,
+			    sizeof(obuf) - 1);
 		OLD_PS_STRINGS->old_ps_nargvstr = 1;
-		OLD_PS_STRINGS->old_ps_argvstr = buf;
+		OLD_PS_STRINGS->old_ps_argvstr = nargvp[0];
 	}
 }