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

[Josh Poimboeuf <jpoimboe@redhat.com>]

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@redhat.com>
---
 arch/x86/include/asm/thread_info.h |  4 +++-
 arch/x86/kernel/signal.c           |  4 ++++
 include/linux/livepatch.h          |  2 ++
 kernel/livepatch/transition.c      | 15 +++++++++++++++
 kernel/sched/idle.c                |  4 ++++
 5 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/arch/x86/include/asm/thread_info.h b/arch/x86/include/asm/thread_info.h
index 547e344..4e46d36 100644
--- a/arch/x86/include/asm/thread_info.h
+++ b/arch/x86/include/asm/thread_info.h
@@ -78,6 +78,7 @@ struct thread_info {
 #define TIF_MCE_NOTIFY		10	/* notify userspace of an MCE */
 #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 */
@@ -102,6 +103,7 @@ struct thread_info {
 #define _TIF_SECCOMP		(1 << TIF_SECCOMP)
 #define _TIF_MCE_NOTIFY		(1 << TIF_MCE_NOTIFY)
 #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)
@@ -141,7 +143,7 @@ struct thread_info {
 /* Only used for 64 bit */
 #define _TIF_DO_NOTIFY_MASK						\
 	(_TIF_SIGPENDING | _TIF_MCE_NOTIFY | _TIF_NOTIFY_RESUME |	\
-	 _TIF_USER_RETURN_NOTIFY | _TIF_UPROBE)
+	 _TIF_USER_RETURN_NOTIFY | _TIF_UPROBE | _TIF_KLP_NEED_UPDATE)
 
 /* flags to check in __switch_to() */
 #define _TIF_WORK_CTXSW							\
diff --git a/arch/x86/kernel/signal.c b/arch/x86/kernel/signal.c
index ed37a76..1d4b8e6 100644
--- a/arch/x86/kernel/signal.c
+++ b/arch/x86/kernel/signal.c
@@ -23,6 +23,7 @@
 #include <linux/user-return-notifier.h>
 #include <linux/uprobes.h>
 #include <linux/context_tracking.h>
+#include <linux/livepatch.h>
 
 #include <asm/processor.h>
 #include <asm/ucontext.h>
@@ -760,6 +761,9 @@ do_notify_resume(struct pt_regs *regs, void *unused, __u32 thread_info_flags)
 	if (thread_info_flags & _TIF_USER_RETURN_NOTIFY)
 		fire_user_return_notifiers();
 
+	if (unlikely(thread_info_flags & _TIF_KLP_NEED_UPDATE))
+		klp_update_task_universe(current);
+
 	user_enter();
 }
 
diff --git a/include/linux/livepatch.h b/include/linux/livepatch.h
index b8c2f15..14f6a96 100644
--- a/include/linux/livepatch.h
+++ b/include/linux/livepatch.h
@@ -134,6 +134,8 @@ extern int klp_universe_goal;
 
 static inline void klp_update_task_universe(struct task_struct *t)
 {
+	clear_tsk_thread_flag(t, TIF_KLP_NEED_UPDATE);
+
 	/* corresponding smp_wmb() is in klp_set_universe_goal() */
 	smp_rmb();
 
diff --git a/kernel/livepatch/transition.c b/kernel/livepatch/transition.c
index 20fafd2..dac8ea5 100644
--- a/kernel/livepatch/transition.c
+++ b/kernel/livepatch/transition.c
@@ -234,6 +234,9 @@ static void klp_transition_work_fn(struct work_struct *work)
  */
 void klp_start_transition(int universe)
 {
+	struct task_struct *g, *t;
+	unsigned int cpu;
+
 	if (WARN_ON(klp_universe_goal == universe))
 		return;
 
@@ -241,6 +244,18 @@ void klp_start_transition(int universe)
 		  universe == KLP_UNIVERSE_NEW ? "patching" : "unpatching");
 
 	klp_set_universe_goal(universe);
+
+	/* mark all normal tasks as needing a universe update */
+	read_lock(&tasklist_lock);
+	for_each_process_thread(g, t)
+		set_tsk_thread_flag(t, TIF_KLP_NEED_UPDATE);
+	read_unlock(&tasklist_lock);
+
+	/* mark all idle "swapper" tasks as needing a universe update */
+	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 c47fce7..c1390b6 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -7,6 +7,7 @@
 #include <linux/tick.h>
 #include <linux/mm.h>
 #include <linux/stackprotector.h>
+#include <linux/livepatch.h>
 
 #include <asm/tlb.h>
 
@@ -250,6 +251,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.1.0