1 files changed, 155 insertions, 95 deletions

diff --git a/runtime/task_finder.c b/runtime/task_finder.c
index d2e57a6b..6d79c98a 100644
--- a/runtime/task_finder.c
+++ b/runtime/task_finder.c
@@ -1,9 +1,16 @@
 #include <linux/list.h>
+#include <linux/binfmts.h>
 
 static LIST_HEAD(__stp_task_finder_list);
 
 struct stap_task_finder_target;
 
+#define __STP_TF_STARTING	0
+#define __STP_TF_RUNNING	1
+#define __STP_TF_STOPPING	2
+#define __STP_TF_STOPPED	3
+atomic_t __stp_task_finder_state = ATOMIC_INIT(__STP_TF_STARTING);
+
 typedef int (*stap_task_finder_callback)(struct task_struct *tsk,
 					 int register_p,
 					 struct stap_task_finder_target *tgt);
@@ -23,6 +30,10 @@ struct stap_task_finder_target {
 	stap_task_finder_callback callback;
 };
 
+static u32
+__stp_utrace_task_finder_target_death(struct utrace_attached_engine *engine,
+				      struct task_struct *tsk);
+
 static int
 stap_register_task_finder_target(struct stap_task_finder_target *new_tgt)
 {
@@ -38,6 +49,11 @@ stap_register_task_finder_target(struct stap_task_finder_target *new_tgt)
 	else
 		new_tgt->pathlen = 0;
 
+	// Make sure everything is initialized properly.
+	new_tgt->engine_attached = 0;
+	memset(&new_tgt->ops, 0, sizeof(new_tgt->ops));
+	new_tgt->ops.report_death = &__stp_utrace_task_finder_target_death;
+
 	// Search the list for an existing entry for pathname/pid.
 	list_for_each(node, &__stp_task_finder_list) {
 		tgt = list_entry(node, struct stap_task_finder_target, list);
@@ -62,7 +78,6 @@ stap_register_task_finder_target(struct stap_task_finder_target *new_tgt)
 	}
 
 	// Add this target to the callback list for this task.
-	new_tgt->engine_attached = 0;
 	list_add_tail(&new_tgt->callback_list, &tgt->callback_list_head);
 	return 0;
 }
@@ -78,6 +93,10 @@ stap_utrace_detach_ops(struct utrace_engine_ops *ops)
 	rcu_read_lock();
 	for_each_process(tsk) {
 		struct mm_struct *mm;
+
+		if (tsk->pid <= 1)
+			continue;
+
 		mm = get_task_mm(tsk);
 		if (mm) {
 			mmput(mm);
@@ -152,11 +171,12 @@ __stp_get_mm_path(struct mm_struct *mm, char *buf, int buflen)
 		vma = vma->vm_next;
 	}
 	if (vma) {
-		struct vfsmount *mnt = mntget(vma->vm_file->f_path.mnt);
-		struct dentry *dentry = dget(vma->vm_file->f_path.dentry);
-		rc = d_path(dentry, mnt, buf, buflen);
-		dput(dentry);
-		mntput(mnt);
+#if LINUX_VERSION_CODE < KERNEL_VERSION(2,6,25)
+		rc = d_path(vma->vm_file->f_dentry, vma->vm_file->f_vfsmnt,
+			    buf, buflen);
+#else
+		rc = d_path(&(vma->vm_file->f_path), buf, buflen);
+#endif
 	}
 	else {
 		*buf = '\0';
@@ -167,76 +187,82 @@ __stp_get_mm_path(struct mm_struct *mm, char *buf, int buflen)
 }
 
 #define __STP_UTRACE_TASK_FINDER_EVENTS (UTRACE_EVENT(CLONE)	\
-					 | UTRACE_EVENT(EXEC))
+					 | UTRACE_EVENT(EXEC)	\
+					 | UTRACE_EVENT(DEATH))
 
 #define __STP_UTRACE_ATTACHED_TASK_EVENTS (UTRACE_EVENT(DEATH))
 
