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