[PATCH 15/19] nitro_enclaves: Set up the CPU pool under the CPU hotplug read lock

[Alexander Graf <graf@amazon.com>]

An admin who writes a CPU list to ne_cpus while a CPU goes offline can
get a pool built from the topology before and after that CPU went.
ne_setup_cpu_pool() reads cpu_is_offline(), cpu_to_node() and
topology_sibling_cpumask() for the list without the CPU hotplug lock.

Take cpus_read_lock() over those reads, and drop it before the CPUs are
offlined, because remove_cpu() takes the write side. The failure paths
share one unwind in ne_free_core_slots(), so that one cpus_read_unlock()
covers all of them.

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 | 73 +++++++++++++++--------
 1 file changed, 48 insertions(+), 25 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index a897036dee2dd..f8b066c990ce1 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -147,6 +147,27 @@ struct ne_phys_contig_mem_regions {
 	struct range  *regions;
 };
 
+/**
+ * ne_free_core_slots() - Free the core slot arrays of the NE CPU pool.
+ * @void:	No parameters provided.
+ *
+ * Context: Process context. This function is called with the ne_cpu_pool mutex
+ *	    held.
+ */
+static void ne_free_core_slots(void)
+{
+	unsigned int i = 0;
+
+	if (!ne_cpu_pool.avail_threads_per_core)
+		return;
+
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
+		free_cpumask_var(ne_cpu_pool.avail_threads_per_core[i]);
+
+	kfree(ne_cpu_pool.avail_threads_per_core);
+	ne_cpu_pool.avail_threads_per_core = NULL;
+}
+
 /**
  * ne_setup_cpu_pool() - Set the NE CPU pool after handling sanity checks such
  *			 as not sharing CPU cores with the primary / parent VM
@@ -155,6 +176,11 @@ struct ne_phys_contig_mem_regions {
  *			 pool after the checks passed.
  * @cpu_pool:	The CPU mask used for setting NE CPU pool.
  *
+ * The topology of the pool - online state, sibling masks, core ids - is read
+ * under the CPU hotplug read lock. The lock is dropped before the pool CPUs are
+ * offlined, because remove_cpu() takes the write side of the same percpu-rwsem.
+ * Nothing reads topology after the drop.
+ *
  * Context: Process context. This function is called with the ne_pci_dev enclave
  *	    mutex held, if there is a NE PCI device.
  * Return:
@@ -171,6 +197,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 	int rc = -EINVAL;
 
 	mutex_lock(&ne_cpu_pool.mutex);
+	cpus_read_lock();
 
 	cpu = cpumask_any(cpu_pool);
 	if (cpu >= nr_cpu_ids) {
@@ -178,7 +205,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 		rc = -EINVAL;
 
-		goto reset_cpu_pool;
+		goto unlock_hotplug;
 	}
 
 	/*
@@ -192,7 +219,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 			rc = -EINVAL;
 
-			goto reset_cpu_pool;
+			goto unlock_hotplug;
 		}
 
 	/*
@@ -207,7 +234,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 				rc = -EINVAL;
 
-				goto reset_cpu_pool;
+				goto unlock_hotplug;
 			}
 		} else {
 			if (numa_node != cpu_to_node(cpu)) {
@@ -216,7 +243,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 				rc = -EINVAL;
 
-				goto reset_cpu_pool;
+				goto unlock_hotplug;
 			}
 		}
 
@@ -229,7 +256,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 		rc = -EINVAL;
 
-		goto reset_cpu_pool;
+		goto unlock_hotplug;
 	}
 
 	for_each_cpu(cpu_sibling, topology_sibling_cpumask(0)) {
@@ -239,7 +266,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 			rc = -EINVAL;
 
-			goto reset_cpu_pool;
+			goto unlock_hotplug;
 		}
 	}
 
@@ -256,7 +283,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 				rc = -EINVAL;
 
-				goto reset_cpu_pool;
+				goto unlock_hotplug;
 			}
 		}
 	}
@@ -273,14 +300,14 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 	if (!ne_cpu_pool.avail_threads_per_core) {
 		rc = -ENOMEM;
 
-		goto reset_cpu_pool;
+		goto unlock_hotplug;
 	}
 
 	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
 		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL)) {
 			rc = -ENOMEM;
 
-			goto free_cores_cpumask;
+			goto free_slots;
 		}
 
 	/*
@@ -295,12 +322,14 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 			rc = -EINVAL;
 
-			goto clear_cpumask;
+			goto free_slots;
 		}
 
 		cpumask_set_cpu(cpu, ne_cpu_pool.avail_threads_per_core[core_id]);
 	}
 
+	cpus_read_unlock();
+
 	/*
 	 * CPUs that are given to enclave(s) should not be considered online
 	 * by Linux anymore, as the hypervisor will degrade them to floating.
@@ -333,14 +362,13 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 online_cpus:
 	for_each_cpu(cpu, cpu_pool)
 		add_cpu(cpu);
-clear_cpumask:
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		cpumask_clear(ne_cpu_pool.avail_threads_per_core[i]);
-free_cores_cpumask:
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		free_cpumask_var(ne_cpu_pool.avail_threads_per_core[i]);
-	kfree(ne_cpu_pool.avail_threads_per_core);
-reset_cpu_pool:
+	ne_free_core_slots();
+	goto reset_pool;
+free_slots:
+	ne_free_core_slots();
+unlock_hotplug:
+	cpus_read_unlock();
+reset_pool:
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = -1;
@@ -371,7 +399,7 @@ static void ne_teardown_cpu_pool(void)
 		return;
 	}
 
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++) {
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
 		for_each_cpu(cpu, ne_cpu_pool.avail_threads_per_core[i]) {
 			rc = add_cpu(cpu);
 			if (rc != 0)
@@ -379,12 +407,7 @@ static void ne_teardown_cpu_pool(void)
 				       ne_misc_dev.name, cpu, rc);
 		}
 
-		cpumask_clear(ne_cpu_pool.avail_threads_per_core[i]);
-
-		free_cpumask_var(ne_cpu_pool.avail_threads_per_core[i]);
-	}
-
-	kfree(ne_cpu_pool.avail_threads_per_core);
+	ne_free_core_slots();
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = -1;
-- 
2.47.1