2007-08-14 David Smith <dsmith@redhat.com>

	Merge from setuid-branch.  Changes also by Martin Hunt
	<hunt@redhat.com>.

	* staprun.c (init_staprun): Drop CAP_SYS_ADMIN when we're done
	with it.
	(main): Calls parse_modpath instead of path_parse_modname.  Just
	call parse_modpath with argv[optind]. Let it allocate and set
	modpath and modname.  If no modulename was given, display usage
	and exit.  Drop CAP_SYS_NICE when we're done with it.  Set
	atexit(exit_cleanup) so cleanup always gets called and modules get
	removed.  Call handle_symbols.
	(run_stapio): Set argv[0] to stapio so that it executes as itself
	instead of staprun.
	(cleanup): Only do cleanups once and only try to remove module
	when appropriate.
	(exit_cleanup): New. Calls cleanup().
	(mountfs): Sets uid to root before making directory and then
	restores uid.
	(setup_ctl_channel): Uses DEBUGFS define and improved
	error message.
	(setup_relayfs): Ditto.
	(setup_oldrelayfs): Uses DEBUGFS and RELAYFS defines.
	(run_stp_check): Replaced by mountfs().
	(mountfs): New function.  Replaces an external script with C code.
	(init_staprun): Calls mountfs() instead of run_stp_check().

	* staprun.h: Renamed path_parse_modname to parse_modpath.  Added
	MODULE_NAME_LEN define.  Added [_][p]err macros.  Removed
	VERSION_CMD.

	* mainloop.c (cleanup_and_exit): Make sure initialized is 2
	before exiting with code 2.
	(stp_main_loop): Set initialized to 2 when STP_TRANSPORT
	is received. Call cleanup_and_exit() with proper status.
	(start_cmd): exit 1 instead of -1.
	(system_cmd): Ditto.
	(init_staprun): Renamed init_stapio.
	(cleanup_and_exit): Set exit status.

	* cap.c: New file.
	* common.c: New file.
	* stapio.c: New file.
	* staprun_funcs.c: New file.
	* Makefile: Removed.

	* symbols.c (get_sections): Move the filter code up so that
	uninteresting section names are filtered out before
	attempting to open them.
	(do_kernel_symbols): Better detect overfow conditions and realloc
	new space.
	(do_module): After sending all modules, send a null message to
	indicate we are finished.

	* ctl.c (init_ctl_channel): When attempting to attach, if the
	control channel doesn't exist, print a better error message.

	* relay_old.c (init_oldrelayfs): Errors out if
	open_relayfs_files() couldn't open any files.

	PR 4795
	* mainloop.c (send_request): Fixed buffer overflow check.
	* staprun.h: Added buffer overflow checking versions of
	strcpy/sprintf/snprintf.
	* common.c (path_parse_modname): Checks for overflows on
	strcpy/sprintf/snprintf.
	(read_buffer_info): Ditto.
	* ctl.c (init_ctl_channel): Ditto.
	* relay.c (init_relayfs): Ditto.
	* relay_old.c (open_relayfs_files): Ditto.
	(init_oldrelayfs): Ditto.
	* staprun_funcs.c (insert_module): Ditto.
	(check_path): Ditto.
	* symbols.c (get_sections): Ditto.

12 files changed, 1596 insertions, 483 deletions

diff --git a/runtime/staprun/ChangeLog b/runtime/staprun/ChangeLog
index aefd87bf..de522180 100644
--- a/runtime/staprun/ChangeLog
+++ b/runtime/staprun/ChangeLog
@@ -1,3 +1,79 @@
+2007-08-14  David Smith  <dsmith@redhat.com>
+
+	Merge from setuid-branch.  Changes also by Martin Hunt
+	<hunt@redhat.com>.
+
+	* staprun.c (init_staprun): Drop CAP_SYS_ADMIN when we're done
+	with it.
+	(main): Calls parse_modpath instead of path_parse_modname.  Just
+	call parse_modpath with argv[optind]. Let it allocate and set
+	modpath and modname.  If no modulename was given, display usage
+	and exit.  Drop CAP_SYS_NICE when we're done with it.  Set
+	atexit(exit_cleanup) so cleanup always gets called and modules get
+	removed.  Call handle_symbols.
+	(run_stapio): Set argv[0] to stapio so that it executes as itself
+	instead of staprun.
+	(cleanup): Only do cleanups once and only try to remove module
+	when appropriate.
+	(exit_cleanup): New. Calls cleanup().
+	(mountfs): Sets uid to root before making directory and then
+	restores uid.
+	(setup_ctl_channel): Uses DEBUGFS define and improved
+	error message.
+	(setup_relayfs): Ditto.
+	(setup_oldrelayfs): Uses DEBUGFS and RELAYFS defines.
+	(run_stp_check): Replaced by mountfs().
+	(mountfs): New function.  Replaces an external script with C code.
+	(init_staprun): Calls mountfs() instead of run_stp_check().
+
+	* staprun.h: Renamed path_parse_modname to parse_modpath.  Added
+	MODULE_NAME_LEN define.  Added [_][p]err macros.  Removed
+	VERSION_CMD.
+
+	* mainloop.c (cleanup_and_exit): Make sure initialized is 2
+	before exiting with code 2.
+	(stp_main_loop): Set initialized to 2 when STP_TRANSPORT
+	is received. Call cleanup_and_exit() with proper status.
+	(start_cmd): exit 1 instead of -1.
+	(system_cmd): Ditto.
+	(init_staprun): Renamed init_stapio.
+	(cleanup_and_exit): Set exit status.
+
+	* cap.c: New file.
+	* common.c: New file.
+	* stapio.c: New file.
+	* staprun_funcs.c: New file.
+	* Makefile: Removed.
+
+	* symbols.c (get_sections): Move the filter code up so that 
+	uninteresting section names are filtered out before
+	attempting to open them.
+	(do_kernel_symbols): Better detect overfow conditions and realloc
+	new space.
+	(do_module): After sending all modules, send a null message to
+	indicate we are finished.
+
+	* ctl.c (init_ctl_channel): When attempting to attach, if the
+	control channel doesn't exist, print a better error message.
+
+	* relay_old.c (init_oldrelayfs): Errors out if
+	open_relayfs_files() couldn't open any files.
+
+	PR 4795
+	* mainloop.c (send_request): Fixed buffer overflow check.
+	* staprun.h: Added buffer overflow checking versions of
+	strcpy/sprintf/snprintf.
+	* common.c (path_parse_modname): Checks for overflows on
+	strcpy/sprintf/snprintf.
+	(read_buffer_info): Ditto.
+	* ctl.c (init_ctl_channel): Ditto.
+	* relay.c (init_relayfs): Ditto.
+	* relay_old.c (open_relayfs_files): Ditto.
+	(init_oldrelayfs): Ditto.
+	* staprun_funcs.c (insert_module): Ditto.
+	(check_path): Ditto.
+	* symbols.c (get_sections): Ditto.
+
 2007-07-09  David Smith  <dsmith@redhat.com>
 
 	* relay.c (init_relayfs): Fixed a buffer size bug introduced by