-static u32
-__stp_utrace_task_finder_clone(struct utrace_attached_engine *engine,
-			       struct task_struct *parent,
-			       unsigned long clone_flags,
-			       struct task_struct *child)
+static int
+__stp_utrace_attach(struct task_struct *tsk,
+		    const struct utrace_engine_ops *ops, void *data,
+		    unsigned long event_flags)
 {
-	struct utrace_attached_engine *child_engine;
+	struct utrace_attached_engine *engine;
 	struct mm_struct *mm;
+	int rc = 0;
 
-	// On clone, attach to the child.  Ignore threads with no mm
-	// (which are kernel threads).
-	mm = get_task_mm(child);
-	if (mm) {
-		mmput(mm);
-		child_engine = utrace_attach(child, UTRACE_ATTACH_CREATE,
-					     engine->ops, 0);
-		if (IS_ERR(child_engine))
-			_stp_error("attach to clone child %d failed: %ld",
-				   (int)child->pid, PTR_ERR(child_engine));
-		else {
-			utrace_set_flags(child, child_engine,
-					 __STP_UTRACE_TASK_FINDER_EVENTS);
+	// Ignore init
+	if (tsk->pid <= 1)
+		return EPERM;
+
+	// Ignore threads with no mm (which are kernel threads).
+	mm = get_task_mm(tsk);
+	if (! mm)
+		return EPERM;
+	mmput(mm);
+
+	engine = utrace_attach(tsk, UTRACE_ATTACH_CREATE, ops, data);
+	if (IS_ERR(engine)) {
+		int error = -PTR_ERR(engine);
+		if (error != ENOENT) {
+			_stp_error("utrace_attach returned error %d on pid %d",
+				   error, (int)tsk->pid);
+			rc = error;
 		}
 	}
-	return UTRACE_ACTION_RESUME;
+	else if (unlikely(engine == NULL)) {
+		_stp_error("utrace_attach returned NULL on pid %d",
+			   (int)tsk->pid);
+		rc = EFAULT;
+	}
+	else {
+		utrace_set_flags(tsk, engine, event_flags);
+	}
+	return rc;
 }
 
 static u32
-__stp_utrace_task_finder_death(struct utrace_attached_engine *engine,
-			       struct task_struct *tsk)
+__stp_utrace_task_finder_report_clone(struct utrace_attached_engine *engine,
+				      struct task_struct *parent,
+				      unsigned long clone_flags,
+				      struct task_struct *child)
 {
-	struct stap_task_finder_target *tgt = engine->data;
-
-	// The first implementation of this added a
-	// UTRACE_EVENT(DEATH) handler to
-	// __stp_utrace_task_finder_ops.  However, dead threads don't
-	// have a mm_struct, so we can't find the exe's path.  So, we
-	// don't know which callback(s) to call.
-	//
-	// So, now when an "interesting" thread is found, we add a
-	// separate UTRACE_EVENT(DEATH) handler for every probe.
+	struct utrace_attached_engine *child_engine;
+	struct mm_struct *mm;
 
-	if (tgt != NULL && tgt->callback != NULL) {
-		int rc;
+	if (atomic_read(&__stp_task_finder_state) != __STP_TF_RUNNING)
+		return UTRACE_ACTION_RESUME;
 
-		// Call the callback
-		rc = tgt->callback(tsk, 0, tgt);
-		if (rc != 0) {
-			_stp_error("death callback for %d failed: %d",
-				   (int)tsk->pid, rc);
-		}
-	}
+	// On clone, attach to the child.
+	(void) __stp_utrace_attach(child, engine->ops, 0,
+				   __STP_UTRACE_TASK_FINDER_EVENTS);
 	return UTRACE_ACTION_RESUME;
 }
 
 static u32
-__stp_utrace_task_finder_exec(struct utrace_attached_engine *engine,
-			      struct task_struct *tsk,
-			      const struct linux_binprm *bprm,
-			      struct pt_regs *regs)
+__stp_utrace_task_finder_report_exec(struct utrace_attached_engine *engine,
+				     struct task_struct *tsk,
+				     const struct linux_binprm *bprm,
+				     struct pt_regs *regs)
 {
 	size_t filelen;
 	struct list_head *tgt_node;
 	struct stap_task_finder_target *tgt;
 	int found_node = 0;
 
+	if (atomic_read(&__stp_task_finder_state) != __STP_TF_RUNNING)
+		return UTRACE_ACTION_RESUME;
+
 	// On exec, check bprm
 	if (bprm->filename == NULL)
 		return UTRACE_ACTION_RESUME;
@@ -258,6 +284,8 @@ __stp_utrace_task_finder_exec(struct utrace_attached_engine *engine,
 		struct list_head *cb_node;
 		list_for_each(cb_node, &tgt->callback_list_head) {
 			struct stap_task_finder_target *cb_tgt;
+			int rc;
+
 			cb_tgt = list_entry(cb_node,
 					    struct stap_task_finder_target,
 					    callback_list);
@@ -274,31 +302,59 @@ __stp_utrace_task_finder_exec(struct utrace_attached_engine *engine,
 			}
 
 			// Set up thread death notification.
-			memset(&cb_tgt->ops, 0, sizeof(cb_tgt->ops));
-			cb_tgt->ops.report_death
-				= &__stp_utrace_task_finder_death;
-
-			engine = utrace_attach(tsk,
-					       UTRACE_ATTACH_CREATE,
-					       &cb_tgt->ops, cb_tgt);
-			if (IS_ERR(engine)) {
-				_stp_error("attach to exec'ed %d failed: %ld",
-					   (int)tsk->pid,
-					   PTR_ERR(engine));
-			}
-			else {
-				utrace_set_flags(tsk, engine,
+			rc = __stp_utrace_attach(tsk, &cb_tgt->ops, cb_tgt,
 						 __STP_UTRACE_ATTACHED_TASK_EVENTS);
-				cb_tgt->engine_attached = 1;
-			}
+			if (rc != 0 && rc != EPERM)
+				break;
+			cb_tgt->engine_attached = 1;
 		}
 	}
 	return UTRACE_ACTION_RESUME;
 }
 
+static u32
+stap_utrace_task_finder_report_death(struct utrace_attached_engine *engine,
+				     struct task_struct *tsk)
+{
+	return UTRACE_ACTION_DETACH;
+}
+
+static u32
+__stp_utrace_task_finder_target_death(struct utrace_attached_engine *engine,
+				      struct task_struct *tsk)
+{
+	struct stap_task_finder_target *tgt = engine->data;
+
+	if (atomic_read(&__stp_task_finder_state) != __STP_TF_RUNNING) {
+		return UTRACE_ACTION_DETACH;
+	}
+
+	// The first implementation of this added a
+	// UTRACE_EVENT(DEATH) handler to
+	// __stp_utrace_task_finder_ops.  However, dead threads don't
+	// have a mm_struct, so we can't find the exe's path.  So, we
+	// don't know which callback(s) to call.
+	//
+	// So, now when an "interesting" thread is found, we add a
+	// separate UTRACE_EVENT(DEATH) handler for every probe.
+
+	if (tgt != NULL && tgt->callback != NULL) {
+		int rc;
+
+		// Call the callback
+		rc = tgt->callback(tsk, 0, tgt);
+		if (rc != 0) {
+			_stp_error("death callback for %d failed: %d",
+				   (int)tsk->pid, rc);
+		}
+	}
+	return UTRACE_ACTION_DETACH;
+}
+
 struct utrace_engine_ops __stp_utrace_task_finder_ops = {
-	.report_clone = __stp_utrace_task_finder_clone,
-	.report_exec = __stp_utrace_task_finder_exec,
+	.report_clone = __stp_utrace_task_finder_report_clone,
+	.report_exec = __stp_utrace_task_finder_report_exec,
+	.report_death = stap_utrace_task_finder_report_death,
 };
 
 int
@@ -314,44 +370,36 @@ stap_start_task_finder(void)
 		return ENOMEM;
 	}
 
+	atomic_set(&__stp_task_finder_state, __STP_TF_RUNNING);
+
 	rcu_read_lock();
 	for_each_process(tsk) {
-		struct utrace_attached_engine *engine;
 		struct mm_struct *mm;
 		char *mmpath;
 		size_t mmpathlen;
 		struct list_head *tgt_node;
 
+		/* Attach to the thread */
+		rc = __stp_utrace_attach(tsk, &__stp_utrace_task_finder_ops, 0,
+					 __STP_UTRACE_TASK_FINDER_EVENTS);
+		if (rc == EPERM) {
+			/* Ignore EPERM errors, which mean this wasn't
+			 * a thread we can attach to. */
+			rc = 0;
+			continue;
+		}
+		else if (rc != 0) {
+			/* If we get a real error, quit. */
+			break;
+		}
+
+		/* Grab the path associated with this task. */
 		mm = get_task_mm(tsk);
 		if (! mm) {
 		    /* If the thread doesn't have a mm_struct, it is
 		     * a kernel thread which we need to skip. */
 		    continue;
 		}
-
-		/* Attach to the thread */
-		engine = utrace_attach(tsk, UTRACE_ATTACH_CREATE,
-				       &__stp_utrace_task_finder_ops, 0);
-		if (IS_ERR(engine)) {
-			int error = -PTR_ERR(engine);
-			if (error != ENOENT) {
-				mmput(mm);
-				_stp_error("utrace_attach returned error %d on pid %d",
-					   error, (int)tsk->pid);
-				rc = error;
-				break;
-			}
-		}
-		else if (unlikely(engine == NULL)) {
-			mmput(mm);
-			_stp_error("utrace_attach returned NULL on pid %d",
-				   (int)tsk->pid);
-			rc = EFAULT;
-			break;
-		}
-		utrace_set_flags(tsk, engine, __STP_UTRACE_TASK_FINDER_EVENTS);
-
-		/* Check the thread's exe's path/pid against our list. */
 		mmpath = __stp_get_mm_path(mm, mmpath_buf, PATH_MAX);
 		mmput(mm);		/* We're done with mm */
 		if (IS_ERR(mmpath)) {
@@ -361,6 +409,7 @@ stap_start_task_finder(void)
 			break;
 		}
 
+		/* Check the thread's exe's path/pid against our list. */
 		mmpathlen = strlen(mmpath);
 		list_for_each(tgt_node, &__stp_task_finder_list) {
 			struct stap_task_finder_target *tgt;
@@ -394,10 +443,19 @@ stap_start_task_finder(void)
 						   (int)tsk->pid, rc);
 					break;
 				}
+
+				// Set up thread death notification.
+				rc = __stp_utrace_attach(tsk, &cb_tgt->ops,
+							 cb_tgt,
+							 __STP_UTRACE_ATTACHED_TASK_EVENTS);
+				if (rc != 0 && rc != EPERM)
+					break;
+				cb_tgt->engine_attached = 1;
 			}
 		}
 	}
 	rcu_read_unlock();
+
 	_stp_kfree(mmpath_buf);
 	return rc;
 }
@@ -405,6 +463,8 @@ stap_start_task_finder(void)
 static void
 stap_stop_task_finder(void)
 {
+	atomic_set(&__stp_task_finder_state, __STP_TF_STOPPING);
 	stap_utrace_detach_ops(&__stp_utrace_task_finder_ops);
 	__stp_task_finder_cleanup();
+	atomic_set(&__stp_task_finder_state, __STP_TF_STOPPED);
 }