diff --git a/.gitignore b/.gitignore
index ef046e0..87945ee 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,7 @@ kpatch-diff-gen/kpatch-diff-gen
 *.ko
 *.ko.cmd
 *.mod.c
+*.swp
 kmod/.tmp_versions/
 kmod/Module.symvers
 kmod/modules.order
diff --git a/kmod/base.c b/kmod/base.c
index 49f528e..f4a124c 100644
--- a/kmod/base.c
+++ b/kmod/base.c
@@ -152,7 +152,7 @@ int kpatch_register(struct module *mod, void *kpatch_relas,
 	int size;
 	int num_patches;
 	struct kpatch_patch *patches;
-	struct kpatch_func *funcs;
+	struct kpatch_func *funcs, *f;
 
 	num_relas = (kpatch_relas_end - kpatch_relas) / sizeof(*relas);
 	relas = kpatch_relas;
@@ -181,6 +181,7 @@ int kpatch_register(struct module *mod, void *kpatch_relas,
 		//printk("%p <- %lx\n", loc, val);
 		//printk("%lx\n", (unsigned long)__va(__pa((unsigned long)loc)));
 		//loc = __va(__pa((unsigned long)loc));
+		/* TODO: safe to assume it was ro to start with? */
 		set_memory_rw((unsigned long)loc & PAGE_MASK, 1);
 		ret = probe_kernel_write(loc, &val, size);
 		set_memory_ro((unsigned long)loc & PAGE_MASK, 1);
@@ -199,7 +200,18 @@ int kpatch_register(struct module *mod, void *kpatch_relas,
 	for (i = 0; i < num_patches; i++) {
 		funcs[i].old_func_addr = patches[i].orig;
 		funcs[i].new_func_addr = patches[i].new;
-		funcs[i].old_func_name = "FIXME";
+		funcs[i].mod = mod;
+		funcs[i].old_func_name = "TODO";
+		/* TODO: need old_func_addr_end too */
+		/* TODO: verify name/address with kallsyms */
+
+		/* Do any needed incremental patching. */
+		for (f = kpatch_funcs; f->old_func_name; f++) {
+			if (funcs[i].old_func_addr == f->old_func_addr) {
+				funcs[i].old_func_addr = f->new_func_addr;
+				ref_module(funcs[i].mod, f->mod);
+			}
+		}
 
 		ret = ftrace_set_filter_ip(&kpatch_ftrace_ops, patches[i].orig,
 					   0, 0);
@@ -226,13 +238,6 @@ int kpatch_register(struct module *mod, void *kpatch_relas,
 			goto out;
 		}
 
-		/* Do any needed incremental patching. */
-		for (g = kpatch_funcs; g->old_func_name; g++)
-			if (f->old_func_addr == g->old_func_addr) {
-				f->old_func_addr = g->new_func_addr;
-				ref_module(f->owner, g->owner);
-			}
-
 
 		if (!kallsyms_lookup_size_offset(f->old_func_addr, &size,
 						 &offset)) {
@@ -289,7 +294,6 @@ out:
 }
 EXPORT_SYMBOL(kpatch_register);
 
-#if 0
 /* Called from stop_machine */
 static int kpatch_remove_patch(void *data)
 {
@@ -300,7 +304,7 @@ static int kpatch_remove_patch(void *data)
 	if (ret)
 		goto out;
 
-	for (i = 0; i < KPATCH_MAX_FUNCS; i++)
+	for (i = 0; i < KPATCH_MAX_FUNCS; i++) /* TODO iterate by pointer */
 		if (kpatch_funcs[i].old_func_addr == funcs->old_func_addr)
 			break;
 
@@ -321,13 +325,21 @@ static int kpatch_remove_patch(void *data)
 out:
 	return ret;
 }
-#endif
 
 int kpatch_unregister(struct module *mod)
 {
 	int ret = 0;
-#if 0
-	struct kpatch_func *f;
+	struct kpatch_func *funcs, *f;
+	int num_funcs, i;
+
+	num_funcs = kpatch_num_funcs(kpatch_funcs);
+
+	funcs = kmalloc((num_funcs + 1) * sizeof(*funcs), GFP_KERNEL);
+
+	for (f = kpatch_funcs, i = 0; f->old_func_name; f++)
+		if (f->mod == mod)
+			memcpy(&funcs[i++], f, sizeof(*funcs));
+	memset(&funcs[i], 0, sizeof(*funcs));
 
 	ret = stop_machine(kpatch_remove_patch, funcs, NULL);
 	if (ret)
@@ -353,7 +365,7 @@ int kpatch_unregister(struct module *mod)
 	}
 
 out:
-#endif
+	kfree(funcs);
 	return ret;
 }
 EXPORT_SYMBOL(kpatch_unregister);
diff --git a/kmod/kpatch.h b/kmod/kpatch.h
index 3216cf1..08b89eb 100644
--- a/kmod/kpatch.h
+++ b/kmod/kpatch.h
@@ -7,7 +7,7 @@ struct kpatch_func {
 	unsigned long new_func_addr;
 	char *old_func_name;
 	unsigned long old_func_addr_end;
-	struct module *owner;
+	struct module *mod;
 };
 
 struct kpatch_rela {
diff --git a/kpatch-create b/kpatch-create
index d4b64f7..319c1c7 100755
--- a/kpatch-create
+++ b/kpatch-create
@@ -9,6 +9,7 @@ set -o pipefail
 CUR_DIR="$PWD"
 SCRIPT="`basename $BASH_SOURCE`"
 SCRIPT_DIR="`dirname $BASH_SOURCE`"
+SCRIPT_DIR="`readlink -f \"$SCRIPT_DIR\"`"
 
 KMOD_DIR="$SCRIPT_DIR/kmod"
 KPATCH_REG="$KMOD_DIR/kpatch-register.o"
@@ -106,7 +107,7 @@ touch "$PROGRESS_FILE"
 $MAKE_CMD > /dev/null
 rm -f "$PROGRESS_FILE"
 
-scriptecho "generating kpatch binaries"
+scriptecho "diffing binaries"
 find . -type f -name '*.o.orig_kpatch' | while read file; do
     origfile="${file#./}"
     newfile="${origfile%.orig_kpatch}"
@@ -118,6 +119,7 @@ find . -type f -name '*.o.orig_kpatch' | while read file; do
     "$KPATCH_GEN" "$origfile" "$newfile" -v vmlinux.orig_kpatch -o "$newfile.kpatch_gen"
 done
 
+scriptecho "generating kpatch modules"
 unset CROSS_COMPILE
 cp -a "$KMOD_DIR" "$TMPDIR/kmod"
 
@@ -131,4 +133,4 @@ make -C "$KERNEL_DIR" M="$TMPDIR/kmod" kpatch-patch.ko > /dev/null
 
 cp kpatch.ko kpatch-patch.ko "$OUTPUT_DIR"
 
-scriptecho Success!
+scriptecho success!