[FEATURE] add instrumentation child threads

diff --git a/driver/helper.c b/driver/helper.c
index 811577e..922e155 100644 (file)
--- a/driver/helper.c
+++ b/driver/helper.c
@@ -35,47 +35,41 @@ static int entry_handler_pf(struct kretprobe_instance *ri, struct pt_regs *regs)
 /* Detects when IPs are really loaded into phy mem and installs probes. */
 static int ret_handler_pf(struct kretprobe_instance *ri, struct pt_regs *regs)
 {
-       struct task_struct *task = current->group_leader;
-       struct mm_struct *mm = task->mm;
+       struct task_struct *task;
+       struct sspt_proc *proc;
+
        /*
         * Because process threads have same address space
         * we instrument only group_leader of all this threads
         */
-       struct pf_data *data;
-       unsigned long addr = 0;
-       int valid_addr;
-
-       if (task->flags & PF_KTHREAD) {
-               goto out;
-       }
-
-       if (!is_us_instrumentation()) {
-               goto out;
-       }
+       task = current->group_leader;
+       if (task->flags & PF_KTHREAD)
+               return 0;
 
-       data = (struct pf_data *)ri->data;
-       addr = data->addr;
+       if (!is_us_instrumentation())
+               return 0;
 
-       valid_addr = mm && page_present(mm, addr);
-       if (!valid_addr) {
-               goto out;
-       }
+       proc = sspt_proc_get_by_task(task);
+       if (proc)
+               goto install_proc;
 
        task = check_task(task);
        if (task) {
-               struct sspt_proc *proc;
-               proc = sspt_proc_get_by_task_or_new(task);
-               if (proc) {
-                       if (proc->first_install) {
-                               unsigned long page = addr & PAGE_MASK;
-                               sspt_proc_install_page(proc, page);
-                       } else {
-                               sspt_proc_install(proc);
-                       }
-               }
+               proc = sspt_proc_get_new(task);
+               goto install_proc;
+       }
+
+       return 0;
+
+install_proc:
+       if (proc->first_install) {
+               unsigned long page;
+               page = ((struct pf_data *)ri->data)->addr & PAGE_MASK;
+               sspt_proc_install_page(proc, page);
+       } else {
+               sspt_proc_install(proc);
        }
 
-out:
        return 0;
 }
 
@@ -110,14 +104,21 @@ static void rm_uprobes_child(struct task_struct *task)
 /* Delete uprobs in children at fork */
 static int ret_handler_cp(struct kretprobe_instance *ri, struct pt_regs *regs)
 {
-       struct task_struct* task = (struct task_struct *)regs_return_value(regs);
+       struct task_struct *task = (struct task_struct *)regs_return_value(regs);
 
        if(!task || IS_ERR(task))
                goto out;
 
-       if(task->mm != current->mm)     /* check flags CLONE_VM */
+       if(task->mm != current->mm) {   /* check flags CLONE_VM */
                rm_uprobes_child(task);
 
+               if (check_task(current)) {
+                       struct sspt_proc *proc;
+
+                       proc = sspt_proc_get_new(task);
+                       sspt_proc_install(proc);
+               }
+       }
 out:
        return 0;
 }

diff --git a/driver/sspt/sspt_file.c b/driver/sspt/sspt_file.c
index 2f5da69..c0d3e69 100644 (file)
--- a/driver/sspt/sspt_file.c
+++ b/driver/sspt/sspt_file.c
@@ -251,7 +251,7 @@ void sspt_file_set_mapping(struct sspt_file *file, struct vm_area_struct *vma)
        file->vm_start = vma->vm_start;
        file->vm_end = vma->vm_end;
 
-       pack_event_info(DYN_LIB_PROBE_ID, RECORD_ENTRY, "dspdd",
-                       task->tgid, file->name, vma->vm_start,
-                       vma->vm_end - vma->vm_start, app_flag);
+       pack_task_event_info(task, DYN_LIB_PROBE_ID, RECORD_ENTRY, "dspdd",
+                            task->tgid, file->name, vma->vm_start,
+                            vma->vm_end - vma->vm_start, app_flag);
 }

diff --git a/driver/sspt/sspt_proc.c b/driver/sspt/sspt_proc.c
index 2e33c68..5d8baae 100644 (file)
--- a/driver/sspt/sspt_proc.c
+++ b/driver/sspt/sspt_proc.c
@@ -104,13 +104,22 @@ static void add_proc_probes(struct sspt_proc *proc)
        list_add_tail(&proc->list, &proc_probes_list);
 }
 
+struct sspt_proc *sspt_proc_get_new(struct task_struct *task)
+{
+       struct sspt_proc *proc;
+
+       proc = sspt_proc_copy(us_proc_info.pp, task);
+       proc->sm = create_sm_us(task);
+       add_proc_probes(proc);
+
+       return proc;
+}
+
 struct sspt_proc *sspt_proc_get_by_task_or_new(struct task_struct *task)
 {
        struct sspt_proc *proc = sspt_proc_get_by_task(task);
        if (proc == NULL) {
-               proc = sspt_proc_copy(us_proc_info.pp, task);
-               proc->sm = create_sm_us(task);
-               add_proc_probes(proc);
+               proc = sspt_proc_get_new(task);
        }
 
        return proc;
@@ -234,8 +243,8 @@ void sspt_proc_install(struct sspt_proc *proc)
                        struct sspt_file *file = sspt_proc_find_file(proc, dentry);
                        if (file) {
                                if (!file->loaded) {
-                                       sspt_file_set_mapping(file, vma);
                                        file->loaded = 1;
+                                       sspt_file_set_mapping(file, vma);
                                }
 
                                sspt_file_install(file);

diff --git a/driver/sspt/sspt_proc.h b/driver/sspt/sspt_proc.h
index 02ceced..453e004 100644 (file)
--- a/driver/sspt/sspt_proc.h
+++ b/driver/sspt/sspt_proc.h
@@ -46,6 +46,7 @@ struct sspt_proc *sspt_proc_create(struct dentry* dentry, struct task_struct *ta
 struct sspt_proc *sspt_proc_copy(struct sspt_proc *proc, struct task_struct *task);
 void sspt_proc_free(struct sspt_proc *proc);
 
+struct sspt_proc *sspt_proc_get_new(struct task_struct *task);
 struct sspt_proc *sspt_proc_get_by_task(struct task_struct *task);
 struct sspt_proc *sspt_proc_get_by_task_or_new(struct task_struct *task);
 void sspt_proc_free_all(void);