Add pledge support to pstat

The filemode(), ttymode() and vnodemode() functions can be split into two
pieces. The *_prep() piece must be hoisted to before pledge and the rest
can run under pledge "stdio rpath vminfo". The magic block that decides
which ones of the *_prep() pieces must be run with which combination of
command line switches is quite impressive.

Patch from Rob Pierce, thanks!

ok deraadt

1 files changed, 118 insertions, 75 deletions

diff --git a/usr.sbin/pstat/pstat.c b/usr.sbin/pstat/pstat.c
index 858569cb488..1c79db35ac0 100644
--- a/usr.sbin/pstat/pstat.c
+++ b/usr.sbin/pstat/pstat.c
@@ -1,4 +1,4 @@
-/*	$OpenBSD: pstat.c,v 1.109 2016/09/19 20:08:12 tb Exp $	*/
+/*	$OpenBSD: pstat.c,v 1.110 2016/09/19 20:10:54 tb Exp $	*/
 /*	$NetBSD: pstat.c,v 1.27 1996/10/23 22:50:06 cgd Exp $	*/
 
 /*-
@@ -82,6 +82,8 @@ struct nlist vnodenl[] = {
 	{ NULL }
 };
 
+struct itty *globalitp;
+struct kinfo_file *kf;
 struct nlist *globalnl;
 
 struct e_vnode {
@@ -89,10 +91,15 @@ struct e_vnode {
 	struct vnode vnode;
 };
 
-int	need_nlist;
-int	usenumflag;
-int	totalflag;
 int	kflag;
+int	totalflag;
+int	usenumflag;
+int	hideroot;
+int	maxfile;
+int	need_nlist;
+int	nfile;
+int	ntty;
+int	numvnodes;
 char	*nlistf	= NULL;
 char	*memf	= NULL;
 kvm_t	*kd = NULL;
@@ -112,27 +119,29 @@ kvm_t	*kd = NULL;
 	}
 
 void	filemode(void);
+void	filemodeprep(void);
 struct mount *
 	getmnt(struct mount *);
 struct e_vnode *
-	kinfo_vnodes(int *);
+	kinfo_vnodes(void);
 void	mount_print(struct mount *);
 void	nfs_header(void);
 int	nfs_print(struct vnode *);
 void	swapmode(void);
 void	ttymode(void);
+void	ttymodeprep(void);
 void	ttyprt(struct itty *);
 void	tty2itty(struct tty *tp, struct itty *itp);
 void	ufs_header(void);
 int	ufs_print(struct vnode *);
 void	ext2fs_header(void);
 int	ext2fs_print(struct vnode *);
-void	usage(void);
+static void __dead	usage(void);
 void	vnode_header(void);
 void	vnode_print(struct vnode *, struct vnode *);
 void	vnodemode(void);
+void	vnodemodeprep(void);
 
-int	hideroot;
 
 int
 main(int argc, char *argv[])
@@ -191,7 +200,7 @@ main(int argc, char *argv[])
 	if ((dformat == NULL && argc > 0) || (dformat && argc == 0))
 		usage();
 
-	need_nlist = vnodeflag || dformat;
+	need_nlist = dformat || vnodeflag;
 
 	if (nlistf != NULL || memf != NULL) {
 		if (fileflag || totalflag)
@@ -203,6 +212,25 @@ main(int argc, char *argv[])
 		    O_RDONLY | (need_nlist ? 0 : KVM_NO_FILES), buf)) == 0)
 			errx(1, "kvm_openfiles: %s", buf);
 
+	if (need_nlist)
+		if (kvm_nlist(kd, vnodenl) == -1)
+			errx(1, "kvm_nlist: %s", kvm_geterr(kd));
+
+	if (!(fileflag | vnodeflag | ttyflag | swapflag | totalflag || dformat))
+		usage();
+
+	if(!dformat) {
+		if (fileflag || totalflag)
+			filemodeprep();
+		if (vnodeflag || totalflag)
+			vnodemodeprep();
+		if (ttyflag)
+			ttymodeprep();
+	}
+
+	if (pledge("stdio rpath vminfo", NULL) == -1)
+		err(1, "pledge");
+
 	if (dformat) {
 		struct nlist *nl;
 		int longformat = 0, stringformat = 0, error = 0, n;
@@ -312,15 +340,9 @@ main(int argc, char *argv[])
 		for (i = 0; i < argc; i++)
 			free(nl[i].n_name);
 		free(nl);
-		exit(error);
+		return error;
 	}
 
-	if (need_nlist)
-		if (kvm_nlist(kd, vnodenl) == -1)
-			errx(1, "kvm_nlist: %s", kvm_geterr(kd));
-
-	if (!(fileflag | vnodeflag | ttyflag | swapflag | totalflag || dformat))
-		usage();
 	if (fileflag || totalflag)
 		filemode();
 	if (vnodeflag || totalflag)
@@ -329,7 +351,7 @@ main(int argc, char *argv[])
 		ttymode();
 	if (swapflag || totalflag)
 		swapmode();
-	exit(0);
+	return 0;
 }
 
 void
@@ -338,11 +360,10 @@ vnodemode(void)
 	struct e_vnode *e_vnodebase, *endvnode, *evp;
 	struct vnode *vp;
 	struct mount *maddr, *mp = NULL;
-	int numvnodes;
 
 	globalnl = vnodenl;
 
-	e_vnodebase = kinfo_vnodes(&numvnodes);
+	e_vnodebase = kinfo_vnodes();
 	if (totalflag) {
 		(void)printf("%7d vnodes\n", numvnodes);
 		return;
@@ -399,6 +420,21 @@ vnodemode(void)
 }
 
 void
+vnodemodeprep(void)
+{
+	int mib[2];
+	size_t num;
+
+	if (kd == NULL) {
+		mib[0] = CTL_KERN;
+		mib[1] = KERN_NUMVNODES;
+		num = sizeof(numvnodes);
+		if (sysctl(mib, 2, &numvnodes, &num, NULL, 0) < 0)
+			err(1, "sysctl(KERN_NUMVNODES) failed");
+	}
+}
+
+void
 vnode_header(void)
 {
 	(void)printf("%*s TYP VFLAG  USE HOLD", 2 * (int)sizeof(long), "ADDR");
@@ -792,24 +828,16 @@ mount_print(struct mount *mp)
  * simulate what a running kernel does in kinfo_vnode
  */
 struct e_vnode *
