[RFC PATCH v1.9 12/14] livepatch: create per-task consistency model

[Josh Poimboeuf]

Add a basic per-task consistency model.  This is the foundation which
will eventually enable us to patch those ~10% of security patches which
change function prototypes and/or data semantics.

When a patch is enabled, livepatch enters into a transition state where
tasks are converging from the old universe to the new universe.  If a
given task isn't using any of the patched functions, it's switched to
the new universe.  Once all the tasks have been converged to the new
universe, patching is complete.

The same sequence occurs when a patch is disabled, except the tasks
converge from the new universe to the old universe.

The /sys/kernel/livepatch/<patch>/transition file shows whether a patch
is in transition.  Only a single patch (the topmost patch on the stack)
can be in transition at a given time.  A patch can remain in the
transition state indefinitely, if any of the tasks are stuck in the
previous universe.

A transition can be reversed and effectively canceled by writing the
opposite value to the /sys/kernel/livepatch/<patch>/enabled file while
the transition is in progress.  Then all the tasks will attempt to
converge back to the original universe.

Signed-off-by: Josh Poimboeuf <jpoimboe@xxxxxxxxxx>
---
 include/linux/livepatch.h     |  27 +++
 include/linux/sched.h         |   3 +
 kernel/fork.c                 |   2 +
 kernel/livepatch/Makefile     |   2 +-
 kernel/livepatch/core.c       |  99 +++++++---
 kernel/livepatch/patch.c      |  43 ++++-
 kernel/livepatch/patch.h      |   1 +
 kernel/livepatch/transition.c | 406 ++++++++++++++++++++++++++++++++++++++++++
 kernel/livepatch/transition.h |  20 +++
 kernel/sched/core.c           |   2 +
 10 files changed, 582 insertions(+), 23 deletions(-)
 create mode 100644 kernel/livepatch/transition.c
 create mode 100644 kernel/livepatch/transition.h

diff --git a/include/linux/livepatch.h b/include/linux/livepatch.h
index dd5db74..4d2e26d 100644
--- a/include/linux/livepatch.h
+++ b/include/linux/livepatch.h
@@ -34,12 +34,14 @@
  * @new_func:	pointer to the patched function code
  * @old_sympos: a hint indicating which symbol position the old function
  *		can be found (optional)
+ * @immediate:  patch the func immediately, bypassing backtrace safety checks
  * @old_addr:	the address of the function being patched
  * @kobj:	kobject for sysfs resources
  * @stack_node:	list node for klp_ops func_stack list
  * @old_size:	size of the old function
  * @new_size:	size of the new function
  * @patched:	the func has been added to the klp_ops list
