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
  89 static atomic_t start_flag = ATOMIC_INIT(0);
  90
  91 static int copy_process_ret_handler(struct kretprobe_instance *ri,
  92                                     struct pt_regs *regs)
  93 {
  94         struct task_struct *task;
  95
  96         if (atomic_read(&start_flag) == 0)
  97                 return 0;
  98
  99         task = (struct task_struct *)regs_return_value(regs);
 100         if (!IS_ERR(task))
 101                 swap_task_data_clean(task);
 102
 103         return 0;
 104 }
 105
 106 static int do_exit_handler(struct kprobe *p, struct pt_regs *regs)
 107 {
 108         struct task_data *td;
 109
 110         if (atomic_read(&start_flag) == 0)
 111                 return 0;
 112
 113         td = __td(current);
 114         __td_free(td);
 115
 116         return 0;
 117 }
 118
 119 static struct kretprobe copy_process_rp = {
 120         .handler = copy_process_ret_handler
 121 };
 122
 123 static struct kprobe do_exit_probe = {
 124         .pre_handler = do_exit_handler
 125 };
 126
 127 static int __task_data_init(void *data)
 128 {
 129         struct task_struct *g, *t;
 130
 131         /* set start_flags */
 132         atomic_set(&start_flag, 1);
 133
 134         do_each_thread(g, t) {
 135                 swap_task_data_clean(t);
 136         } while_each_thread(g, t);
 137
 138         return 0;
 139 }
 140
 141 static int __task_data_exit(void *data)
 142 {
 143         struct task_struct *g, *t;
 144         struct task_data *td;
 145
 146         do_each_thread(g, t) {
 147                 td = __td(t);
 148                 __td_free(td);
 149         } while_each_thread(g, t);
 150
 151         /* reset start_flags */
 152         atomic_set(&start_flag, 0);
 153
 154         return 0;
 155 }
 156
 157 static void task_data_start(void)
 158 {
 159         int ret;
 160
 161         ret = swap_register_kprobe(&do_exit_probe);
 162         if (ret) {
 163                 pr_err(TD_PREFIX "register on 'do_exit' failed: ret=%d\n", ret);
 164                 return;
 165         }
 166
 167         ret = swap_register_kretprobe(&copy_process_rp);
 168         if (ret) {
 169                 swap_unregister_kprobe(&do_exit_probe);
 170                 pr_err(TD_PREFIX "register on 'copy_process' failed: ret=%d\n", ret);
 171                 return;
 172         }
 173
 174         /* stop_machine: cannot get tasklist_lock from module */
 175         ret = stop_machine(__task_data_init, NULL, NULL);
 176         if (ret)
 177                 printk(TD_PREFIX "task data initialization failed: %d\n", ret);
 178 }
 179
 180 static void task_data_stop(void)
 181 {
 182         int ret;
 183
 184         /* stop_machine: the same here */
 185         ret = stop_machine(__task_data_exit, NULL, NULL);
 186         if (ret) {
 187                 printk(TD_PREFIX "task data cleanup failed: %d\n", ret);
 188                 /* something went wrong: at least make sure we unregister
 189                  * all the installed probes */
 190                 swap_unregister_kprobe(&do_exit_probe);
 191         }
 192
 193         swap_unregister_kretprobe(&copy_process_rp);
 194         swap_unregister_kprobe(&do_exit_probe);
 195 }
 196
 197 static int task_data_once(void)
 198 {
 199         const char *sym;
 200
 201         sym = "copy_process";
 202         copy_process_rp.kp.addr = (kprobe_opcode_t *)swap_ksyms_substr(sym);
 203         if (copy_process_rp.kp.addr == NULL)
 204                 goto not_found;
 205
 206         sym = "do_exit";
 207         do_exit_probe.addr = (kprobe_opcode_t *)swap_ksyms_substr(sym);
 208         if (do_exit_probe.addr == NULL)
 209                 goto not_found;
 210
 211         return 0;
 212
 213 not_found:
 214         pr_err(TD_PREFIX "ERROR: symbol %s(...) not found\n", sym);
 215         return -ESRCH;
 216 }
 217
 218 static int task_data_init(void)
 219 {
 220         int ret = 0;
 221
 222         __task_data_cbs_start_h = us_manager_reg_cb(START_CB, task_data_start);
 223         if (__task_data_cbs_start_h < 0) {
 224                 ret = __task_data_cbs_start_h;
 225                 printk(KERN_ERR TD_PREFIX "start_cb registration failed\n");
 226                 goto out;
 227         }
 228
 229         __task_data_cbs_stop_h = us_manager_reg_cb(STOP_CB_TD, task_data_stop);
 230         if (__task_data_cbs_stop_h < 0) {
 231                 ret = __task_data_cbs_stop_h;
 232                 us_manager_unreg_cb(__task_data_cbs_start_h);
 233                 printk(KERN_ERR TD_PREFIX "stop_cb registration failed\n");
 234         }
 235
 236 out:
 237         return ret;
 238 }
 239
 240 static void task_data_exit(void)
 241 {
 242         us_manager_unreg_cb(__task_data_cbs_start_h);
 243         us_manager_unreg_cb(__task_data_cbs_stop_h);
 244 }
 245
 246 SWAP_LIGHT_INIT_MODULE(task_data_once, task_data_init, task_data_exit,
 247                        NULL, NULL);
 248
 249 MODULE_LICENSE("GPL");
 250 MODULE_DESCRIPTION("SWAP Task Data Module");