-kinfo_vnodes(int *avnodes)
+kinfo_vnodes(void)
 {
 	struct mntlist kvm_mountlist;
 	struct mount *mp, mount;
 	struct vnode *vp, vnode;
 	char *vbuf, *evbuf, *bp;
-	int mib[2], numvnodes;
 	size_t num;
 
-	if (kd == 0) {
-		mib[0] = CTL_KERN;
-		mib[1] = KERN_NUMVNODES;
-		num = sizeof(numvnodes);
-		if (sysctl(mib, 2, &numvnodes, &num, NULL, 0) < 0)
-			err(1, "sysctl(KERN_NUMVNODES) failed");
-	} else
+	if (kd != NULL)
 		KGET(V_NUMV, numvnodes);
-	*avnodes = numvnodes;
 	if (totalflag)
 		return NULL;
 	if ((vbuf = calloc(numvnodes + 20,
@@ -837,7 +865,7 @@ kinfo_vnodes(int *avnodes)
 			num++;
 		}
 	}
-	*avnodes = num;
+	numvnodes = num;
 	return ((struct e_vnode *)vbuf);
 }
 
@@ -866,32 +894,17 @@ ttymode(void)
 {
 	struct ttylist_head tty_head;
 	struct tty *tp, tty;
-	int mib[3], ntty, i;
-	struct itty itty, *itp;
-	size_t nlen;
+	int i;
+	struct itty itty;
 
-	if (!need_nlist) {
-		mib[0] = CTL_KERN;
-		mib[1] = KERN_TTYCOUNT;
-		nlen = sizeof(ntty);
-		if (sysctl(mib, 2, &ntty, &nlen, NULL, 0) < 0)
-			err(1, "sysctl(KERN_TTYCOUNT) failed");
-	} else
+	if (need_nlist)
 		KGET(TTY_NTTY, ntty);
 	(void)printf("%d terminal device%s\n", ntty, ntty == 1 ? "" : "s");
 	(void)printf("%s", hdr);
 	if (!need_nlist) {
-		mib[0] = CTL_KERN;
-		mib[1] = KERN_TTY;
-		mib[2] = KERN_TTY_INFO;
-		if ((itp = reallocarray(NULL, ntty, sizeof(struct itty))) == NULL)
-			err(1, "malloc");
-		nlen = ntty * sizeof(struct itty);
-		if (sysctl(mib, 3, itp, &nlen, NULL, 0) < 0)
-			err(1, "sysctl(KERN_TTY_INFO) failed");
 		for (i = 0; i < ntty; i++)
-			ttyprt(&itp[i]);
-		free(itp);
+			ttyprt(&globalitp[i]);
+		free(globalitp);
 	} else {
 		KGET(TTY_TTYLIST, tty_head);
 		for (tp = TAILQ_FIRST(&tty_head); tp;
@@ -903,6 +916,30 @@ ttymode(void)
 	}
 }
 
+void
+ttymodeprep(void)
+{
+	int mib[3];
+	size_t nlen;
+
+	if (!need_nlist) {
+		mib[0] = CTL_KERN;
+		mib[1] = KERN_TTYCOUNT;
+		nlen = sizeof(ntty);
+		if (sysctl(mib, 2, &ntty, &nlen, NULL, 0) < 0)
+			err(1, "sysctl(KERN_TTYCOUNT) failed");
+
+		mib[0] = CTL_KERN;
+		mib[1] = KERN_TTY;
+		mib[2] = KERN_TTY_INFO;
+		if ((globalitp = reallocarray(NULL, ntty, sizeof(struct itty))) == NULL)
+			err(1, "malloc");
+		nlen = ntty * sizeof(struct itty);
+		if (sysctl(mib, 3, globalitp, &nlen, NULL, 0) < 0)
+			err(1, "sysctl(KERN_TTY_INFO) failed");
+	}
+}
+
 struct {
 	int flag;
 	char val;
@@ -976,28 +1013,12 @@ ttyprt(struct itty *tp)
 void
 filemode(void)
 {
-	struct kinfo_file *kf;
 	char flagbuf[16], *fbp;
 	static char *dtypes[] = { "???", "inode", "socket", "pipe", "kqueue", "???", "???" };
-	int mib[2], maxfile, nfile;
-	size_t len;
 
 	globalnl = vnodenl;
 
-	if (nlistf == NULL && memf == NULL) {
-		mib[0] = CTL_KERN;
-		mib[1] = KERN_MAXFILES;
-		len = sizeof(maxfile);
-		if (sysctl(mib, 2, &maxfile, &len, NULL, 0) < 0)
-			err(1, "sysctl(KERN_MAXFILES) failed");
-		if (totalflag) {
-			mib[0] = CTL_KERN;
-			mib[1] = KERN_NFILES;
-			len = sizeof(nfile);
-			if (sysctl(mib, 2, &nfile, &len, NULL, 0) < 0)
-				err(1, "sysctl(KERN_NFILES) failed");
-		}
-	} else {
+	if (nlistf != NULL || memf != NULL) {
 		KGET(FNL_MAXFILE, maxfile);
 		if (totalflag) {
 			KGET(FNL_NFILE, nfile);
@@ -1006,14 +1027,6 @@ filemode(void)
 		}
 	}
 
-	if (!totalflag) {
-		kf = kvm_getfiles(kd, KERN_FILE_BYFILE, 0, sizeof *kf, &nfile);
-		if (kf == NULL) {
-			warnx("kvm_getfiles: %s", kvm_geterr(kd));
-			return;
-		}
-	}
-
 	(void)printf("%d/%d open files\n", nfile, maxfile);
 	if (totalflag)
 		return;
@@ -1059,6 +1072,36 @@ filemode(void)
 	}
 }
 
+void
+filemodeprep(void)
+{
+	int mib[2];
+	size_t len;
+
+	if (nlistf == NULL && memf == NULL) {
+		mib[0] = CTL_KERN;
+		mib[1] = KERN_MAXFILES;
+		len = sizeof(maxfile);
+		if (sysctl(mib, 2, &maxfile, &len, NULL, 0) < 0)
+			err(1, "sysctl(KERN_MAXFILES) failed");
+		if (totalflag) {
+			mib[0] = CTL_KERN;
+			mib[1] = KERN_NFILES;
+			len = sizeof(nfile);
+			if (sysctl(mib, 2, &nfile, &len, NULL, 0) < 0)
+				err(1, "sysctl(KERN_NFILES) failed");
+		}
+	}
+
+	if (!totalflag) {
+		kf = kvm_getfiles(kd, KERN_FILE_BYFILE, 0, sizeof *kf, &nfile);
+		if (kf == NULL) {
+			warnx("kvm_getfiles: %s", kvm_geterr(kd));
+			return;
+		}
+	}
+}
+
 /*
  * swapmode is based on a program called swapinfo written
  * by Kevin Lahey <kml@rokkaku.atl.ga.us>.
@@ -1152,7 +1195,7 @@ swapmode(void)
 	}
 }
 
-void
+static void __dead
 usage(void)
 {
 	(void)fprintf(stderr, "usage: "