[RFC PATCH 34/36] arm64: kprobes: Disable NMIs

[Vladimir Murzin <vladimir.murzin@arm.com>]

Kprobes masks all DAIF exceptions to keep things simple and avoid
nested exceptions. With FEAT_NMI, that is no longer enough. Mask
ALLINT as well to ensure things stay simple.

Signed-off-by: Vladimir Murzin <vladimir.murzin@arm.com>
---
 arch/arm64/kernel/probes/kprobes.c | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/arch/arm64/kernel/probes/kprobes.c b/arch/arm64/kernel/probes/kprobes.c
index d9d25d184309..defbe6061546 100644
--- a/arch/arm64/kernel/probes/kprobes.c
+++ b/arch/arm64/kernel/probes/kprobes.c
@@ -193,17 +193,27 @@ static void __kprobes set_current_kprobe(struct kprobe *p)
  * the kprobe state is per-CPU and doesn't get migrated.
  */
 static void __kprobes kprobes_save_local_irqflag(struct kprobe_ctlblk *kcb,
-						struct pt_regs *regs)
+						 struct pt_regs *regs)
 {
-	kcb->saved_irqflag = regs->pstate & DAIF_MASK;
-	regs->pstate |= DAIF_MASK;
+	arm64_exc_hwstate_t next = arm64_exc_hwstate_of_context(CRITICAL_CONTEXT);
+
+	kcb->saved_irqflag = arm64_inherit_exc_hwstate(regs).flags;
+
+	regs->pstate &= ~(DAIF_MASK | PSR_ALLINT_BIT);
+
+	regs->pstate |= next.daif;
+	regs->pstate |= next.allint;
 }
 
 static void __kprobes kprobes_restore_local_irqflag(struct kprobe_ctlblk *kcb,
-						struct pt_regs *regs)
+						    struct pt_regs *regs)
 {
-	regs->pstate &= ~DAIF_MASK;
-	regs->pstate |= kcb->saved_irqflag;
+	arm64_exc_hwstate_t prev = {.flags = kcb->saved_irqflag};
+
+	regs->pstate &= ~(DAIF_MASK | PSR_ALLINT_BIT);
+
+	regs->pstate |= prev.daif;
+	regs->pstate |= prev.allint;
 }
 
 static void __kprobes setup_singlestep(struct kprobe *p,
-- 
2.34.1