[RFC][PATCH v2 09/11] context_tracking,livepatch: Dont disturb NOHZ_FULL

[Peter Zijlstra <peterz@infradead.org>]

Using the new context_tracking infrastructure, avoid disturbing
userspace tasks when context tracking is enabled.

When context_tracking_set_cpu_work() returns true, we have the
guarantee that klp_update_patch_state() is called from noinstr code
before it runs normal kernel code. This covers
syscall/exceptions/interrupts and NMI entry.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 include/linux/context_tracking.h |    2 +-
 include/linux/livepatch.h        |    2 ++
 kernel/context_tracking.c        |   11 +++++------
 kernel/livepatch/transition.c    |   29 ++++++++++++++++++++++++++---
 4 files changed, 34 insertions(+), 10 deletions(-)

--- a/include/linux/context_tracking.h
+++ b/include/linux/context_tracking.h
@@ -10,7 +10,7 @@
 #include <asm/ptrace.h>
 
 enum ct_work {
-	CT_WORK_n = 0,
+	CT_WORK_KLP = 1,
 };
 
 /*
--- a/include/linux/livepatch.h
+++ b/include/linux/livepatch.h
@@ -201,6 +201,7 @@ void klp_module_going(struct module *mod
 
 void klp_copy_process(struct task_struct *child);
 void klp_update_patch_state(struct task_struct *task);
+void __klp_update_patch_state(struct task_struct *task);
 
 static inline bool klp_patch_pending(struct task_struct *task)
 {
@@ -242,6 +243,7 @@ static inline int klp_module_coming(stru
 static inline void klp_module_going(struct module *mod) {}
 static inline bool klp_patch_pending(struct task_struct *task) { return false; }
 static inline void klp_update_patch_state(struct task_struct *task) {}
+static inline void __klp_update_patch_state(struct task_struct *task) {}
 static inline void klp_copy_process(struct task_struct *child) {}
 
 static inline
--- a/kernel/context_tracking.c
+++ b/kernel/context_tracking.c
@@ -21,6 +21,7 @@
 #include <linux/hardirq.h>
 #include <linux/export.h>
 #include <linux/kprobes.h>
+#include <linux/livepatch.h>
 
 #define CREATE_TRACE_POINTS
 #include <trace/events/context_tracking.h>
@@ -55,15 +56,13 @@ static noinstr void ct_exit_user_work(struct
 {
 	unsigned int work = arch_atomic_read(&ct->work);
 
-#if 0
-	if (work & CT_WORK_n) {
+	if (work & CT_WORK_KLP) {
 		/* NMI happens here and must still do/finish CT_WORK_n */
-		do_work_n();
+		__klp_update_patch_state(current);
 
 		smp_mb__before_atomic();
-		arch_atomic_andnot(CT_WORK_n, &ct->work);
+		arch_atomic_andnot(CT_WORK_KLP, &ct->work);
 	}
-#endif
 
 	smp_mb__before_atomic();
 	arch_atomic_andnot(CT_SEQ_WORK, &ct->seq);
--- a/kernel/livepatch/transition.c
+++ b/kernel/livepatch/transition.c
@@ -10,6 +10,7 @@
 #include <linux/cpu.h>
 #include <linux/stacktrace.h>
 #include <linux/tracehook.h>
+#include <linux/context_tracking.h>
 #include "core.h"
 #include "patch.h"
 #include "transition.h"
@@ -153,6 +154,11 @@ void klp_cancel_transition(void)
 	klp_complete_transition();
 }
 
+noinstr void __klp_update_patch_state(struct task_struct *task)
+{
+	task->patch_state = READ_ONCE(klp_target_state);
+}
+
 /*
  * Switch the patched state of the task to the set of functions in the target
  * patch state.
@@ -180,8 +186,10 @@ void klp_update_patch_state(struct task_
 	 *    of func->transition, if klp_ftrace_handler() is called later on
 	 *    the same CPU.  See __klp_disable_patch().
 	 */
-	if (test_and_clear_tsk_thread_flag(task, TIF_PATCH_PENDING))
+	if (test_tsk_thread_flag(task, TIF_PATCH_PENDING)) {
 		task->patch_state = READ_ONCE(klp_target_state);
+		clear_tsk_thread_flag(task, TIF_PATCH_PENDING);
+	}
 
 	preempt_enable_notrace();
 }
@@ -270,15 +278,30 @@ static int klp_check_and_switch_task(str
 {
 	int ret;
 
-	if (task_curr(task))
+	if (task_curr(task)) {
+		/*
+		 * This only succeeds when the task is in NOHZ_FULL user
+		 * mode, the true return value guarantees any kernel entry
+		 * will call klp_update_patch_state().
+		 *
+		 * XXX: ideally we'd simply return 0 here and leave
+		 * TIF_PATCH_PENDING alone, to be fixed up by
+		 * klp_update_patch_state(), except livepatching goes wobbly
+		 * with 'pending' TIF bits on.
+		 */
+		if (context_tracking_set_cpu_work(task_cpu(task), CT_WORK_KLP))
+			goto clear;
+
 		return -EBUSY;
+	}
 
 	ret = klp_check_stack(task, arg);
 	if (ret)
 		return ret;
 
-	clear_tsk_thread_flag(task, TIF_PATCH_PENDING);
 	task->patch_state = klp_target_state;
+clear:
+	clear_tsk_thread_flag(task, TIF_PATCH_PENDING);
 	return 0;
 }