[IMPROVE] create slot_manager
[kernel/swap-modules.git] / uprobe / swap_uprobes.c
1 /*
2  *  Dynamic Binary Instrumentation Module based on KProbes
3  *  modules/kprobe/dbi_uprobes.h
4  *
5  * This program is free software; you can redistribute it and/or modify
6  * it under the terms of the GNU General Public License as published by
7  * the Free Software Foundation; either version 2 of the License, or
8  * (at your option) any later version.
9  *
10  * This program is distributed in the hope that it will be useful,
11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13  * GNU General Public License for more details.
14  *
15  * You should have received a copy of the GNU General Public License
16  * along with this program; if not, write to the Free Software
17  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
18  *
19  * Copyright (C) Samsung Electronics, 2006-2010
20  *
21  * 2008-2009    Alexey Gerenkov <a.gerenkov@samsung.com> User-Space
22  *              Probes initial implementation; Support x86/ARM/MIPS for both user and kernel spaces.
23  * 2010         Ekaterina Gorelkina <e.gorelkina@samsung.com>: redesign module for separating core and arch parts
24  *
25  */
26
27
28 #include "swap_uprobes.h"
29 #include "dbi_kdebug.h"
30
31 #include <asm/swap_uprobes.h>
32
33 #include <linux/hash.h>
34 #include <linux/mempolicy.h>
35 #include <linux/module.h>
36 #include <dbi_insn_slots.h>
37 #include <dbi_kprobes_deps.h>
38
39 enum {
40         UPROBE_HASH_BITS  = 10,
41         UPROBE_TABLE_SIZE = (1 << UPROBE_HASH_BITS)
42 };
43
44 struct hlist_head uprobe_insn_slot_table[UPROBE_TABLE_SIZE];
45 struct hlist_head uprobe_table[UPROBE_TABLE_SIZE];
46 struct hlist_head uprobe_insn_pages;
47
48 DEFINE_SPINLOCK(uretprobe_lock);        /* Protects uretprobe_inst_table */
49 static struct hlist_head uretprobe_inst_table[UPROBE_TABLE_SIZE];
50
51 #define DEBUG_PRINT_HASH_TABLE 0
52
53 #if DEBUG_PRINT_HASH_TABLE
54 void print_kprobe_hash_table(void)
55 {
56         int i;
57         struct hlist_head *head;
58         struct hlist_node *node;
59         struct kprobe *p;
60
61         // print uprobe table
62         for (i = 0; i < KPROBE_TABLE_SIZE; ++i) {
63                 head = &kprobe_table[i];
64                 hlist_for_each_entry_rcu (p, node, head, is_hlist_arm) {
65                         printk("####### find K tgid=%u, addr=%x\n",
66                                         p->tgid, p->addr);
67                 }
68         }
69 }
70
71 void print_kretprobe_hash_table(void)
72 {
73         int i;
74         struct hlist_head *head;
75         struct hlist_node *node;
76         struct kprobe *p;
77
78         // print uprobe table
79         for (i = 0; i < KPROBE_TABLE_SIZE; ++i) {
80                 head = &kretprobe_inst_table[i];
81                 hlist_for_each_entry_rcu (p, node, head, is_hlist_arm) {
82                         printk("####### find KR tgid=%u, addr=%x\n",
83                                         p->tgid, p->addr);
84                 }
85         }
86 }
87
88 void print_uprobe_hash_table(void)
89 {
90         int i;
91         struct hlist_head *head;
92         struct hlist_node *node;
93         struct kprobe *p;
94
95         // print uprobe table
96         for (i = 0; i < UPROBE_TABLE_SIZE; ++i) {
97                 head = &uprobe_insn_slot_table[i];
98                 hlist_for_each_entry_rcu (p, node, head, is_hlist_arm) {
99                         printk("####### find U tgid=%u, addr=%x\n",
100                                         p->tgid, p->addr);
101                 }
102         }
103 }
104 #endif
105
106 /*
107  * Keep all fields in the uprobe consistent
108  */
109 static inline void copy_uprobe(struct kprobe *old_p, struct kprobe *p)
110 {
111         memcpy(&p->opcode, &old_p->opcode, sizeof(kprobe_opcode_t));
112         memcpy(&p->ainsn, &old_p->ainsn, sizeof(struct arch_specific_insn));
113         p->ss_addr = old_p->ss_addr;
114 #ifdef CONFIG_ARM
115         p->safe_arm = old_p->safe_arm;
116         p->safe_thumb = old_p->safe_thumb;
117 #endif
118 }
119
120 /*
121  * Aggregate handlers for multiple uprobes support - these handlers
122  * take care of invoking the individual uprobe handlers on p->list
123  */
124 static int aggr_pre_uhandler(struct kprobe *p, struct pt_regs *regs)
125 {
126         struct kprobe *kp;
127         int ret;
128
129         list_for_each_entry_rcu(kp, &p->list, list) {
130                 if (kp->pre_handler) {
131                         ret = kp->pre_handler(kp, regs);
132                         if (ret) {
133                                 return ret;
134                         }
135                 }
136         }
137
138         return 0;
139 }
140
141 static void aggr_post_uhandler(struct kprobe *p, struct pt_regs *regs, unsigned long flags)
142 {
143         struct kprobe *kp;
144
145         list_for_each_entry_rcu(kp, &p->list, list) {
146                 if (kp->post_handler) {
147                         kp->post_handler(kp, regs, flags);
148                 }
149         }
150 }
151
152 static int aggr_fault_uhandler(struct kprobe *p, struct pt_regs *regs, int trapnr)
153 {
154         return 0;
155 }
156
157 static int aggr_break_uhandler(struct kprobe *p, struct pt_regs *regs)
158 {
159         return 0;
160 }
161
162 /*
163  * Add the new probe to old_p->list. Fail if this is the
164  * second ujprobe at the address - two ujprobes can't coexist
165  */
166 static int add_new_uprobe(struct kprobe *old_p, struct kprobe *p)
167 {
168         if (p->break_handler) {
169                 if (old_p->break_handler) {
170                         return -EEXIST;
171                 }
172
173                 list_add_tail_rcu(&p->list, &old_p->list);
174                 old_p->break_handler = aggr_break_uhandler;
175         } else {
176                 list_add_rcu (&p->list, &old_p->list);
177         }
178
179         if (p->post_handler && !old_p->post_handler) {
180                 old_p->post_handler = aggr_post_uhandler;
181         }
182
183         return 0;
184 }
185
186 /*
187  * Fill in the required fields of the "manager uprobe". Replace the
188  * earlier uprobe in the hlist with the manager uprobe
189  */
190 static inline void add_aggr_uprobe(struct kprobe *ap, struct kprobe *p)
191 {
192         copy_uprobe(p, ap);
193
194         ap->addr = p->addr;
195         ap->pre_handler = aggr_pre_uhandler;
196         ap->fault_handler = aggr_fault_uhandler;
197
198         if (p->post_handler) {
199                 ap->post_handler = aggr_post_uhandler;
200         }
201
202         if (p->break_handler) {
203                 ap->break_handler = aggr_break_uhandler;
204         }
205
206         INIT_LIST_HEAD(&ap->list);
207         list_add_rcu(&p->list, &ap->list);
208
209         hlist_replace_rcu(&p->hlist, &ap->hlist);
210 }
211
212 /*
213  * This is the second or subsequent uprobe at the address - handle
214  * the intricacies
215  */
216 static int register_aggr_uprobe(struct kprobe *old_p, struct kprobe *p)
217 {
218         int ret = 0;
219         struct kprobe *ap;
220
221         if (old_p->pre_handler == aggr_pre_uhandler) {
222                 copy_uprobe(old_p, p);
223                 ret = add_new_uprobe(old_p, p);
224         } else {
225                 struct uprobe *uap = kzalloc(sizeof(*uap), GFP_KERNEL);
226                 if (!uap) {
227                         return -ENOMEM;
228                 }
229
230                 uap->task = kp2up(p)->task;
231                 ap = up2kp(uap);
232                 add_aggr_uprobe(ap, old_p);
233                 copy_uprobe(ap, p);
234                 ret = add_new_uprobe(ap, p);
235         }
236
237         return ret;
238 }
239
240 static void arm_uprobe(struct uprobe *p)
241 {
242         kprobe_opcode_t insn = BREAKPOINT_INSTRUCTION;
243         int ret = write_proc_vm_atomic(p->task, (unsigned long)p->kp.addr,
244                                        &insn, sizeof(insn));
245         if (!ret) {
246                 panic("arm_uprobe: failed to write memory "
247                       "tgid=%u addr=%p!\n", p->task->tgid, p->kp.addr);
248         }
249 }
250
251 void disarm_uprobe(struct uprobe *p)
252 {
253         int ret = write_proc_vm_atomic(p->task, (unsigned long)p->kp.addr,
254                                        &p->kp.opcode, sizeof(p->kp.opcode));
255         if (!ret) {
256                 panic("disarm_uprobe: failed to write memory "
257                       "tgid=%u, addr=%p!\n", p->task->tgid, p->kp.addr);
258         }
259 }
260 EXPORT_SYMBOL_GPL(disarm_uprobe);
261
262 static void init_uprobes_insn_slots(void)
263 {
264         int i;
265         for (i = 0; i < UPROBE_TABLE_SIZE; ++i) {
266                 INIT_HLIST_HEAD(&uprobe_insn_slot_table[i]);
267         }
268 }
269
270 static void init_uprobe_table(void)
271 {
272         int i;
273         for (i = 0; i < UPROBE_TABLE_SIZE; ++i) {
274                 INIT_HLIST_HEAD(&uprobe_table[i]);
275         }
276 }
277
278 static void init_uretprobe_inst_table(void)
279 {
280         int i;
281         for (i = 0; i < UPROBE_TABLE_SIZE; ++i) {
282                 INIT_HLIST_HEAD (&uretprobe_inst_table[i]);
283         }
284 }
285
286 struct kprobe *get_ukprobe(void *addr, pid_t tgid)
287 {
288         struct hlist_head *head;
289         struct hlist_node *node;
290         struct kprobe *p;
291
292         head = &uprobe_table[hash_ptr(addr, UPROBE_HASH_BITS)];
293         hlist_for_each_entry_rcu(p, node, head, hlist) {
294                 if (p->addr == addr && kp2up(p)->task->tgid == tgid) {
295                         return p;
296                 }
297         }
298
299         return NULL;
300 }
301
302 static void add_uprobe_table(struct kprobe *p)
303 {
304 #ifdef CONFIG_ARM
305         INIT_HLIST_NODE(&p->is_hlist_arm);
306         hlist_add_head_rcu(&p->is_hlist_arm, &uprobe_insn_slot_table[hash_ptr(p->ainsn.insn_arm, UPROBE_HASH_BITS)]);
307         INIT_HLIST_NODE(&p->is_hlist_thumb);
308         hlist_add_head_rcu(&p->is_hlist_thumb, &uprobe_insn_slot_table[hash_ptr(p->ainsn.insn_thumb, UPROBE_HASH_BITS)]);
309 #else /* CONFIG_ARM */
310         INIT_HLIST_NODE(&p->is_hlist);
311         hlist_add_head_rcu(&p->is_hlist, &uprobe_insn_slot_table[hash_ptr(p->ainsn.insn, UPROBE_HASH_BITS)]);
312 #endif /* CONFIG_ARM */
313 }
314
315 #ifdef CONFIG_ARM
316 static struct kprobe *get_ukprobe_bis_arm(void *addr, pid_t tgid)
317 {
318         struct hlist_head *head;
319         struct hlist_node *node;
320         struct kprobe *p;
321
322         /* TODO: test - two processes invokes instrumented function */
323         head = &uprobe_insn_slot_table[hash_ptr(addr, UPROBE_HASH_BITS)];
324         hlist_for_each_entry_rcu(p, node, head, is_hlist_arm) {
325                 if (p->ainsn.insn == addr && kp2up(p)->task->tgid == tgid) {
326                         return p;
327                 }
328         }
329
330         return NULL;
331 }
332
333 static struct kprobe *get_ukprobe_bis_thumb(void *addr, pid_t tgid)
334 {
335         struct hlist_head *head;
336         struct hlist_node *node;
337         struct kprobe *p;
338
339         /* TODO: test - two processes invokes instrumented function */
340         head = &uprobe_insn_slot_table[hash_ptr(addr, UPROBE_HASH_BITS)];
341         hlist_for_each_entry_rcu(p, node, head, is_hlist_thumb) {
342                 if (p->ainsn.insn == addr && kp2up(p)->task->tgid == tgid) {
343                         return p;
344                 }
345         }
346
347         return NULL;
348 }
349
350 struct kprobe *get_ukprobe_by_insn_slot(void *addr, pid_t tgid, struct pt_regs *regs)
351 {
352         return thumb_mode(regs) ?
353                         get_ukprobe_bis_thumb(addr - 0x1a, tgid) :
354                         get_ukprobe_bis_arm(addr - 4 * UPROBES_TRAMP_RET_BREAK_IDX, tgid);
355 }
356 #else /* CONFIG_ARM */
357 struct kprobe *get_ukprobe_by_insn_slot(void *addr, pid_t tgid, struct pt_regs *regs)
358 {
359         struct hlist_head *head;
360         struct hlist_node *node;
361         struct kprobe *p;
362
363         /* TODO: test - two processes invokes instrumented function */
364         head = &uprobe_insn_slot_table[hash_ptr(addr, UPROBE_HASH_BITS)];
365         hlist_for_each_entry_rcu(p, node, head, is_hlist) {
366                 if (p->ainsn.insn == addr && kp2up(p)->task->tgid == tgid) {
367                         return p;
368                 }
369         }
370
371         return NULL;
372 }
373 #endif /* CONFIG_ARM */
374
375
376 static void remove_uprobe(struct uprobe *up)
377 {
378         struct kprobe *p = &up->kp;
379         struct task_struct *task = up->task;
380
381 #ifdef CONFIG_ARM
382         free_insn_slot(up->sm, p->ainsn.insn_arm);
383         free_insn_slot(up->sm, p->ainsn.insn_thumb);
384 #else /* CONFIG_ARM */
385         free_insn_slot(up->sm, p->ainsn.insn);
386 #endif /* CONFIG_ARM */
387 }
388
389 static struct hlist_head *uretprobe_inst_table_head(void *hash_key)
390 {
391         return &uretprobe_inst_table[hash_ptr (hash_key, UPROBE_HASH_BITS)];
392 }
393
394 /* Called with uretprobe_lock held */
395 static void add_urp_inst(struct uretprobe_instance *ri)
396 {
397         /*
398          * Remove rp inst off the free list -
399          * Add it back when probed function returns
400          */
401         hlist_del(&ri->uflist);
402
403         /* Add rp inst onto table */
404         INIT_HLIST_NODE(&ri->hlist);
405         hlist_add_head(&ri->hlist, uretprobe_inst_table_head(ri->task->mm));
406
407         /* Also add this rp inst to the used list. */
408         INIT_HLIST_NODE(&ri->uflist);
409         hlist_add_head(&ri->uflist, &ri->rp->used_instances);
410 }
411
412 /* Called with uretprobe_lock held */
413 static void recycle_urp_inst(struct uretprobe_instance *ri)
414 {
415         if (ri->rp) {
416                 hlist_del(&ri->hlist);
417                 /* remove rp inst off the used list */
418                 hlist_del(&ri->uflist);
419                 /* put rp inst back onto the free list */
420                 INIT_HLIST_NODE(&ri->uflist);
421                 hlist_add_head(&ri->uflist, &ri->rp->free_instances);
422         }
423 }
424
425 /* Called with uretprobe_lock held */
426 static struct uretprobe_instance *get_used_urp_inst(struct uretprobe *rp)
427 {
428         struct hlist_node *node;
429         struct uretprobe_instance *ri;
430
431         hlist_for_each_entry(ri, node, &rp->used_instances, uflist) {
432                 return ri;
433         }
434
435         return NULL;
436 }
437
438 /* Called with uretprobe_lock held */
439 struct uretprobe_instance *get_free_urp_inst_no_alloc(struct uretprobe *rp)
440 {
441         struct hlist_node *node;
442         struct uretprobe_instance *ri;
443
444         hlist_for_each_entry (ri, node, &rp->free_instances, uflist) {
445                 return ri;
446         }
447
448         return NULL;
449 }
450
451 /* Called with uretprobe_lock held */
452 static void free_urp_inst(struct uretprobe *rp)
453 {
454         struct uretprobe_instance *ri;
455         while ((ri = get_free_urp_inst_no_alloc(rp)) != NULL) {
456                 hlist_del(&ri->uflist);
457                 kfree(ri);
458         }
459 }
460
461 #define COMMON_URP_NR 10
462
463 static int alloc_nodes_uretprobe(struct uretprobe *rp)
464 {
465         int alloc_nodes;
466         struct uretprobe_instance *inst;
467         int i;
468
469 #if 1//def CONFIG_PREEMPT
470         rp->maxactive += max(COMMON_URP_NR, 2 * NR_CPUS);
471 #else
472         rp->maxacpptive += NR_CPUS;
473 #endif
474         alloc_nodes = COMMON_URP_NR;
475
476         for (i = 0; i < alloc_nodes; ++i) {
477                 inst = kmalloc(sizeof(*inst), GFP_ATOMIC);
478                 if (inst == NULL) {
479                         free_urp_inst(rp);
480                         return -ENOMEM;
481                 }
482                 INIT_HLIST_NODE(&inst->uflist);
483                 hlist_add_head(&inst->uflist, &rp->free_instances);
484         }
485
486         return 0;
487 }
488
489 /* Called with uretprobe_lock held */
490 static struct uretprobe_instance *get_free_urp_inst(struct uretprobe *rp)
491 {
492         struct hlist_node *node;
493         struct uretprobe_instance *ri;
494
495         hlist_for_each_entry(ri, node, &rp->free_instances, uflist) {
496                 return ri;
497         }
498
499         if (!alloc_nodes_uretprobe(rp)) {
500                 hlist_for_each_entry(ri, node, &rp->free_instances, uflist) {
501                         return ri;
502                 }
503         }
504
505         return NULL;
506 }
507 // ===================================================================
508
509 int dbi_register_uprobe(struct uprobe *up, int atomic)
510 {
511         int ret = 0;
512         struct kprobe *p, *old_p;
513
514         p = &up->kp;
515         if (!p->addr) {
516                 return -EINVAL;
517         }
518
519         DBPRINTF("p->addr = 0x%p p = 0x%p\n", p->addr, p);
520
521 // thumb address = address-1;
522 #if defined(CONFIG_ARM)
523         // TODO: must be corrected in 'bundle'
524         if ((unsigned long) p->addr & 0x01) {
525                 p->addr = (kprobe_opcode_t *)((unsigned long)p->addr & 0xfffffffe);
526         }
527 #endif
528
529         p->mod_refcounted = 0;
530         p->nmissed = 0;
531         INIT_LIST_HEAD(&p->list);
532 #ifdef KPROBES_PROFILE
533         p->start_tm.tv_sec = p->start_tm.tv_usec = 0;
534         p->hnd_tm_sum.tv_sec = p->hnd_tm_sum.tv_usec = 0;
535         p->count = 0;
536 #endif
537
538         // get the first item
539         old_p = get_ukprobe(p->addr, kp2up(p)->task->tgid);
540         if (old_p) {
541 #ifdef CONFIG_ARM
542                 p->safe_arm = old_p->safe_arm;
543                 p->safe_thumb = old_p->safe_thumb;
544 #endif
545                 ret = register_aggr_uprobe(old_p, p);
546                 if (!ret) {
547 //                      atomic_inc(&kprobe_count);
548                         add_uprobe_table(p);
549                 }
550                 DBPRINTF("goto out\n", ret);
551                 goto out;
552         }
553
554         ret = arch_prepare_uprobe(up, &uprobe_insn_pages, atomic);
555         if (ret) {
556                 DBPRINTF("goto out\n", ret);
557                 goto out;
558         }
559
560         DBPRINTF ("before out ret = 0x%x\n", ret);
561
562         // TODO: add uprobe (must be in function)
563         INIT_HLIST_NODE(&p->hlist);
564         hlist_add_head_rcu(&p->hlist, &uprobe_table[hash_ptr(p->addr, UPROBE_HASH_BITS)]);
565         add_uprobe_table(p);
566         arm_uprobe(up);
567
568 out:
569         DBPRINTF("out ret = 0x%x\n", ret);
570         return ret;
571 }
572
573 void dbi_unregister_uprobe(struct uprobe *up, int atomic)
574 {
575         struct kprobe *p, *old_p, *list_p;
576         int cleanup_p;
577
578         p = &up->kp;
579         old_p = get_ukprobe(p->addr, kp2up(p)->task->tgid);
580         if (unlikely(!old_p)) {
581                 return;
582         }
583
584         if (p != old_p) {
585                 list_for_each_entry_rcu(list_p, &old_p->list, list) {
586                         if (list_p == p) {
587                                 /* uprobe p is a valid probe */
588                                 goto valid_p;
589                         }
590                 }
591
592                 return;
593         }
594
595 valid_p:
596         if ((old_p == p) || ((old_p->pre_handler == aggr_pre_uhandler) &&
597             (p->list.next == &old_p->list) && (p->list.prev == &old_p->list))) {
598                 /* Only probe on the hash list */
599                 disarm_uprobe(up);
600                 hlist_del_rcu(&old_p->hlist);
601                 cleanup_p = 1;
602         } else {
603                 list_del_rcu(&p->list);
604                 cleanup_p = 0;
605         }
606
607         if (cleanup_p) {
608                 if (p != old_p) {
609                         list_del_rcu(&p->list);
610                         kfree(old_p);
611                 }
612
613                 if (!in_atomic()) {
614                         synchronize_sched();
615                 }
616
617                 remove_uprobe(up);
618         } else {
619                 if (p->break_handler) {
620                         old_p->break_handler = NULL;
621                 }
622
623                 if (p->post_handler) {
624                         list_for_each_entry_rcu (list_p, &old_p->list, list) {
625                                 if (list_p->post_handler) {
626                                         cleanup_p = 2;
627                                         break;
628                                 }
629                         }
630
631                         if (cleanup_p == 0) {
632                                 old_p->post_handler = NULL;
633                         }
634                 }
635         }
636 }
637
638 int dbi_register_ujprobe(struct ujprobe *jp, int atomic)
639 {
640         int ret = 0;
641
642         /* Todo: Verify probepoint is a function entry point */
643         jp->up.kp.pre_handler = setjmp_upre_handler;
644         jp->up.kp.break_handler = longjmp_break_uhandler;
645
646         ret = dbi_register_uprobe(&jp->up, atomic);
647
648         return ret;
649 }
650
651 void dbi_unregister_ujprobe(struct ujprobe *jp, int atomic)
652 {
653         dbi_unregister_uprobe(&jp->up, atomic);
654         /*
655          * Here is an attempt to unregister even those probes that have not been
656          * installed (hence not added to the hlist).
657          * So if we try to delete them from the hlist we will get NULL pointer
658          * dereference error. That is why we check whether this node
659          * really belongs to the hlist.
660          */
661 #ifdef CONFIG_ARM
662         if (!(hlist_unhashed(&jp->up.kp.is_hlist_arm))) {
663                 hlist_del_rcu(&jp->up.kp.is_hlist_arm);
664         }
665         if (!(hlist_unhashed(&jp->up.kp.is_hlist_thumb))) {
666                 hlist_del_rcu(&jp->up.kp.is_hlist_thumb);
667         }
668 #else /* CONFIG_ARM */
669         if (!(hlist_unhashed(&jp->up.kp.is_hlist))) {
670                 hlist_del_rcu(&jp->up.kp.is_hlist);
671         }
672 #endif /* CONFIG_ARM */
673 }
674
675 int trampoline_uprobe_handler(struct kprobe *p, struct pt_regs *regs)
676 {
677         struct uretprobe_instance *ri = NULL;
678         struct hlist_head *head;
679         struct hlist_node *node, *tmp;
680         unsigned long flags, tramp_addr, orig_ret_addr = 0;
681
682         tramp_addr = arch_get_trampoline_addr(p, regs);
683         spin_lock_irqsave(&uretprobe_lock, flags);
684
685         head = uretprobe_inst_table_head(current->mm);
686
687         /*
688          * It is possible to have multiple instances associated with a given
689          * task either because an multiple functions in the call path
690          * have a return probe installed on them, and/or more then one
691          * return probe was registered for a target function.
692          *
693          * We can handle this because:
694          *     - instances are always inserted at the head of the list
695          *     - when multiple return probes are registered for the same
696          *       function, the first instance's ret_addr will point to the
697          *       real return address, and all the rest will point to
698          *       uretprobe_trampoline
699          */
700         hlist_for_each_entry_safe(ri, node, tmp, head, hlist) {
701                 if (ri->task != current) {
702                         /* another task is sharing our hash bucket */
703                         continue;
704                 }
705
706                 if (ri->rp && ri->rp->handler) {
707                         ri->rp->handler(ri, regs, ri->rp->priv_arg);
708                 }
709
710                 orig_ret_addr = (unsigned long)ri->ret_addr;
711                 recycle_urp_inst(ri);
712
713                 if (orig_ret_addr != tramp_addr) {
714                         /*
715                          * This is the real return address. Any other
716                          * instances associated with this task are for
717                          * other calls deeper on the call stack
718                          */
719                         break;
720                 }
721         }
722
723         spin_unlock_irqrestore(&uretprobe_lock, flags);
724         arch_set_orig_ret_addr(orig_ret_addr, regs);
725
726         return 1;
727 }
728
729 static int pre_handler_uretprobe(struct kprobe *p, struct pt_regs *regs)
730 {
731         struct uprobe *up = container_of(p, struct uprobe, kp);
732         struct uretprobe *rp = container_of(up, struct uretprobe, up);
733         struct uretprobe_instance *ri;
734         unsigned long flags;
735
736         /* TODO: consider to only swap the RA after the last pre_handler fired */
737         spin_lock_irqsave(&uretprobe_lock, flags);
738
739         /* TODO: test - remove retprobe after func entry but before its exit */
740         if ((ri = get_free_urp_inst(rp)) != NULL) {
741                 ri->rp = rp;
742                 ri->task = current;
743
744                 arch_prepare_uretprobe_hl(ri, regs);
745
746                 add_urp_inst(ri);
747         } else {
748                 ++rp->nmissed;
749         }
750
751         spin_unlock_irqrestore(&uretprobe_lock, flags);
752
753         return 0;
754 }
755
756 int dbi_register_uretprobe(struct uretprobe *rp, int atomic)
757 {
758         int i, ret = 0;
759         struct uretprobe_instance *inst;
760
761         DBPRINTF ("START\n");
762
763         rp->up.kp.pre_handler = pre_handler_uretprobe;
764         rp->up.kp.post_handler = NULL;
765         rp->up.kp.fault_handler = NULL;
766         rp->up.kp.break_handler = NULL;
767
768         /* Pre-allocate memory for max kretprobe instances */
769         if (rp->maxactive <= 0) {
770 #if 1//def CONFIG_PREEMPT
771                 rp->maxactive = max(10, 2 * NR_CPUS);
772 #else
773                 rp->maxactive = NR_CPUS;
774 #endif
775         }
776
777         INIT_HLIST_HEAD(&rp->used_instances);
778         INIT_HLIST_HEAD(&rp->free_instances);
779
780         for (i = 0; i < rp->maxactive; i++) {
781                 inst = kmalloc(sizeof(*inst), GFP_KERNEL);
782                 if (inst == NULL) {
783                         free_urp_inst(rp);
784                         ret = -ENOMEM;
785                         goto out;
786                 }
787
788                 INIT_HLIST_NODE(&inst->uflist);
789                 hlist_add_head(&inst->uflist, &rp->free_instances);
790         }
791
792         rp->nmissed = 0;
793
794         /* Establish function entry probe point */
795         ret = dbi_register_uprobe(&rp->up, atomic);
796         if (ret) {
797                 free_urp_inst(rp);
798                 goto out;
799         }
800
801 out:
802         return ret;
803 }
804
805 int dbi_disarm_urp_inst(struct uretprobe_instance *ri, struct task_struct *rm_task)
806 {
807         struct task_struct *task = rm_task ? rm_task : ri->task;
808         kprobe_opcode_t *tramp;
809         kprobe_opcode_t *sp = (kprobe_opcode_t *)((long)ri->sp & ~1);
810         kprobe_opcode_t *stack = sp - RETPROBE_STACK_DEPTH + 1;
811         kprobe_opcode_t *found = NULL;
812         kprobe_opcode_t *buf[RETPROBE_STACK_DEPTH];
813         int i, retval;
814
815         /* Understand function mode */
816         if ((long)ri->sp & 1) {
817                 tramp = (kprobe_opcode_t *)
818                         ((unsigned long)ri->rp->up.kp.ainsn.insn + 0x1b);
819         } else {
820                 tramp = (kprobe_opcode_t *)
821                         (ri->rp->up.kp.ainsn.insn + UPROBES_TRAMP_RET_BREAK_IDX);
822         }
823
824         retval = read_proc_vm_atomic(task, (unsigned long)stack, buf, sizeof(buf));
825         if (retval != sizeof(buf)) {
826                 printk("---> %s (%d/%d): failed to read stack from %08lx",
827                         task->comm, task->tgid, task->pid, (unsigned long)stack);
828                 retval = -EFAULT;
829                 goto out;
830         }
831
832         /* search the stack from the bottom */
833         for (i = RETPROBE_STACK_DEPTH - 1; i >= 0; i--) {
834                 if (buf[i] == tramp) {
835                         found = stack + i;
836                         break;
837                 }
838         }
839
840         if (found) {
841                 printk("---> %s (%d/%d): trampoline found at %08lx (%08lx /%+d) - %p\n",
842                                 task->comm, task->tgid, task->pid,
843                                 (unsigned long)found, (unsigned long)sp,
844                                 found - sp, ri->rp->up.kp.addr);
845                 retval = write_proc_vm_atomic(task, (unsigned long)found, &ri->ret_addr,
846                                 sizeof(ri->ret_addr));
847                 if (retval != sizeof(ri->ret_addr)) {
848                         printk("---> %s (%d/%d): failed to write value to %08lx",
849                                 task->comm, task->tgid, task->pid, (unsigned long)found);
850                         retval = -EFAULT;
851                 } else {
852                         retval = 0;
853                 }
854         } else {
855                 struct pt_regs *uregs = task_pt_regs(ri->task);
856                 unsigned long ra = dbi_get_ret_addr(uregs);
857                 if (ra == (unsigned long)tramp) {
858                         printk("---> %s (%d/%d): trampoline found at lr = %08lx - %p\n",
859                                         task->comm, task->tgid, task->pid, ra, ri->rp->up.kp.addr);
860                         dbi_set_ret_addr(uregs, (unsigned long)tramp);
861                         retval = 0;
862                 } else {
863                         printk("---> %s (%d/%d): trampoline NOT found at sp = %08lx, lr = %08lx - %p\n",
864                                         task->comm, task->tgid, task->pid,
865                                         (unsigned long)sp, ra, ri->rp->up.kp.addr);
866                         retval = -ENOENT;
867                 }
868         }
869
870 out:
871         return retval;
872 }
873
874 /* Called with uretprobe_lock held */
875 int dbi_disarm_urp_inst_for_task(struct task_struct *parent, struct task_struct *task)
876 {
877         struct uretprobe_instance *ri;
878         struct hlist_node *node, *tmp;
879         struct hlist_head *head = uretprobe_inst_table_head(parent->mm);
880
881         hlist_for_each_entry_safe(ri, node, tmp, head, hlist) {
882                 if (parent == ri->task) {
883                         dbi_disarm_urp_inst(ri, task);
884                 }
885         }
886
887         return 0;
888 }
889 EXPORT_SYMBOL_GPL(dbi_disarm_urp_inst_for_task);
890
891 void dbi_unregister_uretprobe(struct uretprobe *rp, int atomic)
892 {
893         unsigned long flags;
894         struct uretprobe_instance *ri;
895
896         spin_lock_irqsave (&uretprobe_lock, flags);
897
898         while ((ri = get_used_urp_inst(rp)) != NULL) {
899                 if (dbi_disarm_urp_inst(ri, NULL) != 0)
900                         /*panic*/printk("%s (%d/%d): cannot disarm urp instance (%08lx)\n",
901                                         ri->task->comm, ri->task->tgid, ri->task->pid,
902                                         (unsigned long)rp->up.kp.addr);
903                 recycle_urp_inst(ri);
904         }
905
906         if (hlist_empty(&rp->used_instances)) {
907                 struct kprobe *p = &rp->up.kp;
908 #ifdef CONFIG_ARM
909                 if (!(hlist_unhashed(&p->is_hlist_arm))) {
910                         hlist_del_rcu(&p->is_hlist_arm);
911                 }
912
913                 if (!(hlist_unhashed(&p->is_hlist_thumb))) {
914                         hlist_del_rcu(&p->is_hlist_thumb);
915                 }
916 #else /* CONFIG_ARM */
917                 if (!(hlist_unhashed(&p->is_hlist))) {
918                         hlist_del_rcu(&p->is_hlist);
919                 }
920 #endif /* CONFIG_ARM */
921         }
922
923         while ((ri = get_used_urp_inst(rp)) != NULL) {
924                 ri->rp = NULL;
925                 hlist_del(&ri->uflist);
926         }
927
928         spin_unlock_irqrestore(&uretprobe_lock, flags);
929         free_urp_inst(rp);
930
931         dbi_unregister_uprobe(&rp->up, atomic);
932 }
933
934 void dbi_unregister_all_uprobes(struct task_struct *task, int atomic)
935 {
936         struct hlist_head *head;
937         struct hlist_node *node, *tnode;
938         struct kprobe *p;
939         int i;
940
941         for (i = 0; i < UPROBE_TABLE_SIZE; ++i) {
942                 head = &uprobe_table[i];
943                 hlist_for_each_entry_safe(p, node, tnode, head, hlist) {
944                         if (kp2up(p)->task->tgid == task->tgid) {
945                                 struct uprobe *up = container_of(p, struct uprobe, kp);
946                                 printk("dbi_unregister_all_uprobes: delete uprobe at %p[%lx] for %s/%d\n",
947                                                 p->addr, (unsigned long)p->opcode, task->comm, task->pid);
948                                 dbi_unregister_uprobe(up, atomic);
949                         }
950                 }
951         }
952 }
953
954 void dbi_uprobe_return(void)
955 {
956         dbi_arch_uprobe_return();
957 }
958
959 static int __init init_uprobes(void)
960 {
961         init_uprobe_table();
962         init_uprobes_insn_slots();
963         init_uretprobe_inst_table();
964
965         return swap_arch_init_uprobes();
966 }
967
968 static void __exit exit_uprobes(void)
969 {
970         swap_arch_exit_uprobes();
971 }
972
973 EXPORT_SYMBOL_GPL(dbi_uprobe_return);
974 EXPORT_SYMBOL_GPL(dbi_register_ujprobe);
975 EXPORT_SYMBOL_GPL(dbi_unregister_ujprobe);
976 EXPORT_SYMBOL_GPL(dbi_register_uretprobe);
977 EXPORT_SYMBOL_GPL(dbi_unregister_uretprobe);
978 EXPORT_SYMBOL_GPL(dbi_unregister_all_uprobes);
979
980 module_init(init_uprobes);
981 module_exit(exit_uprobes);
982
983 MODULE_LICENSE ("GPL");