task_data/task_data.c

   1 #include <linux/kernel.h>
   2 #include <linux/module.h>
   3 #include <linux/sched.h>
   4 #include <linux/stop_machine.h>
   5 #include <linux/slab.h>
   6 #include <kprobe/swap_kprobes.h>
   7 #include <ksyms/ksyms.h>
   8 #include <master/swap_initializer.h>
   9 #include <us_manager/callbacks.h>
  10 #include "task_data.h"
  11
  12 /* lower bits are used as flags */
  13 #define TD_MAGIC_MASK 0xfffffff0
  14 #define TD_FLAGS_MASK (~TD_MAGIC_MASK)
  15
  16 #define __DEFINE_TD_MAGIC(m) ((m) & TD_MAGIC_MASK)
  17
  18 #define TD_MAGIC __DEFINE_TD_MAGIC(0xbebebebe)
  19 #define TD_OFFSET 1  /* skip STACK_END_MAGIC */
  20 #define TD_PREFIX "[TASK_DATA] "
  21
  22 struct task_data {
  23         void *data;
  24         unsigned long magic;
  25 } __attribute__((packed));
  26
  27 #define get_magic(td) ((td)->magic & TD_MAGIC_MASK)
  28 #define get_flags(td) ((td)->magic & TD_FLAGS_MASK)
  29
  30 static int __task_data_cbs_start_h = -1;
  31 static int __task_data_cbs_stop_h = -1;
  32
  33 static inline struct task_data *__td(struct task_struct *task)
  34 {
  35         return (struct task_data *)(end_of_stack(task) + TD_OFFSET);
  36 }
  37
  38 static inline bool __td_check(struct task_data *td)
  39 {
  40         return (get_magic(td) == TD_MAGIC);
  41 }
  42
  43 static inline void __td_init(struct task_data *td, void *data,
  44                              unsigned long flags)
  45 {
  46         td->magic = TD_MAGIC | (flags & TD_FLAGS_MASK);
  47         td->data = data;
  48 }
  49
  50 static inline void __td_free(struct task_data *td)
  51 {
  52         unsigned long flags = get_flags(td);
  53         bool ok = __td_check(td);
  54
  55         /* freeing the data if consistency check fails is dangerous:
  56          * better leave it as a memory leak instead */
  57         if (ok) {
  58                 if ((flags & SWAP_TD_FREE) && td->data)
  59                         kfree(td->data);
  60                 td->magic = 0;
  61                 td->data = NULL;
  62                 return;
  63         }
  64
  65         WARN(!ok, TD_PREFIX "td(%p) check failed: %08lx", td, get_magic(td));
  66 }
  67
  68 void *swap_task_data_get(struct task_struct *task, int *ok)
  69 {
  70         struct task_data *td = __td(task);
  71
  72         if (ok)
  73                 *ok = __td_check(td);
  74
  75         return td->data;
  76 }
  77 EXPORT_SYMBOL_GPL(swap_task_data_get);
  78
  79 void swap_task_data_set(struct task_struct *task, void *data,
  80                         unsigned long flags)
  81 {
  82         struct task_data *td = __td(task);
  83
  84         __td_init(td, data, flags);
  85 }
  86 EXPORT_SYMBOL_GPL(swap_task_data_set);
  87
  88 static int copy_process_ret_handler(struct kretprobe_instance *ri,
  89                                     struct pt_regs *regs)
  90 {
  91         struct task_struct *task;
  92
  93         task = (struct task_struct *)regs_return_value(regs);
  94         if (!IS_ERR(task))
  95                 swap_task_data_clean(task);
  96
  97         return 0;
  98 }
  99
 100 static int do_exit_handler(struct kprobe *p, struct pt_regs *regs)
 101 {
 102         struct task_data *td = __td(current);
 103
 104         __td_free(td);
 105
 106         return 0;
 107 }
 108
 109 static struct kretprobe copy_process_rp = {
 110         .handler = copy_process_ret_handler
 111 };
 112
 113 static struct kprobe do_exit_probe = {
 114         .pre_handler = do_exit_handler
 115 };
 116
 117 static int __set_helper_probes(void)
 118 {
 119         unsigned long addr;
 120         int ret;
 121
 122         addr = swap_ksyms_substr("copy_process");
 123         if (addr == 0) {
 124                 printk(TD_PREFIX "Cannot find address for copy_process\n");
 125                 return -EINVAL;
 126         }
 127         copy_process_rp.kp.addr = (kprobe_opcode_t *)addr;
 128         ret = swap_register_kretprobe(&copy_process_rp);
 129         if (ret)
 130                 goto reg_failed;
 131
 132         addr = swap_ksyms_substr("do_exit");
 133         if (addr == 0) {
 134                 printk(TD_PREFIX "Cannot find address for do_exit\n");
 135                 return -EINVAL;
 136         }
 137         do_exit_probe.addr = (kprobe_opcode_t *)addr;
 138         ret = swap_register_kprobe(&do_exit_probe);
 139         if (ret)
 140                 goto unreg_copy_process;
 141
 142         return 0;
 143
 144 unreg_copy_process:
 145         swap_unregister_kretprobe(&copy_process_rp);
 146
 147 reg_failed:
 148         printk(TD_PREFIX "0x%lx: probe registration failed\n", addr);
 149
 150         return ret;
 151 }
 152
 153 static void __remove_helper_probes(void)
 154 {
 155         swap_unregister_kretprobe(&copy_process_rp);
 156         swap_unregister_kprobe(&do_exit_probe);
 157 }
 158
 159 static int __task_data_init(void *data)
 160 {
 161         struct task_struct *g, *t;
 162
 163         do_each_thread(g, t) {
 164                 swap_task_data_clean(t);
 165         } while_each_thread(g, t);
 166
 167         return 0;
 168
 169
 170 }
 171
 172 static int __task_data_exit(void *data)
 173 {
 174         struct task_struct *g, *t;
 175         struct task_data *td;
 176
 177         do_each_thread(g, t) {
 178                 td = __td(t);
 179                 __td_free(td);
 180         } while_each_thread(g, t);
 181
 182         return 0;
 183 }
 184
 185 static void task_data_start(void)
 186 {
 187         int ret;
 188
 189         ret = __set_helper_probes();
 190         if (ret)
 191                 return;
 192
 193         /* stop_machine: cannot get tasklist_lock from module */
 194         ret = stop_machine(__task_data_init, NULL, NULL);
 195         if (ret)
 196                 printk(TD_PREFIX "task data initialization failed: %d\n", ret);
 197 }
 198
 199 static void task_data_stop(void)
 200 {
 201         int ret;
 202
 203         __remove_helper_probes();
 204
 205         /* stop_machine: the same here */
 206         ret = stop_machine(__task_data_exit, NULL, NULL);
 207         if (ret) {
 208                 printk(TD_PREFIX "task data cleanup failed: %d\n", ret);
 209                 /* something went wrong: at least make sure we unregister
 210                  * all the installed probes */
 211                 swap_unregister_kprobe(&do_exit_probe);
 212         }
 213 }
 214
 215 static int task_data_init(void)
 216 {
 217         int ret = 0;
 218
 219         __task_data_cbs_start_h = us_manager_reg_cb(START_CB, task_data_start);
 220
 221         if (__task_data_cbs_start_h < 0) {
 222                 ret = __task_data_cbs_start_h;
 223                 printk(KERN_ERR TD_PREFIX "start_cb registration failed\n");
 224                 goto out;
 225         }
 226
 227         __task_data_cbs_stop_h = us_manager_reg_cb(STOP_CB_TD, task_data_stop);
 228
 229         if (__task_data_cbs_stop_h < 0) {
 230                 ret = __task_data_cbs_stop_h;
 231                 us_manager_unreg_cb(__task_data_cbs_start_h);
 232                 printk(KERN_ERR TD_PREFIX "stop_cb registration failed\n");
 233         }
 234
 235 out:
 236         return ret;
 237 }
 238
 239 static void task_data_exit(void)
 240 {
 241         us_manager_unreg_cb(__task_data_cbs_start_h);
 242         us_manager_unreg_cb(__task_data_cbs_stop_h);
 243 }
 244
 245 SWAP_LIGHT_INIT_MODULE(NULL, task_data_init, task_data_exit, NULL, NULL);
 246
 247
 248 MODULE_LICENSE("GPL");
 249 MODULE_DESCRIPTION("SWAP Task Data Module");