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 __task_data_init(void *data)
 118 {
 119         struct task_struct *g, *t;
 120         unsigned long addr;
 121         int ret;
 122
 123         addr = swap_ksyms_substr("copy_process");
 124         if (addr == 0) {
 125                 printk(TD_PREFIX "Cannot find address for copy_process\n");
 126                 return -EINVAL;
 127         }
 128         copy_process_rp.kp.addr = (kprobe_opcode_t *)addr;
 129         ret = swap_register_kretprobe(&copy_process_rp);
 130         if (ret)
 131                 goto reg_failed;
 132
 133         addr = swap_ksyms_substr("do_exit");
 134         if (addr == 0) {
 135                 printk(TD_PREFIX "Cannot find address for do_exit\n");
 136                 return -EINVAL;
 137         }
 138         do_exit_probe.addr = (kprobe_opcode_t *)addr;
 139         ret = swap_register_kprobe(&do_exit_probe);
 140         if (ret)
 141                 goto unreg_copy_process;
 142
 143         do_each_thread(g, t) {
 144                 swap_task_data_clean(t);
 145         } while_each_thread(g, t);
 146
 147         return 0;
 148
 149 unreg_copy_process:
 150         swap_unregister_kretprobe(&copy_process_rp);
 151
 152 reg_failed:
 153         printk(TD_PREFIX "0x%lx: probe registration failed\n", addr);
 154
 155         return ret;
 156 }
 157
 158 static int __task_data_exit(void *data)
 159 {
 160         struct task_struct *g, *t;
 161         struct task_data *td;
 162
 163         swap_unregister_kprobe(&do_exit_probe);
 164
 165         do_each_thread(g, t) {
 166                 td = __td(t);
 167                 __td_free(td);
 168         } while_each_thread(g, t);
 169
 170         return 0;
 171 }
 172
 173 static void task_data_start(void)
 174 {
 175         int ret;
 176
 177         /* stop_machine: cannot get tasklist_lock from module */
 178         ret = stop_machine(__task_data_init, NULL, NULL);
 179         if (ret)
 180                 printk(TD_PREFIX "task data initialization failed: %d\n", ret);
 181 }
 182
 183 static void task_data_stop(void)
 184 {
 185         int ret;
 186
 187         swap_unregister_kretprobe(&copy_process_rp);
 188
 189         /* stop_machine: the same here */
 190         ret = stop_machine(__task_data_exit, NULL, NULL);
 191         if (ret) {
 192                 printk(TD_PREFIX "task data cleanup failed: %d\n", ret);
 193                 /* something went wrong: at least make sure we unregister
 194                  * all the installed probes */
 195                 swap_unregister_kprobe(&do_exit_probe);
 196         }
 197 }
 198
 199 static int task_data_init(void)
 200 {
 201         int ret = 0;
 202
 203         __task_data_cbs_start_h = us_manager_reg_cb(START_CB, task_data_start);
 204
 205         if (__task_data_cbs_start_h < 0) {
 206                 ret = __task_data_cbs_start_h;
 207                 printk(KERN_ERR TD_PREFIX "start_cb registration failed\n");
 208                 goto out;
 209         }
 210
 211         __task_data_cbs_stop_h = us_manager_reg_cb(STOP_CB_TD, task_data_stop);
 212
 213         if (__task_data_cbs_stop_h < 0) {
 214                 ret = __task_data_cbs_stop_h;
 215                 us_manager_unreg_cb(__task_data_cbs_start_h);
 216                 printk(KERN_ERR TD_PREFIX "stop_cb registration failed\n");
 217         }
 218
 219 out:
 220         return ret;
 221 }
 222
 223 static void task_data_exit(void)
 224 {
 225         us_manager_unreg_cb(__task_data_cbs_start_h);
 226         us_manager_unreg_cb(__task_data_cbs_stop_h);
 227 }
 228
 229 SWAP_LIGHT_INIT_MODULE(NULL, task_data_init, task_data_exit, NULL, NULL);
 230
 231
 232 MODULE_LICENSE("GPL");
 233 MODULE_DESCRIPTION("SWAP Task Data Module");