Reprivilege CPUs on pkvm initialization failure

arch/x86/include/asm/kvm_pkvm.h

@@ -224,6 +224,7 @@ extern unsigned int pkvm_sym(nr_cpu_ids);
 DECLARE_STATIC_KEY_FALSE(pkvm_sym(switch_vcpu_ibpb));
 extern struct fpu_state_config pkvm_sym(fpu_kernel_cfg);
 extern struct fpu_state_config pkvm_sym(fpu_user_cfg);
+extern struct pkvm_init_ops *pkvm_sym(init_ops);
 
 u64 pkvm_total_reserve_pages(void);
 PKVM_DECLARE(void *, pkvm_early_alloc_page, (void));

arch/x86/include/asm/pkvm_hypercalls.h

@@ -8,7 +8,9 @@ BUILD_BUG_ON(1)
 #endif
 
 /* Hypercalls used only during pKVM initialization */
+PKVM_HC(init)
 PKVM_HC(init_finalize)
+PKVM_HC(reprivilege_cpu)
 
 /* pKVM vmexit tracing/profiling */
 PKVM_HC(enable_vmexit_trace)

arch/x86/kvm/pkvm/Makefile

@@ -13,7 +13,7 @@ asflags-y += -include $(srctree)/arch/x86/kvm/pkvm/undef.h
 cmd_deps			:= $(src)/compiling_cmds.sh
 
 pkvm-hyp-y			:= early_alloc.o pkvm.o memory.o cpu.o \
-				   idt.o entry.o init_finalize.o pgtable.o \
+				   idt.o entry.o init.o pgtable.o \
 				   mmu.o page_alloc.o lapic.o trace.o \
 				   fpu.o
 
@@ -39,7 +39,7 @@ pkvm-hyp-y			+= $(kernel-lib)/sort.o $(kernel-lib)/bsearch.o \
 kvm				:= ..
 pkvm-hyp-y			+= $(kvm)/x86.o $(kvm)/cpuid.o
 
-pkvm-hyp-$(CONFIG_PKVM_INTEL)	+= vmx/host_vmentry.o vmx/host_vmx.o \
+pkvm-hyp-$(CONFIG_PKVM_INTEL)	+= vmx/host_vmentry.o vmx/host_vmx.o vmx/host_repriv.o \
 				   $(kvm)/vmx/vmx.o vmx/idt.o vmx/ept.o \
 				   $(kvm)/vmx/main.o
 

arch/x86/kvm/pkvm/init_finalize.c → arch/x86/kvm/pkvm/init.c

@@ -3,7 +3,7 @@
 #include <asm/kvm_pkvm.h>
 #include "early_alloc.h"
 #include "fpu.h"
-#include "init_finalize.h"
+#include "init.h"
 #include "lapic.h"
 #include "memory.h"
 #include "mmu.h"
@@ -13,7 +13,7 @@
 static void *hyp_pgt_base;
 static void *host_pgt_base;
 static void *pkvm_vmemmap_base;
