[PATCH 12/19] nitro_enclaves: Put the CPU pool back when the new CPU list is refused

[Alexander Graf <graf@amazon.com>]

An admin who writes a CPU list that the driver refuses, for example one
that names CPU 0, leaves out a sibling or spans two NUMA nodes, loses
the pool that was there. The checks need every CPU of the list online
and a CPU the pool holds is offline, so they can only run after the old
pool is torn down.

Snapshot the CPUs of the pool before the teardown, and set the pool up
again from the snapshot when the new list is refused. Where that fails
as well, clear ne_cpus so that it does not name a pool that is gone.

Fixes: ff8a4d3e3a99 ("nitro_enclaves: Add logic for setting an enclave vCPU")
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 35 +++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index de3bf93a4ef0f..70d740d64ee17 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -382,6 +382,24 @@ static void ne_teardown_cpu_pool(void)
 	mutex_unlock(&ne_cpu_pool.mutex);
 }
 
+/**
+ * ne_snapshot_cpu_pool() - Copy out the CPUs the NE CPU pool holds.
+ * @cpu_pool:	CPU mask to fill in with the CPUs from the NE CPU pool.
+ *
+ * Context: Process context.
+ */
+static void ne_snapshot_cpu_pool(struct cpumask *cpu_pool)
+{
+	unsigned int i = 0;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
+		cpumask_or(cpu_pool, cpu_pool, ne_cpu_pool.avail_threads_per_core[i]);
+
+	mutex_unlock(&ne_cpu_pool.mutex);
+}
+
 /**
  * ne_set_kernel_param() - Set the NE CPU pool value via the NE kernel parameter.
  * @val:	NE CPU pool string value.
@@ -396,6 +414,7 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 {
 	char error_val[] = "";
 	cpumask_var_t cpu_pool;
+	cpumask_var_t old_pool;
 	struct ne_pci_dev *ne_pci_dev = NULL;
 	int rc = -EINVAL;
 
@@ -405,6 +424,12 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 	if (!zalloc_cpumask_var(&cpu_pool, GFP_KERNEL))
 		return -ENOMEM;
 
+	if (!zalloc_cpumask_var(&old_pool, GFP_KERNEL)) {
+		free_cpumask_var(cpu_pool);
+
+		return -ENOMEM;
+	}
+
 	rc = down_read_killable(&ne_devs.lock);
 	if (rc)
 		goto free_cpumasks;
@@ -443,13 +468,16 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 		goto unlock;
 	}
 
+	ne_snapshot_cpu_pool(old_pool);
+
 	ne_teardown_cpu_pool();
 
 	rc = ne_setup_cpu_pool(cpu_pool);
 	if (rc < 0) {
 		pr_err("%s: Error in setup CPU pool [rc=%d]\n", ne_misc_dev.name, rc);
 
-		param_set_copystring(error_val, kp);
+		if (cpumask_empty(old_pool) || ne_setup_cpu_pool(old_pool) < 0)
+			param_set_copystring(error_val, kp);
 
 		goto unlock;
 	}
@@ -460,7 +488,8 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 
 		ne_teardown_cpu_pool();
 
-		param_set_copystring(error_val, kp);
+		if (cpumask_empty(old_pool) || ne_setup_cpu_pool(old_pool) < 0)
+			param_set_copystring(error_val, kp);
 	}
 
 unlock:
@@ -471,6 +500,8 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 	up_read(&ne_devs.lock);
 
 free_cpumasks:
+	free_cpumask_var(old_pool);
+
 	free_cpumask_var(cpu_pool);
 
 	return rc;
-- 
2.47.1