kmod/core: Handle registering error and unroll it

Handle registering error to unroll the ftrace filter.
This also introduces get_kpatch_func() and
kpatch_remove_funcs_from_filter() for holding up
redundant loops.

Changes from v2:
 - Rebased on the latest kpatch.

Changes from v1:
 - Rename get_kpatch_func to kpatch_get_func.
 - Fix function definition style issue.
 - Do not jump to a label in "if" block.
 - Rollback the ftrace user counter if we hit an error.

Signed-off-by: Masami Hiramatsu <masami.hiramatsu.pt@hitachi.com>

kmod/core/core.c

@@ -169,6 +169,15 @@ out:
 	return ret;
 }
 
+static struct kpatch_func *kpatch_get_func(unsigned long ip)
+{
+	struct kpatch_func *f;
+
+	hash_for_each_possible(kpatch_func_hash, f, node, ip)
+		if (f->old_addr == ip)
+			return f;
+	return NULL;
+}
 
 void notrace kpatch_ftrace_handler(unsigned long ip, unsigned long parent_ip,
 				   struct ftrace_ops *op, struct pt_regs *regs)
@@ -198,6 +207,35 @@ static struct ftrace_ops kpatch_ftrace_ops __read_mostly = {
 	.flags = FTRACE_OPS_FL_SAVE_REGS,
 };
 
+/* Remove kpatch_funcs from ftrace filter */
+static int kpatch_remove_funcs_from_filter(struct kpatch_func *funcs,
+					   int num_funcs)
+{
+	int i, ret = 0;
+
+	for (i = 0; i < num_funcs; i++) {
+		struct kpatch_func *func = &funcs[i];
+
+		/*
+		 * If any other modules have also patched this function, don't
+		 * remove its ftrace handler.
+		 */
+		if (kpatch_get_func(func->old_addr))
+			continue;
+
+		/* Remove the ftrace handler for this function. */
+		ret = ftrace_set_filter_ip(&kpatch_ftrace_ops, func->old_addr,
+					   1, 0);
+		if (ret) {
+			pr_err("can't remove ftrace filter at address 0x%lx\n",
+			       func->old_addr);
+			break;
+		}
+	}
+
+	return ret;
+}
+
 int kpatch_register(struct module *mod, struct kpatch_func *funcs,
 		    int num_funcs)
 {
@@ -210,8 +248,7 @@ int kpatch_register(struct module *mod, struct kpatch_func *funcs,
 	down(&kpatch_mutex);
 
 	for (i = 0; i < num_funcs; i++) {
-		struct kpatch_func *f, *func = &funcs[i];
+		struct kpatch_func *func = &funcs[i];
-		bool found = false;
 
 		func->mod = mod;
 
@@ -219,14 +256,7 @@ int kpatch_register(struct module *mod, struct kpatch_func *funcs,
 		 * If any other modules have also patched this function, it
 		 * already has an ftrace handler.
 		 */
-		hash_for_each_possible(kpatch_func_hash, f, node,
+		if (kpatch_get_func(func->old_addr))
-				       func->old_addr) {
-			if (f->old_addr == func->old_addr) {
-				found = true;
-				break;
-			}
-		}
-		if (found)
 			continue;
 
 		/* Add an ftrace handler for this function. */
@@ -235,6 +265,7 @@ int kpatch_register(struct module *mod, struct kpatch_func *funcs,
 		if (ret) {
 			pr_err("can't set ftrace filter at address 0x%lx\n",
 			       func->old_addr);
+			num_funcs = i;
 			goto out;
 		}
 	}
@@ -244,6 +275,8 @@ int kpatch_register(struct module *mod, struct kpatch_func *funcs,
 		ret = register_ftrace_function(&kpatch_ftrace_ops);
 		if (ret) {
 			pr_err("can't register ftrace handler\n");
+			/* For the next time, the counter should be unrolled */
+			--kpatch_num_registered;
 			goto out;
 		}
 	}
@@ -259,25 +292,30 @@ int kpatch_register(struct module *mod, struct kpatch_func *funcs,
 			if (ret2)
 				pr_err("ftrace unregister failed (%d)\n", ret2);
 		}
-
-		goto out;
 	}
 
-	/* TODO: need TAINT_KPATCH */
-	pr_notice_once("tainting kernel with TAINT_USER\n");
-	add_taint(TAINT_USER, LOCKDEP_STILL_OK);
-
-	pr_notice("loaded patch module \"%s\"\n", mod->name);
 out:
+	/* Rollback the filter if we get any error */
+	if (ret)
+		kpatch_remove_funcs_from_filter(funcs, num_funcs);
+	else {
+		/* TODO: need TAINT_KPATCH */
+		pr_notice_once("tainting kernel with TAINT_USER\n");
+		add_taint(TAINT_USER, LOCKDEP_STILL_OK);
+
+		pr_notice("loaded patch module \"%s\"\n", mod->name);
+	}
+
 	up(&kpatch_mutex);
 	return ret;
+
 }
 EXPORT_SYMBOL(kpatch_register);
 
 int kpatch_unregister(struct module *mod, struct kpatch_func *funcs,
 		      int num_funcs)
 {
-	int i, ret;
+	int ret;
 	struct kpatch_stop_machine_args args = {
 		.funcs = funcs,
 		.num_funcs = num_funcs,
@@ -297,35 +335,9 @@ int kpatch_unregister(struct module *mod, struct kpatch_func *funcs,
 		}
 	}
 
-	for (i = 0; i < num_funcs; i++) {
+	ret = kpatch_remove_funcs_from_filter(funcs, num_funcs);
-		struct kpatch_func *f, *func = &funcs[i];
+	if (ret == 0)
-		bool found = false;
+		pr_notice("unloaded patch module \"%s\"\n", mod->name);
-
-		/*
-		 * If any other modules have also patched this function, don't
-		 * remove its ftrace handler.
-		 */
-		hash_for_each_possible(kpatch_func_hash, f, node,
-				       func->old_addr) {
-			if (f->old_addr == func->old_addr) {
-				found = true;
-				break;
-			}
-		}
-		if (found)
-			continue;
-
-		/* Remove the ftrace handler for this function. */
-		ret = ftrace_set_filter_ip(&kpatch_ftrace_ops, func->old_addr,
-					   1, 0);
-		if (ret) {
-			pr_err("can't remove ftrace filter at address 0x%lx\n",
-			       func->old_addr);
-			goto out;
-		}
-	}
-
-	pr_notice("unloaded patch module \"%s\"\n", mod->name);
 
 out:
 	up(&kpatch_mutex);