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