diff mbox series

[2/4] arm64: Extract kprobes_save_local_irqflag() and kprobes_restore_local_irqflag()

Message ID 20200509214159.19680-3-liwei391@huawei.com (mailing list archive)
State New, archived
Headers show
Series arm64: kgdb/kdb: Fix single-step debugging issues | expand

Commit Message

Wei Li May 9, 2020, 9:41 p.m. UTC
PSTATE.I and PSTATE.D are very important for single-step working.

Without disabling interrupt on local CPU, there is a chance of
interrupt occurrence in the period of exception return and start of
out-of-line single-step, that result in wrongly single stepping
into the interrupt handler. And if D bit is set then, it results into
undefined exception and when it's handler enables dbg then single-step
exception is generated, not as expected.

As they are maintained well in kprobes_save_local_irqflag() and
kprobes_restore_local_irqflag() for kprobe module, extract them as
kernel_prepare_single_step() and kernel_cleanup_single_step() for
general use.

Signed-off-by: Wei Li <liwei391@huawei.com>
---
 arch/arm64/include/asm/debug-monitors.h |  4 ++++
 arch/arm64/kernel/debug-monitors.c      | 26 +++++++++++++++++++++++
 arch/arm64/kernel/probes/kprobes.c      | 28 ++-----------------------
 3 files changed, 32 insertions(+), 26 deletions(-)

Comments

Masami Hiramatsu (Google) May 10, 2020, 8:59 a.m. UTC | #1
Hi Wei,

On Sun, 10 May 2020 05:41:57 +0800
Wei Li <liwei391@huawei.com> wrote:

> PSTATE.I and PSTATE.D are very important for single-step working.
> 
> Without disabling interrupt on local CPU, there is a chance of
> interrupt occurrence in the period of exception return and start of
> out-of-line single-step, that result in wrongly single stepping
> into the interrupt handler. And if D bit is set then, it results into
> undefined exception and when it's handler enables dbg then single-step
> exception is generated, not as expected.
> 
> As they are maintained well in kprobes_save_local_irqflag() and
> kprobes_restore_local_irqflag() for kprobe module, extract them as
> kernel_prepare_single_step() and kernel_cleanup_single_step() for
> general use.

Agreed, these prepare/cleanup should be generic.

Reviewed-by: Masami Hiramatsu <mhiramat@kernel.org>

Thank you!

