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 "task_data.h"
  10
  11 /* lower bits are used as flags */
  12 #define TD_MAGIC_MASK 0xfffffff0
  13 #define TD_FLAGS_MASK (~TD_MAGIC_MASK)
  14
  15 #define __DEFINE_TD_MAGIC(m) ((m) & TD_MAGIC_MASK)
  16
  17 #define TD_MAGIC __DEFINE_TD_MAGIC(0xbebebebe)
  18 #define TD_OFFSET 1  /* skip STACK_END_MAGIC */
  19 #define TD_PREFIX "[TASK_DATA] "
  20
  21 struct task_data {
  22         void *data;
  23         unsigned long magic;
  24 } __attribute__((packed));
  25
  26 #define get_magic(td) ((td)->magic & TD_MAGIC_MASK)
  27 #define get_flags(td) ((td)->magic & TD_FLAGS_MASK)
  28
  29 static inline struct task_data *__td(struct task_struct *task)
  30 {
  31         return (struct task_data *)(end_of_stack(task) + TD_OFFSET);
  32 }
  33
  34 static inline bool __td_check(struct task_data *td)
  35 {
  36         return (get_magic(td) == TD_MAGIC);
  37 }
  38
  39 static inline void __td_init(struct task_data *td, void *data,
  40                              unsigned long flags)
  41 {
  42         td->magic = TD_MAGIC | (flags & TD_FLAGS_MASK);
  43         td->data = data;
  44 }
  45
  46 static inline void __td_free(struct task_data *td)
  47 {
  48         unsigned long flags = get_flags(td);
  49         bool ok = __td_check(td);
  50
  51         /* freeing the data if consistency check fails is dangerous:
  52          * better leave it as a memory leak instead */
  53         if (ok) {
  54                 if ((flags & SWAP_TD_FREE) && td->data)
  55                         kfree(td->data);
  56                 td->magic = 0;
  57                 td->data = NULL;
  58                 return;
  59         }
  60
  61         WARN(!ok, TD_PREFIX "td(%p) check failed: %08lx", td, get_magic(td));
  62 }
  63
  64 void *swap_task_data_get(struct task_struct *task, int *ok)
  65 {
  66         struct task_data *td = __td(task);
  67
  68         if (ok)
  69                 *ok = __td_check(td);
  70
  71         return td->data;
  72 }
  73 EXPORT_SYMBOL_GPL(swap_task_data_get);
  74
  75 void swap_task_data_set(struct task_struct *task, void *data,
  76                         unsigned long flags)
  77 {
  78         struct task_data *td = __td(task);
  79
  80         __td_init(td, data, flags);
  81 }
  82 EXPORT_SYMBOL_GPL(swap_task_data_set);
  83
  84 static int copy_process_ret_handler(struct kretprobe_instance *ri,
  85                                     struct pt_regs *regs)
  86 {
  87         struct task_struct *task;
  88
  89         task = (struct task_struct *)regs_return_value(regs);
  90         if (task)
  91                 swap_task_data_clean(task);
  92
  93         return 0;
  94 }
  95
  96 static int do_exit_handler(struct kprobe *p, struct pt_regs *regs)
  97 {
  98         struct task_data *td = __td(current);
  99
 100         __td_free(td);
 101
 102         return 0;
 103 }
 104
 105 static struct kretprobe copy_process_rp = {
 106         .handler = copy_process_ret_handler
 107 };
 108
 109 static struct kprobe do_exit_probe = {
 110         .pre_handler = do_exit_handler
 111 };
 112
 113 static int __task_data_init(void *data)
 114 {
 115         struct task_struct *g, *t;
 116         unsigned long addr;
 117         int ret;
 118
 119         addr = swap_ksyms_substr("copy_process");
 120         if (addr == 0) {
 121                 printk(TD_PREFIX "Cannot find address for copy_process\n");
 122                 return -EINVAL;
 123         }
 124         copy_process_rp.kp.addr = (kprobe_opcode_t *)addr;
 125         ret = swap_register_kretprobe(&copy_process_rp);
 126         if (ret)
 127                 goto reg_failed;
 128
 129         addr = swap_ksyms_substr("do_exit");
 130         if (addr == 0) {
 131                 printk(TD_PREFIX "Cannot find address for do_exit\n");
 132                 return -EINVAL;
 133         }
 134         do_exit_probe.addr = (kprobe_opcode_t *)addr;
 135         ret = swap_register_kprobe(&do_exit_probe);
 136         if (ret)
 137                 goto unreg_copy_process;
 138
 139         do_each_thread(g, t) {
 140                 swap_task_data_clean(t);
 141         } while_each_thread(g, t);
 142
 143         return 0;
 144
 145 unreg_copy_process:
 146         swap_unregister_kretprobe(&copy_process_rp);
 147
 148 reg_failed:
 149         printk(TD_PREFIX "0x%lx: probe registration failed\n", addr);
 150
 151         return ret;
 152 }
 153
 154 static int __task_data_exit(void *data)
 155 {
 156         struct task_struct *g, *t;
 157         struct task_data *td;
 158
 159         swap_unregister_kprobe(&do_exit_probe);
 160         swap_unregister_kretprobe(&copy_process_rp);
 161
 162         do_each_thread(g, t) {
 163                 td = __td(t);
 164                 __td_free(td);
 165         } while_each_thread(g, t);
 166
 167         return 0;
 168 }
 169
 170 static int task_data_init(void)
 171 {
 172         int ret;
 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         return ret;
 180 }
 181
 182 static void task_data_exit(void)
 183 {
 184         int ret;
 185
 186         /* stop_machine: the same here */
 187         ret = stop_machine(__task_data_exit, &copy_process_rp, NULL);
 188         if (ret) {
 189                 printk(TD_PREFIX "task data cleanup failed: %d\n", ret);
 190                 /* something went wrong: at least make sure we unregister
 191                  * all the installed probes */
 192                 swap_unregister_kprobe(&do_exit_probe);
 193                 swap_unregister_kretprobe(&copy_process_rp);
 194         }
 195 }
 196
 197 SWAP_LIGHT_INIT_MODULE(NULL, task_data_init, task_data_exit, NULL, NULL);
 198
 199
 200 MODULE_LICENSE("GPL");
 201 MODULE_DESCRIPTION("SWAP Task Data Module");