diff --git a/runtime/staprun/cap.c b/runtime/staprun/cap.c
new file mode 100644
index 00000000..df4a7130
--- /dev/null
+++ b/runtime/staprun/cap.c
@@ -0,0 +1,158 @@
+/* -*- linux-c -*-
+ *
+ * cap.c - staprun capabilities functions
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
+ *
+ * Copyright (C) 2007 Red Hat, Inc.
+ *
+ */
+
+#include "staprun.h"
+#include <sys/prctl.h>
+
+/* like perror, but exits */
+#define ferror(msg) {						  \
+		_perr(msg);					  \
+		exit(1);					  \
+	}							  \
+		
+/*
+ * init_cap() sets up the initial capabilities for staprun. Then
+ * it calls prctl( PR_SET_KEEPCAPS) to arrrange to keep these capabilities
+ * even when not running as root. Next it resets the real, effective, and 
+ * saved uid and gid back to the normal user.
+ *
+ * There are two sets of capabilities we are concerned with; permitted
+ * and effective. The permitted capabilities are all the capabilities
+ * that this process is ever permitted to have. They are defined in init_cap()
+ * and may be permanently removed with drop_cap().
+ *
+ * Effective capabilities are the capabilities from the permitted set
+ * that are currently enabled. A good practice would be to only enable 
+ * capabilities when necessary and to delete or drop them as soon as possible.
+ *
+ * Capabilities we might use include:
+ *
+ * CAP_SYS_MODULE - insert and remove kernel modules
+ * CAP_SYS_ADMIN - misc, including mounting and unmounting
+ * CAP_SYS_NICE - setpriority()
+ * CAP_SETUID - allows setuid
+ * CAP_SETGID - allows setgid
+ * CAP_CHOWN - allows chown
+ */
+
+int init_cap(void)
+{
+	cap_t caps = cap_init();
+	cap_value_t capv[] = {CAP_SYS_MODULE, CAP_SYS_ADMIN, CAP_SYS_NICE, CAP_SETUID, CAP_SETGID};
+	const int numcaps = sizeof(capv) / sizeof(capv[0]);
+	uid_t uid = getuid();
+	gid_t gid = getgid();
+
+	cap_clear(caps);
+	if (caps == NULL)
+		ferror("cap_init");
+
+	if (cap_set_flag(caps, CAP_PERMITTED, numcaps, capv, CAP_SET) < 0)
+		ferror("cap_set_flag");
+
+	if (cap_set_proc(caps) < 0)
+		ferror("cap_set_proc");
+
+	cap_free(caps);
+
+	if (prctl(PR_SET_KEEPCAPS, 1, 0, 0, 0) < 0)
+		ferror("prctl");
+
+	if (setresuid(uid, uid, uid) < 0)
+		ferror("setresuid");
+
+	if (setresgid(gid, gid, gid) < 0)
+		ferror("setresgid");
+
+	return 1;
+}
+
+void print_cap(char *text)
+{
+	int p;
+	cap_t caps = cap_get_proc();
+	uid_t uid, euid, suid;
+	gid_t gid, egid, sgid;
+
+	if (caps == NULL) {
+		perr("cap_get_proc");
+		return;
+	}
+	
+	getresuid(&uid, &euid, &suid);
+	getresgid(&gid, &egid, &sgid);
+
+	printf("***** %s\n", text);
+
+	if ((p=prctl(PR_GET_KEEPCAPS, 0, 0, 0, 0)) < 0)
+		perr("Couldn't get PR_SET_KEEPCAPS flag value");
+	else 
+		printf("KEEPCAPS: %d\n", p);
+
+	printf("uid: %d, euid: %d, suid: %d\ngid: %d. egid: %d, sgid: %d\n", 
+	       uid, euid, suid, gid, egid, sgid );
+	printf("Caps: %s\n", cap_to_text(caps, NULL));
+	cap_free(caps);
+	printf("*****\n\n");
+}
+
+/* drop_cap() permanently removes a capability from the permitted set. There is
+ * no way to recover the capability after this.  You do not need to remove
+ * it from the effective set before calling this.
+ */
+void drop_cap(cap_value_t cap)
+{
+	cap_t caps = cap_get_proc();
+	if (caps == NULL)
+		ferror("cap_get_proc failed");
+	if (cap_set_flag(caps, CAP_PERMITTED, 1, &cap, CAP_CLEAR) < 0)
+		ferror("Could not clear effective capabilities");
+	if (cap_set_proc(caps) < 0)
+		ferror("Could not apply capability set");
+	cap_free(caps);
+}
+
+/* add_cap() adds a permitted capability to the effective set. */
+void add_cap(cap_value_t cap)
+{
+	cap_t caps = cap_get_proc();
+	if (caps == NULL)
+		ferror("cap_get_proc failed");
+	if (cap_set_flag(caps, CAP_EFFECTIVE, 1, &cap, CAP_SET) < 0)
+		ferror("Could not set effective capabilities");
+	if (cap_set_proc(caps) < 0)
+		ferror("Could not apply capability set");
+	cap_free(caps);
+}
+
+/* del_cap() deletes a permitted capability from the effective set. */
+void del_cap(cap_value_t cap)
+{
+	cap_t caps = cap_get_proc();
+	if (caps == NULL)
+		ferror("cap_get_proc failed");
+	if (cap_set_flag(caps, CAP_EFFECTIVE, 1, &cap, CAP_CLEAR) < 0)
+		ferror("Could not clear effective capabilities");
+	if (cap_set_proc(caps) < 0)
+		ferror("Could not apply capability set");
+	cap_free(caps);
+}
diff --git a/runtime/staprun/common.c b/runtime/staprun/common.c
new file mode 100644
index 00000000..cbe88be0
--- /dev/null
+++ b/runtime/staprun/common.c
@@ -0,0 +1,307 @@
+/* -*- linux-c -*-
+ *
+ * common.c - staprun suid/user common code
+ *
+ * This file is part of systemtap, and is free software.  You can
+ * redistribute it and/or modify it under the terms of the GNU General
+ * Public License (GPL); either version 2, or (at your option) any
+ * later version.
+ *
+ * Copyright (C) 2007 Red Hat Inc.
+ */
+
+#include "staprun.h"
+#include <sys/types.h>
+#include <unistd.h>
+#include <sys/utsname.h>
+
+/* variables needed by parse_args() */
+int verbose;
+int target_pid;
+unsigned int buffer_size;
+char *target_cmd;
+char *outfile_name;
+int attach_mod;
+int load_only;
+
+/* module variables */
+char *modname = NULL;
+char *modpath = "";
+char *modoptions[MAXMODOPTIONS];
+
+int initialized = 0;
+int control_channel = 0;
+
+void parse_args(int argc, char **argv)
+{
+	int c;
+
+	/* Initialize option variables. */
+	verbose = 0;
+	target_pid = 0;
+	buffer_size = 0;
+	target_cmd = NULL;
+	outfile_name = NULL;
+	attach_mod = 0;
+	load_only = 0;
+
+	while ((c = getopt(argc, argv, "ALvb:t:d:c:o:x:")) != EOF) {
+		switch (c) {
+		case 'v':
+			verbose++;
+			break;
+		case 'b':
+			buffer_size = (unsigned)atoi(optarg);
+			if (buffer_size < 1 || buffer_size > 64) {
+				err("Invalid buffer size '%d' (should be 1-64).\n", buffer_size);
+				usage(argv[0]);
+			}
+			break;
+		case 't':
+		case 'x':
+			target_pid = atoi(optarg);
+			break;
+		case 'd':
+			/* obsolete internal option used by stap */
+			break;
+		case 'c':
+			target_cmd = optarg;
+			break;
+		case 'o':
+			outfile_name = optarg;
+			break;
+		case 'A':
+			attach_mod = 1;
+			break;
+		case 'L':
+			load_only = 1;
+			break;
+		default:
+			usage(argv[0]);
+		}
+	}
+
+	if (attach_mod && load_only) {
+		err("You can't specify the '-A' and '-L' options together.\n");
+		usage(argv[0]);
+	}
+
+	if (attach_mod && buffer_size) {
+		err("You can't specify the '-A' and '-b' options together.  The '-b'\n"
+		    "buffer size option only has an effect when the module is inserted.\n");
+		usage(argv[0]);
+	}
+
+	if (attach_mod && target_cmd) {
+		err("You can't specify the '-A' and '-c' options together.  The '-c cmd'\n"
+		    "option used to start a command only has an effect when the module\n"
+		    "is inserted.\n");
+		usage(argv[0]);
+	}
+
+	if (attach_mod && target_pid) {
+		err("You can't specify the '-A' and '-x' options together.  The '-x pid'\n"
+		    "option only has an effect when the module is inserted.\n");
+		usage(argv[0]);
+	}
+}
+
+void usage(char *prog)
+{
+	err("\n%s [-v]  [-c cmd ] [-x pid] [-u user]\n"
+                "\t[-A|-L] [-b bufsize] [-o FILE] MODULE [module-options]\n", prog);
+	err("-v              Increase verbosity.\n");
+	err("-c cmd          Command \'cmd\' will be run and staprun will\n");
+	err("                exit when it does.  The '_stp_target' variable\n");
+	err("                will contain the pid for the command.\n");
+	err("-x pid          Sets the '_stp_target' variable to pid.\n");
+	err("-o FILE         Send output to FILE.\n");
+	err("-b buffer size  The systemtap module specifies a buffer size.\n");
+	err("                Setting one here will override that value.  The\n");
+	err("                value should be an integer between 1 and 64\n");
+	err("                which be assumed to be the buffer size in MB.\n");
+	err("                That value will be per-cpu in bulk mode.\n");
+	err("-L              Load module and start probes, then detach.\n");
+	err("-A              Attach to loaded systemtap module.\n");
+	err("MODULE can be either a module name or a module path.  If a\n");
+	err("module name is used, it is looked for in the following\n");
+	err("directory: /lib/modules/`uname -r`/systemtap\n");
+	exit(1);
+}
+
+/*
+ * parse_modpath.  Here's how this code interprets the global modpath:
+ *
+ * (1) If modpath contains a '/', it is assumed to be an absolute or
+ * relative file path to a module (such as "../foo.ko" or
+ * "/tmp/stapXYZ/stap_foo.ko").
+ *
+ * (2) If modpath doesn't contain a '/' and ends in '.ko', it is a file
+ * path to a module in the current directory (such as "foo.ko").
+ *
+ * (3) If modpath doesn't contain a '/' and doesn't end in '.ko', then
+ * it is a module name and the full pathname of the module is
+ * '/lib/modules/`uname -r`/systemtap/PATH.ko'.  For instance, if
+ * modpath was "foo", the full module pathname would be
+ * '/lib/modules/`uname -r`/systemtap/foo.ko'.
+ */
+void parse_modpath(const char *inpath)
+{
+	const char *mptr = rindex(inpath, '/');
+	char *ptr;
+
+	dbug(3, "inpath=%s\n", inpath);
+
+	/* If we couldn't find a '/', ... */
+	if (mptr == NULL) {
+		size_t plen = strlen(inpath);
+
+		/* If the path ends with the '.ko' file extension,
+		 * then we've got a module in the current
+		 * directory. */
+		if (plen > 3 && strcmp(&inpath[plen - 3], ".ko") == 0) {
+			mptr = inpath;
+			modpath = strdup(inpath);
+			if (!modpath) {
+				err("Memory allocation failed. Exiting.\n");
+				exit(1);
+			}
+		} else {
+			/* If we didn't find the '.ko' file extension, then
+			 * we've just got a module name, not a module path.
+			 * Look for the module in /lib/modules/`uname
+			 * -r`/systemtap. */
+
+			struct utsname utsbuf;
+			int len;
+			#define MODULE_PATH "/lib/modules/%s/systemtap/%s.ko"
+
+			/* First, we need to figure out what the
+			 * kernel version. */
+			if (uname(&utsbuf) != 0) {
+				perr("Unable to determine kernel version, uname failed");
+				exit(-1);
+			}
+
+			/* Build the module path, which will look like
+			 * '/lib/modules/KVER/systemtap/{path}.ko'. */
+			len = sizeof(MODULE_PATH) + sizeof(utsbuf.release) + strlen(inpath);
+			modpath = malloc(len);
+			if (!modpath) {
+				err("Memory allocation failed. Exiting.\n");
+				exit(1);
+			}
+			
+			if (snprintf_chk(modpath, len, MODULE_PATH, utsbuf.release, inpath))
+				exit(-1);
+
+			dbug(2, "modpath=\"%s\"\n", modpath);
+
+			mptr = rindex(modpath, '/');
+			mptr++;
+		}
+	} else {
+		/* We found a '/', so the module name starts with the next
+		 * character. */
+		mptr++;
+
+		modpath = strdup(inpath);
+		if (!modpath) {
+			err("Memory allocation failed. Exiting.\n");
+			exit(1);
+		}
+	}
+
+	modname = strdup(mptr);
+	if (!modname) {
+		err("Memory allocation failed. Exiting.\n");
+		exit(1);
+	}
+
+	ptr = rindex(modname, '.');
+	if (ptr)
+		*ptr = '\0';
+
+	/* We've finally got a real modname.  Make sure it isn't too
+	 * long.  If it is too long, init_module() will appear to
+	 * work, but the module can't be removed (because you end up
+	 * with control characters in the module name). */
+	if (strlen(modname) > MODULE_NAME_LEN) {
+		err("ERROR: Module name ('%s') is too long.\n", modname);
+		exit(1);
+	}
+}
+
+#define ERR_MSG "\nUNEXPECTED FATAL ERROR in staprun. Please file a bug report.\n"
+static void fatal_handler (int signum)
+{
+        int rc;
+        char *str = strsignal(signum);
+        rc = write (STDERR_FILENO, ERR_MSG, sizeof(ERR_MSG));
+        rc = write (STDERR_FILENO, str, strlen(str));
+        rc = write (STDERR_FILENO, "\n", 1);
+	if (initialized)
+		_exit(3);
+	else
+		_exit(1);
+}
+
+void setup_signals(void)
+{
+	sigset_t s;
+	struct sigaction a;
+
+	/* blocking all signals while we set things up */
+	sigfillset(&s);
+#ifdef SINGLE_THREADED
+	sigprocmask(SIG_SETMASK, &s, NULL);
+#else
+	pthread_sigmask(SIG_SETMASK, &s, NULL);
+#endif
+	/* set some of them to be ignored */
+	memset(&a, 0, sizeof(a));
+	sigfillset(&a.sa_mask);
+	a.sa_handler = SIG_IGN;
+	sigaction(SIGPIPE, &a, NULL);
+	sigaction(SIGUSR2, &a, NULL);
+
+	/* for serious errors, handle them in fatal_handler */
+	a.sa_handler = fatal_handler;
+	sigaction(SIGBUS, &a, NULL);
+	sigaction(SIGFPE, &a, NULL);
+	sigaction(SIGILL, &a, NULL);
+	sigaction(SIGSEGV, &a, NULL);
+	sigaction(SIGXCPU, &a, NULL);
+	sigaction(SIGXFSZ, &a, NULL);
+
+	/* unblock all signals */
+	sigemptyset(&s);
+
+#ifdef SINGLE_THREADED
+	sigprocmask(SIG_SETMASK, &s, NULL);
+#else
+	pthread_sigmask(SIG_SETMASK, &s, NULL);
+#endif
+}
+
+/**
+ *	send_request - send request to kernel over control channel
+ *	@type: the relay-app command id
+ *	@data: pointer to the data to be sent
+ *	@len: length of the data to be sent
+ *
+ *	Returns 0 on success, negative otherwise.
+ */
+int send_request(int type, void *data, int len)
+{
+	char buf[1024];
+
+	/* Before doing memcpy, make sure 'buf' is big enough. */
+	if ((len + 4) > (int)sizeof(buf)) {
+		_err("exceeded maximum send_request size.\n");
+		return -1;
+	}
+	memcpy(buf, &type, 4);
+	memcpy(&buf[4], data, len);
+	return write(control_channel, buf, len+4);
+}
diff --git a/runtime/staprun/ctl.c b/runtime/staprun/ctl.c
index 53c27190..72592bdf 100644
--- a/runtime/staprun/ctl.c
+++ b/runtime/staprun/ctl.c
@@ -12,62 +12,28 @@
 
 #include "staprun.h"
 
