]> nv-tegra.nvidia Code Review - linux-2.6.git/blobdiff - arch/ia64/kernel/kprobes.c
[PATCH] kprobes: enable booster on the preemptible kernel
[linux-2.6.git] / arch / ia64 / kernel / kprobes.c
index 471086b808a4b942781aa19abe940fb08dc8533b..4d592ee9300b2739d6c82f569f0b1a02254ee09e 100644 (file)
  *              <anil.s.keshavamurthy@intel.com> adapted from i386
  */
 
-#include <linux/config.h>
 #include <linux/kprobes.h>
 #include <linux/ptrace.h>
-#include <linux/spinlock.h>
 #include <linux/string.h>
 #include <linux/slab.h>
 #include <linux/preempt.h>
 #include <asm/pgtable.h>
 #include <asm/kdebug.h>
 #include <asm/sections.h>
+#include <asm/uaccess.h>
 
 extern void jprobe_inst_return(void);
 
-/* kprobe_status settings */
-#define KPROBE_HIT_ACTIVE      0x00000001
-#define KPROBE_HIT_SS          0x00000002
-
-static struct kprobe *current_kprobe, *kprobe_prev;
-static unsigned long kprobe_status, kprobe_status_prev;
-static struct pt_regs jprobe_saved_regs;
+DEFINE_PER_CPU(struct kprobe *, current_kprobe) = NULL;
+DEFINE_PER_CPU(struct kprobe_ctlblk, kprobe_ctlblk);
 
 enum instruction_type {A, I, M, F, B, L, X, u};
 static enum instruction_type bundle_encoding[32][3] = {
@@ -96,7 +90,7 @@ static void __kprobes update_kprobe_inst_flag(uint template, uint  slot,
        p->ainsn.target_br_reg = 0;
 
        /* Check for Break instruction
-        * Bits 37:40 Major opcode to be zero
+        * Bits 37:40 Major opcode to be zero
         * Bits 27:32 X6 to be zero
         * Bits 32:35 X3 to be zero
         */
@@ -110,19 +104,19 @@ static void __kprobes update_kprobe_inst_flag(uint template, uint  slot,
                switch (major_opcode) {
                  case INDIRECT_CALL_OPCODE:
                        p->ainsn.inst_flag |= INST_FLAG_FIX_BRANCH_REG;
-                       p->ainsn.target_br_reg = ((kprobe_inst >> 6) & 0x7);
-                       break;
+                       p->ainsn.target_br_reg = ((kprobe_inst >> 6) & 0x7);
+                       break;
                  case IP_RELATIVE_PREDICT_OPCODE:
                  case IP_RELATIVE_BRANCH_OPCODE:
                        p->ainsn.inst_flag |= INST_FLAG_FIX_RELATIVE_IP_ADDR;
-                       break;
+                       break;
                  case IP_RELATIVE_CALL_OPCODE:
-                       p->ainsn.inst_flag |= INST_FLAG_FIX_RELATIVE_IP_ADDR;
-                       p->ainsn.inst_flag |= INST_FLAG_FIX_BRANCH_REG;
-                       p->ainsn.target_br_reg = ((kprobe_inst >> 6) & 0x7);
-                       break;
+                       p->ainsn.inst_flag |= INST_FLAG_FIX_RELATIVE_IP_ADDR;
+                       p->ainsn.inst_flag |= INST_FLAG_FIX_BRANCH_REG;
+                       p->ainsn.target_br_reg = ((kprobe_inst >> 6) & 0x7);
+                       break;
                }
-       } else if (bundle_encoding[template][slot] == X) {
+       } else if (bundle_encoding[template][slot] == X) {
                switch (major_opcode) {
                  case LONG_CALL_OPCODE:
                        p->ainsn.inst_flag |= INST_FLAG_FIX_BRANCH_REG;
@@ -142,10 +136,8 @@ static void __kprobes update_kprobe_inst_flag(uint template, uint  slot,
 static int __kprobes unsupported_inst(uint template, uint  slot,
                                      uint major_opcode,
                                      unsigned long kprobe_inst,
-                                     struct kprobe *p)
+                                     unsigned long addr)
 {
-       unsigned long addr = (unsigned long)p->addr;
-
        if (bundle_encoding[template][slot] == I) {
                switch (major_opcode) {
                        case 0x0: //I_UNIT_MISC_OPCODE:
@@ -223,7 +215,7 @@ static void __kprobes prepare_break_inst(uint template, uint  slot,
                                         struct kprobe *p)
 {
        unsigned long break_inst = BREAK_INST;
-       bundle_t *bundle = &p->ainsn.insn.bundle;
+       bundle_t *bundle = &p->opcode.bundle;
 
        /*
         * Copy the original kprobe_inst qualifying predicate(qp)
@@ -256,7 +248,7 @@ static void __kprobes prepare_break_inst(uint template, uint  slot,
        update_kprobe_inst_flag(template, slot, major_opcode, kprobe_inst, p);
 }
 
-static inline void get_kprobe_inst(bundle_t *bundle, uint slot,
+static void __kprobes get_kprobe_inst(bundle_t *bundle, uint slot,
                unsigned long *kprobe_inst, uint *major_opcode)
 {
        unsigned long kprobe_inst_p0, kprobe_inst_p1;
@@ -266,24 +258,24 @@ static inline void get_kprobe_inst(bundle_t *bundle, uint slot,
 
        switch (slot) {
          case 0:
-               *major_opcode = (bundle->quad0.slot0 >> SLOT0_OPCODE_SHIFT);
-               *kprobe_inst = bundle->quad0.slot0;
-               break;
+               *major_opcode = (bundle->quad0.slot0 >> SLOT0_OPCODE_SHIFT);
+               *kprobe_inst = bundle->quad0.slot0;
+                 break;
          case 1:
-               *major_opcode = (bundle->quad1.slot1_p1 >> SLOT1_p1_OPCODE_SHIFT);
-               kprobe_inst_p0 = bundle->quad0.slot1_p0;
-               kprobe_inst_p1 = bundle->quad1.slot1_p1;
-               *kprobe_inst = kprobe_inst_p0 | (kprobe_inst_p1 << (64-46));
+               *major_opcode = (bundle->quad1.slot1_p1 >> SLOT1_p1_OPCODE_SHIFT);
+               kprobe_inst_p0 = bundle->quad0.slot1_p0;
+               kprobe_inst_p1 = bundle->quad1.slot1_p1;
+               *kprobe_inst = kprobe_inst_p0 | (kprobe_inst_p1 << (64-46));
                break;
          case 2:
-               *major_opcode = (bundle->quad1.slot2 >> SLOT2_OPCODE_SHIFT);
-               *kprobe_inst = bundle->quad1.slot2;
+               *major_opcode = (bundle->quad1.slot2 >> SLOT2_OPCODE_SHIFT);
+               *kprobe_inst = bundle->quad1.slot2;
                break;
        }
 }
 
 /* Returns non-zero if the addr is in the Interrupt Vector Table */
-static inline int in_ivt_functions(unsigned long addr)
+static int __kprobes in_ivt_functions(unsigned long addr)
 {
        return (addr >= (unsigned long)__start_ivt_text
                && addr < (unsigned long)__end_ivt_text);
@@ -298,11 +290,11 @@ static int __kprobes valid_kprobe_addr(int template, int slot,
                return -EINVAL;
        }
 
-       if (in_ivt_functions(addr)) {
-               printk(KERN_WARNING "Kprobes can't be inserted inside "
+       if (in_ivt_functions(addr)) {
+               printk(KERN_WARNING "Kprobes can't be inserted inside "
                                "IVT functions at 0x%lx\n", addr);
-               return -EINVAL;
-       }
+               return -EINVAL;
+       }
 
        if (slot == 1 && bundle_encoding[template][1] != L) {
                printk(KERN_WARNING "Inserting kprobes on slot #1 "
@@ -313,21 +305,22 @@ static int __kprobes valid_kprobe_addr(int template, int slot,
        return 0;
 }
 
-static inline void save_previous_kprobe(void)
+static void __kprobes save_previous_kprobe(struct kprobe_ctlblk *kcb)
 {
-       kprobe_prev = current_kprobe;
-       kprobe_status_prev = kprobe_status;
+       kcb->prev_kprobe.kp = kprobe_running();
+       kcb->prev_kprobe.status = kcb->kprobe_status;
 }
 
-static inline void restore_previous_kprobe(void)
+static void __kprobes restore_previous_kprobe(struct kprobe_ctlblk *kcb)
 {
-       current_kprobe = kprobe_prev;
-       kprobe_status = kprobe_status_prev;
+       __get_cpu_var(current_kprobe) = kcb->prev_kprobe.kp;
+       kcb->kprobe_status = kcb->prev_kprobe.status;
 }
 
-static inline void set_current_kprobe(struct kprobe *p)
+static void __kprobes set_current_kprobe(struct kprobe *p,
+                       struct kprobe_ctlblk *kcb)
 {
-       current_kprobe = p;
+       __get_cpu_var(current_kprobe) = p;
 }
 
 static void kretprobe_trampoline(void)
@@ -345,13 +338,15 @@ static void kretprobe_trampoline(void)
 int __kprobes trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
 {
        struct kretprobe_instance *ri = NULL;
-       struct hlist_head *head;
+       struct hlist_head *head, empty_rp;
        struct hlist_node *node, *tmp;
-       unsigned long orig_ret_address = 0;
+       unsigned long flags, orig_ret_address = 0;
        unsigned long trampoline_address =
                ((struct fnptr *)kretprobe_trampoline)->ip;
 
-        head = kretprobe_inst_table_head(current);
+       INIT_HLIST_HEAD(&empty_rp);
+       spin_lock_irqsave(&kretprobe_lock, flags);
+       head = kretprobe_inst_table_head(current);
 
        /*
         * It is possible to have multiple instances associated with a given
@@ -367,15 +362,15 @@ int __kprobes trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
         *       kretprobe_trampoline
         */
        hlist_for_each_entry_safe(ri, node, tmp, head, hlist) {
-                if (ri->task != current)
+               if (ri->task != current)
                        /* another task is sharing our hash bucket */
-                        continue;
+                       continue;
 
                if (ri->rp && ri->rp->handler)
                        ri->rp->handler(ri, regs);
 
                orig_ret_address = (unsigned long)ri->ret_addr;
-               recycle_rp_inst(ri);
+               recycle_rp_inst(ri, &empty_rp);
 
                if (orig_ret_address != trampoline_address)
                        /*
@@ -389,17 +384,23 @@ int __kprobes trampoline_probe_handler(struct kprobe *p, struct pt_regs *regs)
        BUG_ON(!orig_ret_address || (orig_ret_address == trampoline_address));
        regs->cr_iip = orig_ret_address;
 
-       unlock_kprobes();
+       reset_current_kprobe();
+       spin_unlock_irqrestore(&kretprobe_lock, flags);
        preempt_enable_no_resched();
 
-        /*
-         * By returning a non-zero value, we are telling
-         * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption.
-         */
-        return 1;
+       hlist_for_each_entry_safe(ri, node, tmp, &empty_rp, hlist) {
+               hlist_del(&ri->hlist);
+               kfree(ri);
+       }
+       /*
+        * By returning a non-zero value, we are telling
+        * kprobe_handler() that we don't want the post_handler
+        * to run (and have re-enabled preemption)
+        */
+       return 1;
 }
 
+/* Called with kretprobe_lock held */
 void __kprobes arch_prepare_kretprobe(struct kretprobe *rp,
                                      struct pt_regs *regs)
 {
@@ -425,26 +426,31 @@ int __kprobes arch_prepare_kprobe(struct kprobe *p)
        unsigned long *kprobe_addr = (unsigned long *)(addr & ~0xFULL);
        unsigned long kprobe_inst=0;
        unsigned int slot = addr & 0xf, template, major_opcode = 0;
-       bundle_t *bundle = &p->ainsn.insn.bundle;
+       bundle_t *bundle;
 
-       memcpy(&p->opcode.bundle, kprobe_addr, sizeof(bundle_t));
-       memcpy(&p->ainsn.insn.bundle, kprobe_addr, sizeof(bundle_t));
-
-       template = bundle->quad0.template;
+       bundle = &((kprobe_opcode_t *)kprobe_addr)->bundle;
+       template = bundle->quad0.template;
 
        if(valid_kprobe_addr(template, slot, addr))
                return -EINVAL;
 
        /* Move to slot 2, if bundle is MLX type and kprobe slot is 1 */
-       if (slot == 1 && bundle_encoding[template][1] == L)
-               slot++;
+       if (slot == 1 && bundle_encoding[template][1] == L)
+               slot++;
 
        /* Get kprobe_inst and major_opcode from the bundle */
        get_kprobe_inst(bundle, slot, &kprobe_inst, &major_opcode);
 
-       if (unsupported_inst(template, slot, major_opcode, kprobe_inst, p))
+       if (unsupported_inst(template, slot, major_opcode, kprobe_inst, addr))
                        return -EINVAL;
 
+
+       p->ainsn.insn = get_insn_slot();
+       if (!p->ainsn.insn)
+               return -ENOMEM;
+       memcpy(&p->opcode, kprobe_addr, sizeof(kprobe_opcode_t));
+       memcpy(p->ainsn.insn, kprobe_addr, sizeof(kprobe_opcode_t));
+
        prepare_break_inst(template, slot, major_opcode, kprobe_inst, p);
 
        return 0;
@@ -455,8 +461,10 @@ void __kprobes arch_arm_kprobe(struct kprobe *p)
        unsigned long addr = (unsigned long)p->addr;
        unsigned long arm_addr = addr & ~0xFULL;
 
-       memcpy((char *)arm_addr, &p->ainsn.insn.bundle, sizeof(bundle_t));
-       flush_icache_range(arm_addr, arm_addr + sizeof(bundle_t));
+       flush_icache_range((unsigned long)p->ainsn.insn,
+                       (unsigned long)p->ainsn.insn + sizeof(kprobe_opcode_t));
+       memcpy((char *)arm_addr, &p->opcode, sizeof(kprobe_opcode_t));
+       flush_icache_range(arm_addr, arm_addr + sizeof(kprobe_opcode_t));
 }
 
 void __kprobes arch_disarm_kprobe(struct kprobe *p)
@@ -464,15 +472,18 @@ void __kprobes arch_disarm_kprobe(struct kprobe *p)
        unsigned long addr = (unsigned long)p->addr;
        unsigned long arm_addr = addr & ~0xFULL;
 
-       /* p->opcode contains the original unaltered bundle */
-       memcpy((char *) arm_addr, (char *) &p->opcode.bundle, sizeof(bundle_t));
-       flush_icache_range(arm_addr, arm_addr + sizeof(bundle_t));
+       /* p->ainsn.insn contains the original unaltered kprobe_opcode_t */
+       memcpy((char *) arm_addr, (char *) p->ainsn.insn,
+                                        sizeof(kprobe_opcode_t));
+       flush_icache_range(arm_addr, arm_addr + sizeof(kprobe_opcode_t));
 }
 
 void __kprobes arch_remove_kprobe(struct kprobe *p)
 {
+       mutex_lock(&kprobe_mutex);
+       free_insn_slot(p->ainsn.insn, 0);
+       mutex_unlock(&kprobe_mutex);
 }
-
 /*
  * We are resuming execution after a single step fault, so the pt_regs
  * structure reflects the register state after we executed the instruction
@@ -483,21 +494,22 @@ void __kprobes arch_remove_kprobe(struct kprobe *p)
  */
 static void __kprobes resume_execution(struct kprobe *p, struct pt_regs *regs)
 {
-       unsigned long bundle_addr = ((unsigned long) (&p->opcode.bundle)) & ~0xFULL;
-       unsigned long resume_addr = (unsigned long)p->addr & ~0xFULL;
-       unsigned long template;
-       int slot = ((unsigned long)p->addr & 0xf);
+       unsigned long bundle_addr = (unsigned long) (&p->ainsn.insn->bundle);
+       unsigned long resume_addr = (unsigned long)p->addr & ~0xFULL;
+       unsigned long template;
+       int slot = ((unsigned long)p->addr & 0xf);
 
-       template = p->opcode.bundle.quad0.template;
+       template = p->ainsn.insn->bundle.quad0.template;
 
-       if (slot == 1 && bundle_encoding[template][1] == L)
-               slot = 2;
+       if (slot == 1 && bundle_encoding[template][1] == L)
+               slot = 2;
 
        if (p->ainsn.inst_flag) {
 
                if (p->ainsn.inst_flag & INST_FLAG_FIX_RELATIVE_IP_ADDR) {
                        /* Fix relative IP address */
-                       regs->cr_iip = (regs->cr_iip - bundle_addr) + resume_addr;
+                       regs->cr_iip = (regs->cr_iip - bundle_addr) +
+                                       resume_addr;
                }
 
                if (p->ainsn.inst_flag & INST_FLAG_FIX_BRANCH_REG) {
@@ -534,23 +546,23 @@ static void __kprobes resume_execution(struct kprobe *p, struct pt_regs *regs)
        }
 
        if (slot == 2) {
-               if (regs->cr_iip == bundle_addr + 0x10) {
-                       regs->cr_iip = resume_addr + 0x10;
-               }
-       } else {
-               if (regs->cr_iip == bundle_addr) {
-                       regs->cr_iip = resume_addr;
-               }
+               if (regs->cr_iip == bundle_addr + 0x10) {
+                       regs->cr_iip = resume_addr + 0x10;
+               }
+       } else {
+               if (regs->cr_iip == bundle_addr) {
+                       regs->cr_iip = resume_addr;
+               }
        }
 
 turn_ss_off:
-       /* Turn off Single Step bit */
-       ia64_psr(regs)->ss = 0;
+       /* Turn off Single Step bit */
+       ia64_psr(regs)->ss = 0;
 }
 
 static void __kprobes prepare_ss(struct kprobe *p, struct pt_regs *regs)
 {
-       unsigned long bundle_addr = (unsigned long) &p->opcode.bundle;
+       unsigned long bundle_addr = (unsigned long) &p->ainsn.insn->bundle;
        unsigned long slot = (unsigned long)p->addr & 0xf;
 
        /* single step inline if break instruction */
@@ -581,7 +593,7 @@ static int __kprobes is_ia64_break_inst(struct pt_regs *regs)
 
        /* Move to slot 2, if bundle is MLX type and kprobe slot is 1 */
        if (slot == 1 && bundle_encoding[template][1] == L)
-               slot++;
+               slot++;
 
        /* Get Kprobe probe instruction at given slot*/
        get_kprobe_inst(&bundle, slot, &kprobe_inst, &major_opcode);
@@ -606,17 +618,22 @@ static int __kprobes pre_kprobes_handler(struct die_args *args)
        int ret = 0;
        struct pt_regs *regs = args->regs;
        kprobe_opcode_t *addr = (kprobe_opcode_t *)instruction_pointer(regs);
+       struct kprobe_ctlblk *kcb;
 
+       /*
+        * We don't want to be preempted for the entire
+        * duration of kprobe processing
+        */
        preempt_disable();
+       kcb = get_kprobe_ctlblk();
 
        /* Handle recursion cases */
        if (kprobe_running()) {
                p = get_kprobe(addr);
                if (p) {
-                       if ( (kprobe_status == KPROBE_HIT_SS) &&
+                       if ((kcb->kprobe_status == KPROBE_HIT_SS) &&
                             (p->ainsn.inst_flag == INST_FLAG_BREAK_INST)) {
-                               ia64_psr(regs)->ss = 0;
-                               unlock_kprobes();
+                               ia64_psr(regs)->ss = 0;
                                goto no_kprobe;
                        }
                        /* We have reentered the pre_kprobe_handler(), since
@@ -625,30 +642,35 @@ static int __kprobes pre_kprobes_handler(struct die_args *args)
                         * just single step on the instruction of the new probe
                         * without calling any user handlers.
                         */
-                       save_previous_kprobe();
-                       set_current_kprobe(p);
-                       p->nmissed++;
+                       save_previous_kprobe(kcb);
+                       set_current_kprobe(p, kcb);
+                       kprobes_inc_nmissed_count(p);
                        prepare_ss(p, regs);
-                       kprobe_status = KPROBE_REENTER;
+                       kcb->kprobe_status = KPROBE_REENTER;
                        return 1;
                } else if (args->err == __IA64_BREAK_JPROBE) {
                        /*
                         * jprobe instrumented function just completed
                         */
-                       p = current_kprobe;
+                       p = __get_cpu_var(current_kprobe);
                        if (p->break_handler && p->break_handler(p, regs)) {
                                goto ss_probe;
                        }
+               } else if (!is_ia64_break_inst(regs)) {
+                       /* The breakpoint instruction was removed by
+                        * another cpu right after we hit, no further
+                        * handling of this interrupt is appropriate
+                        */
+                       ret = 1;
+                       goto no_kprobe;
                } else {
                        /* Not our break */
                        goto no_kprobe;
                }
        }
 
-       lock_kprobes();
        p = get_kprobe(addr);
        if (!p) {
-               unlock_kprobes();
                if (!is_ia64_break_inst(regs)) {
                        /*
                         * The breakpoint instruction was removed right
@@ -665,8 +687,8 @@ static int __kprobes pre_kprobes_handler(struct die_args *args)
                goto no_kprobe;
        }
 
-       kprobe_status = KPROBE_HIT_ACTIVE;
-       set_current_kprobe(p);
+       set_current_kprobe(p, kcb);
+       kcb->kprobe_status = KPROBE_HIT_ACTIVE;
 
        if (p->pre_handler && p->pre_handler(p, regs))
                /*
@@ -678,7 +700,7 @@ static int __kprobes pre_kprobes_handler(struct die_args *args)
 
 ss_probe:
        prepare_ss(p, regs);
-       kprobe_status = KPROBE_HIT_SS;
+       kcb->kprobe_status = KPROBE_HIT_SS;
        return 1;
 
 no_kprobe:
@@ -688,23 +710,25 @@ no_kprobe:
 
 static int __kprobes post_kprobes_handler(struct pt_regs *regs)
 {
-       if (!kprobe_running())
+       struct kprobe *cur = kprobe_running();
+       struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+
+       if (!cur)
                return 0;
 
-       if ((kprobe_status != KPROBE_REENTER) && current_kprobe->post_handler) {
-               kprobe_status = KPROBE_HIT_SSDONE;
-               current_kprobe->post_handler(current_kprobe, regs, 0);
+       if ((kcb->kprobe_status != KPROBE_REENTER) && cur->post_handler) {
+               kcb->kprobe_status = KPROBE_HIT_SSDONE;
+               cur->post_handler(cur, regs, 0);
        }
 
-       resume_execution(current_kprobe, regs);
+       resume_execution(cur, regs);
 
        /*Restore back the original saved kprobes variables and continue. */
-       if (kprobe_status == KPROBE_REENTER) {
-               restore_previous_kprobe();
+       if (kcb->kprobe_status == KPROBE_REENTER) {
+               restore_previous_kprobe(kcb);
                goto out;
        }
-
-       unlock_kprobes();
+       reset_current_kprobe();
 
 out:
        preempt_enable_no_resched();
@@ -713,17 +737,59 @@ out:
 
 static int __kprobes kprobes_fault_handler(struct pt_regs *regs, int trapnr)
 {
-       if (!kprobe_running())
-               return 0;
+       struct kprobe *cur = kprobe_running();
+       struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
 
-       if (current_kprobe->fault_handler &&
-           current_kprobe->fault_handler(current_kprobe, regs, trapnr))
-               return 1;
 
-       if (kprobe_status & KPROBE_HIT_SS) {
-               resume_execution(current_kprobe, regs);
-               unlock_kprobes();
+       switch(kcb->kprobe_status) {
+       case KPROBE_HIT_SS:
+       case KPROBE_REENTER:
+               /*
+                * We are here because the instruction being single
+                * stepped caused a page fault. We reset the current
+                * kprobe and the instruction pointer points back to
+                * the probe address and allow the page fault handler
+                * to continue as a normal page fault.
+                */
+               regs->cr_iip = ((unsigned long)cur->addr) & ~0xFULL;
+               ia64_psr(regs)->ri = ((unsigned long)cur->addr) & 0xf;
+               if (kcb->kprobe_status == KPROBE_REENTER)
+                       restore_previous_kprobe(kcb);
+               else
+                       reset_current_kprobe();
                preempt_enable_no_resched();
+               break;
+       case KPROBE_HIT_ACTIVE:
+       case KPROBE_HIT_SSDONE:
+               /*
+                * We increment the nmissed count for accounting,
+                * we can also use npre/npostfault count for accouting
+                * these specific fault cases.
+                */
+               kprobes_inc_nmissed_count(cur);
+
+               /*
+                * We come here because instructions in the pre/post
+                * handler caused the page_fault, this could happen
+                * if handler tries to access user space by
+                * copy_from_user(), get_user() etc. Let the
+                * user-specified handler try to fix it first.
+                */
+               if (cur->fault_handler && cur->fault_handler(cur, regs, trapnr))
+                       return 1;
+               /*
+                * In case the user-specified fault handler returned
+                * zero, try to fix up.
+                */
+               if (ia64_done_with_exception(regs))
+                       return 1;
+
+               /*
+                * Let ia64_do_page_fault() fix it.
+                */
+               break;
+       default:
+               break;
        }
 
        return 0;
@@ -733,31 +799,90 @@ int __kprobes kprobe_exceptions_notify(struct notifier_block *self,
                                       unsigned long val, void *data)
 {
        struct die_args *args = (struct die_args *)data;
+       int ret = NOTIFY_DONE;
+
+       if (args->regs && user_mode(args->regs))
+               return ret;
+
        switch(val) {
        case DIE_BREAK:
-               if (pre_kprobes_handler(args))
-                       return NOTIFY_STOP;
+               /* err is break number from ia64_bad_break() */
+               if (args->err == 0x80200 || args->err == 0x80300 || args->err == 0)
+                       if (pre_kprobes_handler(args))
+                               ret = NOTIFY_STOP;
                break;
-       case DIE_SS:
-               if (post_kprobes_handler(args->regs))
-                       return NOTIFY_STOP;
+       case DIE_FAULT:
+               /* err is vector number from ia64_fault() */
+               if (args->err == 36)
+                       if (post_kprobes_handler(args->regs))
+                               ret = NOTIFY_STOP;
                break;
        case DIE_PAGE_FAULT:
-               if (kprobes_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
+               /* kprobe_running() needs smp_processor_id() */
+               preempt_disable();
+               if (kprobe_running() &&
+                       kprobes_fault_handler(args->regs, args->trapnr))
+                       ret = NOTIFY_STOP;
+               preempt_enable();
        default:
                break;
        }
-       return NOTIFY_DONE;
+       return ret;
+}
+
+struct param_bsp_cfm {
+       unsigned long ip;
+       unsigned long *bsp;
+       unsigned long cfm;
+};
+
+static void ia64_get_bsp_cfm(struct unw_frame_info *info, void *arg)
+{
+       unsigned long ip;
+       struct param_bsp_cfm *lp = arg;
+
+       do {
+               unw_get_ip(info, &ip);
+               if (ip == 0)
+                       break;
+               if (ip == lp->ip) {
+                       unw_get_bsp(info, (unsigned long*)&lp->bsp);
+                       unw_get_cfm(info, (unsigned long*)&lp->cfm);
+                       return;
+               }
+       } while (unw_unwind(info) >= 0);
+       lp->bsp = 0;
+       lp->cfm = 0;
+       return;
 }
 
 int __kprobes setjmp_pre_handler(struct kprobe *p, struct pt_regs *regs)
 {
        struct jprobe *jp = container_of(p, struct jprobe, kp);
        unsigned long addr = ((struct fnptr *)(jp->entry))->ip;
+       struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+       struct param_bsp_cfm pa;
+       int bytes;
+
+       /*
+        * Callee owns the argument space and could overwrite it, eg
+        * tail call optimization. So to be absolutely safe
+        * we save the argument space before transfering the control
+        * to instrumented jprobe function which runs in
+        * the process context
+        */
+       pa.ip = regs->cr_iip;
+       unw_init_running(ia64_get_bsp_cfm, &pa);
+       bytes = (char *)ia64_rse_skip_regs(pa.bsp, pa.cfm & 0x3f)
+                               - (char *)pa.bsp;
+       memcpy( kcb->jprobes_saved_stacked_regs,
+               pa.bsp,
+               bytes );
+       kcb->bsp = pa.bsp;
+       kcb->cfm = pa.cfm;
 
        /* save architectural state */
-       jprobe_saved_regs = *regs;
+       kcb->jprobe_saved_regs = *regs;
 
        /* after rfi, execute the jprobe instrumented function */
        regs->cr_iip = addr & ~0xFULL;
@@ -768,14 +893,29 @@ int __kprobes setjmp_pre_handler(struct kprobe *p, struct pt_regs *regs)
         * fix the return address to our jprobe_inst_return() function
         * in the jprobes.S file
         */
-       regs->b0 = ((struct fnptr *)(jprobe_inst_return))->ip;
+       regs->b0 = ((struct fnptr *)(jprobe_inst_return))->ip;
 
        return 1;
 }
 
 int __kprobes longjmp_break_handler(struct kprobe *p, struct pt_regs *regs)
 {
-       *regs = jprobe_saved_regs;
+       struct kprobe_ctlblk *kcb = get_kprobe_ctlblk();
+       int bytes;
+
+       /* restoring architectural state */
+       *regs = kcb->jprobe_saved_regs;
+
+       /* restoring the original argument space */
+       flush_register_stack();
+       bytes = (char *)ia64_rse_skip_regs(kcb->bsp, kcb->cfm & 0x3f)
+                               - (char *)kcb->bsp;
+       memcpy( kcb->bsp,
+               kcb->jprobes_saved_stacked_regs,
+               bytes );
+       invalidate_stacked_regs();
+
+       preempt_enable_no_resched();
        return 1;
 }