> 
> Signed-off-by: Wei Li <liwei391@huawei.com>
> ---
>  arch/arm64/include/asm/debug-monitors.h |  4 ++++
>  arch/arm64/kernel/debug-monitors.c      | 26 +++++++++++++++++++++++
>  arch/arm64/kernel/probes/kprobes.c      | 28 ++-----------------------
>  3 files changed, 32 insertions(+), 26 deletions(-)
> 
> diff --git a/arch/arm64/include/asm/debug-monitors.h b/arch/arm64/include/asm/debug-monitors.h
> index 7619f473155f..b62469f3475b 100644
> --- a/arch/arm64/include/asm/debug-monitors.h
> +++ b/arch/arm64/include/asm/debug-monitors.h
> @@ -113,6 +113,10 @@ void user_fastforward_single_step(struct task_struct *task);
>  void kernel_enable_single_step(struct pt_regs *regs);
>  void kernel_disable_single_step(void);
>  int kernel_active_single_step(void);
> +void kernel_prepare_single_step(unsigned long *flags,
> +						struct pt_regs *regs);
> +void kernel_cleanup_single_step(unsigned long flags,
> +						struct pt_regs *regs);
>  
>  #ifdef CONFIG_HAVE_HW_BREAKPOINT
>  int reinstall_suspended_bps(struct pt_regs *regs);
> diff --git a/arch/arm64/kernel/debug-monitors.c b/arch/arm64/kernel/debug-monitors.c
> index 48222a4760c2..25ce6b5a52d2 100644
> --- a/arch/arm64/kernel/debug-monitors.c
> +++ b/arch/arm64/kernel/debug-monitors.c
> @@ -429,6 +429,32 @@ int kernel_active_single_step(void)
>  }
>  NOKPROBE_SYMBOL(kernel_active_single_step);
>  
> +/*
> + * Interrupts need to be disabled before single-step mode is set, and not
> + * reenabled until after single-step mode ends.
> + * Without disabling interrupt on local CPU, there is a chance of
> + * interrupt occurrence in the period of exception return and  start of
> + * out-of-line single-step, that result in wrongly single stepping
> + * into the interrupt handler.
> + */
> +void kernel_prepare_single_step(unsigned long *flags,
> +						struct pt_regs *regs)
> +{
> +	*flags = regs->pstate & DAIF_MASK;
> +	regs->pstate |= PSR_I_BIT;
> +	/* Unmask PSTATE.D for enabling software step exceptions. */
> +	regs->pstate &= ~PSR_D_BIT;
> +}
> +NOKPROBE_SYMBOL(kernel_prepare_single_step);
> +
> +void kernel_cleanup_single_step(unsigned long flags,
> +						struct pt_regs *regs)
> +{
> +	regs->pstate &= ~DAIF_MASK;
> +	regs->pstate |= flags;
> +}
> +NOKPROBE_SYMBOL(kernel_cleanup_single_step);
> +
>  /* ptrace API */
>  void user_enable_single_step(struct task_struct *task)
>  {
> diff --git a/arch/arm64/kernel/probes/kprobes.c b/arch/arm64/kernel/probes/kprobes.c
> index d1c95dcf1d78..c655b6b543e3 100644
> --- a/arch/arm64/kernel/probes/kprobes.c
> +++ b/arch/arm64/kernel/probes/kprobes.c
> @@ -168,30 +168,6 @@ static void __kprobes set_current_kprobe(struct kprobe *p)
>  	__this_cpu_write(current_kprobe, p);
>  }
>  
> -/*
> - * Interrupts need to be disabled before single-step mode is set, and not
> - * reenabled until after single-step mode ends.
> - * Without disabling interrupt on local CPU, there is a chance of
> - * interrupt occurrence in the period of exception return and  start of
> - * out-of-line single-step, that result in wrongly single stepping
> - * into the interrupt handler.
> - */
> -static void __kprobes kprobes_save_local_irqflag(struct kprobe_ctlblk *kcb,
> -						struct pt_regs *regs)
> -{
> -	kcb->saved_irqflag = regs->pstate & DAIF_MASK;
> -	regs->pstate |= PSR_I_BIT;
> -	/* Unmask PSTATE.D for enabling software step exceptions. */
> -	regs->pstate &= ~PSR_D_BIT;
> -}
> -
> -static void __kprobes kprobes_restore_local_irqflag(struct kprobe_ctlblk *kcb,
> -						struct pt_regs *regs)
> -{
> -	regs->pstate &= ~DAIF_MASK;
> -	regs->pstate |= kcb->saved_irqflag;
> -}
> -
>  static void __kprobes
>  set_ss_context(struct kprobe_ctlblk *kcb, unsigned long addr)
>  {
> @@ -227,7 +203,7 @@ static void __kprobes setup_singlestep(struct kprobe *p,
>  		set_ss_context(kcb, slot);	/* mark pending ss */
>  
>  		/* IRQs and single stepping do not mix well. */
> -		kprobes_save_local_irqflag(kcb, regs);
> +		kernel_prepare_single_step(&kcb->saved_irqflag, regs);
>  		kernel_enable_single_step(regs);
>  		instruction_pointer_set(regs, slot);
>  	} else {
> @@ -414,7 +390,7 @@ kprobe_single_step_handler(struct pt_regs *regs, unsigned int esr)
>  	retval = kprobe_ss_hit(kcb, instruction_pointer(regs));
>  
>  	if (retval == DBG_HOOK_HANDLED) {
> -		kprobes_restore_local_irqflag(kcb, regs);
> +		kernel_cleanup_single_step(kcb->saved_irqflag, regs);
>  		kernel_disable_single_step();
>  
>  		post_kprobe_handler(kcb, regs);
> -- 
> 2.17.1
>
Doug Anderson May 14, 2020, 12:21 a.m. UTC | #2
Hi,

On Sat, May 9, 2020 at 6:49 AM Wei Li <liwei391@huawei.com> wrote:
>
> PSTATE.I and PSTATE.D are very important for single-step working.
>
> Without disabling interrupt on local CPU, there is a chance of
> interrupt occurrence in the period of exception return and start of
> out-of-line single-step, that result in wrongly single stepping
> into the interrupt handler. And if D bit is set then, it results into
> undefined exception and when it's handler enables dbg then single-step
> exception is generated, not as expected.
>
> As they are maintained well in kprobes_save_local_irqflag() and
> kprobes_restore_local_irqflag() for kprobe module, extract them as
> kernel_prepare_single_step() and kernel_cleanup_single_step() for
> general use.
>
> Signed-off-by: Wei Li <liwei391@huawei.com>
> ---
>  arch/arm64/include/asm/debug-monitors.h |  4 ++++
>  arch/arm64/kernel/debug-monitors.c      | 26 +++++++++++++++++++++++
>  arch/arm64/kernel/probes/kprobes.c      | 28 ++-----------------------
>  3 files changed, 32 insertions(+), 26 deletions(-)
>
> diff --git a/arch/arm64/include/asm/debug-monitors.h b/arch/arm64/include/asm/debug-monitors.h
> index 7619f473155f..b62469f3475b 100644
> --- a/arch/arm64/include/asm/debug-monitors.h
> +++ b/arch/arm64/include/asm/debug-monitors.h
> @@ -113,6 +113,10 @@ void user_fastforward_single_step(struct task_struct *task);
>  void kernel_enable_single_step(struct pt_regs *regs);
>  void kernel_disable_single_step(void);
>  int kernel_active_single_step(void);
> +void kernel_prepare_single_step(unsigned long *flags,
> +                                               struct pt_regs *regs);
> +void kernel_cleanup_single_step(unsigned long flags,
> +                                               struct pt_regs *regs);
>
>  #ifdef CONFIG_HAVE_HW_BREAKPOINT
>  int reinstall_suspended_bps(struct pt_regs *regs);
> diff --git a/arch/arm64/kernel/debug-monitors.c b/arch/arm64/kernel/debug-monitors.c
> index 48222a4760c2..25ce6b5a52d2 100644
> --- a/arch/arm64/kernel/debug-monitors.c
> +++ b/arch/arm64/kernel/debug-monitors.c
> @@ -429,6 +429,32 @@ int kernel_active_single_step(void)
>  }
>  NOKPROBE_SYMBOL(kernel_active_single_step);
>
> +/*
> + * Interrupts need to be disabled before single-step mode is set, and not
> + * reenabled until after single-step mode ends.
> + * Without disabling interrupt on local CPU, there is a chance of
> + * interrupt occurrence in the period of exception return and  start of
> + * out-of-line single-step, that result in wrongly single stepping
> + * into the interrupt handler.
> + */
> +void kernel_prepare_single_step(unsigned long *flags,
> +                                               struct pt_regs *regs)
> +{
> +       *flags = regs->pstate & DAIF_MASK;
> +       regs->pstate |= PSR_I_BIT;
> +       /* Unmask PSTATE.D for enabling software step exceptions. */
> +       regs->pstate &= ~PSR_D_BIT;
> +}
> +NOKPROBE_SYMBOL(kernel_prepare_single_step);

nit: why not just return unsigned long rather than passing by reference?


> +
> +void kernel_cleanup_single_step(unsigned long flags,
> +                                               struct pt_regs *regs)
> +{
> +       regs->pstate &= ~DAIF_MASK;
> +       regs->pstate |= flags;
> +}
> +NOKPROBE_SYMBOL(kernel_cleanup_single_step);
> +
>  /* ptrace API */
>  void user_enable_single_step(struct task_struct *task)
>  {
> diff --git a/arch/arm64/kernel/probes/kprobes.c b/arch/arm64/kernel/probes/kprobes.c
> index d1c95dcf1d78..c655b6b543e3 100644
> --- a/arch/arm64/kernel/probes/kprobes.c
> +++ b/arch/arm64/kernel/probes/kprobes.c
> @@ -168,30 +168,6 @@ static void __kprobes set_current_kprobe(struct kprobe *p)
>         __this_cpu_write(current_kprobe, p);
>  }
>
> -/*
> - * Interrupts need to be disabled before single-step mode is set, and not
> - * reenabled until after single-step mode ends.
> - * Without disabling interrupt on local CPU, there is a chance of
> - * interrupt occurrence in the period of exception return and  start of
> - * out-of-line single-step, that result in wrongly single stepping
> - * into the interrupt handler.
> - */
> -static void __kprobes kprobes_save_local_irqflag(struct kprobe_ctlblk *kcb,
> -                                               struct pt_regs *regs)
> -{
> -       kcb->saved_irqflag = regs->pstate & DAIF_MASK;
> -       regs->pstate |= PSR_I_BIT;
> -       /* Unmask PSTATE.D for enabling software step exceptions. */
> -       regs->pstate &= ~PSR_D_BIT;
> -}
> -
> -static void __kprobes kprobes_restore_local_irqflag(struct kprobe_ctlblk *kcb,
> -                                               struct pt_regs *regs)
> -{
> -       regs->pstate &= ~DAIF_MASK;
> -       regs->pstate |= kcb->saved_irqflag;
> -}
> -
>  static void __kprobes
>  set_ss_context(struct kprobe_ctlblk *kcb, unsigned long addr)
>  {
> @@ -227,7 +203,7 @@ static void __kprobes setup_singlestep(struct kprobe *p,
>                 set_ss_context(kcb, slot);      /* mark pending ss */
>
>                 /* IRQs and single stepping do not mix well. */
> -               kprobes_save_local_irqflag(kcb, regs);
> +               kernel_prepare_single_step(&kcb->saved_irqflag, regs);

Is there some reason to have two functions?  It seems like every time
you call kernel_enable_single_step() you'd want to call
kernel_prepare_single_step().  ...and every time you call
kernel_disable_single_step() you'd want to call
kernel_cleanup_single_step().

Maybe you can just add the flags parameter to
kernel_enable_single_step() / kernel_disable_single_step() and put the
code in there?


-Doug
Wei Li May 16, 2020, 8:47 a.m. UTC | #3
Hi Douglas,

On 2020/5/14 8:21, Doug Anderson wrote:
(SNIP)
>> +/*
>> + * Interrupts need to be disabled before single-step mode is set, and not
>> + * reenabled until after single-step mode ends.
>> + * Without disabling interrupt on local CPU, there is a chance of
>> + * interrupt occurrence in the period of exception return and  start of
>> + * out-of-line single-step, that result in wrongly single stepping
>> + * into the interrupt handler.
>> + */
>> +void kernel_prepare_single_step(unsigned long *flags,
>> +                                               struct pt_regs *regs)
>> +{
>> +       *flags = regs->pstate & DAIF_MASK;
>> +       regs->pstate |= PSR_I_BIT;
>> +       /* Unmask PSTATE.D for enabling software step exceptions. */
>> +       regs->pstate &= ~PSR_D_BIT;
>> +}
>> +NOKPROBE_SYMBOL(kernel_prepare_single_step);
> 
> nit: why not just return unsigned long rather than passing by reference?
Because i just extract this function from kprobes_save_local_irqflag(), i think
return unsigned long is fine.

> 
>> +
>> +void kernel_cleanup_single_step(unsigned long flags,
>> +                                               struct pt_regs *regs)
>> +{
>> +       regs->pstate &= ~DAIF_MASK;
>> +       regs->pstate |= flags;
>> +}
>> +NOKPROBE_SYMBOL(kernel_cleanup_single_step);
>> +
>>  /* ptrace API */
>>  void user_enable_single_step(struct task_struct *task)
>>  {
>> diff --git a/arch/arm64/kernel/probes/kprobes.c b/arch/arm64/kernel/probes/kprobes.c
>> index d1c95dcf1d78..c655b6b543e3 100644
>> --- a/arch/arm64/kernel/probes/kprobes.c
>> +++ b/arch/arm64/kernel/probes/kprobes.c
>> @@ -168,30 +168,6 @@ static void __kprobes set_current_kprobe(struct kprobe *p)
>>         __this_cpu_write(current_kprobe, p);
>>  }
>>
>> -/*
>> - * Interrupts need to be disabled before single-step mode is set, and not
>> - * reenabled until after single-step mode ends.
>> - * Without disabling interrupt on local CPU, there is a chance of
>> - * interrupt occurrence in the period of exception return and  start of
>> - * out-of-line single-step, that result in wrongly single stepping
>> - * into the interrupt handler.
>> - */
>> -static void __kprobes kprobes_save_local_irqflag(struct kprobe_ctlblk *kcb,
>> -                                               struct pt_regs *regs)
>> -{
>> -       kcb->saved_irqflag = regs->pstate & DAIF_MASK;
>> -       regs->pstate |= PSR_I_BIT;
>> -       /* Unmask PSTATE.D for enabling software step exceptions. */
>> -       regs->pstate &= ~PSR_D_BIT;
>> -}
>> -
>> -static void __kprobes kprobes_restore_local_irqflag(struct kprobe_ctlblk *kcb,
>> -                                               struct pt_regs *regs)
>> -{
>> -       regs->pstate &= ~DAIF_MASK;
>> -       regs->pstate |= kcb->saved_irqflag;
>> -}
>> -
>>  static void __kprobes
>>  set_ss_context(struct kprobe_ctlblk *kcb, unsigned long addr)
>>  {
>> @@ -227,7 +203,7 @@ static void __kprobes setup_singlestep(struct kprobe *p,
>>                 set_ss_context(kcb, slot);      /* mark pending ss */
>>
>>                 /* IRQs and single stepping do not mix well. */
>> -               kprobes_save_local_irqflag(kcb, regs);
>> +               kernel_prepare_single_step(&kcb->saved_irqflag, regs);
> 
> Is there some reason to have two functions?  It seems like every time
> you call kernel_enable_single_step() you'd want to call
> kernel_prepare_single_step().  ...and every time you call
> kernel_disable_single_step() you'd want to call
> kernel_cleanup_single_step().
> 
> Maybe you can just add the flags parameter to
> kernel_enable_single_step() / kernel_disable_single_step() and put the
> code in there?
> 

As kernel_enable_single_step() / kernel_disable_single_step() are also called in
breakpoint_handler() and watchpoint_handler(), i am not sure it's a right thing
to put the daif flag prepare/cleanup into them, especially we don't have a context
to save the flags.

Thanks,
Wei
Doug Anderson May 16, 2020, 4:17 p.m. UTC | #4
Hi,

On Sat, May 16, 2020 at 1:47 AM liwei (GF) <liwei391@huawei.com> wrote:
>
> >> -               kprobes_save_local_irqflag(kcb, regs);
> >> +               kernel_prepare_single_step(&kcb->saved_irqflag, regs);
> >
> > Is there some reason to have two functions?  It seems like every time
> > you call kernel_enable_single_step() you'd want to call
> > kernel_prepare_single_step().  ...and every time you call
> > kernel_disable_single_step() you'd want to call
> > kernel_cleanup_single_step().
> >
> > Maybe you can just add the flags parameter to
> > kernel_enable_single_step() / kernel_disable_single_step() and put the
> > code in there?
> >
>
> As kernel_enable_single_step() / kernel_disable_single_step() are also called in
> breakpoint_handler() and watchpoint_handler(), i am not sure it's a right thing
> to put the daif flag prepare/cleanup into them, especially we don't have a context
> to save the flags.

I think you misunderstood what I was suggesting.  Maybe better with
examples?  I was suggesting doing this:

kcb->saved_irqflag = kernel_enable_single_step(regs);
...
kernel_disable_single_step(kcb->saved_irqflag, regs);

To me that seems better than what you have now:

kcb->saved_irqflag = kernel_prepare_single_step(regs);
kernel_enable_single_step(regs);
...
kernel_cleanup_single_step(kcb->saved_irqflag, regs);
kernel_disable_single_step();

...or am I confused?

-Doug
Masami Hiramatsu (Google) May 18, 2020, 3:14 p.m. UTC | #5
On Sat, 16 May 2020 09:17:21 -0700
Doug Anderson <dianders@chromium.org> wrote:

> Hi,
> 
> On Sat, May 16, 2020 at 1:47 AM liwei (GF) <liwei391@huawei.com> wrote:
> >
> > >> -               kprobes_save_local_irqflag(kcb, regs);
> > >> +               kernel_prepare_single_step(&kcb->saved_irqflag, regs);
> > >
> > > Is there some reason to have two functions?  It seems like every time
> > > you call kernel_enable_single_step() you'd want to call
> > > kernel_prepare_single_step().  ...and every time you call
> > > kernel_disable_single_step() you'd want to call
> > > kernel_cleanup_single_step().
> > >
> > > Maybe you can just add the flags parameter to
> > > kernel_enable_single_step() / kernel_disable_single_step() and put the
> > > code in there?
> > >
> >
> > As kernel_enable_single_step() / kernel_disable_single_step() are also called in
> > breakpoint_handler() and watchpoint_handler(), i am not sure it's a right thing
> > to put the daif flag prepare/cleanup into them, especially we don't have a context
> > to save the flags.
> 
> I think you misunderstood what I was suggesting.  Maybe better with
> examples?  I was suggesting doing this:
> 
> kcb->saved_irqflag = kernel_enable_single_step(regs);
> ...
> kernel_disable_single_step(kcb->saved_irqflag, regs);
> 
> To me that seems better than what you have now:
> 
> kcb->saved_irqflag = kernel_prepare_single_step(regs);
> kernel_enable_single_step(regs);
> ...
> kernel_cleanup_single_step(kcb->saved_irqflag, regs);
> kernel_disable_single_step();
> 
> ...or am I confused?

+1, this sounds good to me. Currently arch/arm64/kernel/probes/kprobes.c
has a code which sololy use kernel_disable_single_step() without regs
restoring, but it looks like a bug there. So maybe you need following patch.

Thank you,

-----
From: Masami Hiramatsu <mhiramat@kernel.org>
Date: Mon, 18 May 2020 23:08:28 +0900
Subject: [PATCH] arm64: kprobes: Restore saved interrupt flag before disabling
 single step

Restore the saved interrupt flag in kprobe_ctlblk to regs->pstate
when a page fault happens on single-stepping instruction.
Without this fix, we will lose the flag if it happens because
kcb->saved_irqflag only knows the previous flag.

Fixes: 2dd0e8d2d2a1 ("arm64: Kprobes with single stepping support")
Signed-off-by: Masami Hiramatsu <mhiramat@kernel.org>
Cc: stable@vger.kernel.org
---
 arch/arm64/kernel/probes/kprobes.c | 1 +
 1 file changed, 1 insertion(+)

diff --git a/arch/arm64/kernel/probes/kprobes.c b/arch/arm64/kernel/probes/kprobes.c
index d1c95dcf1d78..73fb99770f69 100644
--- a/arch/arm64/kernel/probes/kprobes.c
+++ b/arch/arm64/kernel/probes/kprobes.c
@@ -308,6 +308,7 @@ int __kprobes kprobe_fault_handler(struct pt_regs *regs, unsigned int fsr)
 		if (!instruction_pointer(regs))
 			BUG();
 
+		kprobes_restore_local_irqflag(kcb, regs);
 		kernel_disable_single_step();
 
 		if (kcb->kprobe_status == KPROBE_REENTER)
diff mbox series

Patch

diff --git a/arch/arm64/include/asm/debug-monitors.h b/arch/arm64/include/asm/debug-monitors.h
index 7619f473155f..b62469f3475b 100644
--- a/arch/arm64/include/asm/debug-monitors.h
+++ b/arch/arm64/include/asm/debug-monitors.h
@@ -113,6 +113,10 @@  void user_fastforward_single_step(struct task_struct *task);
 void kernel_enable_single_step(struct pt_regs *regs);
 void kernel_disable_single_step(void);
 int kernel_active_single_step(void);
+void kernel_prepare_single_step(unsigned long *flags,
+						struct pt_regs *regs);
+void kernel_cleanup_single_step(unsigned long flags,
+						struct pt_regs *regs);
 
 #ifdef CONFIG_HAVE_HW_BREAKPOINT
 int reinstall_suspended_bps(struct pt_regs *regs);
diff --git a/arch/arm64/kernel/debug-monitors.c b/arch/arm64/kernel/debug-monitors.c
index 48222a4760c2..25ce6b5a52d2 100644
--- a/arch/arm64/kernel/debug-monitors.c
+++ b/arch/arm64/kernel/debug-monitors.c
@@ -429,6 +429,32 @@  int kernel_active_single_step(void)
 }
 NOKPROBE_SYMBOL(kernel_active_single_step);
 
+/*
+ * Interrupts need to be disabled before single-step mode is set, and not
+ * reenabled until after single-step mode ends.
+ * Without disabling interrupt on local CPU, there is a chance of
+ * interrupt occurrence in the period of exception return and  start of
+ * out-of-line single-step, that result in wrongly single stepping
+ * into the interrupt handler.
+ */
+void kernel_prepare_single_step(unsigned long *flags,
+						struct pt_regs *regs)
+{
+	*flags = regs->pstate & DAIF_MASK;
+	regs->pstate |= PSR_I_BIT;
+	/* Unmask PSTATE.D for enabling software step exceptions. */
+	regs->pstate &= ~PSR_D_BIT;
+}
+NOKPROBE_SYMBOL(kernel_prepare_single_step);
+
+void kernel_cleanup_single_step(unsigned long flags,
+						struct pt_regs *regs)
+{
+	regs->pstate &= ~DAIF_MASK;
+	regs->pstate |= flags;
+}
+NOKPROBE_SYMBOL(kernel_cleanup_single_step);
+
 /* ptrace API */
 void user_enable_single_step(struct task_struct *task)
 {
diff --git a/arch/arm64/kernel/probes/kprobes.c b/arch/arm64/kernel/probes/kprobes.c
index d1c95dcf1d78..c655b6b543e3 100644
--- a/arch/arm64/kernel/probes/kprobes.c
+++ b/arch/arm64/kernel/probes/kprobes.c
@@ -168,30 +168,6 @@  static void __kprobes set_current_kprobe(struct kprobe *p)
 	__this_cpu_write(current_kprobe, p);
 }
 
-/*
- * Interrupts need to be disabled before single-step mode is set, and not
- * reenabled until after single-step mode ends.
- * Without disabling interrupt on local CPU, there is a chance of
- * interrupt occurrence in the period of exception return and  start of
- * out-of-line single-step, that result in wrongly single stepping
- * into the interrupt handler.
- */
-static void __kprobes kprobes_save_local_irqflag(struct kprobe_ctlblk *kcb,
-						struct pt_regs *regs)
-{
-	kcb->saved_irqflag = regs->pstate & DAIF_MASK;
-	regs->pstate |= PSR_I_BIT;
-	/* Unmask PSTATE.D for enabling software step exceptions. */
-	regs->pstate &= ~PSR_D_BIT;
-}
-
-static void __kprobes kprobes_restore_local_irqflag(struct kprobe_ctlblk *kcb,
-						struct pt_regs *regs)
-{
-	regs->pstate &= ~DAIF_MASK;
-	regs->pstate |= kcb->saved_irqflag;
-}
-
 static void __kprobes
 set_ss_context(struct kprobe_ctlblk *kcb, unsigned long addr)
 {
@@ -227,7 +203,7 @@  static void __kprobes setup_singlestep(struct kprobe *p,
 		set_ss_context(kcb, slot);	/* mark pending ss */
 
 		/* IRQs and single stepping do not mix well. */
-		kprobes_save_local_irqflag(kcb, regs);
+		kernel_prepare_single_step(&kcb->saved_irqflag, regs);
 		kernel_enable_single_step(regs);
 		instruction_pointer_set(regs, slot);
 	} else {
@@ -414,7 +390,7 @@  kprobe_single_step_handler(struct pt_regs *regs, unsigned int esr)
 	retval = kprobe_ss_hit(kcb, instruction_pointer(regs));
 
 	if (retval == DBG_HOOK_HANDLED) {
-		kprobes_restore_local_irqflag(kcb, regs);
+		kernel_cleanup_single_step(kcb->saved_irqflag, regs);
 		kernel_disable_single_step();
 
 		post_kprobe_handler(kcb, regs);