-/* This is only used in the old relayfs code */
-static void read_buffer_info(void)
-{
-	char buf[PATH_MAX];
-	struct statfs st;
-	int fd, len, ret;
-
-	if (!use_old_transport)
-		return;
-
- 	if (statfs("/sys/kernel/debug", &st) == 0 && (int) st.f_type == (int) DEBUGFS_MAGIC)
-		return;
-
-	sprintf (buf, "/proc/systemtap/%s/bufsize", modname);	
-	fd = open(buf, O_RDONLY);
-	if (fd < 0)
-		return;
-
-	len = read(fd, buf, sizeof(buf));
-	if (len <= 0) {
-		fprintf (stderr, "ERROR: couldn't read bufsize.\n");
-		close(fd);
-		return;
-	}
-	ret = sscanf(buf, "%u,%u", &n_subbufs, &subbuf_size);
-	if (ret != 2)
-		fprintf (stderr, "ERROR: couldn't read bufsize.\n");
-
-	dbug(2, "n_subbufs= %u, size=%u\n", n_subbufs, subbuf_size);
-	close(fd);
-	return;
-}
-
-
 int init_ctl_channel(void)
 {
 	char buf[PATH_MAX];
 	struct statfs st;
 
- 	if (statfs("/sys/kernel/debug", &st) == 0 && (int) st.f_type == (int) DEBUGFS_MAGIC)
- 		sprintf (buf, "/sys/kernel/debug/systemtap/%s/cmd", modname);
-	else
-		sprintf (buf, "/proc/systemtap/%s/cmd", modname);
-
+ 	if (statfs("/sys/kernel/debug", &st) == 0 && (int) st.f_type == (int) DEBUGFS_MAGIC) {
+		if (sprintf_chk(buf, "/sys/kernel/debug/systemtap/%s/cmd", modname))
+			return -1;
+	} else {
+		if (sprintf_chk(buf, "/proc/systemtap/%s/cmd", modname))
+			return -1;
+	}
+	
 	dbug(2, "Opening %s\n", buf); 
 	control_channel = open(buf, O_RDWR);
 	if (control_channel < 0) {
-		if (attach_mod) 
-			fprintf (stderr, "ERROR: Cannot connect to module \"%s\".\n", modname);
+		if (attach_mod && errno == ENOENT)
+			err("ERROR: Can not attach. Module %s not running.\n", modname);
 		else
-			fprintf (stderr, "ERROR: couldn't open control channel %s\n", buf);
-		fprintf (stderr, "errcode = %s\n", strerror(errno));
+			perr("Couldn't open control channel '%s'", buf);
 		return -1;
 	}
-
-	read_buffer_info();
 	return 0;
 }
 
diff --git a/runtime/staprun/mainloop.c b/runtime/staprun/mainloop.c
index 7b8c5e31..4acfb001 100644
--- a/runtime/staprun/mainloop.c
+++ b/runtime/staprun/mainloop.c
@@ -11,23 +11,12 @@
  */
 
 #include "staprun.h"
+#include <sys/utsname.h>
 
 /* globals */
-int control_channel = 0;
 int ncpus;
 int use_old_transport = 0;
 
-#define ERR_MSG "\nUNEXPECTED FATAL ERROR in staprun. Please file a bug report.\n"
-void fatal_handler (int signum)
-{
-        int rc;
-        char *str = strsignal(signum);
-        rc = write (STDERR_FILENO, ERR_MSG, sizeof(ERR_MSG));
-        rc = write (STDERR_FILENO, str, strlen(str));
-        rc = write (STDERR_FILENO, "\n", 1);
-        _exit(-1);
-}
-
 static void sigproc(int signum)
 {
 	dbug(2, "sigproc %d (%s)\n", signum, strsignal(signum));
@@ -60,56 +49,6 @@ static void setup_main_signals(int cleanup)
 	sigaction(SIGQUIT, &a, NULL);
 }
 
-void setup_signals(void)
-{
-	sigset_t s;
-	struct sigaction a;
-
-	/* blocking all signals while we set things up */
-	sigfillset(&s);
-	pthread_sigmask(SIG_SETMASK, &s, NULL);
-
-	/* set some of them to be ignored */
-	memset(&a, 0, sizeof(a));
-	sigfillset(&a.sa_mask);
-	a.sa_handler = SIG_IGN;
-	sigaction(SIGPIPE, &a, NULL);
-	sigaction(SIGUSR2, &a, NULL);
-
-	/* for serious errors, handle them in fatal_handler */
-	a.sa_handler = fatal_handler;
-	sigaction(SIGBUS, &a, NULL);
-	sigaction(SIGFPE, &a, NULL);
-	sigaction(SIGILL, &a, NULL);
-	sigaction(SIGSEGV, &a, NULL);
-	sigaction(SIGXCPU, &a, NULL);
-	sigaction(SIGXFSZ, &a, NULL);
-
-	/* unblock all signals */
-	sigemptyset(&s);
-	pthread_sigmask(SIG_SETMASK, &s, NULL);
-}
-
-
-/**
- *	send_request - send request to kernel over control channel
- *	@type: the relay-app command id
- *	@data: pointer to the data to be sent
- *	@len: length of the data to be sent
- *
- *	Returns 0 on success, negative otherwise.
- */
-int send_request(int type, void *data, int len)
-{
-	char buf[1024];
-	if (len > (int)sizeof(buf)) {
-		err("exceeded maximum send_request size.\n");
-		return -1;
-	}
-	memcpy(buf, &type, 4);
-	memcpy(&buf[4],data,len);
-	return write(control_channel, buf, len+4);
-}
 
 /* 
  * start_cmd forks the command given on the command line
@@ -131,23 +70,17 @@ void start_cmd(void)
 
 	dbug (1, "execing target_cmd %s\n", target_cmd);
 	if ((pid = fork()) < 0) {
-		perror ("fork");
-		exit(-1);
+		_perr("fork");
+		exit(1);
 	} else if (pid == 0) {
 		int signum;
 
-		if (setregid(cmd_gid, cmd_gid) < 0) {
-			perror("setregid");
-		}
-		if (setreuid(cmd_uid, cmd_uid) < 0) {
-			perror("setreuid");
-		}
 		/* wait here until signaled */
 		sigwait(&usrset, &signum);
 
 		if (execl("/bin/sh", "sh", "-c", target_cmd, NULL) < 0)
 			perror(target_cmd);
-		_exit(-1);
+		_exit(1);
 	}
 	target_pid = pid;
 }
@@ -156,8 +89,6 @@ void start_cmd(void)
  * system_cmd() executes system commands in response
  * to an STP_SYSTEM message from the module. These
  * messages are sent by the system() systemtap function.
- * uid and gid are set because staprun is running as root and 
- * it is best to run commands as the real user.
  */
 void system_cmd(char *cmd)
 {
@@ -165,57 +96,104 @@ void system_cmd(char *cmd)
 
 	dbug (2, "system %s\n", cmd);
 	if ((pid = fork()) < 0) {
-		perror ("fork");
+		_perr("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);
+			perr("%s", cmd);
+		_exit(1);
 	}
 }
 
+static int using_old_transport(void)
+{
+	struct utsname utsbuf;
+	int i;
+	long int kver[3];
+	char *start, *end;
 
-/* stp_check script */
-#ifdef PKGLIBDIR
-char *stp_check=PKGLIBDIR "/stp_check";
-#else
-char *stp_check="stp_check";
-#endif
+	if (uname(&utsbuf) != 0) {
+		_perr("Unable to determine kernel version, uname failed");
+		return -1;
+	}
+
+	start = utsbuf.release;
+	for (i = 0; i < 3; i++) {
+		errno = 0;
+		kver[i] = strtol(start, &end, 10);
+		if (errno != 0) {
+			_perr("Unable to parse kernel version, strtol failed");
+			return -1;
+		}
+		start = end;
+		start++;
+	}
 
-static int run_stp_check (void)
+	if (KERNEL_VERSION(kver[0], kver[1], kver[2])
+	    <= KERNEL_VERSION(2, 6, 15)) {
+		dbug(2, "Using OLD TRANSPORT\n");
+		return 1;
+	}
+	return 0;
+}
+
+/* This is only used in the old relayfs code */
+static void read_buffer_info(void)
 {
-	int ret ;
-	/* run the _stp_check script */
-	dbug(2, "executing %s\n", stp_check);
-	ret = system(stp_check);
-	return ret;
+	char buf[PATH_MAX];
+	struct statfs st;
+	int fd, len, ret;
+
+	if (!use_old_transport)
+		return;
+
+ 	if (statfs("/sys/kernel/debug", &st) == 0 && (int) st.f_type == (int) DEBUGFS_MAGIC)
+		return;
+
+	if (sprintf_chk(buf, "/proc/systemtap/%s/bufsize", modname))
+		return;
+	fd = open(buf, O_RDONLY);
+	if (fd < 0)
+		return;
+
+	len = read(fd, buf, sizeof(buf));
+	if (len <= 0) {
+		perr("Couldn't read bufsize");
+		close(fd);
+		return;
+	}
+	ret = sscanf(buf, "%u,%u", &n_subbufs, &subbuf_size);
+	if (ret != 2)
+		perr("Couldn't read bufsize");
+
+	dbug(2, "n_subbufs= %u, size=%u\n", n_subbufs, subbuf_size);
+	close(fd);
+	return;
 }
 
+
 /**
- *	init_stp - initialize the app
+ *	init_stapio - initialize the app
  *	@print_summary: boolean, print summary or not at end of run
  *
  *	Returns 0 on success, negative otherwise.
  */
-int init_staprun(void)
+int init_stapio(void)
 {
-	char bufcmd[128];
-	int rstatus;
-	int pid;
+	dbug(2, "init_stapio\n");
 
-	if (system(VERSION_CMD)) {
-		dbug(2, "Using OLD TRANSPORT\n");
-		use_old_transport = 1;
+	use_old_transport = using_old_transport();
+	if (use_old_transport < 0)
+		return -1;
+
+	/* create control channel */
+	if (init_ctl_channel() < 0) {
+		err("Failed to initialize control channel.\n");
+		return -1;
 	}
+	read_buffer_info();
 
 	if (attach_mod) {
-		if (init_ctl_channel() < 0)
-			return -1;
+		dbug(2, "Attaching\n");
 		if (use_old_transport) {
 			if (init_oldrelayfs() < 0) {
 				close_ctl_channel();
@@ -230,57 +208,22 @@ int init_staprun(void)
 		return 0;
 	}
 
-	if (run_stp_check() < 0)
-		return -1;
-
-	/* insert module */
-	sprintf(bufcmd, "_stp_bufsize=%d", buffer_size);
-        modoptions[0] = "insmod";
-        modoptions[1] = modpath;
-        modoptions[2] = bufcmd;
-        /* modoptions[3...N] set by command line parser. */
-
-	if ((pid = fork()) < 0) {
-		perror ("fork");
-		exit(-1);
-	} else if (pid == 0) {
-		if (execvp("/sbin/insmod",  modoptions) < 0)
-			_exit(-1);
-	}
-	if (waitpid(pid, &rstatus, 0) < 0) {
-		perror("waitpid");
-		exit(-1);
-	}
-	if (WIFEXITED(rstatus) && WEXITSTATUS(rstatus)) {
-		fprintf(stderr, "ERROR, couldn't insmod probe module %s\n", modpath);
-		return -1;
-	}
-	
-	/* create control channel */
-	if (init_ctl_channel() < 0) {
-		err("Failed to initialize control channel.\n");
-		goto exit1;
-	}
-
 	/* fork target_cmd if requested. */
 	/* It will not actually exec until signalled. */
 	if (target_cmd)
 		start_cmd();
 
 	return 0;
-
-exit1:
-	snprintf(bufcmd, sizeof(bufcmd), "/sbin/rmmod -w %s", modname);
-	if (system(bufcmd))
-		fprintf(stderr, "ERROR: couldn't rmmod probe module %s.\n", modname);
-	return -1;
 }
 
-
-
+/* cleanup_and_exit() closed channels and frees memory
+ * then exits with the following status codes:
+ * 1 - failed to initialize.
+ * 2 - disconnected
+ * 3 - initialized
+ */
 void cleanup_and_exit (int closed)
 {
-	char tmpbuf[128];
 	pid_t err;
 	static int exiting = 0;
 
@@ -295,7 +238,7 @@ void cleanup_and_exit (int closed)
 	/* what about child processes? we will wait for them here. */
 	err = waitpid(-1, NULL, WNOHANG);
 	if (err >= 0)
-		fprintf(stderr,"\nWaiting for processes to exit\n");
+		err("\nWaiting for processes to exit\n");
 	while(wait(NULL) > 0) ;
 
 	if (use_old_transport)
@@ -306,22 +249,16 @@ void cleanup_and_exit (int closed)
 	dbug(1, "closing control channel\n");
 	close_ctl_channel();
 
-	if (closed == 0) {
-		dbug(1, "removing module\n");
-		snprintf(tmpbuf, sizeof(tmpbuf), "/sbin/rmmod -w %s", modname);
-		if (system(tmpbuf)) {
-			fprintf(stderr, "ERROR: couldn't rmmod probe module %s.\n", modname);
-			exit(1);
-		}
-	} else if (closed == 2) {
-		fprintf(stderr, "\nDisconnecting from systemtap module.\n");
-		fprintf(stderr, "To reconnect, type \"staprun -A %s\"\n", modname); 
-	}
-
-	exit(0);
+	if (initialized == 2 && closed == 2) {
+		err("\nDisconnecting from systemtap module.\n"		\
+		    "To reconnect, type \"staprun -A %s\"\n", modname);
+	} else if (initialized)
+		closed = 3;
+	else
+		closed = 1;
+	exit(closed);
 }
 
-
 /**
  *	stp_main_loop - loop forever reading data
  */
@@ -342,13 +279,11 @@ int stp_main_loop(void)
 	while (1) { /* handle messages from control channel */
 		nb = read(control_channel, recvbuf, sizeof(recvbuf));
 		if (nb <= 0) {
-			if (errno != EINTR) {
-				perror("recv");
-				fprintf(stderr, "WARNING: unexpected EOF. nb=%ld\n", (long)nb);
-			}
+			if (errno != EINTR)
+				_perr("Unexpected EOF in read (nb=%ld)", (long)nb);
 			continue;
 		}
-
+		
 		type = *(int *)recvbuf;
 		data = (void *)(recvbuf + sizeof(int));
 
@@ -362,10 +297,8 @@ int stp_main_loop(void)
 				bw = write(out_fd[0], data, nb - sizeof(int));
 			}
 			if (bw != (nb - (ssize_t)sizeof(int))) {
-				perror("write");
-				fprintf(stderr,
-					"ERROR: write error. nb=%ld\n", (long)nb);
-				cleanup_and_exit(0);
+				_perr("write error (nb=%ld)", (long)nb);
+				cleanup_and_exit(1);
 			}
                         break;
 		}
@@ -388,7 +321,7 @@ int stp_main_loop(void)
 			if (t->res < 0) {
 				if (target_cmd)
 					kill (target_pid, SIGKILL);
-				cleanup_and_exit(0);
+				cleanup_and_exit(1);
 			} else if (target_cmd)
 				kill (target_pid, SIGUSR1);
 			break;
@@ -405,41 +338,20 @@ int stp_main_loop(void)
 			struct _stp_msg_start ts;
 			if (use_old_transport) {
 				if (init_oldrelayfs() < 0)
-					cleanup_and_exit(0);
+					cleanup_and_exit(1);
 			} else {
 				if (init_relayfs() < 0)
-					cleanup_and_exit(0);
+					cleanup_and_exit(1);
 			}
 			ts.target = target_pid;
+			initialized = 2;
 			send_request(STP_START, &ts, sizeof(ts));
 			if (load_only)
 				cleanup_and_exit(2);
 			break;
 		}
-		case STP_MODULE:
-		{
-			dbug(2, "STP_MODULES request received\n");
-			do_module(data);
-			break;
-		}		
-		case STP_SYMBOLS:
-		{
-			struct _stp_msg_symbol *req = (struct _stp_msg_symbol *)data;
-			dbug(2, "STP_SYMBOLS request received\n");
-			if (req->endian != 0x1234) {
-				fprintf(stderr,"ERROR: staprun is compiled with different endianess than the kernel!\n");
-				cleanup_and_exit(0);
-			}
-			if (req->ptr_size != sizeof(char *)) {
-				fprintf(stderr,"ERROR: staprun is compiled with %d-bit pointers and the kernel uses %d-bit.\n",
-					8*(int)sizeof(char *), 8*req->ptr_size);
-				cleanup_and_exit(0);
-			}
-			do_kernel_symbols();
-			break;
-		}
 		default:
-			fprintf(stderr, "WARNING: ignored message of type %d\n", (type));
+			err("WARNING: ignored message of type %d\n", (type));
 		}
 	}
 	fclose(ofp);