+ * @transition:	the func is currently being applied or reverted
  */
 struct klp_func {
 	/* external */
@@ -53,6 +55,7 @@ struct klp_func {
 	 * in kallsyms for the given object is used.
 	 */
 	unsigned long old_sympos;
+	bool immediate;
 
 	/* internal */
 	unsigned long old_addr;
@@ -60,6 +63,7 @@ struct klp_func {
 	struct list_head stack_node;
 	unsigned long old_size, new_size;
 	bool patched;
+	bool transition;
 };
 
 /**
@@ -106,6 +110,7 @@ struct klp_object {
  * struct klp_patch - patch structure for live patching
  * @mod:	reference to the live patch module
  * @objs:	object entries for kernel objects to be patched
+ * @immediate:  patch all funcs immediately, bypassing safety mechanisms
  * @list:	list node for global list of registered patches
  * @kobj:	kobject for sysfs resources
  * @enabled:	the patch is enabled (but operation may be incomplete)
@@ -114,6 +119,7 @@ struct klp_patch {
 	/* external */
 	struct module *mod;
 	struct klp_object *objs;
+	bool immediate;
 
 	/* internal */
 	struct list_head list;
@@ -136,11 +142,32 @@ int klp_disable_patch(struct klp_patch *);
 int klp_module_coming(struct module *mod);
 void klp_module_going(struct module *mod);
 
+extern int klp_universe_goal;
+/*
+ * klp_update_task_universe() - change the patched state of a task
+ * @task:	The task to change
+ *
+ * Converts the patched state of the task so that it will switch to the set of
+ * functions in the goal universe.
+ */
+static inline void klp_update_task_universe(struct task_struct *task)
+{
+	/*
+	 * The corresponding write barriers are in klp_init_transition() and
+	 * klp_start_transition().  See the comments there for an explanation.
+	 */
+	smp_rmb();
+
+	task->klp_universe = klp_universe_goal;
+}
+
 #else /* !CONFIG_LIVEPATCH */
 
 static inline int klp_module_coming(struct module *mod) { return 0; }
 static inline void klp_module_going(struct module *mod) { }
 
+static inline void klp_update_task_universe(struct task_struct *task) {}
+
 #endif /* CONFIG_LIVEPATCH */
 
 #endif /* _LINUX_LIVEPATCH_H_ */
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 62d0961..c27286f 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1848,6 +1848,9 @@ struct task_struct {
 	unsigned long	task_state_change;
 #endif
 	int pagefault_disabled;
+#ifdef CONFIG_LIVEPATCH
+	int klp_universe;
+#endif
 /* CPU-specific state of this task */
 	struct thread_struct thread;
 /*
diff --git a/kernel/fork.c b/kernel/fork.c
index d277e83..27b181e 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -76,6 +76,7 @@
 #include <linux/compiler.h>
 #include <linux/sysctl.h>
 #include <linux/kcov.h>
+#include <linux/livepatch.h>
 
 #include <asm/pgtable.h>
 #include <asm/pgalloc.h>
@@ -1615,6 +1616,7 @@ static struct task_struct *copy_process(unsigned long clone_flags,
 	total_forks++;
 	spin_unlock(&current->sighand->siglock);
 	syscall_tracepoint_update(p);
+	klp_update_task_universe(p);
 	write_unlock_irq(&tasklist_lock);
 
 	proc_fork_connector(p);
diff --git a/kernel/livepatch/Makefile b/kernel/livepatch/Makefile
index e136dad..2b8bdb1 100644
--- a/kernel/livepatch/Makefile
+++ b/kernel/livepatch/Makefile
@@ -1,3 +1,3 @@
 obj-$(CONFIG_LIVEPATCH) += livepatch.o
 
-livepatch-objs := core.o patch.o
+livepatch-objs := core.o patch.o transition.o
diff --git a/kernel/livepatch/core.c b/kernel/livepatch/core.c
index b0fd31d..19afa9b 100644
--- a/kernel/livepatch/core.c
+++ b/kernel/livepatch/core.c
@@ -27,14 +27,18 @@
 #include <linux/list.h>
 #include <linux/kallsyms.h>
 #include <linux/livepatch.h>
+#include <linux/stacktrace.h>
 #include <asm/cacheflush.h>
 #include "patch.h"
+#include "transition.h"
 
 /*
- * The klp_mutex protects the global lists and state transitions of any
- * structure reachable from them.  References to any structure must be obtained
- * under mutex protection (except in klp_ftrace_handler(), which uses RCU to
- * ensure it gets consistent data).
+ * klp_mutex is a coarse lock which serializes access to klp data.  All
+ * accesses to klp-related variables and structures must have mutex protection,
+ * except within the following functions which carefully avoid the need for it:
+ *
+ * - klp_ftrace_handler()
+ * - klp_update_task_universe()
  */
 static DEFINE_MUTEX(klp_mutex);
 
@@ -42,6 +46,30 @@ static LIST_HEAD(klp_patches);
 
 static struct kobject *klp_root_kobj;
 
+static void klp_work_fn(struct work_struct *work);
+static DECLARE_DELAYED_WORK(klp_work, klp_work_fn);
+
+static void klp_schedule_work(void)
+{
+	if (IS_ENABLED(CONFIG_RELIABLE_STACKTRACE))
+		schedule_delayed_work(&klp_work, round_jiffies_relative(HZ));
+}
+
+/*
+ * This work can be performed periodically to finish patching or unpatching any
+ * "straggler" tasks which failed to transition in klp_enable_patch().
+ */
+static void klp_work_fn(struct work_struct *work)
+{
+	mutex_lock(&klp_mutex);
+
+	if (klp_transition_patch)
+		if (!klp_try_complete_transition())
+			klp_schedule_work();
+
+	mutex_unlock(&klp_mutex);
+}
+
 static bool klp_is_module(struct klp_object *obj)
 {
 	return obj->name;
@@ -80,7 +108,6 @@ static void klp_find_object_module(struct klp_object *obj)
 	mutex_unlock(&module_mutex);
 }
 
-/* klp_mutex must be held by caller */
 static bool klp_is_patch_registered(struct klp_patch *patch)
 {
 	struct klp_patch *mypatch;
@@ -244,19 +271,18 @@ out:
 
 static int __klp_disable_patch(struct klp_patch *patch)
 {
-	struct klp_object *obj;
+	if (klp_transition_patch)
+		return -EBUSY;
 
 	/* enforce stacking: only the last enabled patch can be disabled */
 	if (!list_is_last(&patch->list, &klp_patches) &&
 	    list_next_entry(patch, list)->enabled)
 		return -EBUSY;
 
-	pr_notice("disabling patch '%s'\n", patch->mod->name);
-
-	klp_for_each_object(patch, obj) {
-		if (obj->patched)
-			klp_unpatch_object(obj);
-	}
+	klp_init_transition(patch, KLP_UNIVERSE_NEW);
+	klp_start_transition(KLP_UNIVERSE_OLD);
+	if (!klp_try_complete_transition())
+		klp_schedule_work();
 
 	patch->enabled = false;
 
@@ -300,6 +326,9 @@ static int __klp_enable_patch(struct klp_patch *patch)
 	struct klp_object *obj;
 	int ret;
 
+	if (klp_transition_patch)
+		return -EBUSY;
+
 	if (WARN_ON(patch->enabled))
 		return -EINVAL;
 
@@ -311,24 +340,32 @@ static int __klp_enable_patch(struct klp_patch *patch)
 	pr_notice_once("tainting kernel with TAINT_LIVEPATCH\n");
 	add_taint(TAINT_LIVEPATCH, LOCKDEP_STILL_OK);
 
-	pr_notice("enabling patch '%s'\n", patch->mod->name);
+	klp_init_transition(patch, KLP_UNIVERSE_OLD);
 
 	klp_for_each_object(patch, obj) {
 		if (!klp_is_object_loaded(obj))
 			continue;
 
 		ret = klp_patch_object(obj);
-		if (ret)
-			goto unregister;
+		if (ret) {
+			pr_warn("failed to enable patch '%s'\n",
+				patch->mod->name);
+
+			klp_unpatch_objects(patch);
+			klp_complete_transition();
+
+			return ret;
+		}
 	}
 
+	klp_start_transition(KLP_UNIVERSE_NEW);
+
+	if (!klp_try_complete_transition())
+		klp_schedule_work();
+
 	patch->enabled = true;
 
 	return 0;
-
-unregister:
-	WARN_ON(__klp_disable_patch(patch));
-	return ret;
 }
 
 /**
@@ -365,6 +402,7 @@ EXPORT_SYMBOL_GPL(klp_enable_patch);
  * /sys/kernel/livepatch
  * /sys/kernel/livepatch/<patch>
  * /sys/kernel/livepatch/<patch>/enabled
+ * /sys/kernel/livepatch/<patch>/transition
  * /sys/kernel/livepatch/<patch>/<object>
  * /sys/kernel/livepatch/<patch>/<object>/<function,sympos>
  */
@@ -393,7 +431,9 @@ static ssize_t enabled_store(struct kobject *kobj, struct kobj_attribute *attr,
 		goto err;
 	}
 
-	if (val) {
+	if (klp_transition_patch == patch) {
+		klp_reverse_transition();
+	} else if (val) {
 		ret = __klp_enable_patch(patch);
 		if (ret)
 			goto err;
@@ -421,9 +461,21 @@ static ssize_t enabled_show(struct kobject *kobj,
 	return snprintf(buf, PAGE_SIZE-1, "%d\n", patch->enabled);
 }
 
+static ssize_t transition_show(struct kobject *kobj,
+			       struct kobj_attribute *attr, char *buf)
+{
+	struct klp_patch *patch;
+
+	patch = container_of(kobj, struct klp_patch, kobj);
+	return snprintf(buf, PAGE_SIZE-1, "%d\n",
+			klp_transition_patch == patch);
+}
+
 static struct kobj_attribute enabled_kobj_attr = __ATTR_RW(enabled);
+static struct kobj_attribute transition_kobj_attr = __ATTR_RO(transition);
 static struct attribute *klp_patch_attrs[] = {
 	&enabled_kobj_attr.attr,
+	&transition_kobj_attr.attr,
 	NULL
 };
 
@@ -510,6 +562,7 @@ static int klp_init_func(struct klp_object *obj, struct klp_func *func)
 {
 	INIT_LIST_HEAD(&func->stack_node);
 	func->patched = false;
+	func->transition = false;
 
 	/* The format for the sysfs directory is <function,sympos> where sympos
 	 * is the nth occurrence of this symbol in kallsyms for the patched
@@ -738,7 +791,11 @@ int klp_module_coming(struct module *mod)
 				goto err;
 			}
 
-			if (!patch->enabled)
+			/*
+			 * Only patch the module if the patch is enabled or is
+			 * in transition.
+			 */
+			if (!patch->enabled && klp_transition_patch != patch)
 				break;
 
 			pr_notice("applying patch '%s' to loading module '%s'\n",
diff --git a/kernel/livepatch/patch.c b/kernel/livepatch/patch.c
index 92e9ee0..f0fa6b5 100644
--- a/kernel/livepatch/patch.c
+++ b/kernel/livepatch/patch.c
@@ -29,6 +29,7 @@
 #include <linux/bug.h>
 #include <linux/printk.h>
 #include "patch.h"
+#include "transition.h"
 
 static LIST_HEAD(klp_ops);
 
@@ -58,11 +59,42 @@ static void notrace klp_ftrace_handler(unsigned long ip,
 	ops = container_of(fops, struct klp_ops, fops);
 
 	rcu_read_lock();
+
 	func = list_first_or_null_rcu(&ops->func_stack, struct klp_func,
 				      stack_node);
-	if (WARN_ON_ONCE(!func))
+
+	if (!func)
 		goto unlock;
 
+	/*
+	 * See the comment for the 2nd smp_wmb() in klp_init_transition() for
+	 * an explanation of why this read barrier is needed.
+	 */
+	smp_rmb();
+
+	if (unlikely(func->transition)) {
+
+		/*
+		 * See the comment for the 1st smp_wmb() in
+		 * klp_init_transition() for an explanation of why this read
+		 * barrier is needed.
+		 */
+		smp_rmb();
+
+		if (current->klp_universe == KLP_UNIVERSE_OLD) {
+			/*
+			 * Use the previously patched version of the function.
+			 * If no previous patches exist, use the original
+			 * function.
+			 */
+			func = list_entry_rcu(func->stack_node.next,
+					      struct klp_func, stack_node);
+
+			if (&func->stack_node == &ops->func_stack)
+				goto unlock;
+		}
+	}
+
 	klp_arch_set_pc(regs, (unsigned long)func->new_func);
 unlock:
 	rcu_read_unlock();
@@ -183,3 +215,12 @@ int klp_patch_object(struct klp_object *obj)
 
 	return 0;
 }
+
+void klp_unpatch_objects(struct klp_patch *patch)
+{
+	struct klp_object *obj;
+
+	klp_for_each_object(patch, obj)
+		if (obj->patched)
+			klp_unpatch_object(obj);
+}
diff --git a/kernel/livepatch/patch.h b/kernel/livepatch/patch.h
index 2d0cce0..0db2271 100644
--- a/kernel/livepatch/patch.h
+++ b/kernel/livepatch/patch.h
@@ -28,5 +28,6 @@ struct klp_ops *klp_find_ops(unsigned long old_addr);
 
 int klp_patch_object(struct klp_object *obj);
 void klp_unpatch_object(struct klp_object *obj);
+void klp_unpatch_objects(struct klp_patch *patch);
 
 #endif /* _LIVEPATCH_PATCH_H */
diff --git a/kernel/livepatch/transition.c b/kernel/livepatch/transition.c
new file mode 100644
index 0000000..0609d84
--- /dev/null
+++ b/kernel/livepatch/transition.c
@@ -0,0 +1,406 @@
+/*
+ * transition.c - Kernel Live Patching transition functions
+ *
+ * Copyright (C) 2015 Josh Poimboeuf <jpoimboe@xxxxxxxxxx>
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, see <http://www.gnu.org/licenses/>.
+ */
+
+#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
+
+#include <linux/cpu.h>
+#include <linux/stacktrace.h>
+#include "../sched/sched.h"
+
+#include "patch.h"
+#include "transition.h"
+
+#define MAX_STACK_ENTRIES 100
+
+struct klp_patch *klp_transition_patch;
+
+int klp_universe_goal = KLP_UNIVERSE_UNDEFINED;
+
+/*
+ * The transition to the universe goal is complete.  Clean up the data
+ * structures.
+ */
+void klp_complete_transition(void)
+{
+	struct klp_object *obj;
+	struct klp_func *func;
+
+	if (klp_transition_patch->immediate)
+		goto done;
+
+	klp_for_each_object(klp_transition_patch, obj)
+		klp_for_each_func(obj, func)
+			func->transition = false;
+
+done:
+	klp_transition_patch = NULL;
+}
+
+/*
+ * Determine whether the given stack trace includes any references to a
+ * to-be-patched or to-be-unpatched function.
+ */
+static int klp_check_stack_func(struct klp_func *func,
+				struct stack_trace *trace)
+{
+	unsigned long func_addr, func_size, address;
+	struct klp_ops *ops;
+	int i;
+
+	if (func->immediate)
+		return 0;
+
+	for (i = 0; i < trace->max_entries; i++) {
+		address = trace->entries[i];
+
+		if (klp_universe_goal == KLP_UNIVERSE_OLD) {
+			 /*
+			  * Check for the to-be-unpatched function
+			  * (the func itself).
+			  */
+			func_addr = (unsigned long)func->new_func;
+			func_size = func->new_size;
+		} else {
+			/*
+			 * Check for the to-be-patched function
+			 * (the previous func).
+			 */
+			ops = klp_find_ops(func->old_addr);
+
+			if (list_is_singular(&ops->func_stack)) {
+				/* original function */
+				func_addr = func->old_addr;
+				func_size = func->old_size;
+			} else {
+				/* previously patched function */
+				struct klp_func *prev;
+
+				prev = list_next_entry(func, stack_node);
+				func_addr = (unsigned long)prev->new_func;
+				func_size = prev->new_size;
+			}
+		}
+
+		if (address >= func_addr && address < func_addr + func_size)
+			return -EAGAIN;
+	}
+
+	return 0;
+}
+
+/*
+ * Determine whether it's safe to transition the task to the new universe by
+ * looking for any to-be-patched or to-be-unpatched functions on its stack.
+ */
+static int klp_check_stack(struct task_struct *task)
+{
+	static unsigned long entries[MAX_STACK_ENTRIES];
+	struct stack_trace trace;
+	struct klp_object *obj;
+	struct klp_func *func;
+	int ret;
+
+	trace.skip = 0;
+	trace.nr_entries = 0;
+	trace.max_entries = MAX_STACK_ENTRIES;
+	trace.entries = entries;
+	ret = save_stack_trace_tsk_reliable(task, &trace);
+	WARN_ON_ONCE(ret == -ENOSYS);
+	if (ret) {
+		pr_debug("%s: pid %d (%s) has an unreliable stack\n",
+			 __func__, task->pid, task->comm);
+		return ret;
+	}
+
+	klp_for_each_object(klp_transition_patch, obj) {
+		if (!obj->patched)
+			continue;
+		klp_for_each_func(obj, func) {
+			ret = klp_check_stack_func(func, &trace);
+			if (ret) {
+				pr_debug("%s: pid %d (%s) is sleeping on function %s\n",
+					 __func__, task->pid, task->comm,
+					 func->old_name);
+				return ret;
+			}
+		}
+	}
+
+	return 0;
+}
+
+/*
+ * Try to safely switch a task to the universe goal.  If it's currently
+ * running, or it's sleeping on a to-be-patched or to-be-unpatched function, or
+ * if the stack is unreliable, return false.
+ */
+static bool klp_try_switch_task(struct task_struct *task)
+{
+	struct rq *rq;
+	unsigned long flags;
+	int ret;
+	bool success = false;
+
+	/* check if this task has already switched over */
+	if (task->klp_universe == klp_universe_goal)
+		return true;
+
+	/*
+	 * For arches which don't have reliable stack traces, we have to rely
+	 * on other methods (e.g., switching tasks at the syscall barrier).
+	 */
+	if (!IS_ENABLED(CONFIG_RELIABLE_STACKTRACE))
+		return false;
+
+	/*
+	 * Now try to check the stack for any to-be-patched or to-be-unpatched
+	 * functions.  If all goes well, switch the task to the goal universe.
+	 */
+	rq = task_rq_lock(task, &flags);
+
+	if (task_running(rq, task) && task != current) {
+		pr_debug("%s: pid %d (%s) is running\n", __func__, task->pid,
+			 task->comm);
+		goto done;
+	}
+
+	ret = klp_check_stack(task);
+	if (ret)
+		goto done;
+
+	klp_update_task_universe(task);
+
+	success = true;
+done:
+	task_rq_unlock(rq, task, &flags);
+	return success;
+}
+
+/*
+ * Try to switch all remaining tasks to the goal universe by walking the stacks
+ * of sleeping tasks and looking for any to-be-patched or to-be-unpatched
+ * functions.  If such functions are found, the task can't be switched yet.
+ *
+ * If any tasks are still stuck in the starting universe, schedule a retry.
+ */
+bool klp_try_complete_transition(void)
+{
+	unsigned int cpu;
+	struct task_struct *g, *task;
+	bool complete = true;
+
+	/*
+	 * If the patch can be applied or reverted immediately, skip the
+	 * per-task transitions.
+	 */
+	if (klp_transition_patch->immediate)
+		goto success;
+
+	/*
+	 * Try to switch the tasks to the goal universe by walking their stacks
+	 * and looking for any to-be-patched or to-be-unpatched functions.  If
+	 * such functions are found on a stack, or if the stack is deemed
+	 * unreliable, the task can't be switched yet.
+	 *
+	 * Usually this will transition most (or all) of the tasks on a system
+	 * unless the patch includes changes to a very common function.
+	 */
+	read_lock(&tasklist_lock);
+	for_each_process_thread(g, task)
+		if (!klp_try_switch_task(task))
+			complete = false;
+	read_unlock(&tasklist_lock);
+
+	/*
+	 * Ditto for the idle "swapper" tasks.
+	 */
+	get_online_cpus();
+	for_each_online_cpu(cpu)
+		if (!klp_try_switch_task(idle_task(cpu)))
+			complete = false;
+	put_online_cpus();
+
+	/*
+	 * Some tasks weren't able to be switched over.  Try again later and/or
+	 * wait for other methods like syscall barrier switching.
+	 */
+	if (!complete)
+		return false;
+
+success:
+	/*
+	 * When unpatching, all tasks have transitioned to the old universe so
+	 * we can now remove the new functions from the func_stack.
+	 */
+	if (klp_universe_goal == KLP_UNIVERSE_OLD) {
+		klp_unpatch_objects(klp_transition_patch);
+
+		/*
+		 * Don't allow any existing instances of ftrace handlers to
+		 * access any obsolete funcs before we reset the func
+		 * transition states to false.  Otherwise the handler may see
+		 * the deleted "new" func, see that it's not in transition, and
+		 * wrongly pick the new version of the function.
+		 */
+		synchronize_rcu();
+	}
+
+	pr_notice("'%s': %s complete\n", klp_transition_patch->mod->name,
+		  klp_universe_goal == KLP_UNIVERSE_NEW ? "patching" :
+							  "unpatching");
+
+	/* we're done, now cleanup the data structures */
+	klp_complete_transition();
+
+	return true;
+}
+
+/*
+ * Start the transition to the specified universe goal so tasks can begin
+ * switching to it.
+ */
+void klp_start_transition(int universe)
+{
+	if (WARN_ON(klp_universe_goal == universe))
+		return;
+
+	pr_notice("'%s': %s...\n", klp_transition_patch->mod->name,
+		  universe == KLP_UNIVERSE_NEW ? "patching" : "unpatching");
+
+	/*
+	 * Set the global universe goal which tasks will switch to.
+	 *
+	 * Note that any newly forked tasks after this call will be born in the
+	 * goal universe.  So the transition begins here, even before we start
+	 * switching tasks.
+	 */
+	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().
+	 */
+	smp_wmb();
+}
+
+/*
+ * This function can be called in the middle of an existing transition to
+ * reverse the direction of the universe goal.  This can be done to effectively
+ * cancel an existing enable or disable operation if there are any tasks which
+ * are stuck in the starting universe.
+ */
+void klp_reverse_transition(void)
+{
+	struct klp_patch *patch = klp_transition_patch;
+
+	klp_start_transition(!klp_universe_goal);
+	klp_try_complete_transition();
+
+	patch->enabled = !patch->enabled;
+}
+
+/*
+ * Set the global universe goal and all tasks to the starting universe, and
+ * initialize all function transition states to true in preparation for
+ * patching or unpatching.
+ */
+void klp_init_transition(struct klp_patch *patch, int universe)
+{
+	struct task_struct *g, *task;
+	unsigned int cpu;
+	struct klp_object *obj;
+	struct klp_func *func;
+
+	klp_transition_patch = patch;
+
+	/*
+	 * Initialize the universe goal to the starting universe.
+	 */
+	klp_universe_goal = universe;
+
+	/*
+	 * Ensure that if another CPU forks a task after the below task
+	 * universe writes and calls klp_update_task_universe(), it also sees
+	 * the above write to the universe goal.  Otherwise it may undo the
+	 * task universe writes below and mess up the task's starting universe.
+	 */
+	smp_wmb();
+
+	/*
+	 * If the patch can be applied or reverted immediately, skip the
+	 * per-task transitions.
+	 */
+	if (patch->immediate)
+		return;
+
+	/*
+	 * Initialize the task universes to their starting universe to prepare
+	 * them for switching to the goal universe.
+	 */
+	read_lock(&tasklist_lock);
+	for_each_process_thread(g, task)
+		klp_update_task_universe(task);
+	read_unlock(&tasklist_lock);
+
+	/*
+	 * Ditto for the idle "swapper" tasks.
+	 */
+	get_online_cpus();
+	for_each_online_cpu(cpu)
+		klp_update_task_universe(idle_task(cpu));
+	put_online_cpus();
+
+	/*
+	 * Ensure klp_ftrace_handler() sees the task->klp_universe updates
+	 * before the func->transition updates.  Otherwise it could read an
+	 * out-of-date task universe and pick the wrong function.
+	 */
+	smp_wmb();
+
+	/*
+	 * Set the func transition states so klp_ftrace_handler() will know to
+	 * switch to the transition logic.
+	 *
+	 * When patching, the funcs aren't yet in the func_stack and will be
+	 * made visible to the ftrace handler shortly by the calls to
+	 * klp_patch_object().
+	 *
+	 * When unpatching, the funcs are already in the func_stack and so are
+	 * already visible to the ftrace handler.
+	 */
+	klp_for_each_object(patch, obj)
+		klp_for_each_func(obj, func)
+			func->transition = true;
+
+	/*
+	 * For the enable path, ensure klp_ftrace_handler() will see the
+	 * func->transition updates before the funcs become visible to the
+	 * handler.  Otherwise the handler may wrongly pick the new func before
+	 * the task switches to the new universe.
+	 *
+	 * For the disable path, the funcs are already visible to the handler.
+	 * But we still need to ensure the ftrace handler will see the
+	 * func->transition updates before the tasks start switching to the old
+	 * universe.  Otherwise the handler can miss a universe change which
+	 * would result in it wrongly picking the new function.
+	 */
+	smp_wmb();
+}
diff --git a/kernel/livepatch/transition.h b/kernel/livepatch/transition.h
new file mode 100644
index 0000000..98e9930
--- /dev/null
+++ b/kernel/livepatch/transition.h
@@ -0,0 +1,20 @@
+#ifndef _LIVEPATCH_TRANSITION_H
+#define _LIVEPATCH_TRANSITION_H
+
+#include <linux/livepatch.h>
+
+enum {
+	KLP_UNIVERSE_UNDEFINED = -1,
+	KLP_UNIVERSE_OLD,
+	KLP_UNIVERSE_NEW,
+};
+
+extern struct klp_patch *klp_transition_patch;
+
+extern void klp_init_transition(struct klp_patch *patch, int universe);
+extern void klp_start_transition(int universe);
+extern void klp_reverse_transition(void);
+extern bool klp_try_complete_transition(void);
+extern void klp_complete_transition(void);
+
+#endif /* _LIVEPATCH_TRANSITION_H */
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index be1ef22..431007b 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -74,6 +74,7 @@
 #include <linux/context_tracking.h>
 #include <linux/compiler.h>
 #include <linux/frame.h>
+#include <linux/livepatch.h>
 
 #include <asm/switch_to.h>
 #include <asm/tlb.h>
@@ -5082,6 +5083,7 @@ void init_idle(struct task_struct *idle, int cpu)
 #ifdef CONFIG_SMP
 	sprintf(idle->comm, "%s/%d", INIT_TASK_COMM, cpu);
 #endif
+	klp_update_task_universe(idle);
 }
 
 int cpuset_cpumask_can_shrink(const struct cpumask *cur,
-- 
2.4.3