kprobe_opcode_t *addr = NULL;
        unsigned long *lp;
 
-       /* We're in an interrupt, but this is clear and BUG()-safe. */
-       preempt_disable();
        /* Check if the application is using LDT entry for its code segment and
         * calculate the address by reading the base address from the LDT entry.
         */
                goto no_kprobe;
        }
 
+       /*
+        * This preempt_disable() matches the preempt_enable_no_resched()
+        * in post_kprobe_handler()
+        */
+       preempt_disable();
        kprobe_status = KPROBE_HIT_ACTIVE;
        set_current_kprobe(p, regs);
 
        return 1;
 
 no_kprobe:
-       preempt_enable_no_resched();
        return ret;
 }
 
        unlock_kprobes();
        preempt_enable_no_resched();
 
-        /*
-         * By returning a non-zero value, we are telling
-         * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption.
-         */
+       /*
+        * By returning a non-zero value, we are telling
+        * kprobe_handler() that we have handled unlocking
+        * and re-enabling preemption
+        */
         return 1;
 }
 
                                       unsigned long val, void *data)
 {
        struct die_args *args = (struct die_args *)data;
+       int ret = NOTIFY_DONE;
+
+       preempt_disable();
        switch (val) {
        case DIE_INT3:
                if (kprobe_handler(args->regs))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_DEBUG:
                if (post_kprobe_handler(args->regs))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_GPF:
-               if (kprobe_running() &&
-                   kprobe_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
-               break;
        case DIE_PAGE_FAULT:
                if (kprobe_running() &&
                    kprobe_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        default:
                break;
        }
-       return NOTIFY_DONE;
+       preempt_enable();
+       return ret;
 }
 
 int __kprobes setjmp_pre_handler(struct kprobe *p, struct pt_regs *regs)
 
 void __kprobes jprobe_return(void)
 {
-       preempt_enable_no_resched();
        asm volatile ("       xchgl   %%ebx,%%esp     \n"
                      "       int3                      \n"
                      "       .globl jprobe_return_end  \n"
 
         /*
          * By returning a non-zero value, we are telling
          * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption.
+        * and re-enabling preemption
          */
         return 1;
 }
        struct pt_regs *regs = args->regs;
        kprobe_opcode_t *addr = (kprobe_opcode_t *)instruction_pointer(regs);
 
-       preempt_disable();
-
        /* Handle recursion cases */
        if (kprobe_running()) {
                p = get_kprobe(addr);
                goto no_kprobe;
        }
 
+       /*
+        * This preempt_disable() matches the preempt_enable_no_resched()
+        * in post_kprobes_handler()
+        */
+       preempt_disable();
        kprobe_status = KPROBE_HIT_ACTIVE;
        set_current_kprobe(p);
 
        return 1;
 
 no_kprobe:
-       preempt_enable_no_resched();
        return ret;
 }
 
                                       unsigned long val, void *data)
 {
        struct die_args *args = (struct die_args *)data;
+       int ret = NOTIFY_DONE;
+
+       preempt_disable();
        switch(val) {
        case DIE_BREAK:
                if (pre_kprobes_handler(args))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_SS:
                if (post_kprobes_handler(args->regs))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_PAGE_FAULT:
                if (kprobes_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
        default:
                break;
        }
-       return NOTIFY_DONE;
+       preempt_enable();
+       return ret;
 }
 
 int __kprobes setjmp_pre_handler(struct kprobe *p, struct pt_regs *regs)
 
                goto no_kprobe;
        }
 
+       /*
+        * This preempt_disable() matches the preempt_enable_no_resched()
+        * in post_kprobe_handler().
+        */
+       preempt_disable();
        kprobe_status = KPROBE_HIT_ACTIVE;
        current_kprobe = p;
        kprobe_saved_msr = regs->msr;
 ss_probe:
        prepare_singlestep(p, regs);
        kprobe_status = KPROBE_HIT_SS;
-       /*
-        * This preempt_disable() matches the preempt_enable_no_resched()
-        * in post_kprobe_handler().
-        */
-       preempt_disable();
        return 1;
 
 no_kprobe:
        regs->nip = orig_ret_address;
 
        unlock_kprobes();
+       preempt_enable_no_resched();
 
         /*
          * By returning a non-zero value, we are telling
 
        void *addr = (void *) regs->tpc;
        int ret = 0;
 
-       preempt_disable();
-
        if (kprobe_running()) {
                /* We *are* holding lock here, so this is safe.
                 * Disarm the probe we just hit, and ignore it.
                goto no_kprobe;
        }
 
+       /*
+        * This preempt_disable() matches the preempt_enable_no_resched()
+        * in post_kprobes_handler()
+        */
+       preempt_disable();
        set_current_kprobe(p, regs);
        kprobe_status = KPROBE_HIT_ACTIVE;
        if (p->pre_handler && p->pre_handler(p, regs))
        return 1;
 
 no_kprobe:
-       preempt_enable_no_resched();
        return ret;
 }
 
                                       unsigned long val, void *data)
 {
        struct die_args *args = (struct die_args *)data;
+       int ret = NOTIFY_DONE;
+
+       preempt_disable();
        switch (val) {
        case DIE_DEBUG:
                if (kprobe_handler(args->regs))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_DEBUG_2:
                if (post_kprobe_handler(args->regs))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_GPF:
-               if (kprobe_running() &&
-                   kprobe_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
-               break;
        case DIE_PAGE_FAULT:
                if (kprobe_running() &&
                    kprobe_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        default:
                break;
        }
-       return NOTIFY_DONE;
+       preempt_enable();
+       return ret;
 }
 
 asmlinkage void __kprobes kprobe_trap(unsigned long trap_level,
 
 void __kprobes jprobe_return(void)
 {
-       preempt_enable_no_resched();
        __asm__ __volatile__(
                ".globl jprobe_return_trap_instruction\n"
 "jprobe_return_trap_instruction:\n\t"
 
        int ret = 0;
        kprobe_opcode_t *addr = (kprobe_opcode_t *)(regs->rip - sizeof(kprobe_opcode_t));
 
-       /* We're in an interrupt, but this is clear and BUG()-safe. */
-       preempt_disable();
-
        /* Check we're not actually recursing */
        if (kprobe_running()) {
                /* We *are* holding lock here, so this is safe.
                goto no_kprobe;
        }
 
+       /*
+        * This preempt_disable() matches the preempt_enable_no_resched()
+        * in post_kprobe_handler()
+        */
+       preempt_disable();
        kprobe_status = KPROBE_HIT_ACTIVE;
        set_current_kprobe(p, regs);
 
        return 1;
 
 no_kprobe:
-       preempt_enable_no_resched();
        return ret;
 }
 
         /*
          * By returning a non-zero value, we are telling
          * kprobe_handler() that we have handled unlocking
-         * and re-enabling preemption.
+        * and re-enabling preemption
          */
         return 1;
 }
                                       unsigned long val, void *data)
 {
        struct die_args *args = (struct die_args *)data;
+       int ret = NOTIFY_DONE;
+
+       preempt_disable();
        switch (val) {
        case DIE_INT3:
                if (kprobe_handler(args->regs))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_DEBUG:
                if (post_kprobe_handler(args->regs))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        case DIE_GPF:
-               if (kprobe_running() &&
-                   kprobe_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
-               break;
        case DIE_PAGE_FAULT:
                if (kprobe_running() &&
                    kprobe_fault_handler(args->regs, args->trapnr))
-                       return NOTIFY_STOP;
+                       ret = NOTIFY_STOP;
                break;
        default:
                break;
        }
-       return NOTIFY_DONE;
+       preempt_enable();
+       return ret;
 }
 
 int __kprobes setjmp_pre_handler(struct kprobe *p, struct pt_regs *regs)
 
 void __kprobes jprobe_return(void)
 {
-       preempt_enable_no_resched();
        asm volatile ("       xchg   %%rbx,%%rsp     \n"
                      "       int3                      \n"
                      "       .globl jprobe_return_end  \n"