diff --git a/runtime/staprun/relay.c b/runtime/staprun/relay.c
index 84e80660..30c4ce1e 100644
--- a/runtime/staprun/relay.c
+++ b/runtime/staprun/relay.c
@@ -74,9 +74,8 @@ static void *reader_thread(void *data)
 		cpu_set_t cpu_mask;
 		CPU_ZERO(&cpu_mask);
 		CPU_SET(cpu, &cpu_mask);
-		if( sched_setaffinity( 0, sizeof(cpu_mask), &cpu_mask ) < 0 ) {
-			perror("sched_setaffinity");
-		}
+		if( sched_setaffinity( 0, sizeof(cpu_mask), &cpu_mask ) < 0 )
+			_perr("sched_setaffinity");
 #ifdef NEED_PPOLL
 		/* Without a real ppoll, there is a small race condition that could */
 		/* block ppoll(). So use a timeout to prevent that. */
@@ -95,15 +94,14 @@ static void *reader_thread(void *data)
                 if (rc < 0) {
 			dbug(3, "cpu=%d poll=%d errno=%d\n", cpu, rc, errno);
                         if (errno != EINTR) {
-				fprintf(stderr, "poll error: %s\n",strerror(errno));
+				_perr("poll error");
 				return(NULL);
                         }
 			stop_threads = 1;
                 }
 		while ((rc = read(relay_fd[cpu], buf, sizeof(buf))) > 0) {
 			if (write(out_fd[cpu], buf, rc) != rc) {
-				fprintf(stderr, "Couldn't write to output fd %d for cpu %d, exiting: errcode = %d: %s\n", 
-					out_fd[cpu], cpu, errno, strerror(errno));
+				perr("Couldn't write to output %d for cpu %d, exiting.", out_fd[cpu], cpu);
 				return(NULL);
 			}
 		}
@@ -124,16 +122,21 @@ int init_relayfs(void)
 	char rqbuf[128];
 	char buf[PATH_MAX], relay_filebase[PATH_MAX];
 
-	dbug(1, "initializing relayfs\n");
+	dbug(2, "initializing relayfs\n");
 
 	reader[0] = (pthread_t)0;
 	relay_fd[0] = 0;
 	out_fd[0] = 0;
 
- 	if (statfs("/sys/kernel/debug", &st) == 0 && (int) st.f_type == (int) DEBUGFS_MAGIC)
- 		sprintf(relay_filebase, "/sys/kernel/debug/systemtap/%s", modname);
+ 	if (statfs("/sys/kernel/debug", &st) == 0
+	    && (int) st.f_type == (int) DEBUGFS_MAGIC) {
+		if (sprintf_chk(relay_filebase,
+				"/sys/kernel/debug/systemtap/%s",
+				modname))
+			return -1;
+	}
  	else {
-		fprintf(stderr,"Cannot find relayfs or debugfs mount point.\n");
+		err("Cannot find relayfs or debugfs mount point.\n");
 		return -1;
 	}
 
@@ -141,7 +144,8 @@ int init_relayfs(void)
 		bulkmode = 1;
 
 	for (i = 0; i < NR_CPUS; i++) {
-		sprintf(buf, "%s/trace%d", relay_filebase, i);
+		if (sprintf_chk(buf, "%s/trace%d", relay_filebase, i))
+			return -1;
 		dbug(2, "attempting to open %s\n", buf);
 		relay_fd[i] = open(buf, O_RDONLY | O_NONBLOCK);
 		if (relay_fd[i] < 0)
@@ -151,12 +155,12 @@ int init_relayfs(void)
 	dbug(2, "ncpus=%d, bulkmode = %d\n", ncpus, bulkmode);
 
 	if (ncpus == 0) {
-		err("couldn't open %s.\n", buf);
+		_err("couldn't open %s.\n", buf);
 		return -1;
 	}
 	if (ncpus > 1 && bulkmode == 0) {
-		err("ncpus=%d, bulkmode = %d\n", ncpus, bulkmode);
-		err("This is inconsistent! Please file a bug report. Exiting now.\n");
+		_err("ncpus=%d, bulkmode = %d\n", ncpus, bulkmode);
+		_err("This is inconsistent! Please file a bug report. Exiting now.\n");
 		return -1;
 	}
 
@@ -164,16 +168,20 @@ int init_relayfs(void)
 		for (i = 0; i < ncpus; i++) {
 			if (outfile_name) {
 				/* special case: for testing we sometimes want to write to /dev/null */
-				if (strcmp(outfile_name, "/dev/null") == 0)
-					strcpy(buf, outfile_name);
-				else
-					sprintf(buf, "%s_%d", outfile_name, i);
-			} else
-				sprintf(buf, "stpd_cpu%d", i);
+				if (strcmp(outfile_name, "/dev/null") == 0) {
+					strcpy(buf, "/dev/null");
+				} else {
+					if (sprintf_chk(buf, "%s_%d", outfile_name, i))
+						return -1;
+				}
+			} else {
+				if (sprintf_chk(buf, "stpd_cpu%d", i))
+					return -1;
+			}
 			
 			out_fd[i] = open (buf, O_CREAT|O_TRUNC|O_WRONLY, 0666);
 			if (out_fd[i] < 0) {
-				fprintf(stderr, "ERROR: couldn't open output file %s.\n", buf);
+				perr("Couldn't open output file %s", buf);
 				return -1;
 			}
 		}
@@ -182,7 +190,7 @@ int init_relayfs(void)
 		if (outfile_name) {
 			out_fd[0] = open (outfile_name, O_CREAT|O_TRUNC|O_WRONLY, 0666);
 			if (out_fd[0] < 0) {
-				fprintf(stderr, "ERROR: couldn't open output file %s.\n", outfile_name);
+				perr("Couldn't open output file %s", outfile_name);
 				return -1;
 			}
 		} else
