2 files changed, 30 insertions, 0 deletions

diff --git a/include/linux/context_tracking_state.h b/include/linux/context_tracking_state.h
index 6b7b96a32b75..678ecdf90cf6 100644
--- a/include/linux/context_tracking_state.h
+++ b/include/linux/context_tracking_state.h
@@ -12,6 +12,7 @@ struct context_tracking {
          * may be further optimized using static keys.
          */
         bool active;
+        int recursion;
         enum ctx_state {
                 CONTEXT_KERNEL = 0,
                 CONTEXT_USER,
diff --git a/kernel/context_tracking.c b/kernel/context_tracking.c
index 72d59a1a6eb6..5b11a10e196a 100644
--- a/kernel/context_tracking.c
+++ b/kernel/context_tracking.c
@@ -38,6 +38,25 @@ void context_tracking_cpu_set(int cpu)
         }
 }
 
+static bool context_tracking_recursion_enter(void)
+{
+        int recursion;
+
+        recursion = __this_cpu_inc_return(context_tracking.recursion);
+        if (recursion == 1)
+                return true;
+
+        WARN_ONCE((recursion < 1), "Invalid context tracking recursion value %d\n", recursion);
+        __this_cpu_dec(context_tracking.recursion);
+
+        return false;
+}
+
+static void context_tracking_recursion_exit(void)
+{
+        __this_cpu_dec(context_tracking.recursion);
+}
+
 /**
  * context_tracking_enter - Inform the context tracking that the CPU is going
  *                          enter user or guest space mode.
@@ -75,6 +94,9 @@ void context_tracking_enter(enum ctx_state state)
         WARN_ON_ONCE(!current->mm);
 
         local_irq_save(flags);
+        if (!context_tracking_recursion_enter())
+                goto out_irq_restore;
+
         if ( __this_cpu_read(context_tracking.state) != state) {
                 if (__this_cpu_read(context_tracking.active)) {
                         /*
@@ -105,6 +127,8 @@ void context_tracking_enter(enum ctx_state state)
                  */
                 __this_cpu_write(context_tracking.state, state);
         }
+        context_tracking_recursion_exit();
+out_irq_restore:
         local_irq_restore(flags);
 }
 NOKPROBE_SYMBOL(context_tracking_enter);
@@ -139,6 +163,9 @@ void context_tracking_exit(enum ctx_state state)
                 return;
 
         local_irq_save(flags);
+        if (!context_tracking_recursion_enter())
+                goto out_irq_restore;
+
         if (__this_cpu_read(context_tracking.state) == state) {
                 if (__this_cpu_read(context_tracking.active)) {
                         /*
@@ -153,6 +180,8 @@ void context_tracking_exit(enum ctx_state state)
                 }
                 __this_cpu_write(context_tracking.state, CONTEXT_KERNEL);
         }
+        context_tracking_recursion_exit();
+out_irq_restore:
         local_irq_restore(flags);
 }
 NOKPROBE_SYMBOL(context_tracking_exit);