4 files changed, 90 insertions, 25 deletions

diff --git a/runtime/staprun/common.c b/runtime/staprun/common.c
index 9056c710..d3f8835a 100644
--- a/runtime/staprun/common.c
+++ b/runtime/staprun/common.c
@@ -23,6 +23,7 @@ char *target_cmd;
 char *outfile_name;
 int attach_mod;
 int load_only;
+int need_uprobes;
 
 /* module variables */
 char *modname = NULL;
@@ -44,9 +45,13 @@ void parse_args(int argc, char **argv)
 	outfile_name = NULL;
 	attach_mod = 0;
 	load_only = 0;
+	need_uprobes = 0;
 
-	while ((c = getopt(argc, argv, "ALvb:t:d:c:o:x:")) != EOF) {
+	while ((c = getopt(argc, argv, "ALuvb:t:d:c:o:x:")) != EOF) {
 		switch (c) {
+		case 'u':
+			need_uprobes = 1;
+			break;
 		case 'v':
 			verbose++;
 			break;
diff --git a/runtime/staprun/staprun.c b/runtime/staprun/staprun.c
index 71039e0b..67b01abb 100644
--- a/runtime/staprun/staprun.c
+++ b/runtime/staprun/staprun.c
@@ -35,6 +35,16 @@ run_as(uid_t uid, gid_t gid, const char *path, char *const argv[])
 	pid_t pid;
 	int rstatus;
 
+	if (verbose >= 2) {
+		int i = 0;
+		err("execing: ");
+		while (argv[i]) {
+			err("%s ", argv[i]);
+			i++;
+		}
+		err("\n");		
+	}
+
 	if ((pid = fork()) < 0) {
 		_perr("fork");
 		return -1;
@@ -77,18 +87,63 @@ static int run_stapio(char **argv)
 	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);
 }
 
+/*
+ * Module to be inserted has one or more user-space probes.  Make sure
+ * uprobes is enabled.
+ * If /proc/kallsyms lists a symbol in uprobes (e.g. unregister_uprobe),
+ * we're done.
+ * Else try "modprobe uprobes" to load the uprobes module (if any)
+ * built with the kernel.
+ * If that fails, load the uprobes module built in runtime/uprobes.
+ */
+static int enable_uprobes(void)
+{
+	int i;
+	char *argv[10];
+	uid_t uid = getuid();
+	gid_t gid = getgid();
+
+	i = 0;
+	argv[i++] = "/bin/grep";
+	argv[i++] = "-q";
+	argv[i++] = "unregister_uprobe";
+	argv[i++] = "/proc/kallsyms";
+	argv[i] = NULL;
+	if (run_as(uid, gid, argv[0], argv) == 0)
+		return 0;
+
+	/*
+	 * TODO: If user can't setresuid to root here, staprun will exit.
+	 * Is there a situation where that would fail but the subsequent
+	 * attempt to use CAP_SYS_MODULE privileges (in insert_module())
+	 * would succeed?
+	 */
+	dbug(2, "Inserting uprobes module from /lib/modules, if any.\n");
+	i = 0;
+	argv[i++] = "/sbin/modprobe";
+	argv[i++] = "-q";
+	argv[i++] = "uprobes";
+	argv[i] = NULL;
+	if (run_as(0, 0, argv[0], argv) == 0)
+		return 0;
+
+	dbug(2, "Inserting uprobes module from SystemTap runtime.\n");
+	argv[0] = NULL;
+	return insert_module(PKGDATADIR "/runtime/uprobes/uprobes.ko",
+		NULL, argv);
+}
+
+static int insert_stap_module(void)
+{
+	char bufsize_option[128];
+	if (snprintf_chk(bufsize_option, 128, "_stp_bufsize=%d", buffer_size))
+		return -1;
+	return insert_module(modpath, bufsize_option, modoptions);
+}
+
 
 int init_staprun(void)
 {
@@ -101,7 +156,9 @@ int init_staprun(void)
 	drop_cap(CAP_SYS_ADMIN);
  
 	if (!attach_mod) {
-		if (insert_module() < 0)
+		if (need_uprobes && enable_uprobes() != 0)
+			return -1;
+		if (insert_stap_module() < 0)
 			return -1;
 		else
 			inserted_module = 1;
diff --git a/runtime/staprun/staprun.h b/runtime/staprun/staprun.h
index cde44922..1b0f3221 100644
--- a/runtime/staprun/staprun.h
+++ b/runtime/staprun/staprun.h
@@ -141,7 +141,8 @@ 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 insert_module(const char *path, const char *special_options,
+	char **options);
 int mountfs(void);
 int check_permissions(void);
 void handle_symbols(void);
@@ -171,6 +172,7 @@ extern char *target_cmd;
 extern char *outfile_name;
 extern int attach_mod;
 extern int load_only;
+extern int need_uprobes;
 
 /* getopt variables */
 extern char *optarg;
diff --git a/runtime/staprun/staprun_funcs.c b/runtime/staprun/staprun_funcs.c
index 0747b530..eec4ae63 100644
--- a/runtime/staprun/staprun_funcs.c
+++ b/runtime/staprun/staprun_funcs.c
@@ -47,7 +47,7 @@ const char *moderror(int err)
 	}
 }
 
-int insert_module(void)
+int insert_module(const char *path, const char *special_options, char **options)
 {
 	int i;
 	long ret;
@@ -58,34 +58,35 @@ int insert_module(void)
 		
 	dbug(2, "inserting module\n");
 
-	opts = malloc(128);
+	if (special_options)
+		opts = strdup(special_options);
+	else
+		opts = strdup("");
 	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);
+	for (i = 0; options[i] != NULL; i++) {
+		opts = realloc(opts, strlen(opts) + strlen(options[i]) + 2);
 		if (opts == NULL) {
-			_perr("reallocating memory failed");
+			_perr("[re]allocating memory failed");
 			return -1;
 		}
 		strcat(opts, " ");
-		strcat(opts, modoptions[i]);
+		strcat(opts, options[i]);
 	}
 	dbug(2, "module options: %s\n", opts);
 
 	/* Open the module file. */
-	fd = open(modpath, O_RDONLY);
+	fd = open(path, O_RDONLY);
 	if (fd < 0) {
-		perr("Error opening '%s'", modpath);
+		perr("Error opening '%s'", path);
 		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);
+		_perr("Error stat'ing '%s'", path);
 		close(fd);
 		return -1;
 	}
@@ -93,7 +94,7 @@ int insert_module(void)
 	/* 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);
+		_perr("Error mapping '%s'", path);
 		close(fd);
 		free(opts);
 		return -1;
@@ -109,7 +110,7 @@ int insert_module(void)
 	close(fd);
 
 	if (ret != 0) {
-		err("Error inserting module '%s': %s\n", modpath, moderror(saved_errno));
+		err("Error inserting module '%s': %s\n", path, moderror(saved_errno));
 		return -1; 
 	}
 	return 0;