[PATCH v1 1/2] x86/sev: Use per-CPU PSC structure in prep for unaccepted memory support

[Tom Lendacky <thomas.lendacky@amd.com>]

In advance of providing support for unaccepted memory, switch from using
kmalloc() for allocating the Page State Change structure to using a
per-CPU structure. This is needed to avoid a possible recursive call into
set_pages_state() if the allocation requires (more) memory to be accepted,
which would result in a hang.

Protect the use of the per-CPU structure by disabling interrupts during
memory acceptance. Since the set_pages_state() path is the only path into
vmgexit_psc(), rename vmgexit_psc() to __vmgexit_psc() and remove the
calls to disable interrupts which are now performed by set_pages_state().

Even with interrupts disabled, an NMI can be raised while performing
memory acceptance. The NMI could then cause further memory acceptance to
performed. To prevent corruption to the per-CPU structure, use the PSC
MSR protocol in this situation.

Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>
---
 arch/x86/kernel/sev.c | 60 ++++++++++++++++++++++++++++---------------
 1 file changed, 39 insertions(+), 21 deletions(-)

diff --git a/arch/x86/kernel/sev.c b/arch/x86/kernel/sev.c
index c05f0124c410..1f7f6205c4f6 100644
--- a/arch/x86/kernel/sev.c
+++ b/arch/x86/kernel/sev.c
@@ -104,6 +104,15 @@ struct sev_es_runtime_data {
 	 * is currently unsupported in SEV-ES guests.
 	 */
 	unsigned long dr7;
+
+	/*
+	 * Page State Change structure for use when accepting memory or when
+	 * changing page state. Interrupts are disabled when using the structure
+	 * but an NMI could still be raised, so use a flag to indicate when the
+	 * structure is in use and use the MSR protocol in these cases.
+	 */
+	struct snp_psc_desc psc_desc;
+	bool psc_active;
 };
 
 struct ghcb_state {
@@ -660,7 +669,7 @@ static void pvalidate_pages(unsigned long vaddr, unsigned int npages, bool valid
 	}
 }
 
-static void __init early_set_pages_state(unsigned long paddr, unsigned int npages, enum psc_op op)
+static void early_set_pages_state(unsigned long paddr, unsigned int npages, enum psc_op op)
 {
 	unsigned long paddr_end;
 	u64 val;
@@ -742,26 +751,17 @@ void __init snp_prep_memory(unsigned long paddr, unsigned int sz, enum psc_op op
 		WARN(1, "invalid memory op %d\n", op);
 }
 
-static int vmgexit_psc(struct snp_psc_desc *desc)
+static int __vmgexit_psc(struct snp_psc_desc *desc)
 {
 	int cur_entry, end_entry, ret = 0;
 	struct snp_psc_desc *data;
 	struct ghcb_state state;
 	struct es_em_ctxt ctxt;
-	unsigned long flags;
 	struct ghcb *ghcb;
 
-	/*
-	 * __sev_get_ghcb() needs to run with IRQs disabled because it is using
-	 * a per-CPU GHCB.
-	 */
-	local_irq_save(flags);
-
 	ghcb = __sev_get_ghcb(&state);
-	if (!ghcb) {
-		ret = 1;
-		goto out_unlock;
-	}
+	if (!ghcb)
+		return 1;
 
 	/* Copy the input desc into GHCB shared buffer */
 	data = (struct snp_psc_desc *)ghcb->shared_buffer;
@@ -820,9 +820,6 @@ static int vmgexit_psc(struct snp_psc_desc *desc)
 out:
 	__sev_put_ghcb(&state);
 
-out_unlock:
-	local_irq_restore(flags);
-
 	return ret;
 }
 
@@ -861,18 +858,32 @@ static void __set_pages_state(struct snp_psc_desc *data, unsigned long vaddr,
 		i++;
 	}
 
-	if (vmgexit_psc(data))
+	if (__vmgexit_psc(data))
 		sev_es_terminate(SEV_TERM_SET_LINUX, GHCB_TERM_PSC);
 }
 
 static void set_pages_state(unsigned long vaddr, unsigned int npages, int op)
 {
 	unsigned long vaddr_end, next_vaddr;
+	struct sev_es_runtime_data *data;
 	struct snp_psc_desc *desc;
+	unsigned long flags;
 
-	desc = kmalloc(sizeof(*desc), GFP_KERNEL_ACCOUNT);
-	if (!desc)
-		panic("SNP: failed to allocate memory for PSC descriptor\n");
+	/* Disable interrupts since a per-CPU PSC and per-CPU GHCB are used. */
+	local_irq_save(flags);
+
+	data = this_cpu_read(runtime_data);
+	if (!data || data->psc_active) {
+		/* No per-CPU PSC or it is active, use the MSR protocol. */
+		early_set_pages_state(__pa(vaddr), npages, op);
+		goto out;
+	}
+
+	/* Mark the PSC in use. */
+	data->psc_active = true;
+	barrier();
+
+	desc = &data->psc_desc;
 
 	vaddr = vaddr & PAGE_MASK;
 	vaddr_end = vaddr + (npages << PAGE_SHIFT);
@@ -887,7 +898,12 @@ static void set_pages_state(unsigned long vaddr, unsigned int npages, int op)
 		vaddr = next_vaddr;
 	}
 
-	kfree(desc);
+	/* Mark the PSC no longer in use. */
+	barrier();
+	data->psc_active = false;
+
+out:
+	local_irq_restore(flags);
 }
 
 void snp_set_memory_shared(unsigned long vaddr, unsigned int npages)
@@ -1339,6 +1355,8 @@ static void __init alloc_runtime_data(int cpu)
 		panic("Can't allocate SEV-ES runtime data");
 
 	per_cpu(runtime_data, cpu) = data;
+
+	data->psc_active = false;
 }
 
 static void __init init_ghcb(int cpu)
-- 
2.36.1