2006-03-06 Martin Hunt <hunt@redhat.com>

	* librelay.c (start_cmd): Set proper uid/gid before execing
	command.
	(system_cmd): New function.
	(cleanup_and_exit): Wait for any child processes to complete.
	(stp_main_loop): Recognize STP_SYSTEM message.

	* stpd.c (main): Add support for "-u username".

3 files changed, 100 insertions, 28 deletions

diff --git a/runtime/stpd/ChangeLog b/runtime/stpd/ChangeLog
index 3875aecd..dfaf199b 100644
--- a/runtime/stpd/ChangeLog
+++ b/runtime/stpd/ChangeLog
@@ -1,3 +1,13 @@
+2006-03-06  Martin Hunt  <hunt@redhat.com>
+
+	* librelay.c (start_cmd): Set proper uid/gid before execing
+	command.
+	(system_cmd): New function.
+	(cleanup_and_exit): Wait for any child processes to complete.
+	(stp_main_loop): Recognize STP_SYSTEM message.
+
+	* stpd.c (main): Add support for "-u username".
+
 2006-02-25  Martin Hunt  <hunt@redhat.com>
 
 	* librelay.c (init_stp): Better error handling and cleanup.
diff --git a/runtime/stpd/librelay.c b/runtime/stpd/librelay.c
index 0ef15acd..cbc5b5a5 100644
--- a/runtime/stpd/librelay.c
+++ b/runtime/stpd/librelay.c
@@ -16,7 +16,7 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *
  * Copyright (C) IBM Corporation, 2005
- * Copyright (C) Red Hat Inc, 2005
+ * Copyright (C) Red Hat Inc, 2005, 2006
  *
  */
 
@@ -88,6 +88,10 @@ extern int target_pid;
 extern int driver_pid;
 extern char *target_cmd;
 
+/* uid/gid to use when execing external programs */
+extern uid_t cmd_uid;
+extern gid_t cmd_gid;
+
 /* per-cpu buffer info */
 static struct buf_status
 {
@@ -422,6 +426,12 @@ void start_cmd(void)
 		perror ("fork");
 		exit(-1);
 	} else if (pid == 0) {
+		if (setregid(cmd_gid, cmd_gid) < 0) {
+			perror("setregid");
+		}
+		if (setreuid(cmd_uid, cmd_uid) < 0) {
+			perror("setreuid");
+		}
 		/* wait here until signaled */
 		signal(SIGUSR1, sig_usr);
 		sigemptyset(&nullmask);
@@ -433,12 +443,32 @@ void start_cmd(void)
 		sigprocmask(SIG_SETMASK, &oldmask, NULL);
 		if (execl("/bin/sh", "sh", "-c", target_cmd, NULL) < 0)
 			perror(target_cmd);
-		exit(-1);
+		_exit(-1);
 	}
 
 	target_pid = pid;
 }
 
+void system_cmd(char *cmd)
+{
+	pid_t pid;
+
+	dbug ("system %s\n", cmd);
+	if ((pid = fork()) < 0) {
+		perror ("fork");
+	} else if (pid == 0) {
+		if (setregid(cmd_gid, cmd_gid) < 0) {
+			perror("setregid");
+		}
+		if (setreuid(cmd_uid, cmd_uid) < 0) {
+			perror("setreuid");
+		}
+		if (execl("/bin/sh", "sh", "-c", cmd, NULL) < 0)
+			perror(cmd);
+		_exit(-1);
+	}
+}
+
 #include <sys/wait.h>
 static void cleanup_and_exit (int);
 
@@ -466,12 +496,12 @@ int init_stp(const char *relay_filebase, int print_summary)
         modoptions[2] = buf;
         /* modoptions[3...N] set by command line parser. */
 