@@ -191,9 +199,9 @@ int init_relayfs(void)
 	}
 	dbug(2, "starting threads\n");
 	for (i = 0; i < ncpus; i++) {
-		if (pthread_create(&reader[i], NULL, reader_thread, (void *)(long)i) < 0) {
-			fprintf(stderr, "failed to create thread\n");
-			perror("Error creating thread");
+		if (pthread_create(&reader[i], NULL, reader_thread,
+				   (void *)(long)i) < 0) {
+			_perr("failed to create thread");
 			return -1;
 		}
 	}		
diff --git a/runtime/staprun/relay_old.c b/runtime/staprun/relay_old.c
index 19f57788..3f65acbb 100644
--- a/runtime/staprun/relay_old.c
+++ b/runtime/staprun/relay_old.c
@@ -82,32 +82,40 @@ static int open_relayfs_files(int cpu, const char *relay_filebase, const char *p
 	memset(&status[cpu], 0, sizeof(struct buf_status));
 	status[cpu].info.cpu = cpu;
 
-	sprintf(tmp, "%s%d", relay_filebase, cpu);
+	if (sprintf_chk(tmp, "%s%d", relay_filebase, cpu))
+		return -1;
+	dbug(2, "Opening %s.\n", tmp); 
 	relay_fd[cpu] = open(tmp, O_RDONLY | O_NONBLOCK);
 	if (relay_fd[cpu] < 0) {
 		relay_fd[cpu] = 0;
 		return 0;
 	}
 
-	sprintf(tmp, "%s%d", proc_filebase, cpu);
+	if (sprintf_chk(tmp, "%s%d", proc_filebase, cpu))
+		goto err1;
 	dbug(2, "Opening %s.\n", tmp); 
 	proc_fd[cpu] = open(tmp, O_RDWR | O_NONBLOCK);
 	if (proc_fd[cpu] < 0) {
-		fprintf(stderr, "ERROR: couldn't open proc file %s: errcode = %s\n", tmp, strerror(errno));
+		perr("Couldn't open proc file %s", tmp);
 		goto err1;
 	}
 
 	if (outfile_name) {
-		/* special case: for testing we sometimes want to write to /dev/null */
-		if (strcmp(outfile_name, "/dev/null") == 0)
-			strcpy(tmp, outfile_name);
-		else
-			sprintf(tmp, "%s_%d", outfile_name, cpu);
-	} else
-		sprintf(tmp, "stpd_cpu%d", cpu);	
+		/* special case: for testing we sometimes want to
+		 * write to /dev/null */
+		if (strcmp(outfile_name, "/dev/null") == 0) {
+			strcpy(tmp, "/dev/null");
+		} else {
+			if (sprintf_chk(tmp, "%s_%d", outfile_name, cpu))
+				goto err1;
+		}
+	} else {
+		if (sprintf_chk(tmp, "stpd_cpu%d", cpu))
+			goto err1;
+	}
 
 	if((percpu_tmpfile[cpu] = fopen(tmp, "w+")) == NULL) {
-		fprintf(stderr, "ERROR: Couldn't open output file %s: errcode = %s\n", tmp, strerror(errno));
+		perr("Couldn't open output file %s", tmp);
 		goto err2;
 	}
 
@@ -117,12 +125,14 @@ static int open_relayfs_files(int cpu, const char *relay_filebase, const char *p
 				 0);
 	if(relay_buffer[cpu] == MAP_FAILED)
 	{
-		fprintf(stderr, "ERROR: couldn't mmap relay file, total_bufsize (%d) = subbuf_size (%d) * n_subbufs(%d), error = %s \n", (int)total_bufsize, (int)subbuf_size, (int)n_subbufs, strerror(errno));
+		_perr("Couldn't mmap relay file, total_bufsize (%d)"	\
+		     "= subbuf_size (%d) * n_subbufs(%d)",
+		     (int)total_bufsize, (int)subbuf_size, (int)n_subbufs);
 		goto err3;
 	}
-
+	
 	return 1;
-
+	
 err3:
 	fclose(percpu_tmpfile[cpu]);
 err2:
@@ -157,7 +167,7 @@ static int process_subbufs(struct _stp_buf_info *info)
 		len = (subbuf_size - sizeof(padding)) - padding;
 		if (len) {
 			if (fwrite_unlocked (subbuf_ptr, len, 1, percpu_tmpfile[cpu]) != 1) {
-				fprintf(stderr, "ERROR: couldn't write to output file for cpu %d, exiting: errcode = %d: %s\n", cpu, errno, strerror(errno));
+				_perr("Couldn't write to output file for cpu %d, exiting:", cpu);
 				exit(1);
 			}
 		}
@@ -181,9 +191,8 @@ static void *reader_thread(void *data)
 
 	CPU_ZERO(&cpu_mask);
 	CPU_SET(cpu, &cpu_mask);
-	if( sched_setaffinity( 0, sizeof(cpu_mask), &cpu_mask ) < 0 ) {
-		perror("sched_setaffinity");
-	}
+	if( sched_setaffinity( 0, sizeof(cpu_mask), &cpu_mask ) < 0 )
+		_perr("sched_setaffinity");
 
 	pollfd.fd = relay_fd[cpu];
 	pollfd.events = POLLIN;
@@ -192,12 +201,10 @@ static void *reader_thread(void *data)
 		rc = poll(&pollfd, 1, -1);
 		if (rc < 0) {
 			if (errno != EINTR) {
-				fprintf(stderr, "ERROR: poll error: %s\n",
-					strerror(errno));
+				_perr("poll error");
 				exit(1);
 			}
-			fprintf(stderr, "WARNING: poll warning: %s\n",
-				strerror(errno));
+			err("WARNING: poll warning: %s\n", strerror(errno));
 			rc = 0;
 		}
 
@@ -210,7 +217,7 @@ static void *reader_thread(void *data)
 			consumed_info.cpu = cpu;
 			consumed_info.consumed = subbufs_consumed;
 			if (write (proc_fd[cpu], &consumed_info, sizeof(struct _stp_consumed_info)) < 0)
-				fprintf(stderr,"WARNING: writing consumed info failed.\n");
+				perr("writing consumed info failed");
 		}
 		if (status[cpu].info.flushing)
 			pthread_exit(NULL);
@@ -237,7 +244,7 @@ int init_oldrelayfs(void)
 		if (outfile_name) {
 			out_fd[0] = open (outfile_name, O_CREAT|O_TRUNC|O_WRONLY, 0666);
 			if (out_fd[0] < 0) {
-				fprintf(stderr, "ERROR: couldn't open output file %s.\n", outfile_name);
+				perr("Couldn't open output file '%s'", outfile_name);
 				return -1;
 			}
 		} else
@@ -245,14 +252,23 @@ int init_oldrelayfs(void)
 	  return 0;
 	}
 
- 	if (statfs("/sys/kernel/debug", &st) == 0 && (int) st.f_type == (int) DEBUGFS_MAGIC) {
- 		sprintf(relay_filebase, "/sys/kernel/debug/systemtap/%s/trace", modname);
- 		sprintf(proc_filebase, "/sys/kernel/debug/systemtap/%s/", modname);
-	} else if (statfs("/mnt/relay", &st) == 0 && (int) st.f_type == (int) RELAYFS_MAGIC) {
- 		sprintf(relay_filebase, "/mnt/relay/systemtap/%s/trace", modname);
- 		sprintf(proc_filebase, "/proc/systemtap/%s/", modname);
+ 	if (statfs("/sys/kernel/debug", &st) == 0
+	    && (int) st.f_type == (int) DEBUGFS_MAGIC) {
+		if (sprintf_chk(relay_filebase,
+				"/sys/kernel/debug/systemtap/%s/trace",
+				modname))
+			return -1;
+ 		if (sprintf_chk(proc_filebase, 
+				"/sys/kernel/debug/systemtap/%s/", modname))
+			return -1;
+	} else if (statfs("/mnt/relay", &st) == 0
+		   && (int) st.f_type == (int) RELAYFS_MAGIC) {
+		if (sprintf_chk(relay_filebase, "/mnt/relay/systemtap/%s/trace", modname))
+			return -1;
+		if (sprintf_chk(proc_filebase, "/proc/systemtap/%s/", modname))
+			return -1;
  	} else {
-		fprintf(stderr,"Cannot find relayfs or debugfs mount point.\n");
+		err("Cannot find relayfs or debugfs mount point.\n");
 		return -1;
 	}
 
@@ -266,7 +282,7 @@ int init_oldrelayfs(void)
 		if (ret == 0)
 			break;
 		if (ret < 0) {
-			fprintf(stderr, "ERROR: couldn't open relayfs files, cpu = %d\n", i);
+			err("ERROR: couldn't open relayfs files, cpu = %d\n", i);
 			goto err;
 		}
 	}
@@ -274,11 +290,18 @@ int init_oldrelayfs(void)
 	ncpus = i;
 	dbug(2, "ncpus=%d\n", ncpus);
 
+	if (ncpus == 0) {
+		err("couldn't open relayfs files.\n");
+		return -1;
+	}
+
 	for (i = 0; i < ncpus; i++) {
 		/* create a thread for each per-cpu buffer */
 		if (pthread_create(&reader[i], NULL, reader_thread, (void *)(long)i) < 0) {
+			int saved_errno = errno;
 			close_relayfs_files(i);
-			fprintf(stderr, "ERROR: Couldn't create reader thread, cpu = %d\n", i);
+			err("ERROR: Couldn't create reader thread, cpu = %d: %s\n",
+				i, strerror(saved_errno));
 			goto err;
 		}
 	}
diff --git a/runtime/staprun/stapio.c b/runtime/staprun/stapio.c
new file mode 100644
index 00000000..696167af
--- /dev/null
+++ b/runtime/staprun/stapio.c
@@ -0,0 +1,69 @@
+/* -*- linux-c -*-
+ *
+ * stapio.c - SystemTap module io handler.
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
+ *
+ * Copyright (C) 2005-2007 Red Hat, Inc.
+ *
+ */
+
+#include "staprun.h"
+#include <pwd.h>
+char *__name__ = "stapio";
+
+int main(int argc, char **argv)
+{
+	setup_signals();
+	
+	parse_args(argc, argv);
+
+	if (buffer_size)
+		dbug(1, "Using a buffer of %u bytes.\n", buffer_size);
+
+	if (optind < argc) {
+		parse_modpath(argv[optind++]);
+		dbug(2, "modpath=\"%s\", modname=\"%s\"\n", modpath, modname);
+	}
+
+        if (optind < argc) {
+		if (attach_mod) {
+			err("ERROR: Cannot have module options with attach (-A).\n");
+			usage(argv[0]);
+		} else {
+			unsigned start_idx = 3; /* reserve three slots in modoptions[] */
+			while (optind < argc && start_idx+1 < MAXMODOPTIONS)
+				modoptions[start_idx++] = argv[optind++];
+			modoptions[start_idx] = NULL;
+		}
+	}
+
+	if (modpath == NULL || *modpath == '\0') {
+		err("ERROR: Need a module name or path to load.\n");
+		usage(argv[0]);
+	}
+
+	if (init_stapio())
+		exit(1);
+	
+	initialized = 1;
+
+	if (stp_main_loop()) {
+		err("ERROR: Couldn't enter main loop. Exiting.\n");
+		exit(1);
+	}
+
+	return 0;
+}
diff --git a/runtime/staprun/staprun.c b/runtime/staprun/staprun.c
index f40f676d..d9738cd7 100644
--- a/runtime/staprun/staprun.c
+++ b/runtime/staprun/staprun.c
@@ -21,166 +21,198 @@
  */
 
 #include "staprun.h"
