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