[FEATURE] Implement swap_task_data module

The module adds per-task data storage at the end of the kernel
stack.

Change-Id: Ibaccba06aaa9944eee4607a35be80dd32a83932a
Signed-off-by: Vasiliy Ulyanov <v.ulyanov@samsung.com>

5 files changed, 215 insertions, 1 deletions

diff --git a/Kbuild b/Kbuild
index 54f3c3ec..0fa4d0ca 100644
--- a/Kbuild
+++ b/Kbuild
@@ -13,4 +13,5 @@ obj-m := buffer/ \
          energy/ \
          parser/ \
          retprobe/ \
-         webprobe/
+         webprobe/ \
+         task_data/
diff --git a/build.sh b/build.sh
index 3d21d976..8fcaeef6 100755
--- a/build.sh
+++ b/build.sh
@@ -36,6 +36,7 @@ energy_dir=${modules_dir}/energy
 parser_dir=${modules_dir}/parser
 retprobe_dir=${modules_dir}/retprobe
 webprobe_dir=${modules_dir}/webprobe
+task_data_dir=${modules_dir}/task_data
 
 buffer_module_name=swap_buffer.ko
 driver_module_name=swap_driver.ko
@@ -51,6 +52,7 @@ parser_module_name=swap_message_parser.ko
 ksyms_module_name=swap_ksyms.ko
 retprobe_module_name=swap_retprobe.ko
 webprobe_module_name=swap_webprobe.ko
+task_data_module_name=swap_task_data.ko
 
 install_dir="/opt/swap/sdk"
 