-#include <pwd.h>
-
-extern char *optarg;
-extern int optopt;
-extern int optind;
-
-int verbose = 0;
-int target_pid = 0;
-unsigned int buffer_size = 0;
-char modname[128];
-char *modpath = NULL;
-#define MAXMODOPTIONS 64
-char *modoptions[MAXMODOPTIONS];
-char *target_cmd = NULL;
-char *outfile_name = NULL;
-char *username = NULL;
-uid_t cmd_uid;
-gid_t cmd_gid;
-int attach_mod = 0;
-int load_only = 0;
-
-static void path_parse_modname (char *path)
+
+int inserted_module = 0;
+
+/* used in dbug, _err and _perr */
+char *__name__ = "staprun";
+
+extern long delete_module(const char *, unsigned int);
+
+static int
+run_as(uid_t uid, gid_t gid, const char *path, char *const argv[])
 {
-	char *mptr = rindex (path, '/');
-	if (mptr == NULL) 
-		mptr = path;
-	else
-		mptr++;
+	pid_t pid;
+	int rstatus;
 
-	if (strlen(mptr) >= sizeof(modname)) {
-		err("Module name larger than modname buffer.\n");
-		exit (-1);
+	if ((pid = fork()) < 0) {
+		_perr("fork");
+		return -1;
 	}
-	strcpy(modname, mptr);			
+	else if (pid == 0) {
+		/* Make sure we run as the full user.  If we're
+		 * switching to a non-root user, this won't allow
+		 * that process to switch back to root (since the
+		 * original process is setuid). */
+		if (uid != getuid()) {
+			if (do_cap(CAP_SETGID, setresgid, gid, gid, gid) < 0) {
+				_perr("setresgid");
+				exit(1);
+			}
+			if (do_cap(CAP_SETUID, setresuid, uid, uid, uid) < 0) {
+				_perr("setresuid");
+				exit(1);
+			}
+		}
+
+		/* Actually run the command. */
+		if (execv(path, argv) < 0)
+			perror(path);
+		_exit(1);
+	}
+
+	if (waitpid(pid, &rstatus, 0) < 0)
+		return -1;
+
+	if (WIFEXITED(rstatus))
+		return WEXITSTATUS(rstatus);
+	return -1;
+}
+
+/* Keep the uid and gid settings because we will likely */
+/* conditionally restore "-u" */
+static int run_stapio(char **argv)
+{
+	uid_t uid = getuid();
+	gid_t gid = getgid();
+	argv[0] = PKGLIBDIR "/stapio";
+
+	if (verbose >= 2) {
+		int i = 0;
+		err("execing: ");
+		while (argv[i]) {
+			err("%s ", argv[i]);
+			i++;
+		}
+		err("\n");		
+	}
+	return run_as(uid, gid, argv[0], argv);
+}
+
+
+int init_staprun(void)
+{
+	dbug(2, "init_staprun\n");
+
+	if (mountfs() < 0)
+		return -1;
+
+	/* We're done with CAP_SYS_ADMIN. */
+	drop_cap(CAP_SYS_ADMIN);
+ 
+	if (!attach_mod) {
+		if (insert_module() < 0)
+			return -1;
+		else
+			inserted_module = 1;
+	}
+	
+	return 0;
+}
 	
-	mptr = rindex(modname, '.');
-	if (mptr)
-		*mptr = '\0';
+static void cleanup(int rc)
+{
+	/* Only cleanup once. */
+	static int done = 0;
+	if (done == 0)
+		done = 1;
+	else
+		return;
+
+	dbug(2, "rc=%d, inserted_module=%d\n", rc, inserted_module);
+
+	/* rc == 2 means disconnected */
+	if (rc == 2)
+		return;
+
+	/* If we inserted the module and did not get rc==2, then */
+	/* we really want to remove it. */
+	if (inserted_module || rc == 3) {
+		long ret;
+		dbug(2, "removing module %s\n", modname);
+		ret = do_cap(CAP_SYS_MODULE, delete_module, modname, 0);
+		if (ret != 0)
+			err("Error removing module '%s': %s\n", modname, moderror(errno));
+	}
 }
 
-static void usage(char *prog)
+static void exit_cleanup(void)
 {
-	fprintf(stderr, "\n%s [-v]  [-c cmd ] [-x pid] [-u user]\n"
-                "\t[-A modname]] [-L] [-b bufsize] [-o FILE] kmod-name [kmod-options]\n", prog);
-	fprintf(stderr, "-v  increase Verbosity.\n");
-	fprintf(stderr, "-c cmd.  Command \'cmd\' will be run and staprun will exit when it does.\n");
-	fprintf(stderr, "   _stp_target will contain the pid for the command.\n");
-	fprintf(stderr, "-x pid.  Sets _stp_target to 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");
-	fprintf(stderr, "   buffer size in MB. That value will be per-cpu in bulk mode.\n");
-	fprintf(stderr, "-L  Load module and start probes, then detach.\n");
-	fprintf(stderr, "-A modname.  Attach to systemtap module modname.\n");
-	exit(1);
+	dbug(2, "something exited...\n");
+	cleanup(1);
 }
 
 int main(int argc, char **argv)
 {
-	int c;
-
-	setup_signals();
+	int rc;
 
-	while ((c = getopt(argc, argv, "ALvb:t:d:c:o:u:x:")) != EOF) {
-		switch (c) {
-		case 'v':
-			verbose++;
-			break;
-		case 'b':
-		{
-			int size = (unsigned)atoi(optarg);
-			if (!size)
-				usage(argv[0]);
-			if (size > 64) {
-				fprintf(stderr, "Maximum buffer size is 64 (MB)\n");
-				exit(1);
-			}
-			buffer_size = size;
-			break;
-		}
-		case 't':
-		case 'x':
-			target_pid = atoi(optarg);
-			break;
-		case 'd':
-			/* obsolete internal option used by stap */
-			break;
-		case 'c':
-			target_cmd = optarg;
-			break;
-		case 'o':
-			outfile_name = optarg;
-			break;
-		case 'u':
-			username = optarg;
-			break;
-		case 'A':
-			attach_mod = 1;
-			break;
-		case 'L':
-			load_only = 1;
-			break;
-		default:
-			usage(argv[0]);
-		}
+	if (atexit(exit_cleanup)) {
+		_perr("cannot set exit function");
+		exit(1);
 	}
 
-	if (verbose) {
-		if (buffer_size)
-			printf ("Using a buffer of %u bytes.\n", buffer_size);
+	if (!init_cap())
+		return 1;
+
+	/* Get rid of a few standard environment variables (which */
+	/* might cause us to do unintended things). */
+	rc = unsetenv("IFS") || unsetenv("CDPATH") || unsetenv("ENV")
+		|| unsetenv("BASH_ENV");
+	if (rc) {
+		_perr("unsetenv failed");
+		exit(-1);
 	}
 
+	setup_signals();
+
+	parse_args(argc, argv);
+
+	if (buffer_size)
+		dbug(2, "Using a buffer of %u bytes.\n", buffer_size);
+
 	if (optind < argc) {
-		modpath = argv[optind++];
-		path_parse_modname(modpath);
+		parse_modpath(argv[optind++]);
 		dbug(2, "modpath=\"%s\", modname=\"%s\"\n", modpath, modname);
 	}
 
         if (optind < argc) {
 		if (attach_mod) {
-			fprintf(stderr, "Cannot have module options with attach (-A).\n");
+			err("ERROR: Cannot have module options with attach (-A).\n");
 			usage(argv[0]);
 		} else {
-			unsigned start_idx = 3; /* reserve three slots in modoptions[] */
+			unsigned start_idx = 0;
 			while (optind < argc && start_idx+1 < MAXMODOPTIONS)
 				modoptions[start_idx++] = argv[optind++];
 			modoptions[start_idx] = NULL;
 		}
 	}
 
-	if (!modpath) {
-		fprintf (stderr, "Need a module to load.\n");
+	if (modpath == NULL || *modpath == '\0') {
+		err("ERROR: Need a module name or path to load.\n");
 		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 (check_permissions() != 1)
+		usage(argv[0]);
 
 	/* now bump the priority */
-	setpriority (PRIO_PROCESS, 0, -10);
+	rc = do_cap(CAP_SYS_NICE, setpriority, PRIO_PROCESS, 0, -10);
+	/* failure is not fatal in this case */
+	if (rc < 0)
+		_perr("setpriority");
+
+	/* We're done with CAP_SYS_NICE. */
+	drop_cap(CAP_SYS_NICE);
 
 	if (init_staprun())
 		exit(1);
 
-	if (stp_main_loop()) {
-		fprintf(stderr,"Couldn't enter main loop. Exiting.\n");
-		exit(1);
-	}
+	setup_staprun_signals();
+	if (!attach_mod)
+		handle_symbols();
+
+	rc = run_stapio(argv);
+	cleanup(rc);
 
 	return 0;
 }
diff --git a/runtime/staprun/staprun.h b/runtime/staprun/staprun.h
index f8fcfe66..cde44922 100644
--- a/runtime/staprun/staprun.h
+++ b/runtime/staprun/staprun.h
@@ -1,6 +1,6 @@
 /* -*- linux-c -*-
  *
- * staprun.h - include file for staprun
+ * staprun.h - include file for staprun and stapio
  *
  * This file is part of systemtap, and is free software.  You can
  * redistribute it and/or modify it under the terms of the GNU General
@@ -33,32 +33,93 @@
 #include <sys/wait.h>
 #include <sys/statfs.h>
 #include <linux/version.h>
+#include <sys/capability.h>
 
 #define DEBUG
 #ifdef DEBUG
-#define dbug(level, args...) {if (verbose>=level) {fprintf(stderr,"%s:%d ",__FUNCTION__, __LINE__); fprintf(stderr,args);}}
+#define dbug(level, args...) {if (verbose>=level) {fprintf(stderr,"%s:%s:%d ",__name__,__FUNCTION__, __LINE__); fprintf(stderr,args);}}
 #else
 #define dbug(level, args...) ;
 #endif /* DEBUG */
 
-#define err(args...) {fprintf(stderr,"%s:%d ",__FUNCTION__, __LINE__); fprintf(stderr,args); }
+extern char *__name__;
+
+/* print to stderr */
+#define err(args...) fprintf(stderr,args)
+
+/* better perror() */
+#define perr(args...) do {					\
+		int _errno = errno;				\
+		fputs("ERROR: ", stderr);			\
+		fprintf(stderr, args);				\
+		fprintf(stderr, ": %s\n", strerror(_errno));	\
+	} while (0)
+
+/* Error messages. Use these for serious errors, not informational messages to stderr. */
+#define _err(args...) do {fprintf(stderr,"%s:%s:%d: ERROR: ",__name__, __FUNCTION__, __LINE__); fprintf(stderr,args);} while(0)
+#define _perr(args...) do {					\
+		int _errno = errno;				\
+		_err(args);					\
+		fprintf(stderr, ": %s\n", strerror(_errno));	\
+	} while (0)
+#define overflow_error() _err("Internal buffer overflow. Please file a bug report.\n")
+		
+#define do_cap(cap,func,args...) ({			\
+			int _rc, _saved_errno;		\
+			add_cap(cap);			\
+			_rc = func(args);		\
+			_saved_errno = errno;		\
+			del_cap(cap);			\
+			errno = _saved_errno;		\
+			_rc;				\
+		})					\
+
+
+/* Error checking version of sprintf() - returns 1 if overflow error */
+#define sprintf_chk(str, args...) ({			\
+	int _rc;					\
+	_rc = snprintf(str, sizeof(str), args);		\
+	if (_rc >= (int)sizeof(str)) {			\
+		overflow_error();			\
+		_rc = 1;				\
+	}						\
+	else						\
+		_rc = 0;				\
+	_rc;						\
+})
+
+/* Error checking version of snprintf() - returns 1 if overflow error */
+#define snprintf_chk(str, size, args...) ({		\
+	int _rc;					\
+	_rc = snprintf(str, size, args);		\
+	if (_rc >= (int)size) {				\
+		overflow_error();			\
+		_rc = 1;				\
+	}						\
+	else						\
+		_rc = 0;				\
+	_rc;						\
+})
+
+/* Grabbed from linux/module.h kernel include. */
+#define MODULE_NAME_LEN (64 - sizeof(unsigned long))
 
 /* we define this so we are compatible with old transport, but we don't have to use it. */
 #define STP_OLD_TRANSPORT
 #include "../transport/transport_msgs.h"
 
-/* command to check system's kernel version */
-/* KERNEL_VERSION(2.6.15) = 132623 */
-#define VERSION_CMD "uname -r | awk \'{split($1,a,\".\"); split(a[3],b,\"-\"); exit (a[1]*65536+a[2]*256+b[1] <= 132623)}\'"
 extern int use_old_transport;
 
-#define RELAYFS_MAGIC			0xF0B4A981
-#define DEBUGFS_MAGIC			0x64626720
+#define RELAYFS_MAGIC	0xF0B4A981
+#define DEBUGFS_MAGIC	0x64626720
+#define DEBUGFSDIR	"/sys/kernel/debug"
+#define RELAYFSDIR	"/mnt/relay"
 
 /*
  * function prototypes
  */
 int init_staprun(void);
+int init_stapio(void);
 int stp_main_loop(void);
 int send_request(int type, void *data, int len);
 void cleanup_and_exit (int);
@@ -71,28 +132,50 @@ void close_relayfs(void);
 int init_oldrelayfs(void);
 void close_oldrelayfs(int);
 void setup_signals(void);
+/* cap.c */
+void print_cap(char *text);
+int init_cap(void);
+void add_cap(cap_value_t cap);
+void del_cap(cap_value_t cap);
+void drop_cap(cap_value_t cap);
+/* staprun_funcs.c */
+void setup_staprun_signals(void);
+const char *moderror(int err);
+int insert_module(void);
+int mountfs(void);
+int check_permissions(void);
+void handle_symbols(void);
+
+/* common.c functions */
+void parse_args(int argc, char **argv);
+void usage(char *prog);
+void parse_modpath(const char *);
+void setup_signals(void);
 
 /*
  * variables 
  */
 extern int control_channel;
 extern int ncpus;
+extern int initialized;
 
 /* flags */
 extern int verbose;
 extern unsigned int buffer_size;
-extern char modname[];
+extern char *modname;
 extern char *modpath;
-extern char *modoptions[];
+#define MAXMODOPTIONS 64
+extern char *modoptions[MAXMODOPTIONS];
 extern int target_pid;
 extern char *target_cmd;
 extern char *outfile_name;
 extern int attach_mod;
 extern int load_only;
 
-/* uid/gid to use when execing external programs */
-extern uid_t cmd_uid;
-extern gid_t cmd_gid;
+/* getopt variables */
+extern char *optarg;
+extern int optopt;
+extern int optind;
 
 /* maximum number of CPUs we can handle */
 #define NR_CPUS 256