-static DEFINE_PER_CPU(bool, cpu_finalized);
+static DEFINE_PER_CPU(bool, cpu_initialized);
 
 static int divide_memory_pool(phys_addr_t phys, unsigned long size)
 {
@@ -192,12 +192,14 @@ static int create_host_mmu(const struct pkvm_mem_info infos[], int nr_infos,
 	return 0;
 }
 
+/* Set by the host before deprivilege and used through the initialization process. */
+struct pkvm_init_ops *init_ops;
+
 #define TMP_NR_INFOS	16
-static int finalize_global(struct pkvm_mem_info infos[], int nr_infos,
-			   struct pkvm_init_ops *init_ops)
+static int initialize_global(struct pkvm_mem_info infos[], int nr_infos)
 {
 	host_mmu_init_fn_t host_mmu_init_fn = init_ops ? init_ops->host_mmu_init : NULL;
-	hyp_g_finalize_fn_t hyp_g_finalize = init_ops ? init_ops->hyp_g_finalize : NULL;
+	hyp_global_init_fn_t hyp_global_init = init_ops ? init_ops->hyp_global_init : NULL;
 	struct pkvm_mem_info tmp_infos[TMP_NR_INFOS];
 	phys_addr_t mem_base = INVALID_PAGE;
 	unsigned long mem_size = 0;
@@ -238,33 +240,32 @@ static int finalize_global(struct pkvm_mem_info infos[], int nr_infos,
 	if (ret)
 		return ret;
 
-	return hyp_g_finalize ? hyp_g_finalize() : 0;
+	return hyp_global_init ? hyp_global_init() : 0;
 }
 
-int pkvm_init_finalize(struct pkvm_mem_info infos[], int nr_infos,
-		       struct pkvm_init_ops *init_ops)
+int pkvm_init(struct pkvm_mem_info infos[], int nr_infos)
 {
 	hyp_mmu_finalize_fn_t hyp_mmu_finalize_fn = init_ops ? init_ops->hyp_mmu_finalize :
 							       NULL;
 	host_mmu_finalize_fn_t host_mmu_finalize_fn = init_ops ? init_ops->host_mmu_finalize :
 								 NULL;
-	static bool global_finalized;
+	static bool global_initialized;
 	int ret;
 
-	if (this_cpu_read(cpu_finalized))
+	if (this_cpu_read(cpu_initialized))
 		return -EBUSY;
 
-	if (!global_finalized) {
-		ret = finalize_global(infos, nr_infos, init_ops);
+	if (!global_initialized) {
+		ret = initialize_global(infos, nr_infos);
 		if (ret)
 			return ret;
 
-		global_finalized = true;
+		global_initialized = true;
 	} else {
 		/*
 		 * The pKVM hypervisor's MMU was already loaded on the first
-		 * finalized CPU during the global finalize. Need to load it
-		 * on all the other CPUs as well.
+		 * initialized CPU during the global initialize. Need to load
+		 * it on all the other CPUs as well.
 		 */
 		pkvm_hyp_mmu_load();
 	}
@@ -285,6 +286,49 @@ int pkvm_init_finalize(struct pkvm_mem_info infos[], int nr_infos,
 
 	pkvm_vcpu_perf_init(this_cpu_read(host_vcpu));
 
-	this_cpu_write(cpu_finalized, true);
+	this_cpu_write(cpu_initialized, true);
+	return 0;
+}
+
+/*
+ * Flag indicating if pkvm is initialized successfully.
+ * Used to enforce internal hypercalls to be unavailable
+ * for general use once pkvm is initialized.
+ */
+static bool pkvm_initialized __ro_after_init;
+
+int pkvm_reprivilege_vcpu(struct kvm_vcpu *vcpu)
+{
+	if (READ_ONCE(pkvm_initialized))
+		return -EPERM;
+
+	if (!init_ops || !init_ops->reprivilege_cpu)
+		return -EOPNOTSUPP;
+
+	init_ops->reprivilege_cpu(vcpu->arch.regs);
+
+	/* Reach here only if reprivilege operation fails. */
+	return -EFAULT;
+}
+
+int pkvm_init_finalize(void)
+{
+	int cpu;
+
+	if (READ_ONCE(pkvm_initialized))
+		return -EPERM;
+
+	for_each_possible_cpu(cpu) {
+		if (!per_cpu(cpu_initialized, cpu))
+			return -EAGAIN;
+	}
+	WRITE_ONCE(pkvm_initialized, true);
+	if (init_ops)
+		init_ops->reprivilege_cpu = NULL;
+	/*
+	 * TODO: Move reprivilege logic to a separate
+	 * section and zero it out here.
+	 */
+
 	return 0;
 }

arch/x86/kvm/pkvm/init.h

@@ -0,0 +1,38 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __PKVM_X86_INIT_H
+#define __PKVM_X86_INIT_H
+
+#include <asm/kvm_pkvm.h>
+#include "pgtable.h"
+
+typedef int (*hyp_mmu_finalize_fn_t)(struct pkvm_pgtable *pgt);
+typedef int (*host_mmu_init_fn_t)(struct pkvm_pgtable *pgt, void *pool_base,
+				  unsigned long pool_pages);
+typedef int (*host_mmu_finalize_fn_t)(struct pkvm_pgtable *pgt);
+typedef int (*hyp_global_init_fn_t)(void);
+typedef void (*reprivilege_cpu_fn_t)(unsigned long *vcpu_regs);
+
+/**
+ * pkvm_init_ops - The platform vendor specific pKVM init operations used by the
+ *		   pkvm_init. Some operation could be NULL if it is not necessary.
+ *
+ * @hyp_mmu_finalize:	Finalize the hypervisor mmu.
+ * @host_mmu_init:	Initialize the host mmu.
+ * @host_mmu_finalize:	Finalize the host mmu.
+ * @hyp_global_init:	Initialize the hypervisor globally.
+ * @reprivilege_cpu:	Switch the cpu back to root mode. Called if deprivilege
+ *                  	or pKVM initialization fails.
+ */
+struct pkvm_init_ops {
+	hyp_mmu_finalize_fn_t		hyp_mmu_finalize;
+	host_mmu_init_fn_t		host_mmu_init;
+	host_mmu_finalize_fn_t		host_mmu_finalize;
+	hyp_global_init_fn_t		hyp_global_init;
+	reprivilege_cpu_fn_t		reprivilege_cpu;
+};
+
+int pkvm_init(struct pkvm_mem_info infos[], int nr_info);
+int pkvm_init_finalize(void);
+int pkvm_reprivilege_vcpu(struct kvm_vcpu *vcpu);
+
+#endif /* __PKVM_X86_INIT_H */

arch/x86/kvm/pkvm/lapic.c

@@ -61,9 +61,9 @@ void pkvm_lapic_send_init(int cpu)
 	/*
 	 * Pairs with the smp_store_release() in the setup_lapic().
 	 * If remote lapic is not ready, it means the remote CPU is not
-	 * finalized yet. In this case, it is not necessary to send INIT to kick
-	 * as this remote CPU will handle all the pending requests before being
-	 * finalized.
+	 * initialized yet(by pkvm_init()). In this case, it is not
+	 * necessary to send INIT to kick as this remote CPU will handle
+	 * all the pending requests before being initialized.
 	 */
 	if (unlikely(!smp_load_acquire(&remote->ready)))
 		return;

arch/x86/kvm/pkvm/mmu.h

@@ -3,7 +3,7 @@
 #define __PKVM_X86_MMU_H
 
 #include <asm/pkvm_spinlock.h>
-#include "init_finalize.h"
+#include "init.h"
 
 extern pkvm_spinlock_t host_mmu_lock;
 

arch/x86/kvm/pkvm/pkvm.c

@@ -5,7 +5,7 @@
 #include <asm/pkvm_spinlock.h>
 #include "debug.h"
 #include "fpu.h"
-#include "init_finalize.h"
+#include "init.h"
 #include "lapic.h"
 #include "mem_protect.h"
 #include "memory.h"
@@ -480,10 +480,14 @@ void pkvm_handle_host_hypercall(struct kvm_vcpu *vcpu)
 	int ret = 0;
 
 	switch (hc) {
+	case __pkvm__init:
+		ret = pkvm_init((struct pkvm_mem_info *)pkvm_hc_input1(vcpu), pkvm_hc_input2(vcpu));
+		break;
 	case __pkvm__init_finalize:
-		ret = pkvm_init_finalize((struct pkvm_mem_info *)pkvm_hc_input1(vcpu),
-					 pkvm_hc_input2(vcpu),
-					 (struct pkvm_init_ops *)pkvm_hc_input3(vcpu));
+		ret = pkvm_init_finalize();
+		break;
+	case __pkvm__reprivilege_cpu:
+		ret = pkvm_reprivilege_vcpu(vcpu);
 		break;
 	case __pkvm__enable_vmexit_trace:
 		pkvm_enable_vmexit_trace(pkvm_hc_input1(vcpu));

arch/x86/kvm/pkvm/vmx/ept.c

@@ -281,7 +281,7 @@ int pkvm_host_ept_finalize(struct pkvm_pgtable *pgt)
 	ept_sync_global();
 	/*
 	 * Clear the pending TLB flush request left after updating host EPT
-	 * mappings in finalize_global(), as EPT has just been flushed with
+	 * mappings in initialize_global(), as EPT has just been flushed with
 	 * global context anyway.
 	 */
 	kvm_clear_request(KVM_REQ_TLB_FLUSH_CURRENT, hvcpu);