[FIX] task_data: sync register/unregister probes

Change-Id: I2056a4895f21c3a319e6ce3b4710f5ae5bd47845
Signed-off-by: Vyacheslav Cherkashin <v.cherkashin@samsung.com>

1 files changed, 54 insertions, 53 deletions

diff --git a/task_data/task_data.c b/task_data/task_data.c
index 7dba5c1c..cc327652 100644
--- a/task_data/task_data.c
+++ b/task_data/task_data.c
@@ -85,11 +85,17 @@ void swap_task_data_set(struct task_struct *task, void *data,
 }
 EXPORT_SYMBOL_GPL(swap_task_data_set);
 
+
+static atomic_t start_flag = ATOMIC_INIT(0);
+
 static int copy_process_ret_handler(struct kretprobe_instance *ri,
 				    struct pt_regs *regs)
 {
 	struct task_struct *task;
 
+	if (atomic_read(&start_flag) == 0)
+		return 0;
+
 	task = (struct task_struct *)regs_return_value(regs);
 	if (!IS_ERR(task))
 		swap_task_data_clean(task);
@@ -99,8 +105,12 @@ static int copy_process_ret_handler(struct kretprobe_instance *ri,
 
 static int do_exit_handler(struct kprobe *p, struct pt_regs *regs)
 {
-	struct task_data *td = __td(current);
+	struct task_data *td;
 
+	if (atomic_read(&start_flag) == 0)
+		return 0;
+
+	td = __td(current);
 	__td_free(td);
 
 	return 0;
@@ -114,59 +124,18 @@ static struct kprobe do_exit_probe = {
 	.pre_handler = do_exit_handler
 };
 
-static int __set_helper_probes(void)
-{
-	unsigned long addr;
-	int ret;
-
-	addr = swap_ksyms_substr("copy_process");
-	if (addr == 0) {
-		printk(TD_PREFIX "Cannot find address for copy_process\n");
-		return -EINVAL;
-	}
-	copy_process_rp.kp.addr = (kprobe_opcode_t *)addr;
-	ret = swap_register_kretprobe(&copy_process_rp);
-	if (ret)
-		goto reg_failed;
-
-	addr = swap_ksyms_substr("do_exit");
-	if (addr == 0) {
-		printk(TD_PREFIX "Cannot find address for do_exit\n");
-		return -EINVAL;
-	}
-	do_exit_probe.addr = (kprobe_opcode_t *)addr;
-	ret = swap_register_kprobe(&do_exit_probe);
-	if (ret)
-		goto unreg_copy_process;
-
-	return 0;
-
-unreg_copy_process:
-	swap_unregister_kretprobe(&copy_process_rp);
-
-reg_failed:
-	printk(TD_PREFIX "0x%lx: probe registration failed\n", addr);
-
-	return ret;
-}
-
-static void __remove_helper_probes(void)
-{
-	swap_unregister_kretprobe(&copy_process_rp);
-	swap_unregister_kprobe(&do_exit_probe);
-}
-
 static int __task_data_init(void *data)
 {
 	struct task_struct *g, *t;
 
+	/* set start_flags */
+	atomic_set(&start_flag, 1);
+
 	do_each_thread(g, t) {
 		swap_task_data_clean(t);
 	} while_each_thread(g, t);
 
 	return 0;
-
-
 }
 
 static int __task_data_exit(void *data)
@@ -179,6 +148,9 @@ static int __task_data_exit(void *data)
 		__td_free(td);
 	} while_each_thread(g, t);
 
+	/* reset start_flags */
+	atomic_set(&start_flag, 0);
+
 	return 0;
 }
 
@@ -186,9 +158,18 @@ static void task_data_start(void)
 {
 	int ret;
 
-	ret = __set_helper_probes();
-	if (ret)
+	ret = swap_register_kprobe(&do_exit_probe);
+	if (ret) {
+		pr_err(TD_PREFIX "register on 'do_exit' failed: ret=%d\n", ret);
 		return;
+	}
+
+	ret = swap_register_kretprobe(&copy_process_rp);
+	if (ret) {
+		swap_unregister_kprobe(&do_exit_probe);
+		pr_err(TD_PREFIX "register on 'copy_process' failed: ret=%d\n", ret);
+		return;
+	}
 
 	/* stop_machine: cannot get tasklist_lock from module */
 	ret = stop_machine(__task_data_init, NULL, NULL);
@@ -200,8 +181,6 @@ static void task_data_stop(void)
 {
 	int ret;
 
-	__remove_helper_probes();
-
 	/* stop_machine: the same here */
 	ret = stop_machine(__task_data_exit, NULL, NULL);
 	if (ret) {
@@ -210,6 +189,30 @@ static void task_data_stop(void)
 		 * all the installed probes */
 		swap_unregister_kprobe(&do_exit_probe);
 	}
+
+	swap_unregister_kretprobe(&copy_process_rp);
+	swap_unregister_kprobe(&do_exit_probe);
+}
+
+static int task_data_once(void)
+{
+	const char *sym;
+
+	sym = "copy_process";
+	copy_process_rp.kp.addr = (kprobe_opcode_t *)swap_ksyms_substr(sym);
+	if (copy_process_rp.kp.addr == NULL)
+		goto not_found;
+
+	sym = "do_exit";
+	do_exit_probe.addr = (kprobe_opcode_t *)swap_ksyms_substr(sym);
+	if (do_exit_probe.addr == NULL)
+		goto not_found;
+
+	return 0;
+
+not_found:
+	pr_err(TD_PREFIX "ERROR: symbol %s(...) not found\n", sym);
+	return -ESRCH;
 }
 
 static int task_data_init(void)
@@ -217,7 +220,6 @@ static int task_data_init(void)
 	int ret = 0;
 
 	__task_data_cbs_start_h = us_manager_reg_cb(START_CB, task_data_start);
-
 	if (__task_data_cbs_start_h < 0) {
 		ret = __task_data_cbs_start_h;
 		printk(KERN_ERR TD_PREFIX "start_cb registration failed\n");
@@ -225,7 +227,6 @@ static int task_data_init(void)
 	}
 
 	__task_data_cbs_stop_h = us_manager_reg_cb(STOP_CB_TD, task_data_stop);
-
 	if (__task_data_cbs_stop_h < 0) {
 		ret = __task_data_cbs_stop_h;
 		us_manager_unreg_cb(__task_data_cbs_start_h);
@@ -242,8 +243,8 @@ static void task_data_exit(void)
 	us_manager_unreg_cb(__task_data_cbs_stop_h);
 }
 
-SWAP_LIGHT_INIT_MODULE(NULL, task_data_init, task_data_exit, NULL, NULL);
-
+SWAP_LIGHT_INIT_MODULE(task_data_once, task_data_init, task_data_exit,
+		       NULL, NULL);
 
 MODULE_LICENSE("GPL");
 MODULE_DESCRIPTION("SWAP Task Data Module");