diff --git a/runtime/staprun/staprun_funcs.c b/runtime/staprun/staprun_funcs.c
new file mode 100644
index 00000000..0747b530
--- /dev/null
+++ b/runtime/staprun/staprun_funcs.c
@@ -0,0 +1,443 @@
+/* -*- linux-c -*-
+ *
+ * staprun_funcs.c - staprun functions
+ *
+ * This file is part of systemtap, and is free software.  You can
+ * redistribute it and/or modify it under the terms of the GNU General
+ * Public License (GPL); either version 2, or (at your option) any
+ * later version.
+ *
+ * Copyright (C) 2007 Red Hat Inc.
+ */
+
+#include "staprun.h"
+#include <sys/mount.h>
+#include <sys/utsname.h>
+#include <grp.h>
+#include <pwd.h>
+
+void setup_staprun_signals(void)
+{
+	struct sigaction a;
+	memset(&a, 0, sizeof(a));
+	sigfillset(&a.sa_mask);
+	a.sa_handler = SIG_IGN;
+	sigaction(SIGINT, &a, NULL);
+	sigaction(SIGTERM, &a, NULL);
+	sigaction(SIGHUP, &a, NULL);
+	sigaction(SIGQUIT, &a, NULL);
+}
+
+extern long init_module(void *, unsigned long, const char *);
+
+/* Module errors get translated. */
+const char *moderror(int err)
+{
+	switch (err) {
+	case ENOEXEC:
+		return "Invalid module format";
+	case ENOENT:
+		return "Unknown symbol in module";
+	case ESRCH:
+		return "Module has wrong symbol version";
+	case EINVAL:
+		return "Invalid parameters";
+	default:
+		return strerror(err);
+	}
+}
+
+int insert_module(void)
+{
+	int i;
+	long ret;
+	void *file;
+	char *opts;
+	int fd, saved_errno;
+	struct stat sbuf;
+		
+	dbug(2, "inserting module\n");
+
+	opts = malloc(128);
+	if (opts == NULL) {
+		_perr("allocating memory failed");
+		return -1;
+	}
+	if (snprintf_chk(opts, 128, "_stp_bufsize=%d", buffer_size))
+		return -1;
+	for (i = 0; modoptions[i] != NULL; i++) {
+		opts = realloc(opts, strlen(opts) + strlen(modoptions[i]) + 2);
+		if (opts == NULL) {
+			_perr("reallocating memory failed");
+			return -1;
+		}
+		strcat(opts, " ");
+		strcat(opts, modoptions[i]);
+	}
+	dbug(2, "module options: %s\n", opts);
+
+	/* Open the module file. */
+	fd = open(modpath, O_RDONLY);
+	if (fd < 0) {
+		perr("Error opening '%s'", modpath);
+		return -1;
+	}
+	
+	/* Now that the file is open, figure out how big it is. */
+	if (fstat(fd, &sbuf) < 0) {
+		_perr("Error stat'ing '%s'", modpath);
+		close(fd);
+		return -1;
+	}
+
+	/* mmap in the entire module. */
+	file = mmap(NULL, sbuf.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
+	if (file == MAP_FAILED) {
+		_perr("Error mapping '%s'", modpath);
+		close(fd);
+		free(opts);
+		return -1;
+	}
+	    
+	/* Actually insert the module */
+	ret = do_cap(CAP_SYS_MODULE, init_module, file, sbuf.st_size, opts);
+	saved_errno = errno;
+
+	/* Cleanup. */
+	free(opts);
+	munmap(file, sbuf.st_size);
+	close(fd);
+
+	if (ret != 0) {
+		err("Error inserting module '%s': %s\n", modpath, moderror(saved_errno));
+		return -1; 
+	}
+	return 0;
+}
+
+int mountfs(void)
+{
+	struct stat sb;
+	struct statfs st;
+	int rc;
+
+	/* If the debugfs dir is already mounted correctly, we're done. */
+ 	if (statfs(DEBUGFSDIR, &st) == 0
+	    && (int) st.f_type == (int) DEBUGFS_MAGIC)
+		return 0;
+
+	/* If DEBUGFSDIR exists (and is a directory), try to mount
+	 * DEBUGFSDIR. */
+	rc = stat(DEBUGFSDIR, &sb);
+	if (rc == 0 && S_ISDIR(sb.st_mode)) {
+		/* If we can mount the debugfs dir correctly, we're done. */
+		rc = do_cap(CAP_SYS_ADMIN, mount, "debugfs", DEBUGFSDIR,
+			    "debugfs", 0, NULL); 
+		if (rc == 0)
+			return 0;
+		/* If we got ENODEV, that means that debugfs isn't
+		 * supported, so we'll need try try relayfs.  If we
+		 * didn't get ENODEV, we got a real error. */
+		else if (errno != ENODEV) {
+			perr("Couldn't mount %s", DEBUGFSDIR);
+			return -1;
+		}
+	}
+	
+	/* DEBUGFSDIR couldn't be mounted.  So, try RELAYFSDIR. */
+
+	/* If the relayfs dir is already mounted correctly, we're done. */
+	if (statfs(RELAYFSDIR, &st) == 0
+	    && (int)st.f_type == (int)RELAYFS_MAGIC)
+		return 0;
+
+	/* Ensure that RELAYFSDIR exists and is a directory. */
+	rc = stat(RELAYFSDIR, &sb);
+	if (rc == 0 && ! S_ISDIR(sb.st_mode)) {
+		err("%s exists but isn't a directory.\n", RELAYFSDIR);
+		return -1;
+	}
+	else if (rc < 0) {
+		mode_t old_umask;
+		int saved_errno;
+		gid_t gid = getgid();
+		uid_t uid = getuid();
+
+		/* To ensure the directory gets created with the proper
+		 * permissions, set umask to a known value. */
+		old_umask = umask(0002);
+
+		/* To ensure the directory gets created with the
+		 * proper group, we'll have to temporarily switch to
+		 * root. */
+		if (do_cap(CAP_SETUID, setuid, 0) < 0) {
+			_perr("Couldn't change user while creating %s", RELAYFSDIR);
+			return -1;
+		}
+		if (do_cap(CAP_SETGID, setgid, 0) < 0) {
+			_perr("Couldn't change group while creating %s", RELAYFSDIR);
+			return -1;
+		}
+
+		/* Try to create the directory, saving the return
+		 * status and errno value. */
+		rc = mkdir(RELAYFSDIR, 0755);
+		saved_errno = errno;
+
+		/* Restore everything we changed. */
+		if (do_cap(CAP_SETGID, setgid, gid) < 0) {
+			_perr("Couldn't restore group while creating %s", RELAYFSDIR);
+			return -1;
+		}
+		if (do_cap(CAP_SETUID, setuid, uid) < 0) {
+			_perr("Couldn't restore user while creating %s", RELAYFSDIR);
+			return -1;
+		}
+		umask(old_umask);
+
+		/* If creating the directory failed, error out. */
+		if (rc < 0) {
+			err("Couldn't create %s: %s\n", RELAYFSDIR, strerror(saved_errno));
+			return -1;
+		}
+	}
+
+	/* Now that we're sure the directory exists, try mounting RELAYFSDIR. */
+	if (do_cap(CAP_SYS_ADMIN, mount, "relayfs", RELAYFSDIR, "relayfs", 0, NULL) < 0) {
+		perr("Couldn't mount %s", RELAYFSDIR);
+		return -1;
+	}
+	return 0;
+}
+
+
+/*
+ * Members of the 'stapusr' group can only use "blessed" modules -
+ * ones in the '/lib/modules/KVER/systemtap' directory.  Make sure the
+ * module path is in that directory.
+ *
+ * Returns: -1 on errors, 0 on failure, 1 on success.
+ */
+static int
+check_path(void)
+{
+	struct utsname utsbuf;
+	struct stat sb;
+	char staplib_dir_path[PATH_MAX];
+	char staplib_dir_realpath[PATH_MAX];
+	char module_realpath[PATH_MAX];
+
+	/* First, we need to figure out what the kernel
+	 * version is and build the '/lib/modules/KVER/systemtap' path. */
+	if (uname(&utsbuf) != 0) {
+		_perr("ERROR: Unable to determine kernel version, uname failed");
+		return -1;
+	}
+	if (sprintf_chk(staplib_dir_path, "/lib/modules/%s/systemtap", utsbuf.release))
+		return -1;
+
+	/* Validate /lib/modules/KVER/systemtap. */
+	if (stat(staplib_dir_path, &sb) < 0) {
+		perr("Members of the \"stapusr\" group can only use modules within\n"
+		     "  the \"%s\" directory.\n"
+		     "  Error getting information on that directory", staplib_dir_path);
+		return -1;
+	}
+	/* Make sure it is a directory. */
+	if (! S_ISDIR(sb.st_mode)) {
+		err("ERROR: Members of the \"stapusr\" group can only use modules within\n"
+		    "  the \"%s\" directory.\n"
+		    "  That path must refer to a directory.\n", staplib_dir_path);
+		return -1;
+	}
+	/* Make sure it is owned by root. */
+	if (sb.st_uid != 0) {
+		err("ERROR: Members of the \"stapusr\" group can only use modules within\n"
+		    "  the \"%s\" directory.\n"
+		    "  That directory should be owned by root.\n", staplib_dir_path);
+		return -1;
+	}
+	/* Make sure it isn't world writable. */
+	if (sb.st_mode & S_IWOTH) {
+		err("ERROR: Members of the \"stapusr\" group can only use modules within\n"
+		    "  the \"%s\" directory.\n"
+		    "  That directory should not be world writable.\n", staplib_dir_path);
+		return -1;
+	}
+
+	/* Use realpath() to canonicalize the module directory
+	 * path. */
+	if (realpath(staplib_dir_path, staplib_dir_realpath) == NULL) {
+		perr("Members of the \"stapusr\" group can only use modules within\n"
+		     "  the \"%s\" directory.\n"
+		     "  Unable to canonicalize that directory",	staplib_dir_path);
+		return -1;
+	}
+	
+	/* Use realpath() to canonicalize the module path. */
+	if (realpath(modpath, module_realpath) == NULL) {
+		perr("Unable to canonicalize path \"%s\"",modpath);
+		return -1;
+	}
+	
+	/* Now we've got two canonicalized paths.  Make sure
+	 * module_realpath starts with staplib_dir_realpath. */
+	if (strncmp(staplib_dir_realpath, module_realpath,
+		    strlen(staplib_dir_realpath)) != 0) {
+		err("ERROR: Members of the \"stapusr\" group can only use modules within\n"
+		    "  the \"%s\" directory.\n"
+		    "  Module \"%s\" does not exist within that directory.\n",
+		    staplib_dir_path, modpath);
+		return 0;
+	}
+	return 1;
+}
+
+/*
+ * Check the user's permissions.  Is he allowed to run staprun (or is
+ * he limited to "blessed" modules)?
+ *
+ * Returns: -1 on errors, 0 on failure, 1 on success.
+ */
+int check_permissions(void)
+{
+	gid_t gid, gidlist[NGROUPS_MAX];
+	gid_t stapdev_gid, stapusr_gid;
+	int i, ngids;
+	struct group *stgr;
+	int path_check = 0;
+
+	/* If we're root, we can do anything. */
+	if (geteuid() == 0)
+		return 1;
+
+	/* Lookup the gid for group "stapdev" */
+	errno = 0;
+	stgr = getgrnam("stapdev");
+	/* If we couldn't find the group, just set the gid to an
+	 * invalid number.  Just because this group doesn't exist
+	 * doesn't mean the other group doesn't exist. */
+	if (stgr == NULL)
+		stapdev_gid = (gid_t)-1;
+	else
+		stapdev_gid = stgr->gr_gid;
+
+	/* Lookup the gid for group "stapusr" */
+	errno = 0;
+	stgr = getgrnam("stapusr");
+	/* If we couldn't find the group, just set the gid to an
+	 * invalid number.  Just because this group doesn't exist
+	 * doesn't mean the other group doesn't exist. */
+	if (stgr == NULL)
+		stapusr_gid = (gid_t)-1;
+	else
+		stapusr_gid = stgr->gr_gid;
+
+	/* If neither group was found, just return an error. */
+	if (stapdev_gid == (gid_t)-1 && stapusr_gid == (gid_t)-1) {
+		err("ERROR: unable to find either group \"stapdev\" or group \"stapusr\"\n");
+		return -1;
+	}
+
+	/* According to the getgroups() man page, getgroups() may not
+	 * return the effective gid, so try to match it first. */
+	gid = getegid();
+	if (gid == stapdev_gid)
+		return 1;
+	else if (gid == stapusr_gid)
+		path_check = 1;
+
+	/* Get the list of the user's groups. */
+	ngids = getgroups(NGROUPS_MAX, gidlist);
+	if (ngids < 0) {
+		perr("Unable to retrieve group list");
+		return -1;
+	}
+
+	for (i = 0; i < ngids; i++) {
+		/* If the user is a member of 'stapdev', then we're
+		 *  done, since he can use staprun without any
+		 *  restrictions. */
+		if (gidlist[i] == stapdev_gid)
+			return 1;
+
+		/* If the user is a member of 'stapusr', then we'll
+		 * need to check the module path.  However, we'll keep
+		 * checking groups since it is possible the user is a
+		 * member of both groups and we haven't seen the
+		 * 'stapdev' group yet. */
+		if (gidlist[i] == stapusr_gid)
+			path_check = 1;
+	}
+
+	/* If path_check is 0, then the user isn't a member of either
+	 * group.  Error out. */
+	if (path_check == 0) {
+		err("ERROR: you must be a member of either group \"stapdev\" or group \"stapusr\"\n");
+		return 0;
+	}
+
+	/* At this point the user is only a member of the 'stapusr'
+	 * group.  Members of the 'stapusr' group can only use modules
+	 * in /lib/modules/KVER/systemtap.  Make sure the module path
+	 * is in that directory. */
+	return check_path();
+}
+
+/* wait for symbol requests and reply */
+void handle_symbols(void)
+{
+	ssize_t nb;
+	void *data;
+	int type;
+	char recvbuf[8192];
+
+	dbug(2, "waiting for symbol requests\n");
+
+	/* create control channel */
+	if (init_ctl_channel() < 0) {
+		err("Failed to initialize control channel.\n");
+		exit(1);
+	}
+
+	while (1) { /* handle messages from control channel */
+		nb = read(control_channel, recvbuf, sizeof(recvbuf));
+		if (nb <= 0) {
+			if (errno != EINTR)
+				_perr("Unexpected EOF in read (nb=%ld)", (long)nb);
+			continue;
+		}
+		
+		type = *(int *)recvbuf;
+		data = (void *)(recvbuf + sizeof(int));
+
+		switch (type) { 
+		case STP_MODULE:
+		{
+			dbug(2, "STP_MODULES request received\n");
+			do_module(data);
+			goto done;
+		}		
+		case STP_SYMBOLS:
+		{
+			struct _stp_msg_symbol *req = (struct _stp_msg_symbol *)data;
+			dbug(2, "STP_SYMBOLS request received\n");
+			if (req->endian != 0x1234) {
+				err("ERROR: staprun is compiled with different endianess than the kernel!\n");
+				exit(1);
+			}
+			if (req->ptr_size != sizeof(char *)) {
+				err("ERROR: staprun is compiled with %d-bit pointers and the kernel uses %d-bit.\n",
+					8*(int)sizeof(char *), 8*req->ptr_size);
+				exit(1);
+			}
+			do_kernel_symbols();
+			break;
+		}
+		default:
+			err("WARNING: ignored message of type %d\n", (type));
+		}
+	}
+done:
+	close_ctl_channel();
+}
diff --git a/runtime/staprun/symbols.c b/runtime/staprun/symbols.c
index 49755c5e..61b56b2e 100644
--- a/runtime/staprun/symbols.c
+++ b/runtime/staprun/symbols.c
@@ -40,42 +40,53 @@ static int get_sections(char *name, char *data_start, int datalen)
 
 	res = snprintf(dir, sizeof(dir), SECDIR, name);
 	if (res >= (int)sizeof(dir)) {
-		fprintf(stderr, "ERROR: couldn't fit module \"%s\" into dir buffer.\n", name);
-		fprintf(stderr, "This should never happen. Please file a bug report.\n");
-		cleanup_and_exit(0);
+		_err("Couldn't fit module \"%s\" into dir buffer.\n"	\
+		    "This should never happen. Please file a bug report.\n", name);
+		exit(1);
 	}
-
+	
 	if ((secdir = opendir(dir)) == NULL)
 		return 0;
 
+	/* Initialize mod. */
 	memset(mod, 0, sizeof(struct _stp_module));
+
+	/* Copy name in and check for overflow. */
 	strncpy(mod->name, name, STP_MODULE_NAME_LEN);
+	if (mod->name[STP_MODULE_NAME_LEN - 1] != '\0') {
+		_err("Couldn't fit module \"%s\" into mod->name buffer.\n" \
+		    "This should never happen. Please file a bug report.\n", name);
+		exit(1);
+	}
 
 	while ((d = readdir(secdir))) {
 		char *secname = d->d_name;
+
+		/* Copy filename in and check for overflow. */
 		res = snprintf(filename, sizeof(filename), "/sys/module/%s/sections/%s", name, secname);
 		if (res >= (int)sizeof(filename)) {
-			fprintf(stderr, "ERROR: couldn't fit secname \"%s\" into filename buffer.\n", secname);
-			fprintf(stderr, "This should never happen. Please file a bug report.\n");
+			_err("Couldn't fit secname \"%s\" into filename buffer.\n" \
+			    "This should never happen. Please file a bug report.\n", secname);
 			closedir(secdir);
-			cleanup_and_exit(0);
+			exit(1);
+		}
+		
+		/* filter out some non-useful stuff */
+		if (!strncmp(secname,"__",2) 
+		    || !strcmp(secname,".") 
+		    || !strcmp(secname,"..") 
+		    || !strcmp(secname,".module_sig") 
+		    || !strcmp(secname,".modinfo") 
+		    || !strcmp(secname,".strtab") 
+		    || !strcmp(secname,".symtab") ) {
+			continue;
 		}
+		if (!strncmp(secname, ".gnu.linkonce", 13) 
+		    && strcmp(secname, ".gnu.linkonce.this_module"))
+			continue;
+
 		if ((fd = open(filename,O_RDONLY)) >= 0) {
 			if (read(fd, buf, 32) > 0) {
-
-				/* filter out some non-useful stuff */
-				if (!strncmp(secname,"__",2) 
-				    || !strcmp(secname,".module_sig") 
-				    || !strcmp(secname,".modinfo") 
-				    || !strcmp(secname,".strtab") 
-				    || !strcmp(secname,".symtab") ) {
-					close(fd);
-					continue;
-				}
-				if (!strncmp(secname, ".gnu.linkonce", 13) 
-				    && strcmp(secname, ".gnu.linkonce.this_module"))
-					continue;
-
 				/* create next section */
 				sec = (struct _stp_symbol *)data;
 				if (data - data_start + (int)sizeof(struct _stp_symbol) > datalen)
@@ -85,8 +96,10 @@ static int get_sections(char *name, char *data_start, int datalen)
 				sec->symbol = (char *)(strdata - strdata_start);
 				mod->num_sections++;
 
-				/* now create string data for the section */
-				if (strdata - strdata_start + strlen(strdata) >= sizeof(strdata_start))
+				/* now create string data for the
+				 * section (checking for overflow) */
+				if ((strdata - strdata_start + strlen(strdata))
+				    >= sizeof(strdata_start))
 					goto err1;
 				strcpy(strdata, secname);
 				strdata += strlen(secname) + 1;
@@ -120,8 +133,9 @@ err1:
 	close(fd);
 	closedir(secdir);
 err0:
-	err("overflowed buffers.\n");
-	cleanup_and_exit(0);
+	/* if this happens, something went seriously wrong. */
+	_err("Unexpected error. Overflowed buffers.\n");
+	exit(1);
 	return 0; /* not reached */
 }
 #undef SECDIR
@@ -132,8 +146,8 @@ void send_module (char *mname)
 	int len = get_sections(mname, data, sizeof(data));
 	if (len) {
 		if (send_request(STP_MODULE, data, len) < 0) {
-			err("Loading of module %s failed. Exiting...\n", mname);
-			cleanup_and_exit(0);
+			_err("Loading of module %s failed. Exiting...\n", mname);
+			exit(1);
 		}
 	}
 }
@@ -150,6 +164,7 @@ int do_module (void *data)
 				send_module(d->d_name);
 			closedir(moddir);
 		}
+		send_request(STP_MODULE, data, 0);
 		return 1;
 	}
 