-	if ((pid = vfork()) < 0) {
-		perror ("vfork");
+	if ((pid = fork()) < 0) {
+		perror ("fork");
 		exit(-1);
 	} else if (pid == 0) {
 		if (execvp("/sbin/insmod",  modoptions) < 0)
-			exit(-1);
+			_exit(-1);
 	}
 	if (waitpid(pid, &rstatus, 0) < 0) {
 		perror("waitpid");
@@ -598,14 +628,20 @@ static int merge_output(void)
 static void cleanup_and_exit (int closed)
 {
 	char tmpbuf[128];
+	pid_t err;
 
 	if (exiting)
 		return;
-
 	exiting = 1;
 
 	dbug("CLEANUP AND EXIT  closed=%d mode=%d\n", closed, transport_mode);
 
+	/* what about child processes? we will wait for them here. */
+	err = waitpid(-1, NULL, WNOHANG);
+	if (err >= 0)
+		fprintf(stderr,"\nWaititing for processes to exit\n");
+	while(wait(NULL) > 0);
+
 	if (transport_mode == STP_TRANSPORT_RELAYFS) {
 		kill_percpu_threads(ncpus);
 		while(1) {
@@ -634,7 +670,7 @@ static void cleanup_and_exit (int closed)
 	close(control_channel);
 
 	if (!closed) {
-		snprintf(tmpbuf, sizeof(tmpbuf), "/sbin/rmmod %s", modname);
+		snprintf(tmpbuf, sizeof(tmpbuf), "/sbin/rmmod -w %s", modname);
 		if (system(tmpbuf)) {
 			fprintf(stderr, "ERROR: couldn't rmmod probe module %s.  No output will be written.\n",
 				modname);
@@ -644,27 +680,27 @@ static void cleanup_and_exit (int closed)
 	exit(0);
 }
 
-static void sigproc(int signum __attribute__((unused)))
+static void sigproc(int signum)
 {
-  signal(SIGINT, sigproc);
-  signal(SIGTERM, sigproc);
-  send_request(STP_EXIT, NULL, 0);
+	if (signum == SIGCHLD) {
+		pid_t pid = waitpid(-1, NULL, WNOHANG);
+		if (pid != target_pid)
+			return;
+	}
+	send_request(STP_EXIT, NULL, 0);
 }
 
 static void driver_poll (int signum __attribute__((unused)))
 {
-  /* See if the driver process is still alive.  If not, time to exit.  */
-  if (kill (driver_pid, 0) < 0)
-    {
-      send_request(STP_EXIT, NULL, 0);
-      return;
-    }
-  else
-    {
-      /* Check again later. */
-      signal (SIGALRM, driver_poll);
-      alarm (10); // any reasonable poll interval
-    }
+	/* See if the driver process is still alive.  If not, time to exit.  */
+	if (kill (driver_pid, 0) < 0) {
+		send_request(STP_EXIT, NULL, 0);
+		return;
+	} else  {
+		/* Check again later. Use any reasonable poll interval */
+		signal (SIGALRM, driver_poll);
+		alarm (10); 
+	}
 }
 
 
@@ -770,8 +806,14 @@ int stp_main_loop(void)
 				kill (target_pid, SIGUSR1);
 			break;
 		}
+		case STP_SYSTEM:
+		{
+			struct cmd_info *c = (struct cmd_info *)data;
+			system_cmd(c->cmd);
+			break;
+		}
 		default:
-			fprintf(stderr, "WARNING: ignored netlink message of type %d\n", (type));
+			fprintf(stderr, "WARNING: ignored message of type %d\n", (type));
 		}
 	}
 	fclose(ofp);
diff --git a/runtime/stpd/stpd.c b/runtime/stpd/stpd.c
index c6cc913d..bb1a81e0 100644
--- a/runtime/stpd/stpd.c
+++ b/runtime/stpd/stpd.c
@@ -25,6 +25,7 @@
 #include <unistd.h>
 #include <strings.h>
 #include <sys/wait.h>
+#include <pwd.h>
 #include "librelay.h"
 
 extern char *optarg;
@@ -45,6 +46,9 @@ char *modpath = NULL;
 char *modoptions[MAXMODOPTIONS];
 char *target_cmd = NULL;
 char *outfile_name = NULL;
+char *username = NULL;
+uid_t cmd_uid;
+gid_t cmd_gid;
 
  /* relayfs base file name */
 static char stpd_filebase[1024];
@@ -72,6 +76,7 @@ static void usage(char *prog)
 	fprintf(stderr, "-t pid.  Sets _stp_target to pid.\n");
 	fprintf(stderr, "-d pid.  Pass the systemtap driver's pid.\n");
 	fprintf(stderr, "-o FILE. Send output to FILE.\n");
+	fprintf(stderr, "-u username. Run commands as username.\n");
 	fprintf(stderr, "-b buffer size. The systemtap module will specify a buffer size.\n");
 	fprintf(stderr, "   Setting one here will override that value. The value should be\n");
 	fprintf(stderr, "   an integer between 1 and 64 which be assumed to be the\n");
@@ -84,7 +89,7 @@ int main(int argc, char **argv)
 	int c, status;
 	pid_t pid;
 
-	while ((c = getopt(argc, argv, "mpqrb:n:t:d:c:vo:")) != EOF)
+	while ((c = getopt(argc, argv, "mpqrb:n:t:d:c:vo:u:")) != EOF)
 	{
 		switch (c) {
 		case 'm':
@@ -126,6 +131,9 @@ int main(int argc, char **argv)
 		case 'o':
 			outfile_name = optarg;
 			break;
+		case 'u':
+			username = optarg;
+			break;
 		default:
 			usage(argv[0]);
 		}
@@ -167,15 +175,27 @@ int main(int argc, char **argv)
 		usage(argv[0]);
 	}
 
+	if (username) {
+		struct passwd *pw = getpwnam(username);
+		if (!pw) {
+			fprintf(stderr, "Cannot find user \"%s\".\n", username);
+			exit(1);
+		}
+		cmd_uid = pw->pw_uid;
+		cmd_gid = pw->pw_gid;  
+	} else {
+		cmd_uid = getuid();
+		cmd_gid = getgid();
+	}
 
 	if (enable_relayfs) {
 		/* now run the _stp_check script */
-		if ((pid = vfork()) < 0) {
-			perror ("vfork");
+		if ((pid = fork()) < 0) {
+			perror ("fork of stp_check failed.");
 			exit(-1);
 		} else if (pid == 0) {
 			if (execlp(stp_check, stp_check, NULL) < 0)
-				exit (-1);
+				_exit (-1);
 		}
 		if (waitpid(pid, &status, 0) < 0) {
 			perror("waitpid");