[RFC PATCH v1.9 14/14] livepatch: update task universe when exiting kernel

[Josh Poimboeuf]

Update a tasks's universe when returning from a system call or user
space interrupt, or after handling a signal.

This greatly increases the chances of a patch operation succeeding.  If
a task is I/O bound, it can switch universes when returning from a
system call.  If a task is CPU bound, it can switch universes when
returning from an interrupt.  If a task is sleeping on a to-be-patched
function, the user can send SIGSTOP and SIGCONT to force it to switch.

Since the idle "swapper" tasks don't ever exit the kernel, they're
updated from within the idle loop.

Signed-off-by: Josh Poimboeuf <jpoimboe@xxxxxxxxxx>
---
 arch/x86/entry/common.c            |  6 +++++-
 arch/x86/include/asm/thread_info.h |  2 ++
 include/linux/livepatch.h          |  2 ++
 kernel/livepatch/transition.c      | 37 +++++++++++++++++++++++++++++++++----
 kernel/sched/idle.c                |  4 ++++
 5 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/arch/x86/entry/common.c b/arch/x86/entry/common.c
index e79d93d..94639dd 100644
--- a/arch/x86/entry/common.c
+++ b/arch/x86/entry/common.c
@@ -21,6 +21,7 @@
 #include <linux/context_tracking.h>
 #include <linux/user-return-notifier.h>
 #include <linux/uprobes.h>
+#include <linux/livepatch.h>
 
 #include <asm/desc.h>
 #include <asm/traps.h>
@@ -202,7 +203,7 @@ long syscall_trace_enter(struct pt_regs *regs)
 
 #define EXIT_TO_USERMODE_LOOP_FLAGS				\
 	(_TIF_SIGPENDING | _TIF_NOTIFY_RESUME | _TIF_UPROBE |	\
-	 _TIF_NEED_RESCHED | _TIF_USER_RETURN_NOTIFY)
+	 _TIF_NEED_RESCHED | _TIF_USER_RETURN_NOTIFY | _TIF_KLP_NEED_UPDATE)
 
 static void exit_to_usermode_loop(struct pt_regs *regs, u32 cached_flags)
 {
@@ -236,6 +237,9 @@ static void exit_to_usermode_loop(struct pt_regs *regs, u32 cached_flags)
 		if (cached_flags & _TIF_USER_RETURN_NOTIFY)
 			fire_user_return_notifiers();
 
+		if (unlikely(cached_flags & _TIF_KLP_NEED_UPDATE))
+			klp_update_task_universe(current);
+
 		/* Disable IRQs and retry */
 		local_irq_disable();
 
diff --git a/arch/x86/include/asm/thread_info.h b/arch/x86/include/asm/thread_info.h
index 8286669..4e3ea6f 100644
--- a/arch/x86/include/asm/thread_info.h
+++ b/arch/x86/include/asm/thread_info.h
@@ -97,6 +97,7 @@ struct thread_info {
 #define TIF_SECCOMP		8	/* secure computing */
 #define TIF_USER_RETURN_NOTIFY	11	/* notify kernel of userspace return */
 #define TIF_UPROBE		12	/* breakpointed or singlestepping */
+#define TIF_KLP_NEED_UPDATE	13	/* pending live patching update */
 #define TIF_NOTSC		16	/* TSC is not accessible in userland */
 #define TIF_IA32		17	/* IA32 compatibility process */
 #define TIF_FORK		18	/* ret_from_fork */
@@ -120,6 +121,7 @@ struct thread_info {
 #define _TIF_SYSCALL_AUDIT	(1 << TIF_SYSCALL_AUDIT)
 #define _TIF_SECCOMP		(1 << TIF_SECCOMP)
 #define _TIF_USER_RETURN_NOTIFY	(1 << TIF_USER_RETURN_NOTIFY)
+#define _TIF_KLP_NEED_UPDATE	(1 << TIF_KLP_NEED_UPDATE)
 #define _TIF_UPROBE		(1 << TIF_UPROBE)
 #define _TIF_NOTSC		(1 << TIF_NOTSC)
 #define _TIF_IA32		(1 << TIF_IA32)
diff --git a/include/linux/livepatch.h b/include/linux/livepatch.h
index 4d2e26d..29964ac 100644
--- a/include/linux/livepatch.h
+++ b/include/linux/livepatch.h
@@ -152,6 +152,8 @@ extern int klp_universe_goal;
  */
 static inline void klp_update_task_universe(struct task_struct *task)
 {
+	clear_tsk_thread_flag(task, TIF_KLP_NEED_UPDATE);
+
 	/*
 	 * The corresponding write barriers are in klp_init_transition() and
 	 * klp_start_transition().  See the comments there for an explanation.
diff --git a/kernel/livepatch/transition.c b/kernel/livepatch/transition.c
index 0609d84..8a38247 100644
--- a/kernel/livepatch/transition.c
+++ b/kernel/livepatch/transition.c
@@ -277,6 +277,9 @@ success:
  */
 void klp_start_transition(int universe)
 {
+	struct task_struct *g, *task;
+	unsigned int cpu;
+
 	if (WARN_ON(klp_universe_goal == universe))
 		return;
 
@@ -293,12 +296,38 @@ void klp_start_transition(int universe)
 	klp_universe_goal = universe;
 
 	/*
-	 * Enforce the ordering of the universe goal write with later
-	 * task universe writes which are done via
-	 * klp_try_complete_transition().  The corresponding read barrier is in
-	 * klp_update_task_universe().
+	 * Ensure that if another CPU goes through the syscall barrier, sees
+	 * the TIF_KLP_NEED_UPDATE write below, and calls
+	 * klp_update_task_universe(), it also sees the above write to the
+	 * universe goal.  Otherwise it can put the task in the wrong universe.
 	 */
 	smp_wmb();
+
+	/*
+	 * If the patch can be applied or reverted immediately, skip the
+	 * per-task transitions.
+	 */
+	if (klp_transition_patch->immediate)
+		return;
+
+	/*
+	 * Mark all normal tasks as needing a universe update.  As they pass
+	 * through the syscall barrier they'll switch over to the goal universe
+	 * (unless we switch them in klp_try_complete_transition() first).
+	 */
+	read_lock(&tasklist_lock);
+	for_each_process_thread(g, task)
+		set_tsk_thread_flag(task, TIF_KLP_NEED_UPDATE);
+	read_unlock(&tasklist_lock);
+
+	/*
+	 * Ditto for the idle "swapper" tasks, though they never cross the
+	 * syscall barrier.
+	 */
+	get_online_cpus();
+	for_each_online_cpu(cpu)
+		set_tsk_thread_flag(idle_task(cpu), TIF_KLP_NEED_UPDATE);
+	put_online_cpus();
 }
 
 /*
diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index bd12c6c..94bdad9 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -9,6 +9,7 @@
 #include <linux/mm.h>
 #include <linux/stackprotector.h>
 #include <linux/suspend.h>
+#include <linux/livepatch.h>
 
 #include <asm/tlb.h>
 
@@ -266,6 +267,9 @@ static void cpu_idle_loop(void)
 
 		sched_ttwu_pending();
 		schedule_preempt_disabled();
+
+		if (unlikely(test_thread_flag(TIF_KLP_NEED_UPDATE)))
+			klp_update_task_universe(current);
 	}
 }
 
-- 
2.4.3