@@ -166,34 +181,35 @@ static int compar(const void *p1, const void *p2)
 	return 1;
 }
 
-#define MAX_SYMBOLS 32768
+#define MAX_SYMBOLS 32*1024
 
 void do_kernel_symbols(void)
 {
-	FILE *kallsyms;
-	char *sym_base, *data_base;
+	FILE *kallsyms=NULL;
+	char *sym_base=NULL, *data_base=NULL;
 	char buf[128], *ptr, *name, *data, *dataptr, *datamax, type;
 	unsigned long addr;
 	struct _stp_symbol *syms;
-	int num_syms, i = 0;
+	int num_syms, i = 0, max_syms= MAX_SYMBOLS;
+	int data_basesize = MAX_SYMBOLS*32;
 
-	sym_base = malloc(MAX_SYMBOLS*sizeof(struct _stp_symbol)+sizeof(long));
-	data_base = malloc(MAX_SYMBOLS*32);
+	sym_base = malloc(max_syms*sizeof(struct _stp_symbol)+sizeof(long));
+	data_base = malloc(data_basesize);
 	if (data_base == NULL || sym_base == NULL) {
-		fprintf(stderr,"Failed to allocate memory for symbols\n");
-		cleanup_and_exit(0);
+		_err("Failed to allocate memory for symbols\n");
+		goto err;
 	}
 	*(int *)data_base = STP_SYMBOLS;
 	dataptr = data = data_base + sizeof(long);
-	datamax = dataptr + MAX_SYMBOLS*32 - sizeof(long);
+	datamax = data_base + data_basesize;
 
 	*(int *)sym_base = STP_SYMBOLS;
 	syms = (struct _stp_symbol *)(sym_base + sizeof(long));
 
 	kallsyms = fopen ("/proc/kallsyms", "r");
 	if (!kallsyms) {
-		perror("Fatal error: Unable to open /proc/kallsyms:");
-		cleanup_and_exit(0);
+		_perr("Fatal error: Unable to open /proc/kallsyms");
+		goto err;
 	}
 
 	/* put empty string in data */
@@ -216,8 +232,30 @@ void do_kernel_symbols(void)
 			while (*name) *dataptr++ = *name++;
 			*dataptr++ = 0;
 			i++;
-			if (dataptr > datamax - 1000)
-			  break;
+			if (i >= max_syms) {
+				char *s;
+				max_syms *= 2;
+				s = realloc(sym_base, max_syms*sizeof(struct _stp_symbol)+sizeof(long));
+				if (s == NULL) {
+					_err("Could not allocate enough space for symbols.\n");
+					goto err;
+				}
+				syms = (struct _stp_symbol *)(s + sizeof(long));
+				sym_base = s;
+			}
+			if (dataptr > datamax - 1024) {
+				char *db;
+				data_basesize *= 2;
+				db = realloc(data_base, data_basesize);
+				if (db == NULL) {
+					_err("Could not allocate enough space for symbols.\n");
+					goto err;
+				}
+				dataptr = db + (dataptr - data_base);
+				data = db + sizeof(long);
+				datamax = db + data_basesize;
+				data_base = db;
+			}
 		}
 	}
 	num_syms = i;
@@ -241,18 +279,16 @@ void do_kernel_symbols(void)
 	free(data_base);
 	free(sym_base);
 	fclose(kallsyms);
-
-	if (dataptr >= datamax) {
-		err("Error: overflowed symbol data area.\n");
-		cleanup_and_exit(0);
-	}
 	return;
 
 err:
-	free(data_base);
-	free(sym_base);
-	fclose(kallsyms);
-
-	err("Loading of symbols failed. Exiting...\n");
-	cleanup_and_exit(0);
+	if (data_base)
+		free(data_base);
+	if (sym_base)
+		free(sym_base);
+	if (kallsyms)
+		fclose(kallsyms);
+
+	_err("Loading of symbols failed. Exiting...\n");
+	exit(1);
 }