diff --git a/task_data/Kbuild b/task_data/Kbuild
new file mode 100644
index 00000000..7505fcfe
--- /dev/null
+++ b/task_data/Kbuild
@@ -0,0 +1,4 @@
+EXTRA_CFLAGS := $(extra_cflags)
+
+obj-m := swap_task_data.o
+swap_task_data-y := task_data.o
diff --git a/task_data/task_data.c b/task_data/task_data.c
new file mode 100644
index 00000000..13638325
--- /dev/null
+++ b/task_data/task_data.c
@@ -0,0 +1,190 @@
+#include <linux/kernel.h>
+#include <linux/module.h>
+#include <linux/sched.h>
+#include <linux/stop_machine.h>
+#include <kprobe/swap_kprobes.h>
+#include "task_data.h"
+
+/* lower bits are used as flags */
+#define TD_MAGIC_MASK 0xfffffff0
+#define TD_FLAGS_MASK (~TD_MAGIC_MASK)
+
+#define __DEFINE_TD_MAGIC(m) ((m) & TD_MAGIC_MASK)
+
+#define TD_MAGIC __DEFINE_TD_MAGIC(0xbebebebe)
+#define TD_OFFSET 1  /* skip STACK_END_MAGIC */
+#define TD_PREFIX "[TASK_DATA] "
+
+struct task_data {
+	void *data;
+	unsigned long magic;
+} __attribute__((packed));
+
+#define get_magic(td) ((td)->magic & TD_MAGIC_MASK)
+#define get_flags(td) ((td)->magic & TD_FLAGS_MASK)
+
+static inline struct task_data *__td(struct task_struct *task)
+{
+	return (struct task_data *)(end_of_stack(task) + TD_OFFSET);
+}
+
+static inline bool __td_check(struct task_data *td)
+{
+	return (get_magic(td) == TD_MAGIC);
+}
+
+static inline void __td_init(struct task_data *td, void *data,
+			     unsigned long flags)
+{
+	td->magic = TD_MAGIC | (flags & TD_FLAGS_MASK);
+	td->data = data;
+}
+
+static inline void __td_free(struct task_data *td)
+{
+	unsigned long flags = get_flags(td);
+	bool ok = __td_check(td);
+
+	/* freeing the data if consistency check fails is dangerous:
+	 * better leave it as a memory leak instead */
+	if (ok) {
+		if ((flags & SWAP_TD_FREE) && td->data)
+			kfree(td->data);
+		td->magic = 0;
+		td->data = NULL;
+		return;
+	}
+
+	WARN(!ok, TD_PREFIX "td(%p) check failed: %08lx", td, get_magic(td));
+}
+
+void *swap_task_data_get(struct task_struct *task, int *ok)
+{
+	struct task_data *td = __td(task);
+
+	if (ok)
+		*ok = __td_check(td);
+
+	return td->data;
+}
+EXPORT_SYMBOL_GPL(swap_task_data_get);
+
+void swap_task_data_set(struct task_struct *task, void *data,
+			unsigned long flags)
+{
+	struct task_data *td = __td(task);
+
+	__td_init(td, data, flags);
+}
+EXPORT_SYMBOL_GPL(swap_task_data_set);
+
+static int copy_process_ret_handler(struct kretprobe_instance *ri,
+				    struct pt_regs *regs)
+{
+	struct task_struct *task;
+
+	task = (struct task_struct *)regs_return_value(regs);
+	if (task)
+		swap_task_data_clean(task);
+
+	return 0;
+}
+
+static int do_exit_handler(struct kprobe *p, struct pt_regs *regs)
+{
+	struct task_data *td = __td(current);
+
+	__td_free(td);
+
+	return 0;
+}
+
+static struct kretprobe copy_process_rp = {
+	.kp.symbol_name = "copy_process",
+	.handler = copy_process_ret_handler
+};
+
+static struct kprobe do_exit_probe = {
+	.symbol_name = "do_exit",
+	.pre_handler = do_exit_handler
+};
+
+static int __task_data_init(void *data)
+{
+	struct task_struct *g, *t;
+	const char *sym;
+	int ret;
+
+	sym = copy_process_rp.kp.symbol_name;
+	ret = swap_register_kretprobe(&copy_process_rp);
+	if (ret)
+		goto reg_failed;
+
+	sym = do_exit_probe.symbol_name;
+	ret = swap_register_kprobe(&do_exit_probe);
+	if (ret)
+		goto unreg_copy_process;
+
+	do_each_thread(g, t) {
+		swap_task_data_clean(t);
+	} while_each_thread(g, t);
+
+	return 0;
+
+unreg_copy_process:
+	swap_unregister_kretprobe(&copy_process_rp);
+
+reg_failed:
+	printk(TD_PREFIX "%s: probe registration failed\n", sym);
+
+	return ret;
+}
+
+static int __task_data_exit(void *data)
+{
+	struct task_struct *g, *t;
+	struct task_data *td;
+
+	swap_unregister_kprobe(&do_exit_probe);
+	swap_unregister_kretprobe(&copy_process_rp);
+
+	do_each_thread(g, t) {
+		td = __td(t);
+		__td_free(td);
+	} while_each_thread(g, t);
+
+	return 0;
+}
+
+static int __init task_data_init(void)
+{
+	int ret;
+
+	/* stop_machine: cannot get tasklist_lock from module */
+	ret = stop_machine(__task_data_init, NULL, NULL);
+	if (ret)
+		printk(TD_PREFIX "task data initialization failed: %d\n", ret);
+
+	return ret;
+}
+
+static void __exit task_data_exit(void)
+{
+	int ret;
+
+	/* stop_machine: the same here */
+	ret = stop_machine(__task_data_exit, &copy_process_rp, NULL);
+	if (ret) {
+		printk(TD_PREFIX "task data cleanup failed: %d\n", ret);
+		/* something went wrong: at least make sure we unregister
+		 * all the installed probes */
+		swap_unregister_kprobe(&do_exit_probe);
+		swap_unregister_kretprobe(&copy_process_rp);
+	}
+}
+
+module_init(task_data_init);
+module_exit(task_data_exit);
+
+MODULE_LICENSE("GPL");
+MODULE_DESCRIPTION("SWAP Task Data Module");
diff --git a/task_data/task_data.h b/task_data/task_data.h
new file mode 100644
index 00000000..a5e67cf4
--- /dev/null
+++ b/task_data/task_data.h
@@ -0,0 +1,17 @@
+#ifndef __TASK_DATA__
+#define __TASK_DATA__
+
+#define SWAP_TD_FREE 0x1 /* kfree task data automatically */
+
+struct task_struct;
+
+void *swap_task_data_get(struct task_struct *task, int *ok);
+void swap_task_data_set(struct task_struct *task, void *data,
+			unsigned long flags);
+
+static inline void swap_task_data_clean(struct task_struct *task)
+{
+	swap_task_data_set(task, NULL, 0);
+}
+
+#endif /* __TASK_DATA__ */