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