[PATCH 00/40] nitro_enclaves: Add support for Nitro Sandboxes

[PATCH 00/40] nitro_enclaves: Add support for Nitro Sandboxes

[Alexander Graf]

Nitro Sandboxes [1] are a new EC2 feature that build on top of the
technical foundations of Nitro Enclaves - an EC2 feature that allows
you to spawn a virtio-vsock-only VM sibling constructed out of resources
donated by the VM you execute in.

Nitro Sandboxes build on that concept and the same hypervisor interface,
but add the ability to implement virtio devices in the parent, run vCPUs
oversubscribed, add device hotplug support and allow to spawn a sandbox
without attestation, but shared memory. They also boot a full EC2-like VM
that runs with a UEFI stack and executes the EIF-wrapped kernel as UEFI
executable.

In combination, Nitro Sandboxes provide a way to create a confidential
compute VM for almost all workloads, but without the biggest gaps that
other confidential compute offerings have: The parent has no timing
control and the parent also has no control over the base device model,
such as timers, PCIe, UART, etc. All base devices are implemented in
Nitro.

This series adds parent-side support for Nitro Sandboxes to the
nitro_enclaves driver.

This supports adds a few non-obvious technical foundations:

guest_memfd like CMA allocator - With Nitro Enclaves, memory came from
hugetlbfs and user space maintained visibility into it. That led to a
few problems. With this new enablement, we add a guest_memfd like
allocator that seeds from CMA, allows the parent to seed the EIF before
hand-over and then loses any mechanism to access donated memory. Thanks
to CMA, it also provides much larger contiguous chunks than 2M huge pages.

Virtio notifications - To implement virtio in the parent, we need to
handle notifications in both directions. Parent to sandbox goes through
a BAR write; sandbox to parent through multiplexed IRQs. We provide
eventfd support in both directions, which also enables vhost support.

NUMA - With Nitro Sandboxes, NUMA became a first class citizen and all
APIs now support NUMA targeting.

Memory Donations - Every sandbox requires hypervisor memory. To allow
for a larger scale of sandboxes you can run, we add a way to donate
memory from the parent to the hypervisor. The user visible effect of
that is that there is a slot limit sysfs file which you can increase.
On write, we allocate physically contiguous CMA memory and donate it.

This series depends on "nitro_enclaves: Fix enclave release, command
retries and the CPU pool", which I posted earlier today.

I used Kiro, an AI coding assistant running Claude Opus and GPT models,
to draft the initial code and patch descritptions.

I validated all new functionality on x86_64 and arm64 as well as ensured
that Nitro Enclaves are still fully functional after this patch series.

[1] https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/amazon-ec2-nitro-sandboxes.html

Alex

Alexander Graf (40):
  nitro_enclaves: Allow the CPU pool to span NUMA nodes
  nitro_enclaves: Add NE_ADD_ANY_VCPU
  nitro_enclaves: Expose CPU pool state under sysfs
  nitro_enclaves: Document multi-NUMA CPU pools
  dma-contiguous: add dma_get_numa_cma_area() accessor
  nitro_enclaves: Allocate physically contiguous enclave memory
  nitro_enclaves: Add NE_MEMORY_SHARED
  nitro_enclaves: Add NE_ADD_MEMORY_REGION
  nitro_enclaves: Document the memory pool
  nitro_enclaves: Add NE_CREATE_VM2 with a launch flag word
  nitro_enclaves: Add NE_ENCLAVE_PCIE_MODE launch flag
  nitro_enclaves: Name the device's error codes in the UAPI header
  nitro_enclaves: Log a reached device limit at info level
  nitro_enclaves: Add NE_SET_INFO_PAGE
  nitro_enclaves: Document the info page in the memory pool
  nitro_enclaves: Turn the command reply into a per-command union
  nitro_enclaves: Add the slot pool donation commands
  nitro_enclaves: Donate memory for the slot pool
  nitro_enclaves: Add the slot pool growth commands
  nitro_enclaves: Add a helper to grow the slot pool
  nitro_enclaves: Document the slot pool
  nitro_enclaves: Expose slot pool state under sysfs
  nitro_enclaves: Add the device add command
  nitro_enclaves: Add NE_ADD_DEVICE
  nitro_enclaves: Read the interrupt multiplexer geometry at probe
  nitro_enclaves: Donate a bitmap page per multiplexed vector
  nitro_enclaves: Dispatch multiplexed interrupts to per-queue bindings
  nitro_enclaves: Add NE_SET_NOTIFY_EVENTFD
  nitro_enclaves: Report the device's info notification
  nitro_enclaves: Count multiplexed interrupts under debugfs
  nitro_enclaves: Map the doorbell window at probe
  nitro_enclaves: Add NE_SET_VRING_CALL
  nitro_enclaves: Assert the layout of the ioctl argument structs
  nitro_enclaves: Pass on the Nitro Sandbox launch flags
  nitro_enclaves: Add NE_ENCLAVE_DIRECT_MODE
  nitro_enclaves: Add NE_ENCLAVE_PCIE_HOTPLUG_MODE
  nitro_enclaves: Add NE_ENCLAVE_CPU_OVERCOMMIT_MODE
  nitro_enclaves: Add NE_GET_CAPS
  nitro_enclaves: Add NE_UPDATE_DEVICE_CONFIG
  nitro_enclaves: Document the driver side of Nitro Sandboxes

 .../sysfs-devices-virtual-misc-nitro_enclaves |  232 +
 Documentation/virt/ne_overview.rst            |  256 +-
 MAINTAINERS                                   |    1 +
 drivers/virt/nitro_enclaves/Kconfig           |   19 +
 drivers/virt/nitro_enclaves/ne_misc_dev.c     | 4832 +++++++++++++++--
 drivers/virt/nitro_enclaves/ne_misc_dev.h     |  206 +-
 drivers/virt/nitro_enclaves/ne_pci_dev.c      | 1568 +++++-
 drivers/virt/nitro_enclaves/ne_pci_dev.h      |  599 +-
 include/linux/dma-map-ops.h                   |    7 +
 include/uapi/linux/nitro_enclaves.h           | 1182 +++-
 kernel/dma/contiguous.c                       |   36 +
 samples/nitro_enclaves/ne_ioctl_sample.c      |   10 +-
 12 files changed, 8282 insertions(+), 666 deletions(-)
 create mode 100644 Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves


base-commit: 6628e7d4c6b25c8d18f1345e3a20d21829aae7e4
prerequisite-patch-id: 542b3d78fda912095b75a49becd1b93109befdbc
prerequisite-patch-id: 318368e966c1925652b7976948f435d42b87c72f
prerequisite-patch-id: eb35b5d713e7f50f6d1d02310c9afe10b1628fc6
prerequisite-patch-id: 223bf6e60620ca8fe72f3ff6f3bd3f5839cdd6d5
prerequisite-patch-id: e5305ce0515a02f7f88e72adf02c6a2bee2fd2e0
prerequisite-patch-id: 759b5920f842d70e675770608d9088639c10559e
prerequisite-patch-id: e1837df37663da3d3ce53afa4594828b2d602152
prerequisite-patch-id: f0d4302bd36f94de95951a0db1ae9eb9747076aa
prerequisite-patch-id: 41159b92e250fef705be52f714b6e2761a2cd1f7
prerequisite-patch-id: a0c3a56907acac2297c800209faa2849a8df0e48
prerequisite-patch-id: a80a95bd682b540be5760e8fdf7044379fe32d61
prerequisite-patch-id: 9c01e0fecd6810f3448582715c87cd3dd7a6cdf1
prerequisite-patch-id: b94dfc8bb1039740d88a30535445bb7b3dd1fe74
prerequisite-patch-id: bd76f15728dc0ee2ca3417cd95a6b23a7d63e91f
prerequisite-patch-id: 0c60f1574b590f68885a25760017e9ff87752f66
prerequisite-patch-id: 63dff43da52934a6e6f5e65c3e235a7114b6b4dc
prerequisite-patch-id: af52cf938525874cadd6d879aa76495e563876d2
prerequisite-patch-id: f53ee22764b5f13708cacd3e48a59a82d71a270c
prerequisite-patch-id: 258b8c18168db49aacbda5156e33978b539666ab
-- 
2.47.1

[PATCH 01/40] nitro_enclaves: Allow the CPU pool to span NUMA nodes

[Alexander Graf]

With Nitro Sandboxes, we allow users to span a sandbox across multiple
NUMA nodes and place their sandbox on a NUMA node different from 0.
Allow the ne_cpus= CPU pool to span NUMA nodes.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 Documentation/virt/ne_overview.rst        |   3 +-
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 159 +++++++++++++++-------
 drivers/virt/nitro_enclaves/ne_misc_dev.h |   5 +-
 include/uapi/linux/nitro_enclaves.h       |  22 ++-
 samples/nitro_enclaves/ne_ioctl_sample.c  |  10 +-
 5 files changed, 137 insertions(+), 62 deletions(-)

diff --git a/Documentation/virt/ne_overview.rst b/Documentation/virt/ne_overview.rst
index 74c2f5919c886..aecd3704acaf5 100644
--- a/Documentation/virt/ne_overview.rst
+++ b/Documentation/virt/ne_overview.rst
@@ -47,7 +47,8 @@ The memory regions carved out of the primary VM and given to an enclave need to
 be aligned 2 MiB / 1 GiB physically contiguous memory regions (or multiple of
 this size e.g. 8 MiB). The memory can be allocated e.g. by using hugetlbfs from
 user space [2][3][7]. The memory size for an enclave needs to be at least
-64 MiB. The enclave memory and CPUs need to be from the same NUMA node.
+64 MiB. The enclave memory and CPUs need to be from the same NUMA node,
+unless the CPU pool spans several of them.
 
 An enclave runs on dedicated cores. CPU 0 and its CPU siblings need to remain
 available for the primary VM. A CPU pool has to be set for NE purposes by an
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index cb1efc95a7f34..6e36abeff01a2 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -24,6 +24,7 @@
 #include <linux/module.h>
 #include <linux/mutex.h>
 #include <linux/nitro_enclaves.h>
+#include <linux/numa.h>
 #include <linux/pci.h>
 #include <linux/poll.h>
 #include <linux/range.h>
@@ -36,8 +37,8 @@
 
 /**
  * NE_CPUS_SIZE - Size for max 128 CPUs, for now, in a cpu-list string, comma
- *		  separated. The NE CPU pool includes CPUs from a single NUMA
- *		  node.
+ *		  separated. The NE CPU pool may include CPUs from more than
+ *		  one NUMA node.
  */
 #define NE_CPUS_SIZE		(512)
 
@@ -126,13 +127,27 @@ static const struct kernel_param *ne_cpus_kp;
  *				from the pool rather than derived from a thread
  *				count, so a machine whose cores do not all have
  *				the same number of threads is sized correctly.
- * @numa_node:			NUMA node of the CPUs in the pool.
+ * @numa_node:			NUMA node of the CPUs in the pool, or
+ *				NUMA_NO_NODE if the pool spans multiple nodes or
+ *				no pool is set up.
+ * @core_nid:			The NUMA node of each core slot, indexed the
+ *				same way as @avail_threads_per_core. Every slot
+ *				holds a core, so every entry names a node. A slot
+ *				keeps its node after its core is handed to an
+ *				enclave, because ne_get_cpu_from_cpu_pool()
+ *				still reads it when it reuses a sibling of that
+ *				core. Read while
+ *				the pool CPUs are still online, before
+ *				ne_setup_cpu_pool() offlines them, because
+ *				cpu_to_node() is not documented to answer for an
+ *				offline CPU.
  */
 struct ne_cpu_pool {
 	cpumask_var_t	*avail_threads_per_core;
 	struct mutex	mutex;
 	unsigned int	nr_parent_vm_cores;
 	int		numa_node;
+	int		*core_nid;
 };
 
 static struct ne_cpu_pool ne_cpu_pool = {
@@ -160,14 +175,16 @@ static void ne_free_core_slots(void)
 {
 	unsigned int i = 0;
 
-	if (!ne_cpu_pool.avail_threads_per_core)
-		return;
+	if (ne_cpu_pool.avail_threads_per_core) {
+		for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
+			free_cpumask_var(ne_cpu_pool.avail_threads_per_core[i]);
 
-	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;
+	}
 
-	kfree(ne_cpu_pool.avail_threads_per_core);
-	ne_cpu_pool.avail_threads_per_core = NULL;
+	kfree(ne_cpu_pool.core_nid);
+	ne_cpu_pool.core_nid = NULL;
 }
 
 /**
@@ -218,6 +235,10 @@ static int ne_build_core_slots(const struct cpumask *cpu_pool)
 
 	ne_cpu_pool.nr_parent_vm_cores = nr_cores;
 
+	ne_cpu_pool.core_nid = kmalloc_objs(*ne_cpu_pool.core_nid, nr_cores);
+	if (!ne_cpu_pool.core_nid)
+		goto free_slots;
+
 	for (i = 0; i < nr_cores; i++)
 		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL))
 			goto free_slots;
@@ -237,6 +258,8 @@ static int ne_build_core_slots(const struct cpumask *cpu_pool)
 			cpumask_set_cpu(cpu_sibling, processed);
 		}
 
+		ne_cpu_pool.core_nid[next_core_idx] = cpu_to_node(cpu);
+
 		next_core_idx++;
 	}
 
@@ -275,7 +298,9 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 {
 	unsigned int cpu = 0;
 	unsigned int cpu_sibling = 0;
-	int numa_node = -1;
+	bool have_node = false;
+	bool multi_node = false;
+	int numa_node = NUMA_NO_NODE;
 	int rc = -EINVAL;
 
 	mutex_lock(&ne_cpu_pool.mutex);
@@ -305,29 +330,34 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 		}
 
 	/*
-	 * Check if the CPUs from the NE CPU pool are from the same NUMA node.
+	 * Determine the NUMA node of the CPUs in the pool. If the pool spans
+	 * multiple NUMA nodes, set numa_node to NUMA_NO_NODE; the pool is then
+	 * treated as node-agnostic for placement and memory-locality checks.
+	 * A CPU that reports no node at all is refused, because otherwise such
+	 * a CPU would join a single-node pool and turn it into a spanning one.
 	 */
-	for_each_cpu(cpu, cpu_pool)
-		if (numa_node < 0) {
-			numa_node = cpu_to_node(cpu);
-			if (numa_node < 0) {
-				pr_err("%s: Invalid NUMA node %d\n",
-				       ne_misc_dev.name, numa_node);
+	for_each_cpu(cpu, cpu_pool) {
+		int nid = cpu_to_node(cpu);
 
-				rc = -EINVAL;
+		if (nid < 0) {
+			pr_err("%s: CPU %d reports no NUMA node\n",
+			       ne_misc_dev.name, cpu);
 
-				goto unlock_hotplug;
-			}
-		} else {
-			if (numa_node != cpu_to_node(cpu)) {
-				pr_err("%s: CPUs with different NUMA nodes\n",
-				       ne_misc_dev.name);
+			rc = -EINVAL;
 
-				rc = -EINVAL;
+			goto unlock_hotplug;
+		}
 
-				goto unlock_hotplug;
-			}
+		if (!have_node) {
+			numa_node = nid;
+			have_node = true;
+		} else if (numa_node != nid) {
+			multi_node = true;
 		}
+	}
+
+	if (multi_node)
+		numa_node = NUMA_NO_NODE;
 
 	/*
 	 * Check if CPU 0 and its siblings are included in the provided CPU pool
@@ -414,7 +444,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 	cpus_read_unlock();
 reset_pool:
 	ne_cpu_pool.nr_parent_vm_cores = 0;
-	ne_cpu_pool.numa_node = -1;
+	ne_cpu_pool.numa_node = NUMA_NO_NODE;
 	mutex_unlock(&ne_cpu_pool.mutex);
 
 	return rc;
@@ -452,7 +482,7 @@ static void ne_teardown_cpu_pool(void)
 
 	ne_free_core_slots();
 	ne_cpu_pool.nr_parent_vm_cores = 0;
-	ne_cpu_pool.numa_node = -1;
+	ne_cpu_pool.numa_node = NUMA_NO_NODE;
 
 	mutex_unlock(&ne_cpu_pool.mutex);
 }
@@ -633,7 +663,13 @@ static bool ne_donated_cpu(struct ne_enclave *ne_enclave, unsigned int cpu)
 /**
  * ne_get_unused_core_from_cpu_pool() - Get the id of a full core from the
  *					NE CPU pool.
- * @void:	No parameters provided.
+ * @nid:	Required NUMA node for the core, or NUMA_NO_NODE for no
+ *		node constraint.
+ *
+ * When @nid names a specific node, only a free core whose CPUs are
+ * on that node is returned; if none is available the function returns -1
+ * rather than spilling onto another node. When @nid is NUMA_NO_NODE, any
+ * free core is returned.
  *
  * Context: Process context. This function is called with the ne_enclave and
  *	    ne_cpu_pool mutexes held.
@@ -641,19 +677,19 @@ static bool ne_donated_cpu(struct ne_enclave *ne_enclave, unsigned int cpu)
  * * Core id.
  * * -1 if no CPU core available in the pool.
  */
-static int ne_get_unused_core_from_cpu_pool(void)
+static int ne_get_unused_core_from_cpu_pool(int nid)
 {
-	int core_id = -1;
 	unsigned int i = 0;
 
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		if (!cpumask_empty(ne_cpu_pool.avail_threads_per_core[i])) {
-			core_id = i;
-
-			break;
-		}
+	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++) {
+		if (cpumask_empty(ne_cpu_pool.avail_threads_per_core[i]))
+			continue;
+		if (nid != NUMA_NO_NODE && ne_cpu_pool.core_nid[i] != nid)
+			continue;
+		return i;
+	}
 
-	return core_id;
+	return -1;
 }
 
 /**
@@ -710,39 +746,62 @@ static int ne_set_enclave_threads_per_core(struct ne_enclave *ne_enclave,
  *				sibling of a new CPU core.
  * @ne_enclave :	Private data associated with the current enclave.
  * @vcpu_id:		vCPU to get from the NE CPU pool.
+ * @nid:		Required NUMA node for a newly taken core, or
+ *			NUMA_NO_NODE for no node constraint. A core already
+ *			held by the enclave is subject to the same
+ *			constraint, so that a run of requests for one node
+ *			does not hand back a sibling on another.
  *
- * Context: Process context. This function is called with the ne_enclave mutex held.
+ * The sibling walk reads ne_enclave->threads_per_core[] and
+ * ne_cpu_pool.core_nid[] at the same index, and stays below the smaller of
+ * their two lengths.
+ *
+ * Context: Process context. This function is called with the ne_enclave mutex
+ *	    held, and takes the ne_cpu_pool mutex over the sibling walk and the
+ *	    new-core path.
  * Return:
  * * 0 on success.
  * * Negative return value on failure.
  */
-static int ne_get_cpu_from_cpu_pool(struct ne_enclave *ne_enclave, u32 *vcpu_id)
+static int ne_get_cpu_from_cpu_pool(struct ne_enclave *ne_enclave, u32 *vcpu_id,
+				    int nid)
 {
+	unsigned int nr_cores = 0;
 	int core_id = -1;
 	unsigned int cpu = 0;
 	unsigned int i = 0;
 	int rc = -EINVAL;
 
+	mutex_lock(&ne_cpu_pool.mutex);
+
+	nr_cores = min(ne_enclave->nr_parent_vm_cores,
+		       ne_cpu_pool.nr_parent_vm_cores);
+
 	/*
 	 * If previously allocated a thread of a core to this enclave, first
 	 * check remaining sibling(s) for new CPU allocations, so that full
-	 * CPU cores are used for the enclave.
+	 * CPU cores are used for the enclave. Cores on a node other than the
+	 * requested one are skipped, so a run of requests for one node does
+	 * not hand back a sibling on another.
 	 */
-	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++)
+	for (i = 0; i < nr_cores; i++) {
+		if (nid != NUMA_NO_NODE && ne_cpu_pool.core_nid[i] != nid)
+			continue;
+
 		for_each_cpu(cpu, ne_enclave->threads_per_core[i])
 			if (!ne_donated_cpu(ne_enclave, cpu)) {
 				*vcpu_id = cpu;
+				rc = 0;
 
-				return 0;
+				goto unlock_mutex;
 			}
-
-	mutex_lock(&ne_cpu_pool.mutex);
+	}
 
 	/*
 	 * If no remaining siblings, get a core from the NE CPU pool and keep
 	 * track of all the threads in the enclave threads per core data structure.
 	 */
-	core_id = ne_get_unused_core_from_cpu_pool();
+	core_id = ne_get_unused_core_from_cpu_pool(nid);
 
 	rc = ne_set_enclave_threads_per_core(ne_enclave, core_id, *vcpu_id);
 	if (rc < 0)
@@ -970,7 +1029,8 @@ static int ne_sanity_check_user_mem_region_page(struct ne_enclave *ne_enclave,
 		return -NE_ERR_INVALID_PAGE_SIZE;
 	}
 
-	if (ne_enclave->numa_node != page_to_nid(mem_region_page)) {
+	if (ne_enclave->numa_node != NUMA_NO_NODE &&
+	    ne_enclave->numa_node != page_to_nid(mem_region_page)) {
 		dev_err_ratelimited(ne_misc_dev.this_device,
 				    "Page is not from NUMA node %d\n",
 				    ne_enclave->numa_node);
@@ -1329,7 +1389,8 @@ static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned lo
 
 		if (!vcpu_id) {
 			/* Use the CPU pool for choosing a CPU for the enclave. */
-			rc = ne_get_cpu_from_cpu_pool(ne_enclave, &vcpu_id);
+			rc = ne_get_cpu_from_cpu_pool(ne_enclave, &vcpu_id,
+						      NUMA_NO_NODE);
 			if (rc < 0) {
 				dev_err_ratelimited(ne_misc_dev.this_device,
 						    "Error in get CPU from pool [rc=%d]\n",
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 283c17558f2a3..23a783bab3d2f 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -54,7 +54,10 @@ struct ne_mem_region {
  *				per full CPU core in the NE CPU pool the enclave
  *				was created from.
  * @nr_vcpus:			Number of vcpus associated with the enclave.
- * @numa_node:			NUMA node of the enclave memory and CPUs.
+ * @numa_node:			NUMA node of the enclave memory and CPUs, or
+ *				NUMA_NO_NODE if the CPU pool it was created from
+ *				spans nodes, in which case its memory is not held
+ *				to a node.
  * @slot_uid:			Slot unique id mapped to the enclave.
  * @state:			Enclave state, updated during enclave lifetime.
  * @threads_per_core:		Enclave full CPU cores array. Each cpumask in the
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index e808f5ba124d4..206546f24748e 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -26,8 +26,8 @@
  *		  https://www.kernel.org/doc/html/latest/admin-guide/kernel-parameters.html
  *		  CPU 0 and its siblings have to remain available for the
  *		  primary / parent VM, so they cannot be set for enclaves. Full
- *		  CPU core(s), from the same NUMA node, need(s) to be included
- *		  in the CPU pool.
+ *		  CPU core(s) need(s) to be included in the CPU pool. The
+ *		  pool may span NUMA nodes.
  *
  * Context: Process context.
  * Return:
@@ -50,8 +50,9 @@
  * NE_ADD_VCPU - The command is used to set a vCPU for an enclave. The vCPU can
  *		 be auto-chosen from the NE CPU pool or it can be set by the
  *		 caller, with the note that it needs to be available in the NE
- *		 CPU pool. Full CPU core(s), from the same NUMA node, need(s) to
- *		 be associated with an enclave.
+ *		 CPU pool. Full CPU core(s) need(s) to be associated with an
+ *		 enclave. When the CPU pool spans NUMA nodes, the cores of
+ *		 one enclave can be from different nodes.
  *		 The vCPU id is an input / output parameter. If its value is 0,
  *		 then a CPU is chosen from the enclave CPU pool and returned via
  *		 this parameter.
@@ -108,8 +109,11 @@
 /**
  * NE_SET_USER_MEMORY_REGION - The command is used to set a memory region for an
  *			       enclave, given the allocated memory from the
- *			       userspace. Enclave memory needs to be from the
- *			       same NUMA node as the enclave CPUs.
+ *			       userspace. When the CPU pool sits on one
+ *			       NUMA node, enclave memory needs to be from
+ *			       that node. An enclave built out of a pool
+ *			       that spans nodes has no node of its own, and
+ *			       its memory can come from any node.
  *			       The user memory region is an input parameter. It
  *			       includes info provided by the caller - flags,
  *			       memory size and userspace address.
@@ -138,7 +142,9 @@
  * * NE_ERR_MEM_NOT_HUGE_PAGE		- The memory region is not backed by
  *					  huge pages.
  * * NE_ERR_MEM_DIFFERENT_NUMA_NODE	- The memory region is not from the same
- *					  NUMA node as the CPUs.
+ *					  NUMA node as the CPUs. Not returned
+ *					  for an enclave whose CPU pool spans
+ *					  nodes.
  * * NE_ERR_MEM_MAX_REGIONS		- The number of memory regions set for
  *					  the enclave reached maximum.
  * * NE_ERR_INVALID_PAGE_SIZE		- The memory region is not backed by
@@ -223,6 +229,8 @@
 /**
  * NE_ERR_MEM_DIFFERENT_NUMA_NODE - The user space memory region is backed by
  *				    pages from different NUMA nodes than the CPUs.
+ *				    Not returned for an enclave whose CPU pool
+ *				    spans nodes.
  */
 #define NE_ERR_MEM_DIFFERENT_NUMA_NODE		(264)
 /**
diff --git a/samples/nitro_enclaves/ne_ioctl_sample.c b/samples/nitro_enclaves/ne_ioctl_sample.c
index 765b131c73190..84caf5e66fcbd 100644
--- a/samples/nitro_enclaves/ne_ioctl_sample.c
+++ b/samples/nitro_enclaves/ne_ioctl_sample.c
@@ -11,9 +11,10 @@
  * -----
  *
  * Load the nitro_enclaves module, setting also the enclave CPU pool. The
- * enclave CPUs need to be full cores from the same NUMA node. CPU 0 and its
- * siblings have to remain available for the primary / parent VM, so they
- * cannot be included in the enclave CPU pool.
+ * enclave CPUs need to be full cores. They come from the same NUMA node
+ * unless the CPU pool spans several of them. CPU 0 and its siblings have to
+ * remain available for the primary / parent VM, so they cannot be included
+ * in the enclave CPU pool.
  *
  * See the cpu list section from the kernel documentation.
  * https://www.kernel.org/doc/html/latest/admin-guide/kernel-parameters.html#cpu-lists
@@ -41,7 +42,8 @@
  *	dmesg
  *
  * Setup hugetlbfs huge pages. The memory needs to be from the same NUMA node as
- * the enclave CPUs.
+ * the enclave CPUs, unless the CPU pool spans nodes, in which case the enclave
+ * is not held to a node and its memory can come from any of them.
  *
  * https://www.kernel.org/doc/html/latest/admin-guide/mm/hugetlbpage.html
  *
-- 
2.47.1

[PATCH 02/40] nitro_enclaves: Add NE_ADD_ANY_VCPU

[Alexander Graf]

A Nitro Sandbox can span multiple NUMA nodes. The current NE_ADD_VCPU ioctl
only takes a target CPU id as input or a special ANY id. With NUMA, we want
to be able to take any CPU from a specific NUMA node. Add a new
NE_ADD_ANY_VCPU ioctl with nid parameter to select the target node. The new
ioctl returns the selected CPU id or NE_ERR_NO_CPUS_AVAIL_IN_POOL.

To already prepare for later argument extension, filter out the size field
of the ioctl and handle the fields we know about.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 92 ++++++++++++++++++++++-
 include/uapi/linux/nitro_enclaves.h       | 54 +++++++++++++
 2 files changed, 145 insertions(+), 1 deletion(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 6e36abeff01a2..576f1909c7304 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -9,11 +9,13 @@
  */
 
 #include <linux/anon_inodes.h>
+#include <linux/build_bug.h>
 #include <linux/capability.h>
 #include <linux/cpu.h>
 #include <linux/device.h>
 #include <linux/file.h>
 #include <linux/hugetlb.h>
+#include <linux/ioctl.h>
 #include <linux/kernel.h>
 #include <linux/limits.h>
 #include <linux/list.h>
@@ -24,12 +26,15 @@
 #include <linux/module.h>
 #include <linux/mutex.h>
 #include <linux/nitro_enclaves.h>
+#include <linux/nodemask.h>
 #include <linux/numa.h>
 #include <linux/pci.h>
 #include <linux/poll.h>
 #include <linux/range.h>
 #include <linux/slab.h>
+#include <linux/stddef.h>
 #include <linux/types.h>
+#include <linux/uaccess.h>
 #include <uapi/linux/vm_sockets.h>
 
 #include "ne_misc_dev.h"
@@ -42,6 +47,23 @@
  */
 #define NE_CPUS_SIZE		(512)
 
+/**
+ * NE_EA_IOCTL() - Mask the size and direction bits out of an ioctl command, so
+ *		   that an Extensible Argument ioctl matches whatever size of
+ *		   argument the caller passed. The direction bits go with them,
+ *		   so a caller that spells the command _IOW(), _IOR() or
+ *		   _IOWR() reaches the same handler, and the handler decides
+ *		   what it reads and writes.
+ * @cmd:	The ioctl command to mask.
+ */
+#define NE_EA_IOCTL(cmd)	((cmd) & ~(IOC_INOUT | IOCSIZE_MASK))
+
+/*
+ * NE_NUMA_ANY reaches the core allocator as a node id, so the uapi spelling
+ * and the kernel one have to agree.
+ */
+static_assert(NE_NUMA_ANY == NUMA_NO_NODE);
+
 /**
  * NE_EIF_LOAD_OFFSET - The offset where to copy the Enclave Image Format (EIF)
  *			image in enclave memory.
@@ -334,7 +356,8 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 	 * multiple NUMA nodes, set numa_node to NUMA_NO_NODE; the pool is then
 	 * treated as node-agnostic for placement and memory-locality checks.
 	 * A CPU that reports no node at all is refused, because otherwise such
-	 * a CPU would join a single-node pool and turn it into a spanning one.
+	 * a CPU would join a single-node pool and turn it into a spanning one,
+	 * and its core would be unreachable through a node-named pick.
 	 */
 	for_each_cpu(cpu, cpu_pool) {
 		int nid = cpu_to_node(cpu);
@@ -1592,6 +1615,73 @@ static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned lo
 
 		return 0;
 	}
+	}
+
+	/* Extensible Argument ioctls */
+	switch (NE_EA_IOCTL(cmd)) {
+	case NE_EA_IOCTL(NE_ADD_ANY_VCPU): {
+		struct ne_add_any_vcpu_args __user *uargs = (void __user *)arg;
+		unsigned int usize = _IOC_SIZE(cmd);
+		struct ne_add_any_vcpu_args args;
+		u32 vcpu_id = 0;
+		int rc;
+
+		if (usize < offsetofend(struct ne_add_any_vcpu_args, vcpu_id))
+			return -EINVAL;
+
+		rc = copy_struct_from_user(&args, sizeof(args), uargs, usize);
+		if (rc)
+			return rc;
+
+		if (args.numa_node != NE_NUMA_ANY &&
+		    (args.numa_node < 0 || args.numa_node >= nr_node_ids ||
+		     !node_state(args.numa_node, N_POSSIBLE))) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Invalid NUMA node %d\n", args.numa_node);
+
+			return -EINVAL;
+		}
+
+		rc = mutex_lock_killable(&ne_enclave->enclave_info_mutex);
+		if (rc)
+			return rc;
+
+		if (ne_enclave->state != NE_STATE_INIT) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Enclave is not in init state\n");
+
+			mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+			return -NE_ERR_NOT_IN_INIT_STATE;
+		}
+
+		rc = ne_get_cpu_from_cpu_pool(ne_enclave, &vcpu_id, args.numa_node);
+		if (rc < 0) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Error in get CPU from pool [rc=%d node=%d]\n",
+					    rc, args.numa_node);
+
+			mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+			return rc;
+		}
+
+		rc = ne_add_vcpu_ioctl(ne_enclave, vcpu_id);
+		if (rc < 0) {
+			mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+			return rc;
+		}
+
+		mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+		args.vcpu_id = vcpu_id;
+
+		if (copy_to_user(uargs, &args, min_t(size_t, usize, sizeof(args))))
+			return -EFAULT;
+
+		return 0;
+	}
 
 	default:
 		return -ENOTTY;
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index 206546f24748e..e266b8a611c92 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -83,6 +83,60 @@
  */
 #define NE_ADD_VCPU			_IOWR(0xAE, 0x21, __u32)
 
+/**
+ * struct ne_add_any_vcpu_args - Argument for %NE_ADD_ANY_VCPU.
+ * @numa_node:	The NUMA node the vCPU has to come from, or %NE_NUMA_ANY for
+ *		any node.
+ * @vcpu_id:	The vCPU id the kernel chose, set on success. Ignored on
+ *		input.
+ */
+struct ne_add_any_vcpu_args {
+	__s32	numa_node;
+	__u32	vcpu_id;
+};
+
+/**
+ * NE_NUMA_ANY - Value for &struct ne_add_any_vcpu_args.numa_node that asks for
+ *		 a vCPU from any NUMA node.
+ */
+#define NE_NUMA_ANY			(-1)
+
+/**
+ * NE_ADD_ANY_VCPU - The command is used to set a vCPU for an enclave, chosen
+ *		     from the NE CPU pool. @numa_node confines the choice to
+ *		     one NUMA node; %NE_NUMA_ANY leaves it unconstrained. A
+ *		     caller that spreads an enclave over several nodes asks
+ *		     for each node in turn and stops when a node runs out.
+ *		     @vcpu_id carries back the vCPU id the kernel chose, so
+ *		     a caller collects the enclave's vCPU set as it builds
+ *		     it.
+ *		     The ioctl can be invoked on the enclave fd, before an
+ *		     enclave is started.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0				- Logic successfully completed.
+ * * -1				- There was a failure in the ioctl logic.
+ * On failure, errno is set to:
+ * * ENODEV			- The NE PCI device is not available.
+ * * EFAULT			- Argument read or write-back failure.
+ * * EINVAL			- The argument is shorter than @vcpu_id, or
+ *				  @numa_node is neither %NE_NUMA_ANY nor a
+ *				  possible NUMA node.
+ * * EIO			- Current task mm is not the same as the one
+ *				  that created the enclave.
+ * * E2BIG			- The argument is longer than this kernel knows
+ *				  about and the excess is not all zero.
+ * * NE_ERR_NOT_IN_INIT_STATE	- The enclave is not in init state (init =
+ *				  before being started).
+ * * NE_ERR_NO_CPUS_AVAIL_IN_POOL - No CPUs available in the NE CPU pool on
+ *				  @numa_node.
+ * * NE_ERR_VCPU_INVALID_CPU_CORE - The core the pool handed back is outside
+ *				  the enclave's core slots.
+ * * Error codes from the NE PCI device request.
+ */
+#define NE_ADD_ANY_VCPU			_IOWR(0xAE, 0x2D, struct ne_add_any_vcpu_args)
+
 /**
  * NE_GET_IMAGE_LOAD_INFO - The command is used to get information needed for
  *			    in-memory enclave image loading e.g. offset in
-- 
2.47.1

[PATCH 03/40] nitro_enclaves: Expose CPU pool state under sysfs

[Alexander Graf]

When using the dedicated CPU pool, orchestrators need to know whether
there is actually capacity left to spawn a new enclave or sandbox on
a given target NUMA node. The easiest place to expose that is sysfs.

Report the pool under
/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/: total, used and
avail, once as total and once in a nodeN subdirectory per node.
The files are binary attributes, because a CPU list can outgrow the one
page a device attribute gets.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 .../sysfs-devices-virtual-misc-nitro_enclaves |  75 ++++
 MAINTAINERS                                   |   1 +
 drivers/virt/nitro_enclaves/ne_misc_dev.c     | 388 +++++++++++++++++-
 drivers/virt/nitro_enclaves/ne_misc_dev.h     |   3 +
 drivers/virt/nitro_enclaves/ne_pci_dev.c      |  16 +
 5 files changed, 469 insertions(+), 14 deletions(-)
 create mode 100644 Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves

diff --git a/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves b/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
new file mode 100644
index 0000000000000..596054b75e0cf
--- /dev/null
+++ b/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
@@ -0,0 +1,75 @@
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/total
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/used
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/avail
+Date:		November 2026
+KernelVersion:	7.5
+Contact:	aws-nitro-enclaves-devel@amazon.com
+Description:
+		Read-only. CPU-list strings describing the Nitro Enclaves CPU
+		pool. The pool is handed out a whole core at a time, so these
+		files name threads but move in units of a core:
+
+		total
+			every CPU thread dedicated to the pool
+
+		used
+			every thread of every core an enclave holds, including
+			the ones it has not made vCPUs of yet
+
+		avail
+			every thread of every core no enclave holds, so the set
+			a new enclave can draw from
+
+		A thread leaves used when the enclave holding it is freed.
+		A release that fails part way through keeps the thread in
+		used. total changes when the pool is reconfigured, which
+		requires that no enclave exists.
+
+		A read longer than one page is served in parts, and each part
+		renders the pool as it stands when that part is read. A thread
+		that changes hands during such a read can be missing from every
+		part, and the parts joined together can fail to parse as a CPU
+		list at all - see the warning in the kerneldoc of
+		bitmap_print_bitmask_to_buf() in lib/bitmap-str.c, which
+		bitmap_print_list_to_buf() refers to.
+
+		NE_ADD_VCPU auto-pick and NE_ADD_ANY_VCPU each look for a free
+		thread of a core the enclave already holds before taking a
+		core out of avail, so the thread either returns can be one
+		avail does not list.
+
+		A read takes the pool lock, so it can wait for a write to
+		ne_cpus to finish.
+
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/nodeX/total
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/nodeX/used
+What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/nodeX/avail
+Date:		November 2026
+KernelVersion:	7.5
+Contact:	aws-nitro-enclaves-devel@amazon.com
+Description:
+		Read-only. The cpu_pool files of the same name, over the
+		threads of NUMA node X alone. There is a nodeX directory for
+		each node that was online when the device was probed, and its
+		files read empty for a node the pool holds no thread of.
+
+		A thread appears under exactly one node, the one its core was
+		on while that core was still online. The driver offlines the
+		threads it takes into the pool, so
+		/sys/devices/system/node/nodeX/cpulist stops listing them and
+		is not what these files are cut from.
+
+		nodeX/used and nodeX/avail are cut out of nodeX/total the way
+		used and avail are cut out of total, so a node file and its
+		pool-wide counterpart cannot disagree, and the nodeX/total of
+		every node join to total.
+
+		NE_ADD_ANY_VCPU takes a core from the node it names, so
+		nodeX/avail is what a request for that node can draw on. What
+		is written above about a read longer than one page, and about
+		the free thread of a held core that avail does not list,
+		holds for these files as well.
+
+		The directories appear once the device has been probed and the
+		misc device registered, so a reader that opens the misc device
+		the moment it appears can find cpu_pool empty.
diff --git a/MAINTAINERS b/MAINTAINERS
index 20e4a56a197ee..f946a79b34d0d 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -19839,6 +19839,7 @@ R:	The AWS Nitro Enclaves Team <aws-nitro-enclaves-devel@amazon.com>
 L:	linux-kernel@vger.kernel.org
 S:	Supported
 W:	https://aws.amazon.com/ec2/nitro/nitro-enclaves/
+F:	Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
 F:	Documentation/virt/ne_overview.rst
 F:	drivers/virt/nitro_enclaves/
 F:	include/linux/nitro_enclaves.h
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 576f1909c7304..66bcaec8d3199 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -12,11 +12,14 @@
 #include <linux/build_bug.h>
 #include <linux/capability.h>
 #include <linux/cpu.h>
+#include <linux/cpumask.h>
 #include <linux/device.h>
 #include <linux/file.h>
 #include <linux/hugetlb.h>
 #include <linux/ioctl.h>
 #include <linux/kernel.h>
+#include <linux/kobject.h>
+#include <linux/kstrtox.h>
 #include <linux/limits.h>
 #include <linux/list.h>
 #include <linux/miscdevice.h>
@@ -33,6 +36,8 @@
 #include <linux/range.h>
 #include <linux/slab.h>
 #include <linux/stddef.h>
+#include <linux/string.h>
+#include <linux/sysfs.h>
 #include <linux/types.h>
 #include <linux/uaccess.h>
 #include <uapi/linux/vm_sockets.h>
@@ -95,18 +100,6 @@ static const struct file_operations ne_fops = {
 	.compat_ioctl	= compat_ptr_ioctl,
 };
 
-static struct miscdevice ne_misc_dev = {
-	.minor	= MISC_DYNAMIC_MINOR,
-	.name	= "nitro_enclaves",
-	.fops	= &ne_fops,
-	.mode	= 0660,
-};
-
-struct ne_devs ne_devs = {
-	.ne_misc_dev	= &ne_misc_dev,
-	.lock		= __RWSEM_INITIALIZER(ne_devs.lock),
-};
-
 /*
  * TODO: Update logic to create new sysfs entries instead of using
  * a kernel parameter e.g. if multiple sysfs files needed.
@@ -163,6 +156,19 @@ static const struct kernel_param *ne_cpus_kp;
  *				ne_setup_cpu_pool() offlines them, because
  *				cpu_to_node() is not documented to answer for an
  *				offline CPU.
+ * @node_cpus:			The threads of the pool on each NUMA node, indexed
+ *				by node id and sized by nr_node_ids. It partitions
+ *				@total_cpus: a thread sits in the entry of the node
+ *				its core was on while that core was still online,
+ *				which is the node @core_nid records for the same
+ *				core. NULL if no pool is set up.
+ * @total_cpus:			Every CPU thread in the pool, as parsed from
+ *				ne_cpus. Empty if no pool is set up.
+ * @used_cpus:			The subset of @total_cpus that enclaves hold.
+ *				A whole core enters when an enclave takes it
+ *				and leaves when the enclave is freed, in
+ *				the same two places that move the core out of
+ *				and back into @avail_threads_per_core.
  */
 struct ne_cpu_pool {
 	cpumask_var_t	*avail_threads_per_core;
@@ -170,12 +176,330 @@ struct ne_cpu_pool {
 	unsigned int	nr_parent_vm_cores;
 	int		numa_node;
 	int		*core_nid;
+	cpumask_var_t	*node_cpus;
+	struct cpumask	total_cpus;
+	struct cpumask	used_cpus;
 };
 
 static struct ne_cpu_pool ne_cpu_pool = {
 	.mutex = __MUTEX_INITIALIZER(ne_cpu_pool.mutex),
 };
 
+/*
+ * cpu_pool/ sysfs group on the Nitro Enclaves misc device: report the CPU
+ * pool occupancy.
+ *
+ *   /sys/devices/virtual/misc/nitro_enclaves/cpu_pool/{total,used,avail}
+ *
+ * The three files are CPU-list strings rendered from the pool masks under the
+ * pool mutex. avail is what is left of total once used is taken out of it, so
+ * a read has nothing to work out that the allocator has not already decided.
+ * They are binary attributes sized the way the per-node CPU lists under
+ * /sys/devices/system/node are, because a CPU list naming every thread of a
+ * large pool does not fit in the single page a device attribute is handed.
+ */
+
+static ssize_t total_read(struct file *file, struct kobject *kobj,
+			  const struct bin_attribute *attr, char *buf,
+			  loff_t off, size_t count)
+{
+	ssize_t n;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+	n = cpumap_print_list_to_buf(buf, &ne_cpu_pool.total_cpus, off, count);
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	return n;
+}
+
+static const BIN_ATTR_RO(total, CPULIST_FILE_MAX_BYTES);
+
+static ssize_t used_read(struct file *file, struct kobject *kobj,
+			 const struct bin_attribute *attr, char *buf,
+			 loff_t off, size_t count)
+{
+	ssize_t n;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+	n = cpumap_print_list_to_buf(buf, &ne_cpu_pool.used_cpus, off, count);
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	return n;
+}
+
+static const BIN_ATTR_RO(used, CPULIST_FILE_MAX_BYTES);
+
+static ssize_t avail_read(struct file *file, struct kobject *kobj,
+			  const struct bin_attribute *attr, char *buf,
+			  loff_t off, size_t count)
+{
+	cpumask_var_t avail;
+	ssize_t n;
+
+	if (!alloc_cpumask_var(&avail, GFP_KERNEL))
+		return -ENOMEM;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+	cpumask_andnot(avail, &ne_cpu_pool.total_cpus, &ne_cpu_pool.used_cpus);
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	n = cpumap_print_list_to_buf(buf, avail, off, count);
+
+	free_cpumask_var(avail);
+
+	return n;
+}
+
+static const BIN_ATTR_RO(avail, CPULIST_FILE_MAX_BYTES);
+
+static const struct bin_attribute *const ne_cpu_pool_bin_attrs[] = {
+	&bin_attr_total,
+	&bin_attr_used,
+	&bin_attr_avail,
+	NULL,
+};
+
+static const struct attribute_group ne_cpu_pool_group = {
+	.bin_attrs	= ne_cpu_pool_bin_attrs,
+};
+
+/*
+ * cpu_pool/nodeX/ sysfs groups on the Nitro Enclaves misc device: the three
+ * files above, over the threads of one NUMA node.
+ *
+ *   /sys/devices/virtual/misc/nitro_enclaves/cpu_pool/nodeX/{total,used,avail}
+ *
+ * One directory per node that was online when the device was probed, each
+ * holding the same three attributes. The node id is in the directory's name
+ * and is read back out of it, so the three attributes are shared rather than
+ * made per node. A node the pool holds nothing on reads empty in all three.
+ */
+
+/**
+ * enum ne_cpu_pool_view - Which view of a node's pool threads to render.
+ * @NE_CPU_POOL_TOTAL:	Every pool thread of the node.
+ * @NE_CPU_POOL_USED:	The ones enclaves hold.
+ * @NE_CPU_POOL_AVAIL:	The rest.
+ */
+enum ne_cpu_pool_view {
+	NE_CPU_POOL_TOTAL,
+	NE_CPU_POOL_USED,
+	NE_CPU_POOL_AVAIL,
+};
+
+/**
+ * ne_cpu_pool_node_read() - Render one view of one node's pool threads.
+ * @kobj:	The nodeX directory the read came through.
+ * @view:	Which of the three views to render.
+ * @buf:	Buffer to render the CPU list into.
+ * @off:	Offset into the CPU list to render from.
+ * @count:	Bytes of the CPU list to render.
+ *
+ * The node's threads come from the partition of the pool recorded while the
+ * pool's CPUs were online, so this answers for CPUs the driver has offlined.
+ *
+ * Context: Process context.
+ * Return:
+ * * Number of bytes rendered.
+ * * Negative return value on failure.
+ */
+static ssize_t ne_cpu_pool_node_read(struct kobject *kobj,
+				     enum ne_cpu_pool_view view, char *buf,
+				     loff_t off, size_t count)
+{
+	cpumask_var_t threads;
+	ssize_t n;
+	int nid = 0;
+
+	if (kstrtoint(kobject_name(kobj) + strlen("node"), 10, &nid))
+		return -EINVAL;
+
+	if (nid < 0 || (unsigned int)nid >= nr_node_ids)
+		return -EINVAL;
+
+	if (!alloc_cpumask_var(&threads, GFP_KERNEL))
+		return -ENOMEM;
+
+	mutex_lock(&ne_cpu_pool.mutex);
+
+	if (!ne_cpu_pool.node_cpus) {
+		cpumask_clear(threads);
+	} else {
+		cpumask_copy(threads, ne_cpu_pool.node_cpus[nid]);
+
+		if (view == NE_CPU_POOL_USED)
+			cpumask_and(threads, threads, &ne_cpu_pool.used_cpus);
+		else if (view == NE_CPU_POOL_AVAIL)
+			cpumask_andnot(threads, threads,
+				       &ne_cpu_pool.used_cpus);
+	}
+
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	n = cpumap_print_list_to_buf(buf, threads, off, count);
+
+	free_cpumask_var(threads);
+
+	return n;
+}
+
+static ssize_t node_total_read(struct file *file, struct kobject *kobj,
+			       const struct bin_attribute *attr, char *buf,
+			       loff_t off, size_t count)
+{
+	return ne_cpu_pool_node_read(kobj, NE_CPU_POOL_TOTAL, buf, off, count);
+}
+
+static const struct bin_attribute bin_attr_node_total =
+	__BIN_ATTR(total, 0444, node_total_read, NULL, CPULIST_FILE_MAX_BYTES);
+
+static ssize_t node_used_read(struct file *file, struct kobject *kobj,
+			      const struct bin_attribute *attr, char *buf,
+			      loff_t off, size_t count)
+{
+	return ne_cpu_pool_node_read(kobj, NE_CPU_POOL_USED, buf, off, count);
+}
+
+static const struct bin_attribute bin_attr_node_used =
+	__BIN_ATTR(used, 0444, node_used_read, NULL, CPULIST_FILE_MAX_BYTES);
+
+static ssize_t node_avail_read(struct file *file, struct kobject *kobj,
+			       const struct bin_attribute *attr, char *buf,
+			       loff_t off, size_t count)
+{
+	return ne_cpu_pool_node_read(kobj, NE_CPU_POOL_AVAIL, buf, off, count);
+}
+
+static const struct bin_attribute bin_attr_node_avail =
+	__BIN_ATTR(avail, 0444, node_avail_read, NULL, CPULIST_FILE_MAX_BYTES);
+
+static const struct bin_attribute *const ne_cpu_pool_node_bin_attrs[] = {
+	&bin_attr_node_total,
+	&bin_attr_node_used,
+	&bin_attr_node_avail,
+	NULL,
+};
+
+static const struct attribute_group ne_cpu_pool_node_group = {
+	.bin_attrs	= ne_cpu_pool_node_bin_attrs,
+};
+
+static struct kobject *ne_cpu_pool_kobj;
+static struct kobject **ne_cpu_pool_node_kobj;
+
+static struct miscdevice ne_misc_dev = {
+	.minor	= MISC_DYNAMIC_MINOR,
+	.name	= "nitro_enclaves",
+	.fops	= &ne_fops,
+	.mode	= 0660,
+};
+
+struct ne_devs ne_devs = {
+	.ne_misc_dev	= &ne_misc_dev,
+	.lock		= __RWSEM_INITIALIZER(ne_devs.lock),
+};
+
+/**
+ * ne_cpu_pool_sysfs_exit() - Take down the cpu_pool sysfs tree.
+ * @void:	No parameters provided.
+ *
+ * Idempotent, so the setup path can call it on a tree it built part of.
+ *
+ * Context: Process context.
+ */
+void ne_cpu_pool_sysfs_exit(void)
+{
+	unsigned int nid = 0;
+
+	if (ne_cpu_pool_node_kobj) {
+		for (nid = 0; nid < nr_node_ids; nid++) {
+			if (!ne_cpu_pool_node_kobj[nid])
+				continue;
+
+			sysfs_remove_group(ne_cpu_pool_node_kobj[nid],
+					   &ne_cpu_pool_node_group);
+			kobject_put(ne_cpu_pool_node_kobj[nid]);
+			ne_cpu_pool_node_kobj[nid] = NULL;
+		}
+
+		kfree(ne_cpu_pool_node_kobj);
+		ne_cpu_pool_node_kobj = NULL;
+	}
+
+	if (ne_cpu_pool_kobj) {
+		sysfs_remove_group(ne_cpu_pool_kobj, &ne_cpu_pool_group);
+		kobject_put(ne_cpu_pool_kobj);
+		ne_cpu_pool_kobj = NULL;
+	}
+}
+
+/**
+ * ne_cpu_pool_sysfs_init() - Build the cpu_pool sysfs tree.
+ * @void:	No parameters provided.
+ *
+ * The pool-wide files hang on a kobject of their own rather than on a named
+ * attribute group of the misc device, because the per-node directories are
+ * kobjects and a kobject takes a kobject as its parent. So this runs once the
+ * misc device is registered and its kobject exists.
+ *
+ * A directory is made for each node that is online now. A node that comes
+ * online later has no directory, and the pool cannot take a thread of it
+ * either: ne_setup_cpu_pool() refuses a CPU that is offline.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+int ne_cpu_pool_sysfs_init(void)
+{
+	unsigned int nid = 0;
+	int rc = -ENOMEM;
+
+	ne_cpu_pool_kobj = kobject_create_and_add("cpu_pool",
+						  &ne_misc_dev.this_device->kobj);
+	if (!ne_cpu_pool_kobj)
+		return -ENOMEM;
+
+	rc = sysfs_create_group(ne_cpu_pool_kobj, &ne_cpu_pool_group);
+	if (rc < 0)
+		goto sysfs_exit;
+
+	ne_cpu_pool_node_kobj = kzalloc_objs(*ne_cpu_pool_node_kobj,
+					     nr_node_ids);
+	if (!ne_cpu_pool_node_kobj) {
+		rc = -ENOMEM;
+
+		goto sysfs_exit;
+	}
+
+	for_each_online_node(nid) {
+		char name[16];
+
+		snprintf(name, sizeof(name), "node%u", nid);
+
+		ne_cpu_pool_node_kobj[nid] =
+			kobject_create_and_add(name, ne_cpu_pool_kobj);
+		if (!ne_cpu_pool_node_kobj[nid]) {
+			rc = -ENOMEM;
+
+			goto sysfs_exit;
+		}
+
+		rc = sysfs_create_group(ne_cpu_pool_node_kobj[nid],
+					&ne_cpu_pool_node_group);
+		if (rc < 0)
+			goto sysfs_exit;
+	}
+
+	return 0;
+
+sysfs_exit:
+	ne_cpu_pool_sysfs_exit();
+
+	return rc;
+}
+
 /**
  * struct ne_phys_contig_mem_regions - Contiguous physical memory regions.
  * @num:	The number of regions that currently has.
@@ -196,6 +520,7 @@ struct ne_phys_contig_mem_regions {
 static void ne_free_core_slots(void)
 {
 	unsigned int i = 0;
+	unsigned int nid = 0;
 
 	if (ne_cpu_pool.avail_threads_per_core) {
 		for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
@@ -207,6 +532,14 @@ static void ne_free_core_slots(void)
 
 	kfree(ne_cpu_pool.core_nid);
 	ne_cpu_pool.core_nid = NULL;
+
+	if (ne_cpu_pool.node_cpus) {
+		for (nid = 0; nid < nr_node_ids; nid++)
+			free_cpumask_var(ne_cpu_pool.node_cpus[nid]);
+
+		kfree(ne_cpu_pool.node_cpus);
+		ne_cpu_pool.node_cpus = NULL;
+	}
 }
 
 /**
@@ -234,8 +567,10 @@ static int ne_build_core_slots(const struct cpumask *cpu_pool)
 	unsigned int cpu_sibling = 0;
 	unsigned int i = 0;
 	unsigned int next_core_idx = 0;
+	unsigned int nid = 0;
 	unsigned int nr_cores = 0;
 	cpumask_var_t processed;
+	int core_nid = 0;
 	int rc = -ENOMEM;
 
 	if (!zalloc_cpumask_var(&processed, GFP_KERNEL))
@@ -261,6 +596,14 @@ static int ne_build_core_slots(const struct cpumask *cpu_pool)
 	if (!ne_cpu_pool.core_nid)
 		goto free_slots;
 
+	ne_cpu_pool.node_cpus = kzalloc_objs(*ne_cpu_pool.node_cpus, nr_node_ids);
+	if (!ne_cpu_pool.node_cpus)
+		goto free_slots;
+
+	for (nid = 0; nid < nr_node_ids; nid++)
+		if (!zalloc_cpumask_var(&ne_cpu_pool.node_cpus[nid], GFP_KERNEL))
+			goto free_slots;
+
 	for (i = 0; i < nr_cores; i++)
 		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL))
 			goto free_slots;
@@ -280,7 +623,12 @@ static int ne_build_core_slots(const struct cpumask *cpu_pool)
 			cpumask_set_cpu(cpu_sibling, processed);
 		}
 
-		ne_cpu_pool.core_nid[next_core_idx] = cpu_to_node(cpu);
+		core_nid = cpu_to_node(cpu);
+
+		ne_cpu_pool.core_nid[next_core_idx] = core_nid;
+		cpumask_or(ne_cpu_pool.node_cpus[core_nid],
+			   ne_cpu_pool.node_cpus[core_nid],
+			   ne_cpu_pool.avail_threads_per_core[next_core_idx]);
 
 		next_core_idx++;
 	}
@@ -452,6 +800,9 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 		}
 	}
 
+	cpumask_copy(&ne_cpu_pool.total_cpus, cpu_pool);
+	cpumask_clear(&ne_cpu_pool.used_cpus);
+
 	ne_cpu_pool.numa_node = numa_node;
 
 	mutex_unlock(&ne_cpu_pool.mutex);
@@ -466,6 +817,8 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 unlock_hotplug:
 	cpus_read_unlock();
 reset_pool:
+	cpumask_clear(&ne_cpu_pool.total_cpus);
+	cpumask_clear(&ne_cpu_pool.used_cpus);
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.numa_node = NUMA_NO_NODE;
 	mutex_unlock(&ne_cpu_pool.mutex);
@@ -504,6 +857,8 @@ static void ne_teardown_cpu_pool(void)
 		}
 
 	ne_free_core_slots();
+	cpumask_clear(&ne_cpu_pool.total_cpus);
+	cpumask_clear(&ne_cpu_pool.used_cpus);
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.numa_node = NUMA_NO_NODE;
 
@@ -758,6 +1113,9 @@ static int ne_set_enclave_threads_per_core(struct ne_enclave *ne_enclave,
 	for_each_cpu(cpu, ne_cpu_pool.avail_threads_per_core[core_id])
 		cpumask_set_cpu(cpu, ne_enclave->threads_per_core[core_id]);
 
+	cpumask_or(&ne_cpu_pool.used_cpus, &ne_cpu_pool.used_cpus,
+		   ne_cpu_pool.avail_threads_per_core[core_id]);
+
 	cpumask_clear(ne_cpu_pool.avail_threads_per_core[core_id]);
 
 	return 0;
@@ -1768,9 +2126,11 @@ static void ne_enclave_remove_all_vcpu_id_entries(struct ne_enclave *ne_enclave)
 		       ne_cpu_pool.nr_parent_vm_cores);
 
 	for (i = 0; i < nr_cores; i++)
-		for_each_cpu(cpu, ne_enclave->threads_per_core[i])
+		for_each_cpu(cpu, ne_enclave->threads_per_core[i]) {
 			/* Update the available NE CPU pool. */
 			cpumask_set_cpu(cpu, ne_cpu_pool.avail_threads_per_core[i]);
+			cpumask_clear_cpu(cpu, &ne_cpu_pool.used_cpus);
+		}
 
 	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++)
 		free_cpumask_var(ne_enclave->threads_per_core[i]);
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 23a783bab3d2f..483037061eafe 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -115,4 +115,7 @@ struct ne_devs {
 /* Nitro Enclaves (NE) data structure for keeping refs to the NE misc and PCI devices. */
 extern struct ne_devs ne_devs;
 
+int ne_cpu_pool_sysfs_init(void);
+void ne_cpu_pool_sysfs_exit(void);
+
 #endif /* _NE_MISC_DEV_H_ */
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index edc72d4f623e5..2d44cae4ac0cf 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -695,8 +695,20 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 
 	up_write(&ne_devs.lock);
 
+	rc = ne_cpu_pool_sysfs_init();
+	if (rc < 0) {
+		dev_err(&pdev->dev, "Error in cpu pool sysfs init [rc=%d]\n",
+			rc);
+
+		goto deregister_misc_dev;
+	}
+
 	return 0;
 
+deregister_misc_dev:
+	down_write(&ne_devs.lock);
+	ne_pci_dev_detach_enclaves(ne_pci_dev);
+	misc_deregister(ne_devs.ne_misc_dev);
 disable_ne_pci_dev:
 	ne_devs.ne_pci_dev = NULL;
 unlock_ne_devs:
@@ -727,6 +739,8 @@ static void ne_pci_remove(struct pci_dev *pdev)
 {
 	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
 
+	ne_cpu_pool_sysfs_exit();
+
 	down_write(&ne_devs.lock);
 
 	ne_pci_dev_detach_enclaves(ne_pci_dev);
@@ -765,6 +779,8 @@ static void ne_pci_shutdown(struct pci_dev *pdev)
 	if (!ne_pci_dev)
 		return;
 
+	ne_cpu_pool_sysfs_exit();
+
 	down_write(&ne_devs.lock);
 
 	ne_pci_dev_detach_enclaves(ne_pci_dev);
-- 
2.47.1

[PATCH 04/40] nitro_enclaves: Document multi-NUMA CPU pools

[Alexander Graf]

An admin placing enclaves or sandboxes across NUMA nodes needs to know
how the CPU pool, automatic CPU selection and the memory-locality check
interact.

Document those rules in ne_overview.rst, together with the per-node
sysfs directories, and point the cpu_pool ABI file at that section.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 .../sysfs-devices-virtual-misc-nitro_enclaves |  3 ++-
 Documentation/virt/ne_overview.rst            | 27 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves b/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
index 596054b75e0cf..c549364353186 100644
--- a/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
+++ b/Documentation/ABI/testing/sysfs-devices-virtual-misc-nitro_enclaves
@@ -39,7 +39,8 @@ Description:
 		avail does not list.
 
 		A read takes the pool lock, so it can wait for a write to
-		ne_cpus to finish.
+		ne_cpus to finish. For how an enclave is given cores from a
+		particular node, see Documentation/virt/ne_overview.rst.
 
 What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/nodeX/total
 What:		/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/nodeX/used
diff --git a/Documentation/virt/ne_overview.rst b/Documentation/virt/ne_overview.rst
index aecd3704acaf5..223ee7fdec51b 100644
--- a/Documentation/virt/ne_overview.rst
+++ b/Documentation/virt/ne_overview.rst
@@ -48,7 +48,7 @@ be aligned 2 MiB / 1 GiB physically contiguous memory regions (or multiple of
 this size e.g. 8 MiB). The memory can be allocated e.g. by using hugetlbfs from
 user space [2][3][7]. The memory size for an enclave needs to be at least
 64 MiB. The enclave memory and CPUs need to be from the same NUMA node,
-unless the CPU pool spans several of them.
+unless the CPU pool spans several of them; see `Multi-NUMA CPU pools`_.
 
 An enclave runs on dedicated cores. CPU 0 and its CPU siblings need to remain
 available for the primary VM. A CPU pool has to be set for NE purposes by an
@@ -99,3 +99,28 @@ enclave process can exit.
 [6] https://www.kernel.org/doc/html/latest/x86/boot.html
 [7] https://www.kernel.org/doc/html/latest/arm64/hugetlbpage.html
 [8] https://www.kernel.org/doc/html/latest/arm64/booting.html
+
+Multi-NUMA CPU pools
+--------------------
+
+The CPU pool may span multiple NUMA nodes. NE_ADD_VCPU with vcpu_id == 0
+takes the first free core in the pool whatever node it sits on, so an
+enclave built that way can end up holding cores on several nodes.
+NE_ADD_ANY_VCPU names the node instead: the kernel takes a free core on
+that node, and fails rather than crossing to another node once that node
+has none left. User space builds an enclave across nodes by asking for
+each node in turn, and sets numa_node to NE_NUMA_ANY when it does not
+care. Both start from a free thread of a core the enclave already holds,
+and only take a new core once there is none.
+
+The node says where cores come from. An enclave built out of a pool that
+spans nodes has no node of its own, so the rule that its memory be on
+that node does not apply to it and its memory can come from any node.
+
+The pool occupancy is readable under
+/sys/devices/virtual/misc/nitro_enclaves/cpu_pool/, and per node under
+its ``node0``, ``node1``, ... subdirectories, one for each node that was
+online when the device was probed. A node the pool holds nothing on
+reads empty rather than being absent. The driver offlines pool threads,
+so ``nodeN/cpulist`` under /sys/devices/system/node stops listing them
+and answers nothing about the pool.
-- 
2.47.1

[PATCH 05/40] dma-contiguous: add dma_get_numa_cma_area() accessor

[Alexander Graf]

With Nitro Sandboxes, we want to allocate per-NUMA memory from the DMA
contiguous allocator, but have no per-NUMA struct device to pass in.
dma_alloc_contiguous() takes the node from dev_to_node(dev), and the
per-node areas sit in dma_contiguous_numa_area[], which is static.

So instead, introduce dma_get_numa_cma_area(), which returns the generic
per-NUMA DMA CMA area that cma_alloc() would use. This way, the
nitro_enclaves driver can build its own generic CMA pool per-node
allocation logic on top of those pools.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 include/linux/dma-map-ops.h |  7 +++++++
 kernel/dma/contiguous.c     | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/include/linux/dma-map-ops.h b/include/linux/dma-map-ops.h
index afaf2c12189b8..5f7138ea6899a 100644
--- a/include/linux/dma-map-ops.h
+++ b/include/linux/dma-map-ops.h
@@ -93,6 +93,7 @@ static inline void set_dma_ops(struct device *dev,
 #ifdef CONFIG_DMA_CMA
 struct cma *dev_get_cma_area(struct device *dev);
 struct cma *dma_contiguous_get_area_by_idx(unsigned int idx);
+struct cma *dma_get_numa_cma_area(int nid);
 
 void dma_contiguous_reserve(phys_addr_t addr_limit);
 int __init dma_contiguous_reserve_area(phys_addr_t size, phys_addr_t base,
@@ -115,6 +116,12 @@ static inline struct cma *dma_contiguous_get_area_by_idx(unsigned int idx)
 {
 	return NULL;
 }
+
+static inline struct cma *dma_get_numa_cma_area(int nid)
+{
+	return NULL;
+}
+
 static inline void dma_contiguous_reserve(phys_addr_t limit)
 {
 }
diff --git a/kernel/dma/contiguous.c b/kernel/dma/contiguous.c
index 90937dc7c627d..363b84284d087 100644
--- a/kernel/dma/contiguous.c
+++ b/kernel/dma/contiguous.c
@@ -443,6 +443,42 @@ struct page *dma_alloc_contiguous(struct device *dev, size_t size, gfp_t gfp)
 	return cma_alloc_aligned(dma_contiguous_default_area, size, gfp);
 }
 
+/**
+ * dma_get_numa_cma_area() - get the CMA area a NUMA node can allocate from
+ * @nid: NUMA node id
+ *
+ * dma_alloc_contiguous() picks a per-node CMA area from dev_to_node(dev),
+ * which a caller cannot override on a per-request basis. This accessor
+ * returns the area a specific node can allocate from, so a caller can
+ * cma_alloc() from the node it wants.
+ *
+ * A per-node area comes from numa_cma= or cma_pernuma=, or from
+ * CONFIG_CMA_SIZE_PERNUMA, which reserves the default area's size on every
+ * online node on top of the default area itself, all three of which are
+ * CONFIG_DMA_NUMA_CMA. Where there is one online node this returns the default
+ * area, so a host that reserved only cma= reaches its reservation through this
+ * accessor as well, and that holds however CONFIG_DMA_NUMA_CMA is set.
+ *
+ * Return: the area @nid can allocate from, or NULL if @nid has none or is not
+ * a valid node (including NUMA_NO_NODE).
+ */
+struct cma *dma_get_numa_cma_area(int nid)
+{
+	if ((unsigned int)nid >= MAX_NUMNODES)
+		return NULL;
+
+#ifdef CONFIG_DMA_NUMA_CMA
+	if (dma_contiguous_numa_area[nid])
+		return dma_contiguous_numa_area[nid];
+#endif
+
+	if (nr_online_nodes == 1 && node_online(nid))
+		return dma_contiguous_default_area;
+
+	return NULL;
+}
+EXPORT_SYMBOL_GPL(dma_get_numa_cma_area);
+
 /**
  * dma_free_contiguous() - release allocated pages
  * @dev:   Pointer to device for which the pages were allocated.
-- 
2.47.1

[PATCH 06/40] nitro_enclaves: Add NE_ALLOC_MEMORY ioctl

[Alexander Graf]

When donating memory to the hypervisor for Enclave or Sandbox creation,
the hypervisor only provides a limited number of slots of physically
contiguous memory regions that make up the full slot memory donation.

So far, we used hugetlbfs and mmap based memory donation. But that
scheme has a few issues. It means that user space retains access to
memory that is actually donated. Access to that memory turns into an
exception. It also means that any donating allocation ends up with
page tables that back them, even when nobody really should use them.

So we copy the guest_memfd flow of allocating memory, but not exposing
it as mapping to user space by default. EIF writes are a simple fd write.
Each allocation gets its own fd. And for direct mode, where memory really
is supposed to be accessible to the parent, we still allow mmap.

To reduce the number of separate memory regions these fds span, we use
CMA to allocate very large physically contiguous chunks.

Select CMA and DMA_CMA to always have the foundations available.
DMA_NUMA_CMA stays optional.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/Kconfig       |  18 +
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 730 +++++++++++++++++++++-
 drivers/virt/nitro_enclaves/ne_misc_dev.h |  61 ++
 include/uapi/linux/nitro_enclaves.h       |  90 ++-
 4 files changed, 895 insertions(+), 4 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/Kconfig b/drivers/virt/nitro_enclaves/Kconfig
index dc4d25c26256e..7b2456cdbb9fb 100644
--- a/drivers/virt/nitro_enclaves/Kconfig
+++ b/drivers/virt/nitro_enclaves/Kconfig
@@ -8,7 +8,9 @@
 config NITRO_ENCLAVES
 	tristate "Nitro Enclaves Support"
 	depends on (ARM64 || X86) && HOTPLUG_CPU && PCI && SMP
+	select CMA
+	select DMA_CMA
 	help
 	  This driver consists of support for enclave lifetime management
 	  for Nitro Enclaves (NE).
 
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 66bcaec8d3199..56ca1a9d99d41 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -9,12 +9,17 @@
  */
 
 #include <linux/anon_inodes.h>
+#include <linux/bsearch.h>
 #include <linux/build_bug.h>
 #include <linux/capability.h>
+#include <linux/cma.h>
 #include <linux/cpu.h>
 #include <linux/cpumask.h>
 #include <linux/device.h>
+#include <linux/dma-map-ops.h>
 #include <linux/file.h>
+#include <linux/fs.h>
+#include <linux/highmem.h>
 #include <linux/hugetlb.h>
 #include <linux/ioctl.h>
 #include <linux/kernel.h>
@@ -33,13 +38,16 @@
 #include <linux/numa.h>
 #include <linux/pci.h>
 #include <linux/poll.h>
+#include <linux/preempt.h>
 #include <linux/range.h>
+#include <linux/sched/mm.h>
 #include <linux/slab.h>
 #include <linux/stddef.h>
 #include <linux/string.h>
 #include <linux/sysfs.h>
 #include <linux/types.h>
 #include <linux/uaccess.h>
+#include <linux/uio.h>
 #include <uapi/linux/vm_sockets.h>
 
 #include "ne_misc_dev.h"
@@ -86,6 +94,16 @@ static_assert(NE_NUMA_ANY == NUMA_NO_NODE);
  */
 #define NE_MIN_MEM_REGION_SIZE	(2 * 1024UL * 1024UL)
 
+/**
+ * NE_ZERO_SPAN_PAGES - The most one span of the zeroing walk covers.
+ */
+#define NE_ZERO_SPAN_PAGES	(1024UL)
+
+/**
+ * NE_WRITE_SPAN_PAGES - The most one span of the write walk copies.
+ */
+#define NE_WRITE_SPAN_PAGES	(1024UL)
+
 /**
  * NE_PARENT_VM_CID - The CID for the vsock device of the primary / parent VM.
  */
@@ -93,6 +111,16 @@ static_assert(NE_NUMA_ANY == NUMA_NO_NODE);
 
 static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg);
 
+/*
+ * Handed to an enclave as ne_enclave.id once it has a slot, so that an
+ * NE_ALLOC_MEMORY allocation names the enclave it was made for without holding
+ * a reference on it. A 64-bit counter does not come round again, so an id names
+ * one enclave for the lifetime of the driver.
+ */
+static atomic64_t ne_enclave_seq = ATOMIC64_INIT(0);
+
+static int ne_alloc_memory_ioctl(struct ne_enclave *ne_enclave,
+				 struct ne_alloc_memory_args *args);
 static const struct file_operations ne_fops = {
 	.owner		= THIS_MODULE,
 	.llseek		= noop_llseek,
@@ -1731,7 +1759,7 @@ static int ne_start_enclave_ioctl(struct ne_enclave *ne_enclave,
  * Context: Process context. This function is called with ne_devs.lock held
  *	    for reading, on an enclave that is not detached.
  * Return:
- * * 0 on success.
+ * * Ioctl result (e.g. memory allocation file descriptor) on success.
  * * Negative return value on failure.
  */
 static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
@@ -2041,6 +2069,39 @@ static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned lo
 		return 0;
 	}
 
+	case NE_EA_IOCTL(NE_ALLOC_MEMORY): {
+		struct ne_alloc_memory_args __user *uargs = (void __user *)arg;
+		unsigned int usize = _IOC_SIZE(cmd);
+		struct ne_alloc_memory_args args;
+		int rc;
+
+		if (usize < offsetofend(struct ne_alloc_memory_args, flags))
+			return -EINVAL;
+
+		rc = copy_struct_from_user(&args, sizeof(args), uargs, usize);
+		if (rc)
+			return rc;
+
+		rc = mutex_lock_killable(&ne_enclave->enclave_info_mutex);
+		if (rc)
+			return rc;
+
+		if (ne_enclave->state != NE_STATE_INIT) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Enclave is not in init state\n");
+
+			mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+			return -NE_ERR_NOT_IN_INIT_STATE;
+		}
+
+		rc = ne_alloc_memory_ioctl(ne_enclave, &args);
+
+		mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+		return rc;
+	}
+
 	default:
 		return -ENOTTY;
 	}
@@ -2427,6 +2488,8 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 __user *slot_ui
 		goto put_file;
 	}
 
+	ne_enclave->id = atomic64_inc_return(&ne_enclave_seq);
+
 	init_waitqueue_head(&ne_enclave->eventq);
 	ne_enclave->has_event = false;
 	mutex_init(&ne_enclave->enclave_info_mutex);
@@ -2516,6 +2579,671 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 	return 0;
 }
 
+/*
+ * Physically-contiguous memory for enclaves.
+ *
+ * An enclave's memory is handed to the hypervisor in physically contiguous
+ * 2 MiB (or larger) chunks. NE_ALLOC_MEMORY allocates contiguous memory from a
+ * host NUMA node's generic CMA area and returns a file descriptor for it. A
+ * per-node area comes from numa_cma= or cma_pernuma=, or from
+ * CONFIG_CMA_SIZE_PERNUMA, which reserves the cma= default area's size on
+ * every online node on top of that area itself, none of which a kernel has
+ * without CONFIG_DMA_NUMA_CMA. On a host with one online node the accessor
+ * returns the default area, so cma= on its own is enough there, and that much
+ * needs CONFIG_DMA_CMA alone.
+ *
+ * A request the node's area cannot serve fails. Anyone who can open
+ * /dev/nitro_enclaves can create an enclave and ask it for memory, so
+ * serving such a request out of movable memory would take that memory from
+ * the rest of the host and leave the reservation bounding nothing.
+ *
+ * A request larger than one contiguous run is satisfied by several runs, which
+ * the driver keeps on a list of its own so that an enclave can be given every
+ * one of them.
+ *
+ * An allocation is never mapped, and the enclave image reaches it through
+ * write().
+ *
+ * The runs belong to the file the allocation is handed out on. Zeroing a fresh
+ * run and a write into an allocation go through ne_walk_run(), so that neither
+ * can run for as long as the allocation is large; the release path walks the
+ * runs with a preemption point of its own.
+ */
+
+/**
+ * ne_contig_free_run() - Give a contiguous run of pages back.
+ * @cma:	CMA area the run came from.
+ * @page:	First page of the run.
+ * @nr_pages:	Number of pages in the run.
+ */
+void ne_contig_free_run(struct cma *cma, struct page *page,
+			unsigned long nr_pages)
+{
+	cma_release(cma, page, nr_pages);
+}
+
+static void ne_contig_free_alloc(struct ne_contig_alloc *alloc)
+{
+	ne_contig_free_run(alloc->cma, alloc->page, alloc->nr_pages);
+}
+
+/*
+ * Free every run of an allocation. A 512 GiB allocation made of 2 MiB runs is a
+ * quarter of a million entries and cma_release() does not reschedule, so this
+ * walk needs a preemption point of its own. It runs from ->release, which can
+ * land on a workqueue.
+ */
+static void ne_contig_free_allocs(struct ne_contig_mem *mem)
+{
+	unsigned int i;
+
+	for (i = 0; i < mem->nr_allocs; i++) {
+		cond_resched();
+
+		ne_contig_free_alloc(&mem->allocs[i]);
+	}
+
+	mem->nr_allocs = 0;
+	kvfree(mem->allocs);
+	mem->allocs = NULL;
+}
+
+/**
+ * typedef ne_span_fn - What to do with one span of a contiguous run.
+ * @arg:	Context the caller of ne_walk_run() passed.
+ * @page:	Page the span starts in.
+ * @off:	Byte offset of the span inside @page, less than PAGE_SIZE.
+ * @len:	Length of the span in bytes.
+ *
+ * Return: 0 on success, negative return value on failure.
+ */
+typedef int (*ne_span_fn)(void *arg, struct page *page, u64 off, u64 len);
+
+/**
+ * ne_walk_run() - Hand @fn one span of a run at a time.
+ * @page:	First page of the run.
+ * @off:	Byte offset into the run to start at.
+ * @len:	Number of bytes to cover.
+ * @span_pages:	The most one call to @fn covers.
+ * @fn:		What to do with a span.
+ * @arg:	Context for @fn.
+ *
+ * A run can be as large as the whole allocation and none of the things done to
+ * one reschedules on its own, so this is where the driver gets its
+ * preemption points and its bail-out on a fatal signal. @fn gets at most
+ * @span_pages worth of bytes at a time and nothing else bounds it, which is
+ * what keeps the work between two cond_resched() calls from growing with the
+ * allocation. An unaligned @off spreads those bytes over one page more.
+ *
+ * Return: 0 on success, -EINTR when a fatal signal arrived, or whatever @fn
+ *	   returned.
+ */
+static int ne_walk_run(struct page *page, u64 off, u64 len,
+		       unsigned long span_pages, ne_span_fn fn, void *arg)
+{
+	u64 span = (u64)span_pages << PAGE_SHIFT;
+	u64 done = 0;
+
+	while (done < len) {
+		u64 pos = off + done;
+		u64 this = min(len - done, span);
+		int rc;
+
+		rc = fn(arg, page + (pos >> PAGE_SHIFT), pos & ~PAGE_MASK,
+			this);
+		if (rc)
+			return rc;
+
+		done += this;
+
+		if (fatal_signal_pending(current))
+			return -EINTR;
+
+		cond_resched();
+	}
+
+	return 0;
+}
+
+/**
+ * ne_contig_alloc_cmp() - Order a byte offset against one run.
+ * @key:	Byte offset into the allocation, as a u64.
+ * @elem:	The run.
+ *
+ * Return: 0 when @key falls inside @elem, negative when it lies before it and
+ *	   positive when it lies after it.
+ */
+static int ne_contig_alloc_cmp(const void *key, const void *elem)
+{
+	const struct ne_contig_alloc *alloc = elem;
+	u64 off = *(const u64 *)key;
+
+	if (off < alloc->off)
+		return -1;
+
+	if (off >= alloc->off + ((u64)alloc->nr_pages << PAGE_SHIFT))
+		return 1;
+
+	return 0;
+}
+
+/**
+ * ne_find_alloc() - The run a byte offset falls in.
+ * @mem:	The allocation.
+ * @off:	Byte offset into the allocation.
+ *
+ * The runs are ordered by offset and cover the allocation back to back, so a
+ * binary search finds one. A 512 GiB allocation made of 2 MiB runs is a quarter
+ * of a million of them and a fault resolves an offset every time it is taken,
+ * which is why this is not a walk. The array is not mutated after the file is
+ * published, so this needs no lock.
+ *
+ * Return: the run, or NULL when @off is past the end of the allocation.
+ */
+static struct ne_contig_alloc *ne_find_alloc(struct ne_contig_mem *mem, u64 off)
+{
+	return bsearch(&off, mem->allocs, mem->nr_allocs, sizeof(*mem->allocs),
+		       ne_contig_alloc_cmp);
+}
+
+/**
+ * ne_walk_allocs() - Hand @fn the part of an allocation at @off, run by run.
+ * @mem:	The allocation.
+ * @off:	Byte offset into the allocation to start at.
+ * @len:	Number of bytes to cover, which has to fit inside the allocation.
+ * @span_pages:	The most one call to @fn covers.
+ * @fn:		What to do with a span.
+ * @arg:	Context for @fn.
+ *
+ * The runs are handed over in order and back to back, so a caller that only
+ * advances an output cursor sees one flat region.
+ *
+ * Return: 0 on success, negative return value on failure.
+ */
+static int ne_walk_allocs(struct ne_contig_mem *mem, u64 off, u64 len,
+			  unsigned long span_pages, ne_span_fn fn, void *arg)
+{
+	struct ne_contig_alloc *end = mem->allocs + mem->nr_allocs;
+	struct ne_contig_alloc *alloc = ne_find_alloc(mem, off);
+	u64 done = 0;
+
+	if (!alloc)
+		return 0;
+
+	for (; done < len && alloc < end; alloc++) {
+		u64 run = (u64)alloc->nr_pages << PAGE_SHIFT;
+		u64 skip = off > alloc->off ? off - alloc->off : 0;
+		u64 this = min(run - skip, len - done);
+		int rc;
+
+		cond_resched();
+
+		rc = ne_walk_run(alloc->page, skip, this, span_pages, fn, arg);
+		if (rc)
+			return rc;
+
+		done += this;
+	}
+
+	return 0;
+}
+
+static struct cma *ne_numa_cma_area(int nid)
+{
+	struct cma *cma = dma_get_numa_cma_area(nid);
+
+	/*
+	 * A node with no area of its own can serve nothing, however much free
+	 * memory it has, and that is a reservation the host was booted without
+	 * rather than one it has used up. Say so once, since the areas are
+	 * reserved at boot and no node acquires one later.
+	 */
+	if (!cma)
+		pr_warn_once("%s: node %d has no CMA area to serve enclave memory from, reserve one with numa_cma=, cma_pernuma= or cma=\n",
+			     ne_misc_dev.name, nid);
+
+	return cma;
+}
+
+static struct page *ne_alloc_contig_node(unsigned long nr_pages, int nid,
+					 struct cma **out_cma)
+{
+	unsigned int align = ilog2(NE_MIN_MEM_REGION_SIZE >> PAGE_SHIFT);
+	struct cma *cma = ne_numa_cma_area(nid);
+	struct page *page;
+
+	if (!cma)
+		return NULL;
+
+	/*
+	 * The chunk loop in ne_alloc_memory_ioctl() walks sizes downwards and
+	 * expects most of these to fail, so keep a failure cheap and quiet.
+	 */
+	page = cma_alloc(cma, nr_pages, align, true);
+	if (page)
+		*out_cma = cma;
+
+	return page;
+}
+
+/*
+ * Zero one span of a freshly allocated run. CMA pages can carry residual data
+ * from earlier movable use, so nothing reaches an enclave until this has run
+ * over all of it.
+ *
+ * The run is on its way to an enclave and may be mapped into the caller, so this
+ * is clear_user_highpages()' job rather than clear_pages()'. It maps a highmem
+ * page before clearing it and falls back to a batched clear_pages() where there
+ * is no highmem, which is every configuration this driver builds for bar i386
+ * with highmem.
+ * The user address it is given is zero. Neither x86 nor arm64 supplies its own
+ * clear_user_page() or clear_user_highpage(), so the generic clear_user_page()
+ * runs and does not read the address.
+ *
+ * @off is zero and @len a multiple of PAGE_SIZE here, because a run is walked
+ * whole and the span size is a page count.
+ */
+static int ne_zero_span(void *arg, struct page *page, u64 off, u64 len)
+{
+	clear_user_highpages(page, 0, len >> PAGE_SHIFT);
+
+	return 0;
+}
+
+/*
+ * How much of a run one span of the zeroing walk covers. clear_user_highpages()
+ * does not reschedule, so under a cooperatively scheduled model a caller has to
+ * keep a batch within PROCESS_PAGES_NON_PREEMPT_BATCH, which is a single page on
+ * every architecture that does not supply its own clear_pages(). Under
+ * PREEMPT_DYNAMIC the model comes from the command line, so this asks at run
+ * time rather than testing a config symbol. Either way NE_ZERO_SPAN_PAGES stays
+ * the upper bound, so a fatal signal is still noticed within a few megabytes of
+ * clearing.
+ */
+static unsigned long ne_zero_span_pages(void)
+{
+	if (preempt_model_preemptible())
+		return NE_ZERO_SPAN_PAGES;
+
+	return min_t(unsigned long, NE_ZERO_SPAN_PAGES,
+		     PROCESS_PAGES_NON_PREEMPT_BATCH);
+}
+
+/**
+ * ne_alloc_contig() - Allocate a physically-contiguous, zeroed run of pages.
+ * @nr_pages:	Number of pages to allocate.
+ * @nid:	NUMA node to allocate from.
+ * @out_cma:	Set to the CMA area the run came from.
+ *
+ * The run is drawn from @nid's generic CMA area and from nowhere else: an
+ * enclave's memory has to sit on the node the caller asked for, so a run from
+ * another node is no more use to it than no run at all, and a run from outside
+ * the reservation is memory the host budgeted for something else.
+ *
+ * Zeroing a large run takes a while, so a fatal signal abandons it rather than
+ * making the task wait to be killed.
+ *
+ * Return: the first page of the run, NULL on allocation failure, or
+ *	   ERR_PTR(-EINTR) when a fatal signal arrived while zeroing.
+ */
+struct page *ne_alloc_contig(unsigned long nr_pages, int nid,
+			     struct cma **out_cma)
+{
+	struct page *page;
+	int rc;
+
+	*out_cma = NULL;
+
+	page = ne_alloc_contig_node(nr_pages, nid, out_cma);
+	if (!page)
+		return NULL;
+
+	rc = ne_walk_run(page, 0, (u64)nr_pages << PAGE_SHIFT,
+			 ne_zero_span_pages(), ne_zero_span, NULL);
+	if (rc) {
+		struct ne_contig_alloc discard = {
+			.page = page,
+			.nr_pages = nr_pages,
+			.cma = *out_cma,
+		};
+
+		ne_contig_free_alloc(&discard);
+
+		return ERR_PTR(rc);
+	}
+
+	return page;
+}
+
+/**
+ * struct ne_write_ctx - Where the write walk has got to.
+ * @iter:	Where the bytes come from.
+ * @copied:	Bytes copied so far.
+ */
+struct ne_write_ctx {
+	struct iov_iter	*iter;
+	size_t		copied;
+};
+
+/*
+ * Copy one span of a run out of the iterator. Where the kernel has no highmem a
+ * run is contiguous in the direct map, so a span goes in one call. Where it may
+ * have some, page_address() is not enough on its own, so every page goes through
+ * copy_page_from_iter(), which maps it for the copy.
+ */
+static int ne_write_span(void *arg, struct page *page, u64 off, u64 len)
+{
+	struct ne_write_ctx *ctx = arg;
+	size_t n = 0;
+
+	if (!IS_ENABLED(CONFIG_HIGHMEM)) {
+		n = copy_from_iter(page_address(page) + off, len, ctx->iter);
+	} else {
+		while (n < len) {
+			u64 pos = off + n;
+			size_t this;
+
+			this = min_t(size_t, len - n,
+				     PAGE_SIZE - (pos & ~PAGE_MASK));
+			this = copy_page_from_iter(page + (pos >> PAGE_SHIFT),
+						   pos & ~PAGE_MASK, this,
+						   ctx->iter);
+			if (!this)
+				break;
+
+			n += this;
+		}
+	}
+
+	ctx->copied += n;
+
+	return n < len ? -EFAULT : 0;
+}
+
+/**
+ * ne_contig_write_iter() - Place bytes in an allocation.
+ * @iocb:	Holds the allocation's file and the offset to write at.
+ * @from:	Where the bytes come from.
+ *
+ * An enclave image has to be in enclave memory before the enclave starts, and a
+ * private allocation is never mapped, so this is how it gets there. The write
+ * has to fit inside the allocation: a partly written image is of no use to
+ * anyone, so a write running past the end is refused rather than shortened. The
+ * allocation has a fixed maximum offset and no notion of free space, so that
+ * refusal is -EFBIG, the same as generic_write_check_limits() gives for a write
+ * beyond s_maxbytes. There is no RLIMIT_FSIZE in it, so no SIGXFSZ either.
+ *
+ * Return: the number of bytes written, or a negative return value on failure.
+ */
+static ssize_t ne_contig_write_iter(struct kiocb *iocb, struct iov_iter *from)
+{
+	struct ne_contig_mem *mem = iocb->ki_filp->private_data;
+	struct ne_write_ctx ctx = { .iter = from };
+	size_t count = iov_iter_count(from);
+	u64 pos = iocb->ki_pos;
+	int rc;
+
+	if (!count)
+		return 0;
+
+	if (pos >= mem->size || count > mem->size - pos)
+		return -EFBIG;
+
+	rc = ne_walk_allocs(mem, pos, count, NE_WRITE_SPAN_PAGES, ne_write_span,
+			    &ctx);
+
+	if (ctx.copied) {
+		iocb->ki_pos = pos + ctx.copied;
+
+		return ctx.copied;
+	}
+
+	return rc;
+}
+
+/**
+ * ne_contig_release() - Free an allocation once nothing holds it any more.
+ * @inode:	The allocation's inode, unused.
+ * @file:	The allocation's file.
+ *
+ * The last reference on the file is what gets here, and a mapping holds one, so
+ * by this point no page table entry points at the runs.
+ *
+ * Return: 0.
+ */
+static int ne_contig_release(struct inode *inode, struct file *file)
+{
+	struct ne_contig_mem *mem = file->private_data;
+
+	ne_contig_free_allocs(mem);
+
+	if (mem->nr_pinned)
+		atomic64_sub(mem->nr_pinned, &mem->mm->pinned_vm);
+
+	mmdrop(mem->mm);
+
+	kfree(mem);
+
+	return 0;
+}
+
+static const struct file_operations ne_contig_fops = {
+	.owner			= THIS_MODULE,
+	.llseek			= generic_file_llseek,
+	.write_iter		= ne_contig_write_iter,
+	.release		= ne_contig_release,
+};
+
+/**
+ * ne_contig_create_file() - Build a file for an allocation.
+ * @mem:	The allocation, whose runs the file takes over on success.
+ *
+ * On success the file owns @mem, and the last fput() runs ne_contig_release()
+ * to free the runs. On failure @mem is left to the caller.
+ *
+ * Return: the new file, or an ERR_PTR.
+ */
+static struct file *ne_contig_create_file(struct ne_contig_mem *mem)
+{
+	struct file *file;
+
+	/*
+	 * An inode of its own is what carries the size, so that fstat() reports
+	 * it and it bounds the offsets a mapping or a write may use. It also
+	 * gets a security label rather than staying S_PRIVATE, which is what
+	 * user space reaching an inode through a descriptor wants.
+	 */
+	file = anon_inode_create_getfile("ne-mem", &ne_contig_fops, mem, O_RDWR,
+					 NULL);
+	if (IS_ERR(file))
+		return file;
+
+	i_size_write(file_inode(file), mem->size);
+	inode_set_bytes(file_inode(file), mem->size);
+
+	/*
+	 * A file that came through open() would have these from
+	 * do_dentry_open(); one built here does not, and FMODE_PWRITE is what
+	 * lets a positional write reach an allocation that is never mapped.
+	 * FMODE_PREAD gates pread() rather than read(), and there is no read
+	 * operation here for either to reach, so all it buys is pread()
+	 * failing with EINVAL like read() does instead of with ESPIPE.
+	 * Nothing else has this file yet, so there is nobody to see the
+	 * intermediate state.
+	 */
+	file->f_mode |= FMODE_PREAD | FMODE_PWRITE;
+	file->f_flags |= O_LARGEFILE;
+
+	return file;
+}
+
+/**
+ * ne_alloc_memory_ioctl() - Allocate contiguous memory for an enclave.
+ * @ne_enclave:	Enclave the allocation is made for.
+ * @args:	Requested size, node and flags.
+ *
+ * Context: Process context. This function is called with the enclave
+ *	    information mutex held.
+ * Return: a file descriptor on success, negative return value on failure.
+ */
+static int ne_alloc_memory_ioctl(struct ne_enclave *ne_enclave,
+				 struct ne_alloc_memory_args *args)
+{
+	struct ne_contig_mem *mem;
+	unsigned long remaining;
+	u64 pinned, lock_limit;
+	unsigned long nr_pages;
+	struct file *file;
+	int fd, rc;
+
+	if (ne_enclave->mm != current->mm)
+		return -EIO;
+
+	if (args->flags >= NE_ALLOC_MEMORY_MAX_FLAG_VAL)
+		return -NE_ERR_INVALID_FLAG_VALUE;
+
+	if (!args->size || (args->size & (NE_MIN_MEM_REGION_SIZE - 1)))
+		return -EINVAL;
+
+	if (args->size > ULONG_MAX)
+		return -EINVAL;
+
+	if (args->nid >= MAX_NUMNODES || !node_state(args->nid, N_MEMORY))
+		return -EINVAL;
+
+	/*
+	 * An enclave whose cores all sit on one node takes its memory from that
+	 * node, which is what ne_sanity_check_user_mem_region() asks of a
+	 * hugetlbfs region. An enclave built from a CPU pool that spans nodes
+	 * has no node of its own and may take memory from any of them. Refuse
+	 * here, so that a request naming the wrong node does not spend minutes
+	 * zeroing first.
+	 */
+	if (ne_enclave->numa_node != NUMA_NO_NODE &&
+	    ne_enclave->numa_node != (int)args->nid)
+		return -NE_ERR_MEM_DIFFERENT_NUMA_NODE;
+
+	/*
+	 * A node without an area serves nothing, and no area serves more than
+	 * the host has, so refuse those requests before they are charged and
+	 * before the run array is sized from them.
+	 */
+	if (!ne_numa_cma_area(args->nid) ||
+	    args->size > (u64)totalram_pages() << PAGE_SHIFT)
+		return -ENOMEM;
+
+	nr_pages = args->size >> PAGE_SHIFT;
+
+	/*
+	 * Account the pages against RLIMIT_MEMLOCK, the way a driver that pins
+	 * pages for a device accounts them. pinned_vm is an atomic64, so it can
+	 * be uncharged from ne_contig_release() under whatever locks that path
+	 * happens to hold.
+	 */
+	lock_limit = rlimit(RLIMIT_MEMLOCK) >> PAGE_SHIFT;
+	pinned = atomic64_add_return(nr_pages, &current->mm->pinned_vm);
+	if (pinned > lock_limit && !capable(CAP_IPC_LOCK)) {
+		atomic64_sub(nr_pages, &current->mm->pinned_vm);
+
+		return -ENOMEM;
+	}
+
+	mem = kzalloc_obj(*mem);
+	if (!mem) {
+		atomic64_sub(nr_pages, &current->mm->pinned_vm);
+
+		return -ENOMEM;
+	}
+
+	/*
+	 * A run is at least NE_MIN_MEM_REGION_SIZE, so that bounds how many of
+	 * them an allocation can be made of and the array can be sized for the
+	 * worst case here. At 512 GiB of 2 MiB runs it is 8 MiB, next to nothing
+	 * beside the memory about to be allocated, and it keeps the allocation
+	 * loop from having to grow anything.
+	 */
+	mem->allocs = kvmalloc_array(args->size / NE_MIN_MEM_REGION_SIZE,
+				     sizeof(*mem->allocs),
+				     GFP_KERNEL | __GFP_NOWARN);
+	if (!mem->allocs) {
+		atomic64_sub(nr_pages, &current->mm->pinned_vm);
+		kfree(mem);
+
+		return -ENOMEM;
+	}
+
+	mmgrab(current->mm);
+	mem->mm = current->mm;
+	mem->nr_pinned = nr_pages;
+	mem->size = args->size;
+	mem->enclave_id = ne_enclave->id;
+
+	/*
+	 * Take the runs largest power-of-two chunk first, down to the 2 MiB
+	 * minimum. cma_alloc() is given the 2 MiB alignment the hypervisor asks
+	 * for whatever the size, so every run meets it.
+	 */
+	remaining = args->size;
+	while (remaining) {
+		struct ne_contig_alloc *alloc;
+		struct cma *cma = NULL;
+		struct page *page = NULL;
+		unsigned long chunk;
+
+		for (chunk = remaining; chunk >= NE_MIN_MEM_REGION_SIZE;
+		     chunk >>= 1) {
+			chunk = rounddown_pow_of_two(chunk);
+			page = ne_alloc_contig(chunk >> PAGE_SHIFT, args->nid,
+					       &cma);
+			if (IS_ERR(page)) {
+				rc = PTR_ERR(page);
+				goto err_free;
+			}
+			if (page)
+				break;
+		}
+		if (!page) {
+			rc = -ENOMEM;
+			goto err_free;
+		}
+
+		alloc = &mem->allocs[mem->nr_allocs++];
+		alloc->page = page;
+		alloc->nr_pages = chunk >> PAGE_SHIFT;
+		alloc->cma = cma;
+		alloc->off = args->size - remaining;
+
+		remaining -= chunk;
+		cond_resched();
+	}
+
+	fd = get_unused_fd_flags(O_CLOEXEC);
+	if (fd < 0) {
+		rc = fd;
+		goto err_free;
+	}
+
+	file = ne_contig_create_file(mem);
+	if (IS_ERR(file)) {
+		rc = PTR_ERR(file);
+		goto err_put_fd;
+	}
+
+	fd_install(fd, file);
+
+	return fd;
+
+err_put_fd:
+	put_unused_fd(fd);
+err_free:
+	ne_contig_free_allocs(mem);
+	atomic64_sub(nr_pages, &mem->mm->pinned_vm);
+	mmdrop(mem->mm);
+	kfree(mem);
+
+	return rc;
+}
+
 #if defined(CONFIG_NITRO_ENCLAVES_MISC_DEV_TEST)
 #include "ne_misc_dev_test.c"
 #endif
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 483037061eafe..1a76c12815298 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -17,6 +17,48 @@
 
 #include "ne_pci_dev.h"
 
+struct cma;
+
+/**
+ * struct ne_contig_alloc - One physically-contiguous run backing a mapping.
+ * @page:	First page of the run.
+ * @nr_pages:	Number of pages in the run.
+ * @off:	Byte offset of the run inside the allocation.
+ * @cma:	CMA area the run came from.
+ */
+struct ne_contig_alloc {
+	struct page		*page;
+	unsigned long		nr_pages;
+	u64			off;
+	struct cma		*cma;
+};
+
+/**
+ * struct ne_contig_mem - Backing state for one contiguous allocation.
+ * @allocs:	Contiguous runs backing the allocation, ordered by offset.
+ * @nr_allocs:	Number of runs in @allocs.
+ * @size:	Size of the allocation in bytes, mirrored in i_size.
+ * @enclave_id:	Id of the enclave the allocation was made for.
+ * @mm:		The mm charged for the pinned pages, held via mmgrab().
+ * @nr_pinned:	Pages charged to @mm->pinned_vm.
+ *
+ * The runs hang off the file, and a mapping's page table entries cannot outlive
+ * them because a VMA holds vm_file for as long as it exists.
+ *
+ * The array is built before the file is published and not mutated afterwards, so
+ * any holder of a reference can read it without a lock of its own, which is what
+ * lets a page fault resolve an offset to a run. @size is set in the same place
+ * and is read the same way.
+ */
+struct ne_contig_mem {
+	struct ne_contig_alloc	*allocs;
+	unsigned int		nr_allocs;
+	u64			size;
+	u64			enclave_id;
+	struct mm_struct	*mm;
+	unsigned long		nr_pinned;
+};
+
 /**
  * struct ne_mem_region - Entry in the enclave user space memory regions list.
  * @mem_region_list_entry:	Entry in the list of enclave memory regions.
@@ -44,6 +86,9 @@ struct ne_mem_region {
  *				the enclave process via the poll function.
  * @has_event:			Variable used to determine if the out-of-band event
  *				was triggered.
+ * @id:				Id of this enclave, unique for the lifetime of the
+ *				driver, recorded in an %NE_ALLOC_MEMORY allocation
+ *				as the enclave it was made for.
  * @max_mem_regions:		The maximum number of memory regions that can be
  *				handled by the hypervisor.
  * @mem_regions_list:		Enclave user space memory regions list.
@@ -72,6 +117,7 @@ struct ne_enclave {
 	struct list_head	enclave_list_entry;
 	wait_queue_head_t	eventq;
 	bool			has_event;
+	u64			id;
 	u64			max_mem_regions;
 	struct list_head	mem_regions_list;
 	u64			mem_size;
@@ -112,6 +158,21 @@ struct ne_devs {
 	struct rw_semaphore	lock;
 };
 
+/**
+ * ne_alloc_contig() - Allocate a physically-contiguous, zeroed run of pages.
+ * @nr_pages:	Number of pages to allocate.
+ * @nid:	NUMA node to allocate from.
+ * @out_cma:	Set to the CMA area the run came from.
+ *
+ * Return: the first page of the run, NULL on allocation failure, or
+ *	   ERR_PTR(-EINTR) when a fatal signal arrived while zeroing.
+ */
+struct page *ne_alloc_contig(unsigned long nr_pages, int nid,
+			     struct cma **out_cma);
+
+void ne_contig_free_run(struct cma *cma, struct page *page,
+			unsigned long nr_pages);
+
 /* Nitro Enclaves (NE) data structure for keeping refs to the NE misc and PCI devices. */
 extern struct ne_devs ne_devs;
 
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index e266b8a611c92..a2919565bc883 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -209,6 +209,76 @@ struct ne_add_any_vcpu_args {
  */
 #define NE_SET_USER_MEMORY_REGION	_IOW(0xAE, 0x23, struct ne_user_memory_region)
 
+/**
+ * struct ne_alloc_memory_args - Request for physically contiguous enclave
+ *				 memory (in).
+ * @size:	Number of bytes to allocate, a multiple of 2 MiB.
+ * @nid:	Host NUMA node to allocate from. One node, named; unlike
+ *		%NE_ADD_ANY_VCPU it does not take %NE_NUMA_ANY.
+ * @flags:	%NE_MEMORY_PRIVATE.
+ */
+struct ne_alloc_memory_args {
+	__u64	size;
+	__u32	nid;
+	__u32	flags;
+};
+
+/**
+ * NE_ALLOC_MEMORY - The command is used to allocate physically contiguous
+ *		     memory to give to an enclave, drawn from one host NUMA
+ *		     node. The argument is input only; on success the return
+ *		     value is a file descriptor for the allocation.
+ *		     Writing to the descriptor places bytes in the allocation,
+ *		     which is how an enclave image gets there.
+ *		     O_APPEND has no effect. There is no read operation, so
+ *		     read() and pread() both fail with EINVAL and a caller that
+ *		     wants to check what it wrote keeps its own copy.
+ *		     A private allocation is never mapped, so an enclave image
+ *		     goes into one through write() or pwrite() and not through
+ *		     a copy into a mapping, and every write is a call the
+ *		     kernel handles and can act on or refuse.
+ *		     The allocation lives until the descriptor is gone.
+ *		     The ioctl can be invoked on the enclave fd, before an
+ *		     enclave is started.
+ *
+ * Context: Process context.
+ * Return:
+ * * >= 0			- File descriptor for the allocation.
+ * * -1				- There was a failure in the ioctl logic.
+ * On failure, errno is set to:
+ * * EFAULT			- Argument read failure.
+ * * EINVAL			- The argument is shorter than @flags, or @size
+ *				  is zero, not a multiple of 2 MiB or larger than
+ *				  this kernel can represent, or @nid is not a node
+ *				  with memory.
+ * * EIO			- Current task mm is not the same as the one
+ *				  that created the enclave.
+ * * E2BIG			- The argument is longer than this kernel knows
+ *				  about and the excess is not all zero.
+ * * EINTR			- A fatal signal arrived while zeroing.
+ * * ENODEV			- The kernel has no contiguous allocator.
+ * * ENOMEM			- No contiguous run of the requested size is
+ *				  available on that node, or %RLIMIT_MEMLOCK
+ *				  would be exceeded.
+ * * NE_ERR_NOT_IN_INIT_STATE	- The enclave is not in init state (init =
+ *				  before being started).
+ * * NE_ERR_INVALID_FLAG_VALUE	- The value of the provided flag is invalid.
+ * * NE_ERR_MEM_DIFFERENT_NUMA_NODE
+ *				- @nid is not the node the enclave's CPUs are on.
+ *				  Not returned for an enclave whose CPU pool spans
+ *				  nodes.
+ * * Error codes from get_unused_fd_flags() and anon_inode_create_getfile().
+ *
+ * Writing to the descriptor sets errno to:
+ * * EFBIG			- The write does not fit inside the allocation.
+ * * EFAULT			- Part of the buffer is not readable. Any bytes
+ *				  copied before that are reported as a short
+ *				  write, so a caller that has to know how far it
+ *				  got looks at the return value. A fatal signal
+ *				  part way through is reported the same way.
+ */
+#define NE_ALLOC_MEMORY			_IOW(0xAE, 0x2E, struct ne_alloc_memory_args)
+
 /**
  * NE_START_ENCLAVE - The command is used to trigger enclave start after the
  *		      enclave resources, such as memory and CPU, have been set.
@@ -282,9 +352,10 @@ struct ne_add_any_vcpu_args {
 #define NE_ERR_MEM_NOT_HUGE_PAGE		(263)
 /**
  * NE_ERR_MEM_DIFFERENT_NUMA_NODE - The user space memory region is backed by
- *				    pages from different NUMA nodes than the CPUs.
- *				    Not returned for an enclave whose CPU pool
- *				    spans nodes.
+ *				    pages from different NUMA nodes than the CPUs,
+ *				    or a requested node is not the one the CPUs
+ *				    are on. Not returned for an enclave whose CPU
+ *				    pool spans nodes.
  */
 #define NE_ERR_MEM_DIFFERENT_NUMA_NODE		(264)
 /**
@@ -391,6 +462,19 @@ struct ne_user_memory_region {
 	__u64	userspace_addr;
 };
 
+/**
+ * DOC: Enclave memory allocation flags
+ */
+
+/**
+ * NE_MEMORY_PRIVATE - An allocation the parent has no address for, because it
+ *		       cannot be mapped. Writing to the descriptor is the only
+ *		       way to put an enclave image in one.
+ */
+#define NE_MEMORY_PRIVATE		(0x00)
+
+#define NE_ALLOC_MEMORY_MAX_FLAG_VAL	(0x01)
+
 /**
  * DOC: Enclave start info flags
  */
-- 
2.47.1

[PATCH 07/40] nitro_enclaves: Add NE_MEMORY_SHARED

[Alexander Graf]

In preparation to support Nitro Sandboxes, whose parent keeps using some
of the memory it hands over, add NE_MEMORY_SHARED, a flag for
NE_ALLOC_MEMORY that makes an allocation mappable. The parent needs that
to read state the device writes for it, and for virtqueues and buffers
that both the parent and the sandbox access.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 160 +++++++++++++++++++++-
 drivers/virt/nitro_enclaves/ne_misc_dev.h |   6 +-
 include/uapi/linux/nitro_enclaves.h       |  36 ++++-
 3 files changed, 190 insertions(+), 12 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 56ca1a9d99d41..55dbeb5b14d29 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -2597,17 +2597,21 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
  * serving such a request out of movable memory would take that memory from
  * the rest of the host and leave the reservation bounding nothing.
  *
- * A request larger than one contiguous run is satisfied by several runs, which
- * the driver keeps on a list of its own so that an enclave can be given every
- * one of them.
+ * A request larger than one contiguous run is satisfied by several runs, and
+ * mmap() maps them back to back into one VM_PFNMAP VMA, so user space sees a
+ * single region. The driver keeps the runs on a list of its own, so that an
+ * enclave can be given them without anyone reading them back out of the page
+ * tables.
  *
- * An allocation is never mapped, and the enclave image reaches it through
- * write().
+ * An allocation is either private or shared. A private one is never mapped, and
+ * the enclave image reaches it through write() instead. A shared one can also be
+ * mapped, at an offset and a length its own size bounds.
  *
  * The runs belong to the file the allocation is handed out on. Zeroing a fresh
  * run and a write into an allocation go through ne_walk_run(), so that neither
- * can run for as long as the allocation is large; the release path walks the
- * runs with a preemption point of its own.
+ * can run for as long as the allocation is large. A fault resolves the run its
+ * offset falls in with a binary search instead of walking, and the release path
+ * walks the runs with a preemption point of its own.
  */
 
 /**
@@ -2788,6 +2792,82 @@ static int ne_walk_allocs(struct ne_contig_mem *mem, u64 off, u64 len,
 	return 0;
 }
 
+/**
+ * ne_contig_fault_order() - Install one page table entry for a fault.
+ * @vmf:	The fault.
+ * @order:	Order of the entry to install, zero for a PTE.
+ *
+ * A mapping is populated as it is touched, so a caller that maps a whole
+ * allocation and builds an enclave image at the front of it pays page tables
+ * for the part it touched.
+ *
+ * An entry of @order needs that much of one run behind it, at a matching
+ * physical alignment. A run is a power-of-two multiple of 2 MiB aligned to at
+ * least 2 MiB, so a PMD sized entry fits inside one whenever the address it
+ * lands at is aligned too, and anything that does not fit falls back to the
+ * next order down.
+ *
+ * Return: a VM_FAULT_* value.
+ */
+static vm_fault_t ne_contig_fault_order(struct vm_fault *vmf, unsigned int order)
+{
+	struct vm_area_struct *vma = vmf->vma;
+	struct ne_contig_mem *mem = vma->vm_file->private_data;
+	unsigned long len = PAGE_SIZE << order;
+	unsigned long addr = ALIGN_DOWN(vmf->address, len);
+	struct ne_contig_alloc *alloc;
+	unsigned long pfn;
+	u64 off;
+
+	if (addr < vma->vm_start || addr + len > vma->vm_end)
+		return VM_FAULT_FALLBACK;
+
+	off = ((u64)vma->vm_pgoff << PAGE_SHIFT) + (addr - vma->vm_start);
+
+	alloc = ne_find_alloc(mem, off);
+	if (!alloc)
+		return VM_FAULT_SIGBUS;
+
+	if (off + len > alloc->off + ((u64)alloc->nr_pages << PAGE_SHIFT))
+		return VM_FAULT_FALLBACK;
+
+	pfn = page_to_pfn(alloc->page) + ((off - alloc->off) >> PAGE_SHIFT);
+	if (!IS_ALIGNED(pfn, len >> PAGE_SHIFT))
+		return VM_FAULT_FALLBACK;
+
+	if (!order)
+		return vmf_insert_pfn(vma, addr, pfn);
+
+	if (IS_ENABLED(CONFIG_ARCH_SUPPORTS_PMD_PFNMAP) && order == PMD_ORDER)
+		return vmf_insert_pfn_pmd(vmf, pfn,
+					  vmf->flags & FAULT_FLAG_WRITE);
+
+	if (IS_ENABLED(CONFIG_ARCH_SUPPORTS_PUD_PFNMAP) && order == PUD_ORDER)
+		return vmf_insert_pfn_pud(vmf, pfn,
+					  vmf->flags & FAULT_FLAG_WRITE);
+
+	return VM_FAULT_FALLBACK;
+}
+
+static vm_fault_t ne_contig_vma_fault(struct vm_fault *vmf)
+{
+	return ne_contig_fault_order(vmf, 0);
+}
+
+#ifdef CONFIG_ARCH_SUPPORTS_HUGE_PFNMAP
+static vm_fault_t ne_contig_huge_fault(struct vm_fault *vmf, unsigned int order)
+{
+	return ne_contig_fault_order(vmf, order);
+}
+#endif
+
+static const struct vm_operations_struct ne_contig_vm_ops = {
+	.fault		= ne_contig_vma_fault,
+#ifdef CONFIG_ARCH_SUPPORTS_HUGE_PFNMAP
+	.huge_fault	= ne_contig_huge_fault,
+#endif
+};
+
 static struct cma *ne_numa_cma_area(int nid)
 {
 	struct cma *cma = dma_get_numa_cma_area(nid);
@@ -2915,6 +2995,69 @@ struct page *ne_alloc_contig(unsigned long nr_pages, int nid,
 	return page;
 }
 
+/**
+ * ne_contig_supports_mmap() - Whether an allocation may be mapped at all.
+ * @file:	The allocation's file.
+ *
+ * A private allocation is the one an enclave gets to itself, and it is never
+ * mapped. The attribute is fixed when the allocation is made and the file is
+ * not published until it is set, so this reads it without a lock.
+ *
+ * Return: true when the caller may map the allocation.
+ */
+static bool ne_contig_supports_mmap(struct file *file)
+{
+	struct ne_contig_mem *mem = file->private_data;
+
+	return mem->shared;
+}
+
+/**
+ * ne_contig_mmap_prepare() - Set a mapping of part or all of an allocation up.
+ * @desc:	The mapping to set up.
+ *
+ * The offset and the length both have to be multiples of 2 MiB, which is the
+ * unit the allocation itself is made in, and the mapping has to stay inside
+ * the allocation. A caller may map part of an allocation, and may map one
+ * allocation more than once; the runs are not mutated after the file is
+ * published and the file holds them for as long as any mapping exists.
+ *
+ * ne_contig_fault_order() populates the mapping as it is touched. The mapping
+ * asks for VM_HUGEPAGE because a fault on a VM_PFNMAP mapping without it only
+ * reaches ->huge_fault where the transparent_hugepage policy is "always".
+ *
+ * Return: 0 on success, negative return value on failure.
+ */
+static int ne_contig_mmap_prepare(struct vm_area_desc *desc)
+{
+	struct ne_contig_mem *mem = desc->file->private_data;
+	unsigned long size = vma_desc_size(desc);
+	u64 offset = (u64)desc->pgoff << PAGE_SHIFT;
+
+	if (!ne_contig_supports_mmap(desc->file))
+		return -ENODEV;
+
+	/*
+	 * MAP_PRIVATE of an allocation has no meaning: the memory exists to be
+	 * given to an enclave, and a copy that only this process can see never
+	 * reaches one.
+	 */
+	if (!vma_desc_test(desc, VMA_SHARED_BIT))
+		return -EINVAL;
+
+	if ((size | offset) & (NE_MIN_MEM_REGION_SIZE - 1))
+		return -EINVAL;
+
+	if (offset > mem->size || size > mem->size - offset)
+		return -EINVAL;
+
+	desc->vm_ops = &ne_contig_vm_ops;
+	vma_desc_set_flags(desc, VMA_IO_BIT, VMA_PFNMAP_BIT, VMA_DONTEXPAND_BIT,
+			   VMA_DONTDUMP_BIT, VMA_DONTCOPY_BIT, VMA_HUGEPAGE_BIT);
+
+	return 0;
+}
+
 /**
  * struct ne_write_ctx - Where the write walk has got to.
  * @iter:	Where the bytes come from.
@@ -3030,6 +3173,8 @@ static int ne_contig_release(struct inode *inode, struct file *file)
 static const struct file_operations ne_contig_fops = {
 	.owner			= THIS_MODULE,
 	.llseek			= generic_file_llseek,
+	.mmap_prepare		= ne_contig_mmap_prepare,
+	.get_unmapped_area	= thp_get_unmapped_area,
 	.write_iter		= ne_contig_write_iter,
 	.release		= ne_contig_release,
 };
@@ -3176,6 +3321,7 @@ static int ne_alloc_memory_ioctl(struct ne_enclave *ne_enclave,
 	mem->mm = current->mm;
 	mem->nr_pinned = nr_pages;
 	mem->size = args->size;
+	mem->shared = args->flags == NE_MEMORY_SHARED;
 	mem->enclave_id = ne_enclave->id;
 
 	/*
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 1a76c12815298..fcdf457af5800 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -38,6 +38,7 @@ struct ne_contig_alloc {
  * @allocs:	Contiguous runs backing the allocation, ordered by offset.
  * @nr_allocs:	Number of runs in @allocs.
  * @size:	Size of the allocation in bytes, mirrored in i_size.
+ * @shared:	Whether the allocation may be mapped by the caller.
  * @enclave_id:	Id of the enclave the allocation was made for.
  * @mm:		The mm charged for the pinned pages, held via mmgrab().
  * @nr_pinned:	Pages charged to @mm->pinned_vm.
@@ -47,13 +48,14 @@ struct ne_contig_alloc {
  *
  * The array is built before the file is published and not mutated afterwards, so
  * any holder of a reference can read it without a lock of its own, which is what
- * lets a page fault resolve an offset to a run. @size is set in the same place
- * and is read the same way.
+ * lets a page fault resolve an offset to a run. @shared and @size are set in the
+ * same place and are read the same way.
  */
 struct ne_contig_mem {
 	struct ne_contig_alloc	*allocs;
 	unsigned int		nr_allocs;
 	u64			size;
+	bool			shared;
 	u64			enclave_id;
 	struct mm_struct	*mm;
 	unsigned long		nr_pinned;
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index a2919565bc883..f1315c9c876de 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -215,7 +215,7 @@ struct ne_add_any_vcpu_args {
  * @size:	Number of bytes to allocate, a multiple of 2 MiB.
  * @nid:	Host NUMA node to allocate from. One node, named; unlike
  *		%NE_ADD_ANY_VCPU it does not take %NE_NUMA_ANY.
- * @flags:	%NE_MEMORY_PRIVATE.
+ * @flags:	%NE_MEMORY_PRIVATE or %NE_MEMORY_SHARED.
  */
 struct ne_alloc_memory_args {
 	__u64	size;
@@ -237,7 +237,16 @@ struct ne_alloc_memory_args {
  *		     goes into one through write() or pwrite() and not through
  *		     a copy into a mapping, and every write is a call the
  *		     kernel handles and can act on or refuse.
- *		     The allocation lives until the descriptor is gone.
+ *		     A %NE_MEMORY_SHARED allocation can also be mapped, at
+ *		     offset 0 or at any other multiple of 2 MiB, for a multiple
+ *		     of 2 MiB, with %MAP_SHARED. A caller maps one to keep an
+ *		     address for memory it goes on using after it has handed the
+ *		     allocation over. Such a mapping is readable too, which is
+ *		     the only way to read allocated memory back. The mapping is
+ *		     not inherited across fork(); a child that needs one maps
+ *		     the descriptor itself.
+ *		     The allocation lives until the last of the descriptor and
+ *		     its mappings is gone.
  *		     The ioctl can be invoked on the enclave fd, before an
  *		     enclave is started.
  *
@@ -276,6 +285,17 @@ struct ne_alloc_memory_args {
  *				  write, so a caller that has to know how far it
  *				  got looks at the return value. A fatal signal
  *				  part way through is reported the same way.
+ *
+ * An mmap() of the descriptor sets errno to:
+ * * ENODEV			- The allocation is %NE_MEMORY_PRIVATE.
+ * * EINVAL			- The mapping is %MAP_PRIVATE, or its offset or
+ *				  length is not a multiple of 2 MiB, or it does
+ *				  not fit inside the allocation.
+ * A mapping is populated as it is touched, so the page tables for a page are
+ * allocated at first touch and not by mmap(). A touch by the process itself
+ * that fails to get them is retried until it succeeds or the process is
+ * killed. A touch the kernel makes for the process, as read() into the
+ * mapping does, reports EFAULT or a short transfer instead.
  */
 #define NE_ALLOC_MEMORY			_IOW(0xAE, 0x2E, struct ne_alloc_memory_args)
 
@@ -472,8 +492,18 @@ struct ne_user_memory_region {
  *		       way to put an enclave image in one.
  */
 #define NE_MEMORY_PRIVATE		(0x00)
+/**
+ * NE_MEMORY_SHARED - An allocation the parent may also map. A mapping is what a
+ *		      caller wants for memory it goes on reading or writing
+ *		      after it has handed the allocation over. It costs page
+ *		      tables for the part of it a caller touches, and it leaves
+ *		      the memory reachable from the parent for as long as it
+ *		      lives, so an allocation that is only being handed over
+ *		      does not ask for one.
+ */
+#define NE_MEMORY_SHARED		(0x01)
 
-#define NE_ALLOC_MEMORY_MAX_FLAG_VAL	(0x01)
+#define NE_ALLOC_MEMORY_MAX_FLAG_VAL	(0x02)
 
 /**
  * DOC: Enclave start info flags
-- 
2.47.1

[PATCH 08/40] nitro_enclaves: Add NE_ADD_MEMORY_REGION

[Alexander Graf]

The new allocation scheme also requires a new memory add ioctl. Introduce
NE_ADD_MEMORY_REGION which uses the fd that NE_ALLOC_MEMORY returns to
back a slot's memory.

Once memory is donated using this mechanism, disable the ability for user
space to write() into it. It is fully considered donated at this point.

For now, consider slot memory donation as mandatorily donated. You can
not use a SHARED mapping to back it until we introduce DIRECT mode later.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 262 ++++++++++++++++++++++
 drivers/virt/nitro_enclaves/ne_misc_dev.h |  32 ++-
 include/uapi/linux/nitro_enclaves.h       | 104 +++++++--
 3 files changed, 379 insertions(+), 19 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 55dbeb5b14d29..38345ad7c7985 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -119,6 +119,8 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg);
  */
 static atomic64_t ne_enclave_seq = ATOMIC64_INIT(0);
 
+static int ne_add_memory_region_ioctl(struct ne_enclave *ne_enclave,
+				      struct ne_add_memory_region_args *args);
 static int ne_alloc_memory_ioctl(struct ne_enclave *ne_enclave,
 				 struct ne_alloc_memory_args *args);
 static const struct file_operations ne_fops = {
@@ -2069,6 +2071,39 @@ static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned lo
 		return 0;
 	}
 
+	case NE_EA_IOCTL(NE_ADD_MEMORY_REGION): {
+		struct ne_add_memory_region_args __user *uargs = (void __user *)arg;
+		unsigned int usize = _IOC_SIZE(cmd);
+		struct ne_add_memory_region_args args;
+		int rc;
+
+		if (usize < offsetofend(struct ne_add_memory_region_args, flags))
+			return -EINVAL;
+
+		rc = copy_struct_from_user(&args, sizeof(args), uargs, usize);
+		if (rc)
+			return rc;
+
+		rc = mutex_lock_killable(&ne_enclave->enclave_info_mutex);
+		if (rc)
+			return rc;
+
+		if (ne_enclave->state != NE_STATE_INIT) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Enclave is not in init state\n");
+
+			mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+			return -NE_ERR_NOT_IN_INIT_STATE;
+		}
+
+		rc = ne_add_memory_region_ioctl(ne_enclave, &args);
+
+		mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+		return rc;
+	}
+
 	case NE_EA_IOCTL(NE_ALLOC_MEMORY): {
 		struct ne_alloc_memory_args __user *uargs = (void __user *)arg;
 		unsigned int usize = _IOC_SIZE(cmd);
@@ -2139,6 +2174,33 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 	return rc;
 }
 
+/**
+ * ne_enclave_remove_all_contig_regions() - Drop the enclave's references to the
+ *					    allocations it was given.
+ * @ne_enclave :	Private data associated with the current enclave.
+ *
+ * The last reference on an allocation's file frees its runs and gives back its
+ * %RLIMIT_MEMLOCK charge, so an allocation whose descriptor and mappings are
+ * already gone is freed here.
+ *
+ * Context: Process context. This function is called on an enclave that is on
+ *	    no list and has no file left.
+ */
+static void ne_enclave_remove_all_contig_regions(struct ne_enclave *ne_enclave)
+{
+	struct ne_contig_region *ne_contig_region = NULL;
+	struct ne_contig_region *tmp = NULL;
+
+	list_for_each_entry_safe(ne_contig_region, tmp,
+				 &ne_enclave->contig_regions_list, list) {
+		list_del(&ne_contig_region->list);
+
+		fput(ne_contig_region->file);
+
+		kfree(ne_contig_region);
+	}
+}
+
 /**
  * ne_enclave_remove_all_mem_region_entries() - Remove all memory region entries
  *						from the enclave data structure.
@@ -2278,6 +2340,7 @@ static int ne_enclave_free_slot(struct ne_pci_dev *ne_pci_dev,
  */
 static void ne_enclave_free(struct ne_enclave *ne_enclave)
 {
+	ne_enclave_remove_all_contig_regions(ne_enclave);
 	ne_enclave_remove_all_mem_region_entries(ne_enclave);
 	ne_enclave_remove_all_vcpu_id_entries(ne_enclave);
 
@@ -2494,6 +2557,7 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 __user *slot_ui
 	ne_enclave->has_event = false;
 	mutex_init(&ne_enclave->enclave_info_mutex);
 	ne_enclave->max_mem_regions = cmd_reply.mem_regions;
+	INIT_LIST_HEAD(&ne_enclave->contig_regions_list);
 	INIT_LIST_HEAD(&ne_enclave->mem_regions_list);
 	ne_enclave->mm = current->mm;
 	ne_enclave->slot_uid = cmd_reply.slot_uid;
@@ -3116,6 +3180,12 @@ static int ne_write_span(void *arg, struct page *page, u64 off, u64 len)
  * refusal is -EFBIG, the same as generic_write_check_limits() gives for a write
  * beyond s_maxbytes. There is no RLIMIT_FSIZE in it, so no SIGXFSZ either.
  *
+ * Donating an allocation closes this path, and @lock is what makes the two
+ * orderly. It is taken here and in ne_add_memory_region_ioctl() and nowhere
+ * else; setting a mapping up and faulting one in both read only fields that
+ * never change, so nothing takes it under mmap_write_lock or in a fault, and
+ * copying from user space under it cannot invert against the fault path.
+ *
  * Return: the number of bytes written, or a negative return value on failure.
  */
 static ssize_t ne_contig_write_iter(struct kiocb *iocb, struct iov_iter *from)
@@ -3132,6 +3202,11 @@ static ssize_t ne_contig_write_iter(struct kiocb *iocb, struct iov_iter *from)
 	if (pos >= mem->size || count > mem->size - pos)
 		return -EFBIG;
 
+	guard(mutex)(&mem->lock);
+
+	if (mem->donated)
+		return -EPERM;
+
 	rc = ne_walk_allocs(mem, pos, count, NE_WRITE_SPAN_PAGES, ne_write_span,
 			    &ctx);
 
@@ -3165,6 +3240,8 @@ static int ne_contig_release(struct inode *inode, struct file *file)
 
 	mmdrop(mem->mm);
 
+	mutex_destroy(&mem->lock);
+
 	kfree(mem);
 
 	return 0;
@@ -3317,6 +3394,7 @@ static int ne_alloc_memory_ioctl(struct ne_enclave *ne_enclave,
 		return -ENOMEM;
 	}
 
+	mutex_init(&mem->lock);
 	mmgrab(current->mm);
 	mem->mm = current->mm;
 	mem->nr_pinned = nr_pages;
@@ -3385,11 +3463,195 @@ static int ne_alloc_memory_ioctl(struct ne_enclave *ne_enclave,
 	ne_contig_free_allocs(mem);
 	atomic64_sub(nr_pages, &mem->mm->pinned_vm);
 	mmdrop(mem->mm);
+	mutex_destroy(&mem->lock);
 	kfree(mem);
 
 	return rc;
 }
 
+/**
+ * ne_add_memory_region_ioctl() - Give the slot associated with the current
+ *				  enclave a whole allocation.
+ * @ne_enclave:	Private data associated with the current enclave.
+ * @args:	Descriptor for the allocation, and flags.
+ *
+ * Every run the allocation is made of becomes a memory region, except that runs
+ * which turn out to be physically adjacent are merged into one, so a request
+ * that took several runs takes up several of the enclave's regions or fewer.
+ * The runs are 2 MiB aligned and a multiple of 2 MiB long as they come out of
+ * NE_ALLOC_MEMORY, which is what the hypervisor asks of a region, and no
+ * virtual address is involved on this path at all.
+ *
+ * The reference on the file is what keeps the runs and their %RLIMIT_MEMLOCK
+ * charge alive for as long as the enclave has them, whatever user space does
+ * with the descriptor and its mappings afterwards.
+ *
+ * Context: Process context. This function is called with the ne_enclave mutex held.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+static int ne_add_memory_region_ioctl(struct ne_enclave *ne_enclave,
+				      struct ne_add_memory_region_args *args)
+{
+	struct ne_phys_contig_mem_regions phys_contig_mem_regions = {};
+	struct pci_dev *pdev = ne_devs.ne_pci_dev->pdev;
+	struct ne_contig_region *ne_contig_region = NULL;
+	struct ne_contig_region *held = NULL;
+	struct ne_contig_mem *mem = NULL;
+	unsigned long i = 0;
+	int rc = -EINVAL;
+
+	if (ne_enclave->mm != current->mm)
+		return -EIO;
+
+	if (args->flags) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Incorrect flag for memory region\n");
+
+		return -NE_ERR_INVALID_FLAG_VALUE;
+	}
+
+	CLASS(fd, f)(args->fd);
+
+	if (fd_empty(f))
+		return -EBADF;
+
+	if (fd_file(f)->f_op != &ne_contig_fops) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Not an enclave memory allocation\n");
+
+		return -EINVAL;
+	}
+
+	/*
+	 * Every file with these operations came from ne_contig_create_file(),
+	 * which hands the allocation to anon_inode_create_getfile() before there
+	 * is a descriptor to pass in here, so it is there to be read.
+	 */
+	mem = fd_file(f)->private_data;
+
+	if (mem->enclave_id != ne_enclave->id) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Allocation was made for another enclave\n");
+
+		return -EINVAL;
+	}
+
+	/*
+	 * The parent may map a shared allocation and keeps any mapping it has
+	 * for as long as the allocation lives. Memory a region is made of is
+	 * the enclave's alone.
+	 */
+	if (mem->shared) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "The memory region allocation is shared\n");
+
+		return -EINVAL;
+	}
+
+	list_for_each_entry(held, &ne_enclave->contig_regions_list, list) {
+		if (held->file == fd_file(f)) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "The memory region is already used\n");
+
+			return -NE_ERR_MEM_REGION_ALREADY_USED;
+		}
+	}
+
+	ne_contig_region = kzalloc_obj(*ne_contig_region);
+	if (!ne_contig_region)
+		return -ENOMEM;
+
+	/*
+	 * One entry per run is the worst case, since merging only ever reduces
+	 * the count. That is 4 MiB of entries for a 512 GiB allocation the
+	 * allocator had to satisfy in 2 MiB runs, so this uses kvmalloc like
+	 * the run array.
+	 */
+	phys_contig_mem_regions.regions = kvzalloc_objs(*phys_contig_mem_regions.regions,
+							mem->nr_allocs);
+	if (!phys_contig_mem_regions.regions) {
+		rc = -ENOMEM;
+
+		goto free_contig_region;
+	}
+
+	for (i = 0; i < mem->nr_allocs; i++) {
+		struct ne_contig_alloc *alloc = &mem->allocs[i];
+
+		rc = ne_merge_phys_contig_memory_regions(&phys_contig_mem_regions,
+							 page_to_phys(alloc->page),
+							 (u64)alloc->nr_pages << PAGE_SHIFT);
+		if (rc < 0)
+			goto free_regions;
+	}
+
+	if ((ne_enclave->nr_mem_regions + phys_contig_mem_regions.num) >
+	    ne_enclave->max_mem_regions) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Reached max memory regions %lld\n",
+				    ne_enclave->max_mem_regions);
+
+		rc = -NE_ERR_MEM_MAX_REGIONS;
+
+		goto free_regions;
+	}
+
+	/*
+	 * From here the allocation is the enclave's. Closing the write path under
+	 * the same lock a write holds means a write either finished before this
+	 * point or is refused after it, so a donation cannot land in the middle of
+	 * one. Whether the whole image got there is the caller's to check: a write
+	 * cut short by an unreadable buffer reports how far it got.
+	 */
+	scoped_guard(mutex, &mem->lock)
+		mem->donated = true;
+
+	ne_contig_region->file = get_file(fd_file(f));
+
+	list_add(&ne_contig_region->list, &ne_enclave->contig_regions_list);
+
+	for (i = 0; i < phys_contig_mem_regions.num; i++) {
+		struct ne_pci_dev_cmd_reply cmd_reply = {};
+		struct slot_add_mem_req slot_add_mem_req = {};
+
+		slot_add_mem_req.slot_uid = ne_enclave->slot_uid;
+		slot_add_mem_req.paddr = phys_contig_mem_regions.regions[i].start;
+		slot_add_mem_req.size = range_len(&phys_contig_mem_regions.regions[i]);
+
+		rc = ne_do_request(pdev, SLOT_ADD_MEM,
+				   &slot_add_mem_req, sizeof(slot_add_mem_req),
+				   &cmd_reply, sizeof(cmd_reply));
+		if (rc < 0) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Error in slot add mem [rc=%d]\n", rc);
+
+			kvfree(phys_contig_mem_regions.regions);
+
+			/*
+			 * Exit here without dropping the reference on the file
+			 * as memory regions may already been added.
+			 */
+			return rc;
+		}
+
+		ne_enclave->mem_size += slot_add_mem_req.size;
+		ne_enclave->nr_mem_regions++;
+	}
+
+	kvfree(phys_contig_mem_regions.regions);
+
+	return 0;
+
+free_regions:
+	kvfree(phys_contig_mem_regions.regions);
+free_contig_region:
+	kfree(ne_contig_region);
+
+	return rc;
+}
+
 #if defined(CONFIG_NITRO_ENCLAVES_MISC_DEV_TEST)
 #include "ne_misc_dev_test.c"
 #endif
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index fcdf457af5800..2502cca3f3f64 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -42,14 +42,23 @@ struct ne_contig_alloc {
  * @enclave_id:	Id of the enclave the allocation was made for.
  * @mm:		The mm charged for the pinned pages, held via mmgrab().
  * @nr_pinned:	Pages charged to @mm->pinned_vm.
+ * @lock:	Serialises a write against the donation that ends writing.
+ * @donated:	Whether the allocation has been given to its enclave.
  *
- * The runs hang off the file, and a mapping's page table entries cannot outlive
- * them because a VMA holds vm_file for as long as it exists.
+ * The runs hang off the file, so that closing the last descriptor does not free
+ * memory an enclave is still using: a donated region holds the file with
+ * get_file(). That is also what keeps a mapping's page table entries from
+ * outliving the runs, since a VMA holds vm_file for as long as it exists.
  *
  * The array is built before the file is published and not mutated afterwards, so
  * any holder of a reference can read it without a lock of its own, which is what
  * lets a page fault resolve an offset to a run. @shared and @size are set in the
- * same place and are read the same way.
+ * same place and are read the same way. @donated is the one field that changes
+ * after publication, and @lock covers it.
+ *
+ * @enclave_id is compared, never dereferenced, and the counter it comes from is
+ * 64 bits wide, so an allocation cannot be donated to a later enclave that
+ * happens to have been placed at the address of the one it was made for.
  */
 struct ne_contig_mem {
 	struct ne_contig_alloc	*allocs;
@@ -59,6 +68,20 @@ struct ne_contig_mem {
 	u64			enclave_id;
 	struct mm_struct	*mm;
 	unsigned long		nr_pinned;
+	/* Serialises a write against the donation that ends writing. */
+	struct mutex		lock;
+	bool			donated;
+};
+
+/**
+ * struct ne_contig_region - Enclave memory given from one allocation.
+ * @list:	Entry in the enclave's list of allocations given to it.
+ * @file:	File holding the runs, held with get_file() so that they outlive the
+ *		descriptor the allocation was handed out on and its mappings.
+ */
+struct ne_contig_region {
+	struct list_head	list;
+	struct file		*file;
 };
 
 /**
@@ -82,6 +105,8 @@ struct ne_mem_region {
  * @enclave_info_mutex :	Mutex for accessing this internal state.
  * @detached:			The NE PCI device the enclave was created on is
  *				gone. Written with ne_devs.lock held for writing.
+ * @contig_regions_list:	List of %NE_ALLOC_MEMORY allocations given to the
+ *				enclave with %NE_ADD_MEMORY_REGION.
  * @enclave_list_entry :	Entry in the list of created enclaves.
  * @eventq:			Wait queue used for out-of-band event notifications
  *				triggered from the PCI device event handler to
@@ -116,6 +141,7 @@ struct ne_mem_region {
 struct ne_enclave {
 	struct mutex		enclave_info_mutex;
 	bool			detached;
+	struct list_head	contig_regions_list;
 	struct list_head	enclave_list_entry;
 	wait_queue_head_t	eventq;
 	bool			has_event;
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index f1315c9c876de..16515fe27b692 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -225,11 +225,13 @@ struct ne_alloc_memory_args {
 
 /**
  * NE_ALLOC_MEMORY - The command is used to allocate physically contiguous
- *		     memory to give to an enclave, drawn from one host NUMA
- *		     node. The argument is input only; on success the return
- *		     value is a file descriptor for the allocation.
+ *		     memory to give to an enclave with %NE_ADD_MEMORY_REGION,
+ *		     drawn from one host NUMA node. The argument is input only;
+ *		     on success the return value is a file descriptor for the
+ *		     allocation.
  *		     Writing to the descriptor places bytes in the allocation,
- *		     which is how an enclave image gets there.
+ *		     which is how an enclave image gets there, and stops
+ *		     working once the allocation has been donated.
  *		     O_APPEND has no effect. There is no read operation, so
  *		     read() and pread() both fail with EINVAL and a caller that
  *		     wants to check what it wrote keeps its own copy.
@@ -241,12 +243,19 @@ struct ne_alloc_memory_args {
  *		     offset 0 or at any other multiple of 2 MiB, for a multiple
  *		     of 2 MiB, with %MAP_SHARED. A caller maps one to keep an
  *		     address for memory it goes on using after it has handed the
- *		     allocation over. Such a mapping is readable too, which is
- *		     the only way to read allocated memory back. The mapping is
- *		     not inherited across fork(); a child that needs one maps
- *		     the descriptor itself.
- *		     The allocation lives until the last of the descriptor and
- *		     its mappings is gone.
+ *		     allocation over, including once the allocation has been
+ *		     donated and writing to the descriptor has stopped working.
+ *		     Such a mapping is readable too, which is the only way to
+ *		     read allocated memory back. The mapping is not inherited
+ *		     across fork(); a child that needs one maps the descriptor
+ *		     itself.
+ *		     Donating does not take a mapping away, so a parent that
+ *		     mapped a shared allocation can still write to it through
+ *		     that mapping afterwards.
+ *		     The allocation lives until the last of the descriptor, its
+ *		     mappings and the enclave regions donated from it is gone,
+ *		     so a caller can close the descriptor as soon as it has
+ *		     donated.
  *		     The ioctl can be invoked on the enclave fd, before an
  *		     enclave is started.
  *
@@ -279,6 +288,8 @@ struct ne_alloc_memory_args {
  * * Error codes from get_unused_fd_flags() and anon_inode_create_getfile().
  *
  * Writing to the descriptor sets errno to:
+ * * EPERM			- The allocation has been donated with
+ *				  %NE_ADD_MEMORY_REGION.
  * * EFBIG			- The write does not fit inside the allocation.
  * * EFAULT			- Part of the buffer is not readable. Any bytes
  *				  copied before that are reported as a short
@@ -299,6 +310,66 @@ struct ne_alloc_memory_args {
  */
 #define NE_ALLOC_MEMORY			_IOW(0xAE, 0x2E, struct ne_alloc_memory_args)
 
+/**
+ * struct ne_add_memory_region_args - Allocation to give to an enclave (in).
+ * @fd:		Descriptor from %NE_ALLOC_MEMORY.
+ * @flags:	Must be zero.
+ */
+struct ne_add_memory_region_args {
+	__s32	fd;
+	__u32	flags;
+};
+
+/**
+ * NE_ADD_MEMORY_REGION - The command is used to give an enclave the whole of an
+ *			  %NE_ALLOC_MEMORY allocation, as one memory region per
+ *			  contiguous run it is made of, or fewer where runs turn
+ *			  out to be physically adjacent. The argument is input
+ *			  only. The allocation has to have been made for this
+ *			  enclave, and can be given away once.
+ *			  Writing to the descriptor stops working here, so an
+ *			  enclave image goes in before this call. The allocation
+ *			  has to be %NE_MEMORY_PRIVATE, since memory the parent
+ *			  can still reach is not memory the enclave has to
+ *			  itself.
+ *			  An enclave's memory is laid out in the order its
+ *			  allocations are given to it, and the offset from
+ *			  %NE_GET_IMAGE_LOAD_INFO is an offset into that layout,
+ *			  so a caller building an enclave out of several
+ *			  allocations picks the order first, writes at the
+ *			  offsets that order implies, and gives them away in it.
+ *			  The ioctl can be invoked on the enclave fd, before an
+ *			  enclave is started.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0					- Logic successfully completed.
+ * * -1					- There was a failure in the ioctl logic.
+ * On failure, errno is set to:
+ * * ENODEV				- The NE PCI device is not available.
+ * * EFAULT				- Argument read failure.
+ * * EBADF				- @fd is not an open file descriptor.
+ * * EINVAL				- The argument is shorter than @flags, or @fd
+ *					  is not a descriptor from %NE_ALLOC_MEMORY, or
+ *					  the allocation was made for another enclave,
+ *					  or the allocation is %NE_MEMORY_SHARED.
+ * * EIO				- Current task mm is not the same as the one
+ *					  that created the enclave.
+ * * E2BIG				- The argument is longer than this kernel knows
+ *					  about and the excess is not all zero.
+ * * ENOMEM				- Memory allocation failure for internal
+ *					  bookkeeping variables.
+ * * NE_ERR_NOT_IN_INIT_STATE		- The enclave is not in init state (init =
+ *					  before being started).
+ * * NE_ERR_INVALID_FLAG_VALUE		- The value of the provided flag is invalid.
+ * * NE_ERR_MEM_REGION_ALREADY_USED	- The allocation was already given to the
+ *					  enclave.
+ * * NE_ERR_MEM_MAX_REGIONS		- The number of memory regions set for the
+ *					  enclave reached maximum.
+ * * Error codes from the NE PCI device request.
+ */
+#define NE_ADD_MEMORY_REGION		_IOW(0xAE, 0x2F, struct ne_add_memory_region_args)
+
 /**
  * NE_START_ENCLAVE - The command is used to trigger enclave start after the
  *		      enclave resources, such as memory and CPU, have been set.
@@ -362,7 +433,8 @@ struct ne_alloc_memory_args {
  */
 #define NE_ERR_UNALIGNED_MEM_REGION_ADDR	(261)
 /**
- * NE_ERR_MEM_REGION_ALREADY_USED - The user space memory region is already used.
+ * NE_ERR_MEM_REGION_ALREADY_USED - The user space memory region, or the
+ *				      %NE_ALLOC_MEMORY allocation, is already used.
  */
 #define NE_ERR_MEM_REGION_ALREADY_USED		(262)
 /**
@@ -495,11 +567,11 @@ struct ne_user_memory_region {
 /**
  * NE_MEMORY_SHARED - An allocation the parent may also map. A mapping is what a
  *		      caller wants for memory it goes on reading or writing
- *		      after it has handed the allocation over. It costs page
- *		      tables for the part of it a caller touches, and it leaves
- *		      the memory reachable from the parent for as long as it
- *		      lives, so an allocation that is only being handed over
- *		      does not ask for one.
+ *		      after it has handed the allocation over, including after
+ *		      it has been donated. It costs page tables for the part of
+ *		      it a caller touches, and it leaves the memory reachable
+ *		      from the parent for as long as it lives, so an allocation
+ *		      that is only being handed over does not ask for one.
  */
 #define NE_MEMORY_SHARED		(0x01)
 
-- 
2.47.1

[PATCH 09/40] nitro_enclaves: Document the memory pool

[Alexander Graf]

An admin reserving memory for enclaves needs to pick the CMA area and
the kernel parameters that fill it. Document the pool: allocating with
NE_ALLOC_MEMORY, writing the image through the returned descriptor and
adding it with NE_ADD_MEMORY_REGION, the config options a reservation
needs to be reachable, and the difference between private and shared
allocations and how long an allocation lives.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 Documentation/virt/ne_overview.rst | 83 ++++++++++++++++++++++++++++--
 1 file changed, 80 insertions(+), 3 deletions(-)

diff --git a/Documentation/virt/ne_overview.rst b/Documentation/virt/ne_overview.rst
index 223ee7fdec51b..4c4b7febaff0a 100644
--- a/Documentation/virt/ne_overview.rst
+++ b/Documentation/virt/ne_overview.rst
@@ -113,9 +113,10 @@ each node in turn, and sets numa_node to NE_NUMA_ANY when it does not
 care. Both start from a free thread of a core the enclave already holds,
 and only take a new core once there is none.
 
-The node says where cores come from. An enclave built out of a pool that
-spans nodes has no node of its own, so the rule that its memory be on
-that node does not apply to it and its memory can come from any node.
+The node says where cores come from. NE_ALLOC_MEMORY takes its own node;
+see `Memory pool`_. An enclave built out of a pool that spans
+nodes has no node of its own, so the rule that its memory be on that node
+does not apply to it and its memory can come from any node.
 
 The pool occupancy is readable under
 /sys/devices/virtual/misc/nitro_enclaves/cpu_pool/, and per node under
@@ -124,3 +125,79 @@ online when the device was probed. A node the pool holds nothing on
 reads empty rather than being absent. The driver offlines pool threads,
 so ``nodeN/cpulist`` under /sys/devices/system/node stops listing them
 and answers nothing about the pool.
+
+Memory pool
+-----------
+
+Enclave memory must be physically contiguous, in chunks of at least 2 MiB,
+and, for an enclave whose cores all sit on one node, on that node. It has
+traditionally come from hugetlbfs [2], which provides the hugepage sizes the
+architecture has -- 2 MiB and 1 GiB on x86, and 32 MiB, 512 MiB or 16 GiB on
+arm64 depending on the page size -- and needs the hugepage count managed out
+of band.
+
+The driver can hand that memory out itself, backed per NUMA node by the
+kernel's Contiguous Memory Allocator (CMA). Reserve an area on every node
+you intend to run enclaves on, with the numa_cma= or cma_pernuma= kernel
+parameters documented alongside the rest of them [4]. On a host with one
+online node, cma= reserves the one area there is to reserve and the driver
+uses it.
+
+Reaching a reserved area takes CONFIG_DMA_CMA, which NITRO_ENCLAVES selects,
+so the driver always has it. numa_cma= and cma_pernuma= need
+CONFIG_DMA_NUMA_CMA on top of that, which a host with more than one online
+node enables to name the node a reservation lands on. cma= needs neither, so
+on a host with one online node the driver's own select is the whole
+requirement.
+
+NE_ALLOC_MEMORY on the enclave file descriptor allocates 2 MiB aligned,
+physically contiguous memory from one host NUMA node's CMA area. A request
+that area cannot serve fails with -ENOMEM, so the reservation bounds what the
+enclaves on a node can take between them.
+It returns a file descriptor for the allocation. For an enclave whose cores
+all sit on one node the ioctl refuses any other node, so name the same node
+here and in NE_ADD_ANY_VCPU. An enclave built out of a pool that spans nodes
+has no node of its own and takes memory from any of them.
+
+An allocation is NE_MEMORY_PRIVATE unless the flags say NE_MEMORY_SHARED.
+A private one cannot be mapped, and mmap() on it returns ENODEV. Writing to
+the descriptor places bytes in either kind, which is how an enclave image
+gets into one. There is no read operation on the descriptor.
+
+A shared allocation can also be mapped, at any 2 MiB aligned offset for a
+multiple of 2 MiB, with MAP_SHARED. A caller maps one to keep an address for
+memory it goes on using after handing it over, including after the allocation
+has been donated, which is where writing to the descriptor stops working. Such
+a mapping is readable too, and is the only way to read allocated memory back.
+
+The mapping is VM_PFNMAP, so it holds no struct page references and any DMA pin
+over it returns EFAULT, long-term or not; io_uring buffer registration and RDMA
+memory registration are two that a caller might expect to work and that do not.
+Reaching enclave memory from a device in the primary VM is not something this
+interface offers.
+
+The allocation is accounted against RLIMIT_MEMLOCK. NE_ADD_MEMORY_REGION
+takes the descriptor and gives the enclave the whole allocation, as one
+memory region per contiguous run it is made of, or fewer where runs turn out
+to be physically adjacent. It also ends writing: the memory is the enclave's
+from that point, so an image goes in first. It takes private allocations
+only, since memory the parent can still reach is not memory the enclave has
+to itself.
+
+An enclave's memory is laid out in the order its allocations are given to it,
+and the offset NE_GET_IMAGE_LOAD_INFO returns is an offset into that layout.
+A caller building an enclave out of several allocations therefore picks the
+order first, writes at the offsets that order implies, and gives them away in
+it.
+
+A mapping is populated as it is touched, so it costs page tables for the part
+a caller reaches and not for all of it. It keeps the memory reachable from the
+parent for as long as it lives, so an allocation that is only being handed to
+an enclave stays private.
+
+NE_SET_USER_MEMORY_REGION is unchanged and still takes the address of a
+hugetlbfs region.
+
+The allocation lives until the last of its descriptor, its mappings and the
+regions given from it is gone, so neither closing the descriptor nor
+unmapping takes memory away from a running enclave.
-- 
2.47.1

[PATCH 10/40] nitro_enclaves: Add NE_CREATE_VM2 with a launch flag word

[Alexander Graf]

In preparation to support Nitro Sandboxes, whose properties the
hypervisor needs before it allocates a slot, add NE_CREATE_VM2. It takes
a word of launch flags for the slot. The slot id comes back in the
same struct ne_create_vm_args.

No flag is defined in this patch, so any set bit fails with
-NE_ERR_INVALID_FLAG_VALUE.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 53 +++++++++++++++++++++--
 include/uapi/linux/nitro_enclaves.h       | 48 ++++++++++++++++++++
 2 files changed, 97 insertions(+), 4 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 38345ad7c7985..68c497df4f14b 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -2452,6 +2452,8 @@ static const struct file_operations ne_enclave_fops = {
  *			  enclave file descriptor to be further used for enclave
  *			  resources handling e.g. memory regions and CPUs.
  * @ne_pci_dev :	Private data associated with the PCI device.
+ * @flags:		Launch-time flag word from NE_CREATE_VM2, or zero when
+ *			the caller used the flag-less NE_CREATE_VM.
  * @slot_uid:		User pointer to store the generated unique slot id
  *			associated with an enclave to.
  *
@@ -2461,7 +2463,8 @@ static const struct file_operations ne_enclave_fops = {
  * * Enclave fd on success.
  * * Negative return value on failure.
  */
-static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 __user *slot_uid)
+static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 flags,
+			      u64 __user *slot_uid)
 {
 	struct ne_pci_dev_cmd_reply cmd_reply = {};
 	int enclave_fd = -1;
@@ -2472,6 +2475,13 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 __user *slot_ui
 	int rc = -EINVAL;
 	struct slot_alloc_req slot_alloc_req = {};
 
+	if (flags) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Unknown launch flag bit set 0x%llx\n", flags);
+
+		return -NE_ERR_INVALID_FLAG_VALUE;
+	}
+
 	mutex_lock(&ne_cpu_pool.mutex);
 
 	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
@@ -2626,7 +2636,44 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 		if (ne_pci_dev) {
 			rc = mutex_lock_killable(&ne_pci_dev->enclaves_list_mutex);
 			if (!rc) {
-				rc = ne_create_vm_ioctl(ne_pci_dev, slot_uid);
+				rc = ne_create_vm_ioctl(ne_pci_dev, 0, slot_uid);
+				mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+			}
+		}
+
+		up_read(&ne_devs.lock);
+
+		return rc;
+	}
+	}
+
+	/* Extensible Argument ioctls */
+	switch (NE_EA_IOCTL(cmd)) {
+	case NE_EA_IOCTL(NE_CREATE_VM2): {
+		struct ne_create_vm_args args;
+		struct ne_create_vm_args __user *uargs = (void __user *)arg;
+		unsigned int usize = _IOC_SIZE(cmd);
+		struct ne_pci_dev *ne_pci_dev;
+		int rc;
+
+		if (usize < offsetofend(struct ne_create_vm_args, slot_uid))
+			return -EINVAL;
+
+		rc = copy_struct_from_user(&args, sizeof(args), uargs, usize);
+		if (rc)
+			return rc;
+
+		rc = down_read_killable(&ne_devs.lock);
+		if (rc)
+			return rc;
+
+		rc = -ENODEV;
+		ne_pci_dev = ne_devs.ne_pci_dev;
+		if (ne_pci_dev) {
+			rc = mutex_lock_killable(&ne_pci_dev->enclaves_list_mutex);
+			if (!rc) {
+				rc = ne_create_vm_ioctl(ne_pci_dev, args.flags,
+							&uargs->slot_uid);
 				mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
 			}
 		}
@@ -2639,8 +2686,6 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 	default:
 		return -ENOTTY;
 	}
-
-	return 0;
 }
 
 /*
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index 16515fe27b692..8e43be9a0a1fa 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -46,6 +46,54 @@
  */
 #define NE_CREATE_VM			_IOR(0xAE, 0x20, __u64)
 
+/**
+ * struct ne_create_vm_args - NE_CREATE_VM2 argument (in / out).
+ * @flags:	Launch-time flag word (in). Flags select properties of the
+ *		enclave slot that the hypervisor has to know at allocation
+ *		time. NE_ENCLAVE_DEBUG_MODE stays a start-time property and
+ *		is rejected here.
+ * @slot_uid:	Generated unique slot id associated with an enclave (out).
+ */
+struct ne_create_vm_args {
+	__u64	flags;
+	__u64	slot_uid;
+};
+
+/**
+ * NE_CREATE_VM2 - As NE_CREATE_VM, additionally taking a launch-time flag
+ *		   word. The flags carry properties of the slot that have to
+ *		   be decided when the slot is allocated rather than when the
+ *		   enclave starts.
+ *
+ *		   NE_CREATE_VM continues to work; calling it is equivalent
+ *		   to calling NE_CREATE_VM2 with a zero flag word.
+ *
+ *		   The argument is read with copy_struct_from_user(), so a
+ *		   caller built against a header that predates a field this
+ *		   kernel knows leaves that field zero, and a caller that
+ *		   sets a field this kernel does not know is refused.
+ *
+ * Context: Process context.
+ * Return:
+ * * Enclave file descriptor		- As for NE_CREATE_VM.
+ * *  -1				- There was a failure in the ioctl logic.
+ * On failure, errno is set to:
+ * * ENODEV				- The NE PCI device is not available.
+ * * EFAULT				- Argument read or write failure.
+ * * NE_ERR_INVALID_FLAG_VALUE		- Unknown flag bit set, or
+ *					  NE_ENCLAVE_DEBUG_MODE set.
+ * * EINVAL				- The argument stops before @slot_uid.
+ * * E2BIG				- The argument carries a field this kernel
+ *					  does not know, set to a non-zero value.
+ * * ENOMEM				- Memory allocation failure for internal
+ *					  bookkeeping variables.
+ * * NE_ERR_NO_CPUS_AVAIL_IN_POOL	- No NE CPU pool set / no CPUs available
+ *					  in the pool.
+ * * Error codes from get_unused_fd_flags() and anon_inode_getfile().
+ * * Error codes from the NE PCI device request.
+ */
+#define NE_CREATE_VM2			_IOWR(0xAE, 0x30, struct ne_create_vm_args)
+
 /**
  * NE_ADD_VCPU - The command is used to set a vCPU for an enclave. The vCPU can
  *		 be auto-chosen from the NE CPU pool or it can be set by the
-- 
2.47.1

[PATCH 11/40] nitro_enclaves: Add NE_ENCLAVE_PCIE_MODE launch flag

[Alexander Graf]

Nitro Sandboxes give their guest virtio devices that the parent
provides, on a PCIe bus. The hypervisor sets up that bus when it
allocates the slot, so user space asks for it with a launch flag.

Define NE_ENCLAVE_PCIE_MODE, the first launch flag, and forward it in
SLOT_ALLOC. struct slot_alloc_req gains a flag word, and at offset 0,
over the unused field, a payload_version byte that tells the device to
read it. We send version 1 when a flag is set, and the all-zero request
of today otherwise.

A device that takes launch flags publishes a vendor-specific PCI
capability. A device that predates them ignores the flag word and
allocates a plain slot, so refuse a launch flag with -EOPNOTSUPP on a
device without the capability.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 19 ++++++++++++++++++-
 drivers/virt/nitro_enclaves/ne_pci_dev.h  | 23 +++++++++++++++++++++--
 include/uapi/linux/nitro_enclaves.h       | 14 ++++++++++++++
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 68c497df4f14b..0bf62f703806e 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -2475,13 +2475,30 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 flags,
 	int rc = -EINVAL;
 	struct slot_alloc_req slot_alloc_req = {};
 
-	if (flags) {
+	if (flags & ~(u64)NE_ENCLAVE_PCIE_MODE) {
 		dev_err_ratelimited(ne_misc_dev.this_device,
 				    "Unknown launch flag bit set 0x%llx\n", flags);
 
 		return -NE_ERR_INVALID_FLAG_VALUE;
 	}
 
+	/*
+	 * A device that takes launch flags publishes a vendor-specific
+	 * capability. One without it takes no launch flags and would allocate a
+	 * slot without the properties they ask for.
+	 */
+	if (flags && !pci_find_capability(ne_pci_dev->pdev, PCI_CAP_ID_VNDR)) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Device takes no launch flags\n");
+
+		return -EOPNOTSUPP;
+	}
+
+	if (flags) {
+		slot_alloc_req.payload_version = 1;
+		slot_alloc_req.flags = flags;
+	}
+
 	mutex_lock(&ne_cpu_pool.mutex);
 
 	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index 19e129d14a3a0..b0957e2606ec3 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -7,10 +7,12 @@
 #define _NE_PCI_DEV_H_
 
 #include <linux/atomic.h>
+#include <linux/build_bug.h>
 #include <linux/list.h>
 #include <linux/mutex.h>
 #include <linux/pci.h>
 #include <linux/pci_ids.h>
+#include <linux/stddef.h>
 #include <linux/wait.h>
 
 /**
@@ -174,12 +176,29 @@ struct enclave_stop_req {
 
 /**
  * struct slot_alloc_req - SLOT_ALLOC request.
- * @unused:	In order to avoid weird sizeof edge cases.
+ * @payload_version:	Wire version at offset 0. Zero is the value a driver
+ *			without launch flags leaves here, and tells the device
+ *			to disregard @reserved and @flags. 1 means @flags is
+ *			valid.
+ * @reserved:		Alignment padding between @payload_version and
+ *			@flags. Zero.
+ * @flags:		Launch flag word from NE_CREATE_VM2, valid when
+ *			@payload_version >= 1.
  */
 struct slot_alloc_req {
-	u8	unused;
+	u8	payload_version;
+	u8	reserved[3];
+	u32	flags;
 };
 
+/*
+ * The device picks these fields out of the payload by offset, so a change
+ * to the layout here changes the SLOT_ALLOC wire format.
+ */
+static_assert(sizeof(struct slot_alloc_req) == 8);
+static_assert(offsetof(struct slot_alloc_req, reserved) == 1);
+static_assert(offsetof(struct slot_alloc_req, flags) == 4);
+
 /**
  * struct slot_free_req - SLOT_FREE request.
  * @slot_uid:	Slot unique id mapped to the slot to free.
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index 8e43be9a0a1fa..8ccf16e9cf5fb 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -46,6 +46,18 @@
  */
 #define NE_CREATE_VM			_IOR(0xAE, 0x20, __u64)
 
+/**
+ * DOC: NE_CREATE_VM2 launch flags
+ */
+
+/**
+ * NE_ENCLAVE_PCIE_MODE - Allocate a slot with a PCIe bus. The
+ *			  hypervisor selects PCIe backing for the slot at
+ *			  SLOT_ALLOC time, so a caller that wants it asks
+ *			  for it here and not at enclave start.
+ */
+#define NE_ENCLAVE_PCIE_MODE		(1u << 6)
+
 /**
  * struct ne_create_vm_args - NE_CREATE_VM2 argument (in / out).
  * @flags:	Launch-time flag word (in). Flags select properties of the
@@ -82,6 +94,8 @@ struct ne_create_vm_args {
  * * EFAULT				- Argument read or write failure.
  * * NE_ERR_INVALID_FLAG_VALUE		- Unknown flag bit set, or
  *					  NE_ENCLAVE_DEBUG_MODE set.
+ * * EOPNOTSUPP			- A launch flag is set and the NE PCI
+ *					  device takes no launch flags.
  * * EINVAL				- The argument stops before @slot_uid.
  * * E2BIG				- The argument carries a field this kernel
  *					  does not know, set to a non-zero value.
-- 
2.47.1

[PATCH 12/40] nitro_enclaves: Name the device's error codes in the UAPI header

[Alexander Graf]

The nitro_enclaves device refuses a command with an errno or with an
NE_ERR_ code, which ne_do_request() passes on to user space as it is.
The codes of its own that the device returns start at 276, after the
driver's, and we want the device to be able to add more of them later.

Name the device's codes in the header, so that a caller can tell for
example a refused flag from an invalid slot, and document three ranges:
the driver's original codes from 256, the device's codes from 276 to
511, and codes only the driver returns from 1024. That way new codes on
either side never collide. 511 stays below the kernel-internal error
codes, and 1024 well below MAX_ERRNO, so user space reads every code as
an error.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 include/uapi/linux/nitro_enclaves.h | 142 ++++++++++++++++++++++++++++
 1 file changed, 142 insertions(+)

diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index 8ccf16e9cf5fb..6dd00c41c07d5 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -567,6 +567,148 @@ struct ne_add_memory_region_args {
  */
 #define NE_ERR_INVALID_ENCLAVE_CID		(275)
 
+/**
+ * DOC: NE error code ranges
+ *
+ * The NE_ERR_ codes come from three ranges:
+ *
+ * * 256 to 275 - The driver's original codes. The NE PCI device also
+ *   returns NE_ERR_MEM_MAX_REGIONS, with the same meaning.
+ * * 276 to 511 - Codes the NE PCI device returns. The driver passes them on
+ *   to the caller unchanged, and the device adds new codes upward from the
+ *   last one defined here. The range ends below the kernel-internal error
+ *   codes, which start at 512.
+ * * 1024 and above - Codes only the driver returns, added upward from 1024.
+ */
+
+/**
+ * NE_ERR_INVALID_FLAGS - The command's reserved flags field was not zero.
+ */
+#define NE_ERR_INVALID_FLAGS			(276)
+/**
+ * NE_ERR_INVALID_ALIGNMENT - An address or a size the command carries is
+ *			      misaligned or too small.
+ */
+#define NE_ERR_INVALID_ALIGNMENT		(277)
+/**
+ * NE_ERR_INVALID_SIZE - A size or count the command carries is zero, too large
+ *			 or too short.
+ */
+#define NE_ERR_INVALID_SIZE			(278)
+/**
+ * NE_ERR_INVALID_SLOT - The slot the command names does not exist.
+ */
+#define NE_ERR_INVALID_SLOT			(279)
+/**
+ * NE_ERR_INVALID_DEVICE - The command names a device, an address or a memory
+ *			   range the device does not accept.
+ */
+#define NE_ERR_INVALID_DEVICE			(280)
+/**
+ * NE_ERR_TOO_MANY_DEVICES - The enclave already has as many devices as it may.
+ */
+#define NE_ERR_TOO_MANY_DEVICES			(281)
+/**
+ * NE_ERR_TOO_MANY_VQS - The device already has as many queues as it may.
+ */
+#define NE_ERR_TOO_MANY_VQS			(282)
+/**
+ * NE_ERR_NO_RESOURCES - The device could not put together something of its own
+ *			 that the command needs.
+ */
+#define NE_ERR_NO_RESOURCES			(283)
+/**
+ * NE_ERR_ALREADY_EXISTS - The slot already has what the command gives it.
+ */
+#define NE_ERR_ALREADY_EXISTS			(284)
+/**
+ * NE_ERR_MISSING_INFO_PAGE - The slot has no info page, or the one it has is too
+ *			      small for what the command writes there.
+ */
+#define NE_ERR_MISSING_INFO_PAGE		(285)
+/**
+ * NE_ERR_INVALID_FLAG - A flag bit the device does not know, or does not allow
+ *			 here, was set.
+ */
+#define NE_ERR_INVALID_FLAG			(286)
+/**
+ * NE_ERR_RESERVED_FLAG - A reserved flag bit was set.
+ */
+#define NE_ERR_RESERVED_FLAG			(287)
+/**
+ * NE_ERR_MISSING_PCIE - A flag that depends on the PCIe flag was set without it.
+ */
+#define NE_ERR_MISSING_PCIE			(288)
+/**
+ * NE_ERR_INCOMPATIBLE_FLAGS - The flags asked for cannot be had together.
+ */
+#define NE_ERR_INCOMPATIBLE_FLAGS		(289)
+/**
+ * NE_ERR_EIF_FLAG_MISMATCH - The image's matched flags are not the ones the
+ *			      enclave was started with.
+ */
+#define NE_ERR_EIF_FLAG_MISMATCH		(290)
+/**
+ * NE_ERR_DEVICE_DISALLOWED - The device type asked for is not allowed.
+ */
+#define NE_ERR_DEVICE_DISALLOWED		(291)
+/**
+ * NE_ERR_HOTPLUG_DISABLED - The enclave does not take devices once it is
+ *			     running.
+ */
+#define NE_ERR_HOTPLUG_DISABLED			(292)
+/**
+ * NE_ERR_BAD_EIF - The enclave image is missing, or did not pass validation when
+ *		    the enclave was started.
+ */
+#define NE_ERR_BAD_EIF				(293)
+/**
+ * NE_ERR_MISSING_PCIE_VIRTIO - A flag that depends on the virtio-PCIe flag was
+ *				set without it.
+ */
+#define NE_ERR_MISSING_PCIE_VIRTIO		(294)
+/**
+ * NE_ERR_BAD_PAYLOAD_VERSION - The request's payload version is one the device
+ *				does not take.
+ */
+#define NE_ERR_BAD_PAYLOAD_VERSION		(295)
+/**
+ * NE_ERR_INFO_VECTORS_EXHAUSTED - No notification vector is left.
+ */
+#define NE_ERR_INFO_VECTORS_EXHAUSTED		(296)
+/**
+ * NE_ERR_TOO_MANY_ENCLAVES - The instance already has as many enclaves as its
+ *			      quota allows. A terminal answer: neither
+ *			      resending the command nor donating more memory
+ *			      changes it.
+ */
+#define NE_ERR_TOO_MANY_ENCLAVES		(297)
+/**
+ * NE_ERR_NO_MEM_SLOTS_ON_NODE - The memory budget of one NUMA node is used up,
+ *				 which -ENOSPC does not distinguish from the
+ *				 budget as a whole.
+ */
+#define NE_ERR_NO_MEM_SLOTS_ON_NODE		(298)
+/**
+ * NE_ERR_NUMA_SPAN - The enclave's CPUs and memory span more than one host NUMA
+ *		      node, and the enclave was not set up to span nodes.
+ */
+#define NE_ERR_NUMA_SPAN			(299)
+/**
+ * NE_ERR_NUMA_ALIGNMENT - The memory of a guest NUMA node other than the last
+ *			   is not a multiple of the guest's memory section size.
+ */
+#define NE_ERR_NUMA_ALIGNMENT			(300)
+/**
+ * NE_ERR_LOW_RAM_TOO_SMALL - The lowest contiguous run of enclave memory is too
+ *			      small for the boot layout the enclave needs.
+ */
+#define NE_ERR_LOW_RAM_TOO_SMALL		(301)
+/**
+ * NE_ERR_ZERO_REGION - The enclave's memory layout has a region of size zero.
+ */
+#define NE_ERR_ZERO_REGION			(302)
+
 /**
  * DOC: Image load info flags
  */
-- 
2.47.1

[PATCH 13/40] nitro_enclaves: Log a reached device limit at info level

[Alexander Graf]

When a parent asks for more enclave slots than the device has room for,
or for more enclaves than its quota allows, the device answers -ENOSPC
or NE_ERR_TOO_MANY_ENCLAVES and the driver logs an error for each
refused command, although user space gets the code back. Log the two at
info level.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_pci_dev.c | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index 2d44cae4ac0cf..f258e04ee9e47 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -261,8 +261,18 @@ int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 	if (cmd_reply->rc < 0) {
 		rc = cmd_reply->rc;
 
-		dev_err_ratelimited(&pdev->dev, "Error in cmd process logic, cmd type=%u [rc=%d]\n",
-				    cmd_type, rc);
+		if (rc == -ENOSPC) {
+			dev_info_ratelimited(&pdev->dev,
+					     "No room for cmd type=%u [rc=%d]\n",
+					     cmd_type, rc);
+		} else if (rc == -NE_ERR_TOO_MANY_ENCLAVES) {
+			dev_info_ratelimited(&pdev->dev,
+					     "Enclave quota reached for cmd type=%u [rc=%d]\n",
+					     cmd_type, rc);
+		} else {
+			dev_err_ratelimited(&pdev->dev, "Error in cmd process logic, cmd type=%u [rc=%d]\n",
+					    cmd_type, rc);
+		}
 
 		goto unlock_mutex;
 	}
-- 
2.47.1

[PATCH 14/40] nitro_enclaves: Add NE_SET_INFO_PAGE

[Alexander Graf]

Nitro Sandboxes introduce the info page, a page of parent memory that
the device writes and the parent reads, to share various information about
the state of devices inside the slot.

Add the NE_SET_INFO_PAGE ioctl to register a NE_MEMORY_SHARED allocation
as info page. This page is always read-only to user space.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 276 +++++++++++++++++++++-
 drivers/virt/nitro_enclaves/ne_misc_dev.h |  35 ++-
 drivers/virt/nitro_enclaves/ne_pci_dev.c  |  25 ++
 drivers/virt/nitro_enclaves/ne_pci_dev.h  |  21 ++
 include/uapi/linux/nitro_enclaves.h       | 106 ++++++++-
 5 files changed, 435 insertions(+), 28 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 0bf62f703806e..610d98264e5a1 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -123,6 +123,8 @@ static int ne_add_memory_region_ioctl(struct ne_enclave *ne_enclave,
 				      struct ne_add_memory_region_args *args);
 static int ne_alloc_memory_ioctl(struct ne_enclave *ne_enclave,
 				 struct ne_alloc_memory_args *args);
+static int ne_set_info_page_ioctl(struct ne_enclave *ne_enclave,
+				  struct ne_set_info_page_args *args);
 static const struct file_operations ne_fops = {
 	.owner		= THIS_MODULE,
 	.llseek		= noop_llseek,
@@ -2137,6 +2139,39 @@ static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned lo
 		return rc;
 	}
 
+	case NE_EA_IOCTL(NE_SET_INFO_PAGE): {
+		struct ne_set_info_page_args __user *uargs = (void __user *)arg;
+		unsigned int usize = _IOC_SIZE(cmd);
+		struct ne_set_info_page_args args;
+		int rc;
+
+		if (usize < offsetofend(struct ne_set_info_page_args, flags))
+			return -EINVAL;
+
+		rc = copy_struct_from_user(&args, sizeof(args), uargs, usize);
+		if (rc)
+			return rc;
+
+		rc = mutex_lock_killable(&ne_enclave->enclave_info_mutex);
+		if (rc)
+			return rc;
+
+		if (ne_enclave->state != NE_STATE_INIT) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Enclave is not in init state\n");
+
+			mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+			return -NE_ERR_NOT_IN_INIT_STATE;
+		}
+
+		rc = ne_set_info_page_ioctl(ne_enclave, &args);
+
+		mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+		return rc;
+	}
+
 	default:
 		return -ENOTTY;
 	}
@@ -2344,6 +2379,9 @@ static void ne_enclave_free(struct ne_enclave *ne_enclave)
 	ne_enclave_remove_all_mem_region_entries(ne_enclave);
 	ne_enclave_remove_all_vcpu_id_entries(ne_enclave);
 
+	if (ne_enclave->info_file)
+		fput(ne_enclave->info_file);
+
 	kfree(ne_enclave);
 }
 
@@ -3152,6 +3190,9 @@ static bool ne_contig_supports_mmap(struct file *file)
  * asks for VM_HUGEPAGE because a fault on a VM_PFNMAP mapping without it only
  * reaches ->huge_fault where the transparent_hugepage policy is "always".
  *
+ * An allocation that is a slot's info page is mapped without write access, and
+ * without the right to gain it with mprotect().
+ *
  * Return: 0 on success, negative return value on failure.
  */
 static int ne_contig_mmap_prepare(struct vm_area_desc *desc)
@@ -3177,6 +3218,13 @@ static int ne_contig_mmap_prepare(struct vm_area_desc *desc)
 	if (offset > mem->size || size > mem->size - offset)
 		return -EINVAL;
 
+	if (READ_ONCE(mem->read_only)) {
+		if (vma_desc_test(desc, VMA_WRITE_BIT))
+			return -EACCES;
+
+		vma_desc_clear_flags(desc, VMA_MAYWRITE_BIT);
+	}
+
 	desc->vm_ops = &ne_contig_vm_ops;
 	vma_desc_set_flags(desc, VMA_IO_BIT, VMA_PFNMAP_BIT, VMA_DONTEXPAND_BIT,
 			   VMA_DONTDUMP_BIT, VMA_DONTCOPY_BIT, VMA_HUGEPAGE_BIT);
@@ -3242,11 +3290,12 @@ static int ne_write_span(void *arg, struct page *page, u64 off, u64 len)
  * refusal is -EFBIG, the same as generic_write_check_limits() gives for a write
  * beyond s_maxbytes. There is no RLIMIT_FSIZE in it, so no SIGXFSZ either.
  *
- * Donating an allocation closes this path, and @lock is what makes the two
- * orderly. It is taken here and in ne_add_memory_region_ioctl() and nowhere
- * else; setting a mapping up and faulting one in both read only fields that
- * never change, so nothing takes it under mmap_write_lock or in a fault, and
- * copying from user space under it cannot invert against the fault path.
+ * Donating an allocation closes this path, and so does making it an info page.
+ * This function, ne_add_memory_region_ioctl(), ne_contig_make_read_only() and
+ * ne_contig_allow_writes() take @lock, which orders those state changes against
+ * a write. Setting a mapping up and faulting one in do not take it, so nothing
+ * takes it under mmap_write_lock or in a fault, and copying from user space
+ * under it cannot invert against the fault path.
  *
  * Return: the number of bytes written, or a negative return value on failure.
  */
@@ -3266,7 +3315,7 @@ static ssize_t ne_contig_write_iter(struct kiocb *iocb, struct iov_iter *from)
 
 	guard(mutex)(&mem->lock);
 
-	if (mem->donated)
+	if (mem->donated || mem->read_only)
 		return -EPERM;
 
 	rc = ne_walk_allocs(mem, pos, count, NE_WRITE_SPAN_PAGES, ne_write_span,
@@ -3281,6 +3330,60 @@ static ssize_t ne_contig_write_iter(struct kiocb *iocb, struct iov_iter *from)
 	return rc;
 }
 
+/**
+ * ne_contig_make_read_only() - Take write access to an allocation away from
+ *				user space.
+ * @file:	The allocation's file.
+ *
+ * From here on a write to the descriptor fails and a new mapping has no write
+ * access, so only the device and the kernel change the memory. The descriptor
+ * is open for writing, so a mapping made before could be made writable with
+ * mprotect(), and one that exists makes this fail.
+ *
+ * mmap_region() raises i_mmap_writable with full ordering before it calls
+ * ne_contig_mmap_prepare() for a mapping that may write, and the mapping keeps
+ * it raised for as long as it exists. @read_only is set before the count is
+ * read, so a mapping that ne_contig_mmap_prepare() set up without seeing it is
+ * counted here. A mapping that is being set up as this runs counts too, even
+ * one that ends up without write access, so a caller sees -EBUSY for it and
+ * can try again.
+ *
+ * Return: 0 on success, -EBUSY if the allocation is mapped.
+ */
+static int ne_contig_make_read_only(struct file *file)
+{
+	struct ne_contig_mem *mem = file->private_data;
+
+	guard(mutex)(&mem->lock);
+
+	WRITE_ONCE(mem->read_only, true);
+
+	/* Pairs with the full ordering of mapping_map_writable(). */
+	smp_mb();
+
+	if (mapping_writably_mapped(file->f_mapping)) {
+		WRITE_ONCE(mem->read_only, false);
+
+		return -EBUSY;
+	}
+
+	return 0;
+}
+
+/**
+ * ne_contig_allow_writes() - Give write access to an allocation back to user
+ *			      space.
+ * @file:	The allocation's file.
+ */
+static void ne_contig_allow_writes(struct file *file)
+{
+	struct ne_contig_mem *mem = file->private_data;
+
+	guard(mutex)(&mem->lock);
+
+	WRITE_ONCE(mem->read_only, false);
+}
+
 /**
  * ne_contig_release() - Free an allocation once nothing holds it any more.
  * @inode:	The allocation's inode, unused.
@@ -3612,6 +3715,13 @@ static int ne_add_memory_region_ioctl(struct ne_enclave *ne_enclave,
 		return -EINVAL;
 	}
 
+	if (fd_file(f) == ne_enclave->info_file) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "The allocation is the enclave's info page\n");
+
+		return -NE_ERR_MEM_REGION_ALREADY_USED;
+	}
+
 	list_for_each_entry(held, &ne_enclave->contig_regions_list, list) {
 		if (held->file == fd_file(f)) {
 			dev_err_ratelimited(ne_misc_dev.this_device,
@@ -3714,6 +3824,160 @@ static int ne_add_memory_region_ioctl(struct ne_enclave *ne_enclave,
 	return rc;
 }
 
+/**
+ * ne_set_info_page_ioctl() - Give the slot associated with the current enclave
+ *			      the info page.
+ * @ne_enclave:	Private data associated with the current enclave.
+ * @args:	Descriptor for the allocation, and flags.
+ *
+ * The device writes the info page, so the run behind it has to stay where the
+ * device was told it is for as long as the enclave lives. The runs belong to the
+ * file, so a reference on it keeps them whatever user space does with the
+ * descriptor and its mappings, and ne_enclave_free() drops that reference
+ * after SLOT_FREE and after the enclave's memory regions. The reference is taken
+ * before the request, since a request that timed out may still be completed by
+ * the device. A rejected request drops it again.
+ *
+ * Only the device writes the page, and the parent reads back from it what the
+ * device reports about the slot, so user space loses write access to the
+ * allocation before the device is told about it and gets it back only if the
+ * device rejects the request.
+ *
+ * Context: Process context. This function is called with the ne_enclave mutex held.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ * * -ETIMEDOUT if the device did not answer, in which case the slot keeps the
+ *   allocation until the enclave goes away.
+ */
+static int ne_set_info_page_ioctl(struct ne_enclave *ne_enclave,
+				  struct ne_set_info_page_args *args)
+{
+	struct ne_pci_dev_cmd_reply cmd_reply = {};
+	struct slot_add_info_page_req req = {};
+	struct pci_dev *pdev = ne_devs.ne_pci_dev->pdev;
+	struct ne_contig_mem *mem = NULL;
+	int rc = -EINVAL;
+
+	if (ne_enclave->mm != current->mm)
+		return -EIO;
+
+	if (args->flags) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Incorrect flag for info page\n");
+
+		return -NE_ERR_INVALID_FLAG_VALUE;
+	}
+
+	if (ne_enclave->info_file) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "The slot already has an info page\n");
+
+		return -EEXIST;
+	}
+
+	CLASS(fd, f)(args->fd);
+
+	if (fd_empty(f))
+		return -EBADF;
+
+	if (fd_file(f)->f_op != &ne_contig_fops) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Not an enclave memory allocation\n");
+
+		return -EINVAL;
+	}
+
+	mem = fd_file(f)->private_data;
+
+	if (mem->enclave_id != ne_enclave->id) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Allocation was made for another enclave\n");
+
+		return -EINVAL;
+	}
+
+	if (!mem->shared) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "The info page allocation is not shared\n");
+
+		return -EINVAL;
+	}
+
+	/*
+	 * The device takes one address and one length, so an allocation the
+	 * allocator had to satisfy in several runs has no single address to
+	 * name. A request of NE_MIN_MEM_REGION_SIZE cannot be split, since the
+	 * allocator has nothing smaller to fall back to.
+	 */
+	if (mem->nr_allocs != 1) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "The info page allocation is not one run\n");
+
+		return -EINVAL;
+	}
+
+	/*
+	 * An allocation the enclave already has as memory is not available to
+	 * be its info page as well, since the device would be handed the same
+	 * run twice in two roles. @donated is only ever set from
+	 * ne_add_memory_region_ioctl(), which runs under the same enclave mutex
+	 * as this function for the one enclave the allocation was made for, so
+	 * the read needs no lock of its own.
+	 */
+	if (mem->donated) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "The info page allocation was donated\n");
+
+		return -NE_ERR_MEM_REGION_ALREADY_USED;
+	}
+
+	rc = ne_contig_make_read_only(fd_file(f));
+	if (rc < 0) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "The info page allocation is mapped\n");
+
+		return rc;
+	}
+
+	req.slot_uid = ne_enclave->slot_uid;
+	req.info_phys = page_to_phys(mem->allocs[0].page);
+	req.info_size = mem->size;
+
+	/*
+	 * Take the reference before the request, so that the run stays where the
+	 * device is told it is from the moment the device is told, as the memory
+	 * region path does.
+	 */
+	ne_enclave->info_file = get_file(fd_file(f));
+
+	rc = ne_do_request(pdev, SLOT_ADD_INFO_PAGE,
+			   &req, sizeof(req), &cmd_reply, sizeof(cmd_reply));
+	if (rc < 0) {
+		dev_err_ratelimited(ne_misc_dev.this_device,
+				    "Error in slot add info page [rc=%d]\n", rc);
+
+		/*
+		 * A request the device never answered was submitted with the
+		 * run's physical address in it, so the device may write the page
+		 * yet. Keep the reference and let ne_enclave_free() drop it
+		 * after SLOT_FREE. Every other error means the device either
+		 * never saw the request or answered to reject it, and the run is
+		 * the parent's again. ne_request_unanswered() tells the two
+		 * apart.
+		 */
+		if (!ne_request_unanswered(rc)) {
+			ne_contig_allow_writes(ne_enclave->info_file);
+			fput(ne_enclave->info_file);
+			ne_enclave->info_file = NULL;
+		}
+
+		return rc;
+	}
+
+	return 0;
+}
+
 #if defined(CONFIG_NITRO_ENCLAVES_MISC_DEV_TEST)
 #include "ne_misc_dev_test.c"
 #endif
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 2502cca3f3f64..d0d1973bfbe9d 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -38,27 +38,37 @@ struct ne_contig_alloc {
  * @allocs:	Contiguous runs backing the allocation, ordered by offset.
  * @nr_allocs:	Number of runs in @allocs.
  * @size:	Size of the allocation in bytes, mirrored in i_size.
- * @shared:	Whether the allocation may be mapped by the caller.
+ * @shared:	Whether the caller asked for a shared allocation. A shared one may
+ *		be mapped by the caller and may be a slot's info page; a private
+ *		one is neither.
  * @enclave_id:	Id of the enclave the allocation was made for.
  * @mm:		The mm charged for the pinned pages, held via mmgrab().
  * @nr_pinned:	Pages charged to @mm->pinned_vm.
- * @lock:	Serialises a write against the donation that ends writing.
+ * @lock:	Serialises a write against the donation or the info page that
+ *		ends writing.
  * @donated:	Whether the allocation has been given to its enclave.
+ * @read_only:	Whether user space has lost write access to the allocation,
+ *		which it does when the allocation becomes a slot's info page.
  *
  * The runs hang off the file, so that closing the last descriptor does not free
  * memory an enclave is still using: a donated region holds the file with
- * get_file(). That is also what keeps a mapping's page table entries from
- * outliving the runs, since a VMA holds vm_file for as long as it exists.
+ * get_file(), and so does the enclave the allocation is the info page of. That
+ * is also what keeps a mapping's page table entries from outliving the runs,
+ * since a VMA holds vm_file for as long as it exists.
  *
  * The array is built before the file is published and not mutated afterwards, so
  * any holder of a reference can read it without a lock of its own, which is what
  * lets a page fault resolve an offset to a run. @shared and @size are set in the
- * same place and are read the same way. @donated is the one field that changes
- * after publication, and @lock covers it.
+ * same place and are read the same way. @donated and @read_only are the fields
+ * that change after publication. @lock covers both on the write path;
+ * ne_set_info_page_ioctl() reads @donated under the enclave mutex its writer also
+ * holds, and ne_contig_mmap_prepare() reads @read_only without a lock, which
+ * ne_contig_make_read_only() orders against.
  *
  * @enclave_id is compared, never dereferenced, and the counter it comes from is
- * 64 bits wide, so an allocation cannot be donated to a later enclave that
- * happens to have been placed at the address of the one it was made for.
+ * 64 bits wide, so an allocation cannot be donated to, or become the info page
+ * of, a later enclave that happens to have been placed at the address of the one
+ * it was made for.
  */
 struct ne_contig_mem {
 	struct ne_contig_alloc	*allocs;
@@ -68,9 +78,10 @@ struct ne_contig_mem {
 	u64			enclave_id;
 	struct mm_struct	*mm;
 	unsigned long		nr_pinned;
-	/* Serialises a write against the donation that ends writing. */
+	/* Serialises a write against what ends writing. */
 	struct mutex		lock;
 	bool			donated;
+	bool			read_only;
 };
 
 /**
@@ -116,6 +127,11 @@ struct ne_mem_region {
  * @id:				Id of this enclave, unique for the lifetime of the
  *				driver, recorded in an %NE_ALLOC_MEMORY allocation
  *				as the enclave it was made for.
+ * @info_file:			File of the %NE_ALLOC_MEMORY allocation the slot
+ *				was given as its info page, held with get_file()
+ *				so that the device keeps the page after user space
+ *				has closed the descriptor. NULL until
+ *				%NE_SET_INFO_PAGE is called.
  * @max_mem_regions:		The maximum number of memory regions that can be
  *				handled by the hypervisor.
  * @mem_regions_list:		Enclave user space memory regions list.
@@ -146,6 +162,7 @@ struct ne_enclave {
 	wait_queue_head_t	eventq;
 	bool			has_event;
 	u64			id;
+	struct file		*info_file;
 	u64			max_mem_regions;
 	struct list_head	mem_regions_list;
 	u64			mem_size;
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index f258e04ee9e47..a07dbc6d8a9b7 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -128,6 +128,31 @@ static int ne_wait_for_reply(struct pci_dev *pdev, unsigned int timeout_msecs)
 	return 0;
 }
 
+/**
+ * ne_request_unanswered() - Whether a failure left a submitted request in
+ *			     flight.
+ * @rc:		Failure a command request came back with.
+ *
+ * A request the device never answered was submitted with whatever the caller
+ * put in it, and there is nothing to cancel it with, so a resource whose
+ * address is in it may still be written by the device. A caller that has to
+ * decide who owns such a resource asks here instead of naming an errno of its
+ * own.
+ *
+ * Every code the wait for a reply can fail a submitted request with belongs
+ * here. ne_wait_for_reply() times out, so -ETIMEDOUT is named. A code raised
+ * before ne_submit_request(), and a code the device answered with, do not:
+ * nothing is in flight in either case. A
+ * device that answers one of the codes named here of its own accord lands here
+ * too and costs one held resource, which is the side to be wrong on.
+ *
+ * Return: true if the request may still be in flight.
+ */
+bool ne_request_unanswered(int rc)
+{
+	return rc == -ETIMEDOUT;
+}
+
 /**
  * ne_cmd_frees_slot() - Whether a command is one of the two that free a slot.
  * @cmd_type:	Command type of the request.
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index b0957e2606ec3..4f670c53dd3c0 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -123,6 +123,8 @@
  * @NEXT_SLOT:			Get the next slot in the list of allocated slots.
  * @SLOT_INFO:			Get the info for a slot e.g. slot uid, vCPUs count.
  * @SLOT_ADD_BULK_VCPUS:	Add a number of vCPUs, not providing CPU ids.
+ * @SLOT_ADD_INFO_PAGE:		Give a slot the info page, a memory region the
+ *				device writes.
  * @MAX_CMD:			A gatekeeper for max possible command type.
  */
 enum ne_pci_dev_cmd_type {
@@ -138,6 +140,8 @@ enum ne_pci_dev_cmd_type {
 	NEXT_SLOT		= 9,
 	SLOT_INFO		= 10,
 	SLOT_ADD_BULK_VCPUS	= 11,
+	/* 12 is a command this driver does not implement. */
+	SLOT_ADD_INFO_PAGE	= 13,
 	MAX_CMD,
 };
 
@@ -266,6 +270,20 @@ struct slot_add_bulk_vcpus_req {
 	u64	nr_vcpus;
 };
 
+/**
+ * struct slot_add_info_page_req - SLOT_ADD_INFO_PAGE request.
+ * @slot_uid:	Slot unique id mapped to the slot to give the info page to.
+ * @flags:	Flags for the info page. Zero is the only value the driver sends.
+ * @info_phys:	Physical address of the info page.
+ * @info_size:	Size, in bytes, of the info page.
+ */
+struct slot_add_info_page_req {
+	u64	slot_uid;
+	u64	flags;
+	u64	info_phys;
+	u64	info_size;
+};
+
 /**
  * struct ne_pci_dev_cmd_reply - NE PCI device command reply.
  * @rc :		Return code of the logic that processed the request.
@@ -327,6 +345,9 @@ struct ne_pci_dev {
 	struct pci_dev		*pdev;
 };
 
+/* Whether a failure left a submitted request in flight. */
+bool ne_request_unanswered(int rc);
+
 /**
  * ne_do_request() - Submit command request to the PCI device based on the command
  *		     type and retrieve the associated reply.
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index 6dd00c41c07d5..4f604712983f1 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -287,13 +287,15 @@ struct ne_alloc_memory_args {
 
 /**
  * NE_ALLOC_MEMORY - The command is used to allocate physically contiguous
- *		     memory to give to an enclave with %NE_ADD_MEMORY_REGION,
- *		     drawn from one host NUMA node. The argument is input only;
+ *		     memory to give to an enclave with %NE_ADD_MEMORY_REGION or
+ *		     to provide to the device with %NE_SET_INFO_PAGE, drawn from
+ *		     one host NUMA node. The argument is input only;
  *		     on success the return value is a file descriptor for the
  *		     allocation.
  *		     Writing to the descriptor places bytes in the allocation,
  *		     which is how an enclave image gets there, and stops
- *		     working once the allocation has been donated.
+ *		     working once the allocation has been donated or provided
+ *		     with %NE_SET_INFO_PAGE.
  *		     O_APPEND has no effect. There is no read operation, so
  *		     read() and pread() both fail with EINVAL and a caller that
  *		     wants to check what it wrote keeps its own copy.
@@ -308,16 +310,18 @@ struct ne_alloc_memory_args {
  *		     allocation over, including once the allocation has been
  *		     donated and writing to the descriptor has stopped working.
  *		     Such a mapping is readable too, which is the only way to
- *		     read allocated memory back. The mapping is not inherited
- *		     across fork(); a child that needs one maps the descriptor
- *		     itself.
+ *		     read allocated memory back, and is how a parent reads an
+ *		     info page provided with %NE_SET_INFO_PAGE, which it maps
+ *		     without %PROT_WRITE. The mapping is not inherited across
+ *		     fork(); a child that needs one maps the descriptor itself.
  *		     Donating does not take a mapping away, so a parent that
  *		     mapped a shared allocation can still write to it through
  *		     that mapping afterwards.
  *		     The allocation lives until the last of the descriptor, its
- *		     mappings and the enclave regions donated from it is gone,
- *		     so a caller can close the descriptor as soon as it has
- *		     donated.
+ *		     mappings, the enclave regions donated from it and the
+ *		     enclave it is the info page of is gone, so a caller can
+ *		     close the descriptor as soon as it has donated it or
+ *		     provided it to the device.
  *		     The ioctl can be invoked on the enclave fd, before an
  *		     enclave is started.
  *
@@ -351,7 +355,8 @@ struct ne_alloc_memory_args {
  *
  * Writing to the descriptor sets errno to:
  * * EPERM			- The allocation has been donated with
- *				  %NE_ADD_MEMORY_REGION.
+ *				  %NE_ADD_MEMORY_REGION or provided with
+ *				  %NE_SET_INFO_PAGE.
  * * EFBIG			- The write does not fit inside the allocation.
  * * EFAULT			- Part of the buffer is not readable. Any bytes
  *				  copied before that are reported as a short
@@ -361,6 +366,9 @@ struct ne_alloc_memory_args {
  *
  * An mmap() of the descriptor sets errno to:
  * * ENODEV			- The allocation is %NE_MEMORY_PRIVATE.
+ * * EACCES			- The mapping asks for %PROT_WRITE and the
+ *				  allocation has been provided with
+ *				  %NE_SET_INFO_PAGE.
  * * EINVAL			- The mapping is %MAP_PRIVATE, or its offset or
  *				  length is not a multiple of 2 MiB, or it does
  *				  not fit inside the allocation.
@@ -388,7 +396,8 @@ struct ne_add_memory_region_args {
  *			  contiguous run it is made of, or fewer where runs turn
  *			  out to be physically adjacent. The argument is input
  *			  only. The allocation has to have been made for this
- *			  enclave, and can be given away once.
+ *			  enclave, cannot be its info page, and can be given away
+ *			  once.
  *			  Writing to the descriptor stops working here, so an
  *			  enclave image goes in before this call. The allocation
  *			  has to be %NE_MEMORY_PRIVATE, since memory the parent
@@ -425,13 +434,82 @@ struct ne_add_memory_region_args {
  *					  before being started).
  * * NE_ERR_INVALID_FLAG_VALUE		- The value of the provided flag is invalid.
  * * NE_ERR_MEM_REGION_ALREADY_USED	- The allocation was already given to the
- *					  enclave.
+ *					  enclave, or is its info page.
  * * NE_ERR_MEM_MAX_REGIONS		- The number of memory regions set for the
  *					  enclave reached maximum.
  * * Error codes from the NE PCI device request.
  */
 #define NE_ADD_MEMORY_REGION		_IOW(0xAE, 0x2F, struct ne_add_memory_region_args)
 
+/**
+ * struct ne_set_info_page_args - Allocation to provide to the device (in).
+ * @fd:		Descriptor from %NE_ALLOC_MEMORY, made for this enclave, asked
+ *		for as %NE_MEMORY_SHARED, and made of one contiguous run.
+ * @flags:	Must be zero.
+ */
+struct ne_set_info_page_args {
+	__s32	fd;
+	__u32	flags;
+};
+
+/**
+ * NE_SET_INFO_PAGE - The command is used to give an enclave's slot the info
+ *		      page, a %NE_ALLOC_MEMORY allocation the device writes. The
+ *		      argument is input only. The allocation has to be
+ *		      %NE_MEMORY_SHARED, since a mapping is the only way for
+ *		      the parent to read an allocation back. It also has to be
+ *		      one contiguous run, since the device is given one address
+ *		      and one length for it, and a request of 2 MiB is the one
+ *		      size the allocator cannot come back from in several runs.
+ *		      The info page is not enclave memory: it does not count
+ *		      towards the memory regions %NE_ADD_MEMORY_REGION fills.
+ *		      Only the device writes the page, so from this call on the
+ *		      allocation is read-only to user space: writing to the
+ *		      descriptor fails, and the descriptor is mapped without
+ *		      %PROT_WRITE. A mapping made before the call could still be
+ *		      made writable, so the call fails while one exists, and a
+ *		      caller maps the page after the call. A call the device
+ *		      rejects leaves the allocation writable again.
+ *		      A slot takes one info page, and an allocation is either an
+ *		      enclave's memory or its info page. The driver holds the
+ *		      allocation until the enclave is gone, so closing the
+ *		      descriptor does not free the run the device was given.
+ *		      The ioctl can be invoked on the enclave fd, before an
+ *		      enclave is started.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0					- Logic successfully completed.
+ * * -1					- There was a failure in the ioctl logic.
+ * On failure, errno is set to:
+ * * ENODEV				- The NE PCI device is not available.
+ * * EFAULT				- Argument read failure.
+ * * EBADF				- @fd is not an open file descriptor.
+ * * EINVAL				- The argument is shorter than @flags, or @fd
+ *					  is not a descriptor from %NE_ALLOC_MEMORY,
+ *					  or the allocation was made for another
+ *					  enclave, or it is %NE_MEMORY_PRIVATE, or it
+ *					  is made of more than one contiguous run.
+ * * EIO				- Current task mm is not the same as the one
+ *					  that created the enclave.
+ * * EEXIST				- The slot already has an info page.
+ * * EBUSY				- The allocation has a mapping made before the
+ *					  call, or one is being set up.
+ * * E2BIG				- The argument is longer than this kernel knows
+ *					  about and the excess is not all zero.
+ * * NE_ERR_NOT_IN_INIT_STATE		- The enclave is not in init state (init =
+ *					  before being started).
+ * * NE_ERR_INVALID_FLAG_VALUE		- The value of the provided flag is invalid.
+ * * NE_ERR_MEM_REGION_ALREADY_USED	- The allocation was already given to the
+ *					  enclave with %NE_ADD_MEMORY_REGION.
+ * * ETIMEDOUT				- The device did not answer. The slot keeps
+ *					  the allocation, read-only, until the
+ *					  enclave is gone, since the device may
+ *					  still write the page.
+ * * Error codes from the NE PCI device request.
+ */
+#define NE_SET_INFO_PAGE		_IOW(0xAE, 0x25, struct ne_set_info_page_args)
+
 /**
  * NE_START_ENCLAVE - The command is used to trigger enclave start after the
  *		      enclave resources, such as memory and CPU, have been set.
@@ -775,7 +853,9 @@ struct ne_user_memory_region {
  *		      it has been donated. It costs page tables for the part of
  *		      it a caller touches, and it leaves the memory reachable
  *		      from the parent for as long as it lives, so an allocation
- *		      that is only being handed over does not ask for one.
+ *		      that is only being handed over does not ask for one. It is
+ *		      also what an allocation given to %NE_SET_INFO_PAGE has to
+ *		      be.
  */
 #define NE_MEMORY_SHARED		(0x01)
 
-- 
2.47.1

[PATCH 15/40] nitro_enclaves: Document the info page in the memory pool

[Alexander Graf]

An info page takes one contiguous allocation from the same CMA
reservation as enclave memory, so an admin reserving the pool needs to
leave room for it. Document its size, that it is registered before it is
mapped and is read-only to user space from then on, and that it keeps
its allocation alive like a memory region does. Also say that a larger
NE_ALLOC_MEMORY request comes back in several runs where it has to.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 Documentation/virt/ne_overview.rst | 28 +++++++++++++++++++++-------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/Documentation/virt/ne_overview.rst b/Documentation/virt/ne_overview.rst
index 4c4b7febaff0a..3cfc50f788433 100644
--- a/Documentation/virt/ne_overview.rst
+++ b/Documentation/virt/ne_overview.rst
@@ -150,10 +150,11 @@ node enables to name the node a reservation lands on. cma= needs neither, so
 on a host with one online node the driver's own select is the whole
 requirement.
 
-NE_ALLOC_MEMORY on the enclave file descriptor allocates 2 MiB aligned,
-physically contiguous memory from one host NUMA node's CMA area. A request
-that area cannot serve fails with -ENOMEM, so the reservation bounds what the
-enclaves on a node can take between them.
+NE_ALLOC_MEMORY on the enclave file descriptor allocates 2 MiB aligned memory
+from one host NUMA node's CMA area, in one physically contiguous run where it
+can and several where it cannot. A request that area cannot serve fails with
+-ENOMEM, so the reservation bounds what the enclaves on a node can take
+between them.
 It returns a file descriptor for the allocation. For an enclave whose cores
 all sit on one node the ioctl refuses any other node, so name the same node
 here and in NE_ADD_ANY_VCPU. An enclave built out of a pool that spans nodes
@@ -195,9 +196,22 @@ a caller reaches and not for all of it. It keeps the memory reachable from the
 parent for as long as it lives, so an allocation that is only being handed to
 an enclave stays private.
 
+An enclave can also have an info page: an allocation the NE PCI device writes,
+provided to the device by NE_SET_INFO_PAGE. It does not become the enclave's
+memory. Only the device writes it, so from NE_SET_INFO_PAGE on it is read-only
+to user space: write() on the descriptor fails and a mapping has no write
+access. The call fails while a mapping made before it exists, so a caller
+maps the page after it. It is asked for as NE_MEMORY_SHARED, since a mapping is
+the only way for the parent to read back what the device wrote. The device is
+given one address and one length for it, so it has to be one contiguous run,
+and 2 MiB is the only size that cannot come back split. Ask for 2 MiB per
+enclave that will have one, on the node its cores sit on, or on every node it
+may draw from where the pool spans nodes.
+
 NE_SET_USER_MEMORY_REGION is unchanged and still takes the address of a
 hugetlbfs region.
 
-The allocation lives until the last of its descriptor, its mappings and the
-regions given from it is gone, so neither closing the descriptor nor
-unmapping takes memory away from a running enclave.
+An allocation lives until the last of the things holding it is gone: its
+descriptor, its mappings, the regions given from it and the enclave it is the
+info page of. So neither closing the descriptor nor unmapping takes memory away
+from a running enclave.
-- 
2.47.1

[PATCH 16/40] nitro_enclaves: Turn the command reply into a per-command union

[Alexander Graf]

In preparation to support Nitro Sandboxes, whose commands carry
unrelated payloads in their replies, convert struct ne_pci_dev_cmd_reply
into a per-command union. A command with fields of its own in the reply
gets an arm holding them, and a caller reads the reply through the arm
for the command it issued.

The device writes every reply into the same 72 byte window of the PCI
BAR and picks the offset of each field itself, so we place an arm's
fields at the offsets the flat struct gave them and name the holes where
another command owns the bytes. static_assert() on every offset and on
the total size pins each arm to the layout the device uses.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c |  6 +-
 drivers/virt/nitro_enclaves/ne_pci_dev.c  |  4 +-
 drivers/virt/nitro_enclaves/ne_pci_dev.h  | 86 +++++++++++++++++------
 3 files changed, 71 insertions(+), 25 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 610d98264e5a1..77b15eda8d120 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -1749,7 +1749,7 @@ static int ne_start_enclave_ioctl(struct ne_enclave *ne_enclave,
 
 	ne_enclave->state = NE_STATE_RUNNING;
 
-	enclave_start_info->enclave_cid = cmd_reply.enclave_cid;
+	enclave_start_info->enclave_cid = cmd_reply.start.enclave_cid;
 
 	return 0;
 }
@@ -2621,11 +2621,11 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 flags,
 	init_waitqueue_head(&ne_enclave->eventq);
 	ne_enclave->has_event = false;
 	mutex_init(&ne_enclave->enclave_info_mutex);
-	ne_enclave->max_mem_regions = cmd_reply.mem_regions;
+	ne_enclave->max_mem_regions = cmd_reply.alloc.mem_regions;
 	INIT_LIST_HEAD(&ne_enclave->contig_regions_list);
 	INIT_LIST_HEAD(&ne_enclave->mem_regions_list);
 	ne_enclave->mm = current->mm;
-	ne_enclave->slot_uid = cmd_reply.slot_uid;
+	ne_enclave->slot_uid = cmd_reply.alloc.slot_uid;
 	ne_enclave->state = NE_STATE_INIT;
 
 	list_add(&ne_enclave->enclave_list_entry, &ne_pci_dev->enclaves_list);
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index a07dbc6d8a9b7..8ed2349dea878 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -404,8 +404,8 @@ static void ne_event_work_handler(struct work_struct *work)
 		}
 
 		/* Notify enclave process that the enclave state changed. */
-		if (ne_enclave->state != cmd_reply.state) {
-			ne_enclave->state = cmd_reply.state;
+		if (ne_enclave->state != cmd_reply.info.state) {
+			ne_enclave->state = cmd_reply.info.state;
 
 			ne_enclave->has_event = true;
 
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index 4f670c53dd3c0..835413fca8040 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -286,32 +286,78 @@ struct slot_add_info_page_req {
 
 /**
  * struct ne_pci_dev_cmd_reply - NE PCI device command reply.
- * @rc :		Return code of the logic that processed the request.
- * @padding0:		Padding for the overall data structure.
- * @slot_uid:		Valid for all commands except SLOT_COUNT.
- * @enclave_cid:	Valid for ENCLAVE_START command.
- * @slot_count :	Valid for SLOT_COUNT command.
- * @mem_regions:	Valid for SLOT_ALLOC and SLOT_INFO commands.
- * @mem_size:		Valid for SLOT_INFO command.
- * @nr_vcpus:		Valid for SLOT_INFO command.
- * @flags:		Valid for SLOT_INFO command.
- * @state:		Valid for SLOT_INFO command.
- * @padding1:		Padding for the overall data structure.
+ *
+ * The device writes the reply of every command into the same 72-byte
+ * window. A command whose reply carries fields of its own has an arm of
+ * the union below holding them, at the offsets the device writes them
+ * to, and a caller reads through the arm for the command it issued.
+ *
+ * @rc :			Return code of the logic that processed the
+ *				request.
+ * @padding0:			Padding for the overall data structure.
+ * @start:			Fields ENCLAVE_START fills.
+ * @start.slot_uid:		ENCLAVE_START. Slot of the started enclave.
+ * @start.enclave_cid:		ENCLAVE_START. Context ID of the enclave
+ *				vsock device.
+ * @alloc:			Fields SLOT_ALLOC fills.
+ * @alloc.slot_uid:		SLOT_ALLOC. Slot unique id of the new slot.
+ * @alloc.reserved8:		SLOT_ALLOC. Bytes owned by other commands.
+ * @alloc.mem_regions:		SLOT_ALLOC. Maximum number of memory
+ *				regions the slot accepts.
+ * @info:			Fields SLOT_INFO fills.
+ * @info.slot_uid:		SLOT_INFO. Slot the reply describes.
+ * @info.enclave_cid:		SLOT_INFO. Context ID of the enclave vsock
+ *				device.
+ * @info.reserved24:		SLOT_INFO. Bytes owned by other commands.
+ * @info.mem_regions:		SLOT_INFO. Maximum number of memory regions
+ *				the slot accepts.
+ * @info.mem_size:		SLOT_INFO. Memory size added to the slot.
+ * @info.nr_vcpus:		SLOT_INFO. Number of vCPUs added to the slot.
+ * @info.flags:			SLOT_INFO. Slot flags.
+ * @info.state:			SLOT_INFO. Slot state.
+ * @info.padding1:		SLOT_INFO. Padding for the overall data
+ *				structure.
  */
 struct ne_pci_dev_cmd_reply {
 	s32	rc;
 	u8	padding0[4];
-	u64	slot_uid;
-	u64	enclave_cid;
-	u64	slot_count;
-	u64	mem_regions;
-	u64	mem_size;
-	u64	nr_vcpus;
-	u64	flags;
-	u16	state;
-	u8	padding1[6];
+	union {
+		struct {
+			u64	slot_uid;
+			u64	enclave_cid;
+		} start;
+		struct {
+			u64	slot_uid;
+			u64	reserved8[2];
+			u64	mem_regions;
+		} alloc;
+		struct {
+			u64	slot_uid;
+			u64	enclave_cid;
+			u64	reserved24;
+			u64	mem_regions;
+			u64	mem_size;
+			u64	nr_vcpus;
+			u64	flags;
+			u16	state;
+			u8	padding1[6];
+		} info;
+	};
 };
 
+static_assert(sizeof(struct ne_pci_dev_cmd_reply) == 72);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, start.slot_uid) == 8);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, start.enclave_cid) == 16);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, alloc.slot_uid) == 8);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, alloc.mem_regions) == 32);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.slot_uid) == 8);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.enclave_cid) == 16);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.mem_regions) == 32);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.mem_size) == 40);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.nr_vcpus) == 48);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.flags) == 56);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.state) == 64);
+
 /**
  * struct ne_pci_dev - Nitro Enclaves (NE) PCI device.
  * @cmd_reply_avail:		Variable set if a reply has been sent by the
-- 
2.47.1

[PATCH 17/40] nitro_enclaves: Add the slot pool donation commands

[Alexander Graf]

Nitro Sandboxes introduce the concept of memory donations. In order to
launch more sandboxes, the parent instance needs to donate some of its
memory. A device that supports Nitro Sandboxes keeps its slot pool in
that memory: the descriptors of its slots and the memory the slots run
on. In preparation to donate that memory, add the commands that donate
memory for the slot descriptors and the registers that say how much of
it the device wants.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_pci_dev.h | 60 ++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index 835413fca8040..e651e74ab656e 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -95,6 +95,42 @@
  */
 #define NE_RECV_DATA_SIZE	(240)
 
+/**
+ * DOC: Slot pool registers in the NE PCI device MMIO BAR
+ */
+
+/**
+ * NE_SLOT_POOL_MAX - (2 bytes) Register reporting the highest number of
+ *		      slots the device lets the pool grow to, or zero on a
+ *		      device that needs no donation (Read-Only).
+ */
+#define NE_SLOT_POOL_MAX	(0x01f2)
+
+/**
+ * NE_SLOT_POOL_LIMIT - (2 bytes) Register reporting how many slots the
+ *			pool has room for with the pages donated so far
+ *			(Read-Only).
+ */
+#define NE_SLOT_POOL_LIMIT	(0x01f6)
+
+/**
+ * NE_SLOT_POOL_GRAN_SIZE - (4 bytes) Register reporting the size, in bytes,
+ *			    of one donation to the pool (Read-Only).
+ */
+#define NE_SLOT_POOL_GRAN_SIZE	(0x0200)
+
+/**
+ * NE_SLOT_POOL_GRAN_ALIGN - (4 bytes) Register reporting the alignment a
+ *			     donation to the pool has to meet (Read-Only).
+ */
+#define NE_SLOT_POOL_GRAN_ALIGN	(0x0204)
+
+/**
+ * NE_SLOT_POOL_BASE_SIZE - (4 bytes) Register reporting the size, in bytes,
+ *			    of the base donation (Read-Only).
+ */
+#define NE_SLOT_POOL_BASE_SIZE	(0x0208)
+
 /**
  * DOC: MSI-X interrupt vectors
  */
@@ -125,6 +161,10 @@
  * @SLOT_ADD_BULK_VCPUS:	Add a number of vCPUs, not providing CPU ids.
  * @SLOT_ADD_INFO_PAGE:		Give a slot the info page, a memory region the
  *				device writes.
+ * @SLOT_POOL_DONATE:		Donate one page of parent memory to the slot
+ *				pool.
+ * @SLOT_POOL_SHUTDOWN:		Give the whole slot pool back.
+ * @SLOT_POOL_BASE_DONATE:	Donate the base run of the slot pool.
  * @MAX_CMD:			A gatekeeper for max possible command type.
  */
 enum ne_pci_dev_cmd_type {
@@ -142,6 +182,9 @@ enum ne_pci_dev_cmd_type {
 	SLOT_ADD_BULK_VCPUS	= 11,
 	/* 12 is a command this driver does not implement. */
 	SLOT_ADD_INFO_PAGE	= 13,
+	SLOT_POOL_DONATE	= 16,
+	SLOT_POOL_SHUTDOWN	= 17,
+	SLOT_POOL_BASE_DONATE	= 25,
 	MAX_CMD,
 };
 
@@ -284,6 +327,23 @@ struct slot_add_info_page_req {
 	u64	info_size;
 };
 
+/**
+ * struct slot_pool_donate_req - SLOT_POOL_DONATE and SLOT_POOL_BASE_DONATE
+ *				 request.
+ * @paddr:	Physical address of the donated run of pages.
+ */
+struct slot_pool_donate_req {
+	u64	paddr;
+};
+
+/**
+ * struct slot_pool_shutdown_req - SLOT_POOL_SHUTDOWN request.
+ * @unused:	In order to avoid weird sizeof edge cases.
+ */
+struct slot_pool_shutdown_req {
+	u8	unused;
+};
+
 /**
  * struct ne_pci_dev_cmd_reply - NE PCI device command reply.
  *
-- 
2.47.1

[PATCH 18/40] nitro_enclaves: Donate memory for the slot pool

[Alexander Graf]

To launch the first Nitro Sandbox, the parent must donate memory to the
device. So at probe, on a device whose vendor-specific PCI capability
reports a slot pool, donate the base run and one donation of
NE_SLOT_POOL_GRAN_SIZE bytes from ne_alloc_contig().

The device holds a donated page until SLOT_POOL_SHUTDOWN, which it
refuses while a slot is in use. So remove, shutdown and a probe that
fails after donating first run ne_enclave_free_slot() on every enclave
and then take the pool back, before any page goes to the allocator. That
way the device stays usable after an unbind or a kexec with enclaves
running. Like the two free commands, SLOT_POOL_SHUTDOWN waits for
pci_dev_mutex even for a process that is being killed. This teardown
runs on every device, and a slot the device confirms as free lets the
driver release the memory and CPUs of a classic enclave too, once its
file is closed. ENCLAVE_STOP now goes out in any state, because an
enclave that exited may still be scrubbed, and the device answers a stop
for a slot that never started as well.

An enclave whose slot was freed that way gives its CPUs back on close,
so count the enclaves created from the CPU pool and refuse to replace
the pool while the count is not zero. A donation that got no answer
stays recorded as donated, because the device may take it.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c |  53 +++-
 drivers/virt/nitro_enclaves/ne_misc_dev.h |   7 +
 drivers/virt/nitro_enclaves/ne_pci_dev.c  | 334 +++++++++++++++++++++-
 drivers/virt/nitro_enclaves/ne_pci_dev.h  |  25 ++
 4 files changed, 401 insertions(+), 18 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 77b15eda8d120..2fae731acb977 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -201,6 +201,11 @@ static const struct kernel_param *ne_cpus_kp;
  *				and leaves when the enclave is freed, in
  *				the same two places that move the core out of
  *				and back into @avail_threads_per_core.
+ * @nr_enclaves:			How many enclaves were created from the pool and
+ *				have not been freed yet, whether or not they are
+ *				on the list of a device. The pool is not replaced
+ *				while any is, so the cores an enclave gives back
+ *				go back to the pool they came from.
  */
 struct ne_cpu_pool {
 	cpumask_var_t	*avail_threads_per_core;
@@ -211,6 +216,7 @@ struct ne_cpu_pool {
 	cpumask_var_t	*node_cpus;
 	struct cpumask	total_cpus;
 	struct cpumask	used_cpus;
+	unsigned int	nr_enclaves;
 };
 
 static struct ne_cpu_pool ne_cpu_pool = {
@@ -931,6 +937,7 @@ static int ne_replace_cpu_pool(const char *val, const struct kernel_param *kp)
 	cpumask_var_t cpu_pool;
 	cpumask_var_t old_pool;
 	struct ne_pci_dev *ne_pci_dev = NULL;
+	unsigned int nr_enclaves = 0;
 	int rc = -EINVAL;
 
 	if (!zalloc_cpumask_var(&cpu_pool, GFP_KERNEL))
@@ -961,7 +968,11 @@ static int ne_replace_cpu_pool(const char *val, const struct kernel_param *kp)
 			goto unlock_devs;
 	}
 
-	if (ne_pci_dev && !list_empty(&ne_pci_dev->enclaves_list)) {
+	mutex_lock(&ne_cpu_pool.mutex);
+	nr_enclaves = ne_cpu_pool.nr_enclaves;
+	mutex_unlock(&ne_cpu_pool.mutex);
+
+	if (nr_enclaves) {
 		pr_err("%s: The CPU pool is used by enclave(s)\n", ne_misc_dev.name);
 
 		rc = -EPERM;
@@ -2293,6 +2304,8 @@ static void ne_enclave_remove_all_vcpu_id_entries(struct ne_enclave *ne_enclave)
 	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++)
 		free_cpumask_var(ne_enclave->threads_per_core[i]);
 
+	ne_cpu_pool.nr_enclaves--;
+
 	mutex_unlock(&ne_cpu_pool.mutex);
 
 	kfree(ne_enclave->threads_per_core);
@@ -2337,8 +2350,8 @@ static void ne_pci_dev_remove_enclave_entry(struct ne_enclave *ne_enclave,
  * * 0 once the device has freed the slot.
  * * Negative return value on failure.
  */
-static int ne_enclave_free_slot(struct ne_pci_dev *ne_pci_dev,
-				struct ne_enclave *ne_enclave)
+int ne_enclave_free_slot(struct ne_pci_dev *ne_pci_dev,
+			 struct ne_enclave *ne_enclave)
 {
 	struct ne_pci_dev_cmd_reply cmd_reply = {};
 	struct enclave_stop_req enclave_stop_request = {};
@@ -2346,17 +2359,20 @@ static int ne_enclave_free_slot(struct ne_pci_dev *ne_pci_dev,
 	struct pci_dev *pdev = ne_pci_dev->pdev;
 	int rc;
 
-	if (ne_enclave->state != NE_STATE_INIT && ne_enclave->state != NE_STATE_STOPPED) {
-		enclave_stop_request.slot_uid = ne_enclave->slot_uid;
+	/*
+	 * An enclave that exited may still be scrubbed, so the stop goes out
+	 * whatever the state, and the device answers it once the slot is
+	 * neither running nor being scrubbed.
+	 */
+	enclave_stop_request.slot_uid = ne_enclave->slot_uid;
 
-		rc = ne_do_request(pdev, ENCLAVE_STOP,
-				   &enclave_stop_request, sizeof(enclave_stop_request),
-				   &cmd_reply, sizeof(cmd_reply));
-		if (rc < 0)
-			return rc;
+	rc = ne_do_request(pdev, ENCLAVE_STOP,
+			   &enclave_stop_request, sizeof(enclave_stop_request),
+			   &cmd_reply, sizeof(cmd_reply));
+	if (rc < 0)
+		return rc;
 
-		memset(&cmd_reply, 0, sizeof(cmd_reply));
-	}
+	memset(&cmd_reply, 0, sizeof(cmd_reply));
 
 	slot_free_req.slot_uid = ne_enclave->slot_uid;
 
@@ -2416,10 +2432,13 @@ static int ne_enclave_release(struct inode *inode, struct file *file)
 	/*
 	 * The device the enclave was created on is gone, and with it the only
 	 * party that can confirm that the slot no longer uses its memory and
-	 * its CPUs, so the enclave keeps both.
+	 * its CPUs, so the enclave keeps both, unless the device freed the
+	 * slot before it went away.
 	 */
 	if (ne_enclave->detached) {
-		rc = -ENODEV;
+		rc = ne_enclave->slot_freed ? 0 : -ENODEV;
+		if (!rc)
+			ne_enclave_free(ne_enclave);
 
 		goto unlock_devs;
 	}
@@ -2525,7 +2544,7 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 flags,
 	 * capability. One without it takes no launch flags and would allocate a
 	 * slot without the properties they ask for.
 	 */
-	if (flags && !pci_find_capability(ne_pci_dev->pdev, PCI_CAP_ID_VNDR)) {
+	if (flags && !ne_pci_dev->has_vendor_cap) {
 		dev_err_ratelimited(ne_misc_dev.this_device,
 				    "Device takes no launch flags\n");
 
@@ -2630,6 +2649,10 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 flags,
 
 	list_add(&ne_enclave->enclave_list_entry, &ne_pci_dev->enclaves_list);
 
+	mutex_lock(&ne_cpu_pool.mutex);
+	ne_cpu_pool.nr_enclaves++;
+	mutex_unlock(&ne_cpu_pool.mutex);
+
 	if (copy_to_user(slot_uid, &ne_enclave->slot_uid, sizeof(ne_enclave->slot_uid))) {
 		/*
 		 * As we're holding the only reference to 'enclave_file', fput()
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index d0d1973bfbe9d..59374a35c3c7d 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -116,6 +116,9 @@ struct ne_mem_region {
  * @enclave_info_mutex :	Mutex for accessing this internal state.
  * @detached:			The NE PCI device the enclave was created on is
  *				gone. Written with ne_devs.lock held for writing.
+ * @slot_freed:			The device freed the slot of a detached enclave
+ *				before it went away. Written with ne_devs.lock
+ *				held for writing.
  * @contig_regions_list:	List of %NE_ALLOC_MEMORY allocations given to the
  *				enclave with %NE_ADD_MEMORY_REGION.
  * @enclave_list_entry :	Entry in the list of created enclaves.
@@ -157,6 +160,7 @@ struct ne_mem_region {
 struct ne_enclave {
 	struct mutex		enclave_info_mutex;
 	bool			detached;
+	bool			slot_freed;
 	struct list_head	contig_regions_list;
 	struct list_head	enclave_list_entry;
 	wait_queue_head_t	eventq;
@@ -221,6 +225,9 @@ void ne_contig_free_run(struct cma *cma, struct page *page,
 /* Nitro Enclaves (NE) data structure for keeping refs to the NE misc and PCI devices. */
 extern struct ne_devs ne_devs;
 
+int ne_enclave_free_slot(struct ne_pci_dev *ne_pci_dev,
+			 struct ne_enclave *ne_enclave);
+
 int ne_cpu_pool_sysfs_init(void);
 void ne_cpu_pool_sysfs_exit(void);
 
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index 8ed2349dea878..e42be1742a39c 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -7,14 +7,20 @@
  * DOC: Nitro Enclaves (NE) PCI device driver.
  */
 
+#include <linux/cma.h>
 #include <linux/delay.h>
 #include <linux/device.h>
+#include <linux/err.h>
 #include <linux/iopoll.h>
+#include <linux/jiffies.h>
 #include <linux/list.h>
+#include <linux/minmax.h>
 #include <linux/module.h>
 #include <linux/mutex.h>
 #include <linux/nitro_enclaves.h>
 #include <linux/pci.h>
+#include <linux/sizes.h>
+#include <linux/sysfs.h>
 #include <linux/types.h>
 #include <linux/wait.h>
 
@@ -211,10 +217,11 @@ int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 	/*
 	 * Use this mutex so that the PCI device handles one command request at
 	 * a time. A slot being freed can hold it for minutes, so let a task
-	 * that waits for it be killed, unless the task frees a slot itself:
-	 * release has to do that for a process that is being killed too.
+	 * that waits for it be killed, unless the task frees a slot or takes
+	 * the slot pool back: release and unbind have to finish those for a
+	 * process that is being killed too.
 	 */
-	if (ne_cmd_frees_slot(cmd_type)) {
+	if (ne_cmd_frees_slot(cmd_type) || cmd_type == SLOT_POOL_SHUTDOWN) {
 		mutex_lock(&ne_pci_dev->pci_dev_mutex);
 	} else {
 		rc = mutex_lock_killable(&ne_pci_dev->pci_dev_mutex);
@@ -609,6 +616,14 @@ static void ne_pci_dev_disable(struct pci_dev *pdev)
  *				  goes away.
  * @ne_pci_dev:	Private data associated with the PCI device.
  *
+ * A device with a slot pool refuses SLOT_POOL_SHUTDOWN while one of its slots
+ * is in use, and the pages it was donated would stay with it. Each enclave is
+ * therefore stopped and its slot freed here, before ne_slot_pool_release()
+ * runs, the same way as on release. An enclave gives its memory and CPUs back
+ * when its file is released, if the device freed its slot here. An enclave
+ * whose file is already gone is one the device did not free on release, and
+ * it keeps both.
+ *
  * Context: Process context. This function is called with ne_devs.lock held
  *	    for writing.
  */
@@ -625,6 +640,7 @@ static void ne_pci_dev_detach_enclaves(struct ne_pci_dev *ne_pci_dev)
 
 		mutex_lock(&ne_enclave->enclave_info_mutex);
 
+		ne_enclave->slot_freed = !ne_enclave_free_slot(ne_pci_dev, ne_enclave);
 		ne_enclave->detached = true;
 		ne_enclave->has_event = true;
 
@@ -636,6 +652,281 @@ static void ne_pci_dev_detach_enclaves(struct ne_pci_dev *ne_pci_dev)
 	mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
 }
 
+/**
+ * ne_slot_pool_read_geometry() - Read how much memory the slot pool wants.
+ * @pdev:	PCI device to read the slot pool registers of.
+ *
+ * The slot pool registers only exist on a device that publishes a
+ * vendor-specific capability, so a device without one keeps a zero ceiling and
+ * reads no slot pool register at all. A device that stores its slot
+ * descriptors itself reports a zero ceiling and the driver donates nothing to
+ * it. A device that asks for parent memory is only usable if it asks in units
+ * the driver can allocate, so a granularity other than 2 MiB leaves the
+ * ceiling at zero as well.
+ *
+ * Context: Process context.
+ */
+static void ne_slot_pool_read_geometry(struct pci_dev *pdev)
+{
+	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
+	u32 size, align;
+
+	if (!ne_pci_dev->has_vendor_cap)
+		return;
+
+	ne_pci_dev->slot_pool_max = ioread16(ne_pci_dev->iomem_base + NE_SLOT_POOL_MAX);
+	if (!ne_pci_dev->slot_pool_max)
+		return;
+
+	size = ioread32(ne_pci_dev->iomem_base + NE_SLOT_POOL_GRAN_SIZE);
+	align = ioread32(ne_pci_dev->iomem_base + NE_SLOT_POOL_GRAN_ALIGN);
+	if (size != SZ_2M || align != SZ_2M) {
+		dev_dbg(&pdev->dev, "Slot pool granularity size=%u align=%u unusable\n",
+			size, align);
+
+		ne_pci_dev->slot_pool_max = 0;
+
+		return;
+	}
+
+	ne_pci_dev->slot_pool_gran = size;
+}
+
+/**
+ * ne_slot_pool_donate_base() - Donate the base run of the slot pool.
+ * @pdev:	PCI device to donate to.
+ *
+ * The device turns down a donation to the slot pool until it has the base
+ * run, so the base run is donated before any other donation. Donating a
+ * second time is a no-op.
+ *
+ * A donation the device never answered was submitted with the run's physical
+ * address in it, so the run is recorded as donated and left to
+ * ne_slot_pool_release() rather than given back.
+ *
+ * Context: Process context. This function is called with the enclaves list
+ *	    mutex held.
+ * Return:
+ * * 0 on success or if the base run has already been donated.
+ * * Negative return value on failure.
+ * * -ETIMEDOUT, or any other code ne_request_unanswered() names, if the
+ *   request went unanswered, in which case the run stays recorded as donated
+ *   because the driver cannot establish that the device refused it. A device
+ *   that answers such a code itself lands here too and costs one held run.
+ */
+static int ne_slot_pool_donate_base(struct pci_dev *pdev)
+{
+	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
+	struct ne_pci_dev_cmd_reply cmd_reply = {};
+	struct slot_pool_donate_req donate_req = {};
+	struct cma *cma = NULL;
+	unsigned long nr_pages;
+	struct page *page;
+	u32 size;
+	int rc;
+
+	if (ne_pci_dev->slot_pool_base_page)
+		return 0;
+
+	size = ioread32(ne_pci_dev->iomem_base + NE_SLOT_POOL_BASE_SIZE);
+	if (!size || !IS_ALIGNED(size, PAGE_SIZE)) {
+		dev_warn(&pdev->dev, "Slot pool base size=%u unusable\n", size);
+
+		return -EOPNOTSUPP;
+	}
+
+	nr_pages = size >> PAGE_SHIFT;
+
+	page = ne_alloc_contig(nr_pages, numa_node_id(), &cma);
+	if (IS_ERR_OR_NULL(page)) {
+		rc = page ? PTR_ERR(page) : -ENOMEM;
+
+		dev_warn(&pdev->dev,
+			 "Slot pool base size=%u not allocated on node %d [rc=%d]\n",
+			 size, numa_node_id(), rc);
+
+		return rc;
+	}
+
+	donate_req.paddr = page_to_phys(page);
+
+	rc = ne_do_request(pdev, SLOT_POOL_BASE_DONATE,
+			   &donate_req, sizeof(donate_req),
+			   &cmd_reply, sizeof(cmd_reply));
+	if (rc < 0) {
+		dev_err_ratelimited(&pdev->dev, "Error in slot pool base donate [rc=%d]\n",
+				    rc);
+
+		/*
+		 * A request the device never answered was submitted with the
+		 * run's physical address in it, and there is nothing to cancel
+		 * it with, so the device may take the run yet.
+		 * ne_request_unanswered() picks out that case. Record the run as
+		 * donated, so that ne_slot_pool_release() owns it from here the
+		 * way it owns a donation the device accepted. The errors left
+		 * over say the device answered to refuse the donation, or never
+		 * saw it, and the run is the parent's again.
+		 */
+		if (ne_request_unanswered(rc)) {
+			ne_pci_dev->slot_pool_base_page = page;
+			ne_pci_dev->slot_pool_base_cma = cma;
+			ne_pci_dev->slot_pool_base_nr_pages = nr_pages;
+		} else {
+			ne_contig_free_run(cma, page, nr_pages);
+		}
+
+		return rc;
+	}
+
+	ne_pci_dev->slot_pool_base_page = page;
+	ne_pci_dev->slot_pool_base_cma = cma;
+	ne_pci_dev->slot_pool_base_nr_pages = nr_pages;
+
+	return 0;
+}
+
+/**
+ * ne_slot_pool_donate() - Donate one granularity of pages to the slot pool.
+ * @pdev:	PCI device to donate to.
+ *
+ * A donation raises the number of slots the pool has room for, up to the
+ * ceiling the device reports, and a pool that has reached that ceiling takes no
+ * further donation.
+ *
+ * Context: Process context. This function is called with the enclaves list
+ *	    mutex held.
+ * Return:
+ * * 0 on success.
+ * * -ENOSPC if the pool has reached the ceiling the device reports.
+ * * Negative return value on failure.
+ * * -ETIMEDOUT, or any other code ne_request_unanswered() names, if the
+ *   request went unanswered, in which case the run stays recorded as donated
+ *   because the driver cannot establish that the device refused it. A device
+ *   that answers such a code itself lands here too and costs one held run.
+ */
+static int ne_slot_pool_donate(struct pci_dev *pdev)
+{
+	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
+	struct ne_pci_dev_cmd_reply cmd_reply = {};
+	struct slot_pool_donate_req donate_req = {};
+	struct cma *cma = NULL;
+	unsigned long nr_pages;
+	struct page *page;
+	int rc;
+
+	if (ne_pci_dev->slot_pool_page)
+		return 0;
+
+	if (!ne_pci_dev->slot_pool_gran)
+		return -EOPNOTSUPP;
+
+	if (ioread16(ne_pci_dev->iomem_base + NE_SLOT_POOL_LIMIT) >=
+	    ne_pci_dev->slot_pool_max)
+		return -ENOSPC;
+
+	nr_pages = ne_pci_dev->slot_pool_gran >> PAGE_SHIFT;
+
+	page = ne_alloc_contig(nr_pages, numa_node_id(), &cma);
+	if (IS_ERR_OR_NULL(page)) {
+		rc = page ? PTR_ERR(page) : -ENOMEM;
+
+		dev_warn(&pdev->dev,
+			 "Slot pool granularity size=%u not allocated on node %d [rc=%d]\n",
+			 ne_pci_dev->slot_pool_gran, numa_node_id(), rc);
+
+		return rc;
+	}
+
+	donate_req.paddr = page_to_phys(page);
+
+	rc = ne_do_request(pdev, SLOT_POOL_DONATE,
+			   &donate_req, sizeof(donate_req),
+			   &cmd_reply, sizeof(cmd_reply));
+	if (rc < 0) {
+		dev_err_ratelimited(&pdev->dev, "Error in slot pool donate [rc=%d]\n", rc);
+
+		/*
+		 * Record a run whose donation went unanswered as donated
+		 * instead of giving it back, as ne_slot_pool_donate_base()
+		 * does and for the same reason.
+		 */
+		if (ne_request_unanswered(rc)) {
+			ne_pci_dev->slot_pool_page = page;
+			ne_pci_dev->slot_pool_cma = cma;
+		} else {
+			ne_contig_free_run(cma, page, nr_pages);
+		}
+
+		return rc;
+	}
+
+	ne_pci_dev->slot_pool_page = page;
+	ne_pci_dev->slot_pool_cma = cma;
+
+	return 0;
+}
+
+/**
+ * ne_slot_pool_release() - Take the whole slot pool back from the device.
+ * @pdev:	PCI device to take the pool back from.
+ *
+ * The device keeps a slot's descriptor in a page the parent donated, and
+ * SLOT_POOL_SHUTDOWN names no slot, so a pool taken back while the device still
+ * has a slot is parent memory the device goes on writing after the allocator has
+ * handed it to somebody else. ne_pci_dev_detach_enclaves() frees the slot of
+ * every enclave before this runs, and the device refuses SLOT_POOL_SHUTDOWN
+ * while a slot it would not free is still in use. A SLOT_POOL_SHUTDOWN that
+ * fails leaves the pages with the device whatever the failure was: a refusal
+ * means the device still has the pool, and an answer that never came means it
+ * may still have it.
+ *
+ * A donation that was never answered is recorded as donated, and the pages stay
+ * with the device: ne_do_request() sends no command after one that went
+ * unanswered, so SLOT_POOL_SHUTDOWN fails as well.
+ *
+ * Context: Process context.
+ */
+static void ne_slot_pool_release(struct pci_dev *pdev)
+{
+	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
+	struct ne_pci_dev_cmd_reply cmd_reply = {};
+	struct slot_pool_shutdown_req shutdown_req = {};
+	int rc;
+
+	if (!ne_pci_dev->slot_pool_base_page)
+		return;
+
+	mutex_lock(&ne_pci_dev->enclaves_list_mutex);
+
+	rc = ne_do_request(pdev, SLOT_POOL_SHUTDOWN,
+			   &shutdown_req, sizeof(shutdown_req),
+			   &cmd_reply, sizeof(cmd_reply));
+	if (rc < 0) {
+		dev_err(&pdev->dev, "Error in slot pool shutdown [rc=%d]\n", rc);
+
+		goto unlock_mutex;
+	}
+
+	if (ne_pci_dev->slot_pool_page) {
+		ne_contig_free_run(ne_pci_dev->slot_pool_cma,
+				   ne_pci_dev->slot_pool_page,
+				   ne_pci_dev->slot_pool_gran >> PAGE_SHIFT);
+
+		ne_pci_dev->slot_pool_page = NULL;
+		ne_pci_dev->slot_pool_cma = NULL;
+	}
+
+	ne_contig_free_run(ne_pci_dev->slot_pool_base_cma,
+			   ne_pci_dev->slot_pool_base_page,
+			   ne_pci_dev->slot_pool_base_nr_pages);
+
+	ne_pci_dev->slot_pool_base_page = NULL;
+	ne_pci_dev->slot_pool_base_cma = NULL;
+
+unlock_mutex:
+	mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+}
+
 /**
  * ne_pci_probe() - Probe function for the NE PCI device.
  * @pdev:	PCI device to match with the NE PCI driver.
@@ -689,6 +980,15 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 	mutex_init(&ne_pci_dev->pci_dev_mutex);
 	ne_pci_dev->pdev = pdev;
 
+	/*
+	 * The registers past the command interface only exist on a device that
+	 * publishes a vendor-specific capability, and a read of a register the
+	 * device does not implement has no defined value. The answer holds for
+	 * the life of the device, so look it up once here.
+	 */
+	ne_pci_dev->has_vendor_cap =
+		pci_find_capability(pdev, PCI_CAP_ID_VNDR) != 0;
+
 	rc = ne_setup_msix(pdev);
 	if (rc < 0) {
 		dev_err(&pdev->dev, "Error in pci dev msix setup [rc=%d]\n", rc);
@@ -721,6 +1021,29 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 
 	ne_devs.ne_pci_dev = ne_pci_dev;
 
+	ne_slot_pool_read_geometry(pdev);
+
+	if (ne_pci_dev->slot_pool_max) {
+		const char *stage;
+
+		mutex_lock(&ne_pci_dev->enclaves_list_mutex);
+
+		rc = ne_slot_pool_donate_base(pdev);
+		if (rc < 0) {
+			stage = "base";
+		} else {
+			rc = ne_slot_pool_donate(pdev);
+			stage = "granularity";
+		}
+
+		mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+
+		if (rc < 0)
+			dev_warn(&pdev->dev,
+				 "No slot donated to the slot pool [stage=%s rc=%d]\n",
+				 stage, rc);
+	}
+
 	rc = misc_register(ne_devs.ne_misc_dev);
 	if (rc < 0) {
 		dev_err(&pdev->dev, "Error in misc dev register [rc=%d]\n", rc);
@@ -745,6 +1068,7 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 	ne_pci_dev_detach_enclaves(ne_pci_dev);
 	misc_deregister(ne_devs.ne_misc_dev);
 disable_ne_pci_dev:
+	ne_slot_pool_release(pdev);
 	ne_devs.ne_pci_dev = NULL;
 unlock_ne_devs:
 	up_write(&ne_devs.lock);
@@ -782,6 +1106,8 @@ static void ne_pci_remove(struct pci_dev *pdev)
 
 	misc_deregister(ne_devs.ne_misc_dev);
 
+	ne_slot_pool_release(pdev);
+
 	ne_devs.ne_pci_dev = NULL;
 
 	ne_pci_dev_disable(pdev);
@@ -822,6 +1148,8 @@ static void ne_pci_shutdown(struct pci_dev *pdev)
 
 	misc_deregister(ne_devs.ne_misc_dev);
 
+	ne_slot_pool_release(pdev);
+
 	ne_devs.ne_pci_dev = NULL;
 
 	ne_pci_dev_disable(pdev);
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index e651e74ab656e..4e991c4b0e70f 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -433,10 +433,27 @@ static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.state) == 64);
  *				triggered by the Nitro Hypervisor which require
  *				enclave state scanning and propagation to the
  *				enclave process.
+ * @has_vendor_cap:		Whether the PCI device publishes a vendor-specific
+ *				capability, and with it the registers past the
+ *				command interface.
  * @iomem_base :		MMIO region of the PCI device.
  * @notify_work:		Work item for every received out-of-band event.
  * @pci_dev_mutex:		Mutex for accessing the PCI device MMIO space.
  * @pdev:			PCI device data structure.
+ * @slot_pool_max:		Highest number of slots the device lets the slot
+ *				pool grow to, zero on a device that takes no
+ *				donation.
+ * @slot_pool_gran:		Size, in bytes, of one donation to the slot pool.
+ *				Written once at probe, before any donation is
+ *				made, so a later reader of it sees the size the
+ *				donation was allocated with.
+ * @slot_pool_base_page:	First page of the base donation, NULL until it
+ *				has been made.
+ * @slot_pool_base_cma:		CMA area the base donation came from.
+ * @slot_pool_base_nr_pages:	Number of pages in the base donation.
+ * @slot_pool_page:		First page of the donation to the slot pool, NULL
+ *				until it has been made.
+ * @slot_pool_cma:		CMA area the donation came from.
  */
 struct ne_pci_dev {
 	atomic_t		cmd_reply_avail;
@@ -445,10 +462,18 @@ struct ne_pci_dev {
 	struct list_head	enclaves_list;
 	struct mutex		enclaves_list_mutex;
 	struct workqueue_struct	*event_wq;
+	bool			has_vendor_cap;
 	void __iomem		*iomem_base;
 	struct work_struct	notify_work;
 	struct mutex		pci_dev_mutex;
 	struct pci_dev		*pdev;
+	u16			slot_pool_max;
+	u32			slot_pool_gran;
+	struct page		*slot_pool_base_page;
+	struct cma		*slot_pool_base_cma;
+	unsigned long		slot_pool_base_nr_pages;
+	struct page		*slot_pool_page;
+	struct cma		*slot_pool_cma;
 };
 
 /* Whether a failure left a submitted request in flight. */
-- 
2.47.1

[PATCH 19/40] nitro_enclaves: Add the slot pool growth commands

[Alexander Graf]

In preparation to grow the slot pool at runtime, add
SLOT_POOL_MEM_DONATE and SLOT_POOL_GET_LIMITS.

The slot pool has two halves: the slot descriptors, which
SLOT_POOL_DONATE grows, and the memory the device runs slots out of,
which SLOT_POOL_MEM_DONATE grows. Its reply names the part of the
offered run that the device took. A slot needs room in both halves.
SLOT_POOL_GET_LIMITS reports both counts, along with how many slots one
memory donation adds, so that the parent knows which half to grow next.
Both replies get an arm of the reply union.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_pci_dev.h | 50 ++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index 4e991c4b0e70f..3718f893e4f65 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -164,6 +164,10 @@
  * @SLOT_POOL_DONATE:		Donate one page of parent memory to the slot
  *				pool.
  * @SLOT_POOL_SHUTDOWN:		Give the whole slot pool back.
+ * @SLOT_POOL_MEM_DONATE:	Donate a run of parent memory to the memory the
+ *				device runs a slot out of.
+ * @SLOT_POOL_GET_LIMITS:	Report how many slots each half of the pool has
+ *				room for.
  * @SLOT_POOL_BASE_DONATE:	Donate the base run of the slot pool.
  * @MAX_CMD:			A gatekeeper for max possible command type.
  */
@@ -185,6 +189,8 @@ enum ne_pci_dev_cmd_type {
 	SLOT_POOL_DONATE	= 16,
 	SLOT_POOL_SHUTDOWN	= 17,
 	SLOT_POOL_BASE_DONATE	= 25,
+	SLOT_POOL_MEM_DONATE	= 26,
+	SLOT_POOL_GET_LIMITS	= 27,
 	MAX_CMD,
 };
 
@@ -344,6 +350,16 @@ struct slot_pool_shutdown_req {
 	u8	unused;
 };
 
+/**
+ * struct slot_pool_mem_donate_req - SLOT_POOL_MEM_DONATE request.
+ * @paddr:	Physical address of the donated run of memory.
+ * @size:	Size, in bytes, of the donated run.
+ */
+struct slot_pool_mem_donate_req {
+	u64	paddr;
+	u64	size;
+};
+
 /**
  * struct ne_pci_dev_cmd_reply - NE PCI device command reply.
  *
@@ -377,6 +393,21 @@ struct slot_pool_shutdown_req {
  * @info.state:			SLOT_INFO. Slot state.
  * @info.padding1:		SLOT_INFO. Padding for the overall data
  *				structure.
+ * @mem_donate:			Fields SLOT_POOL_MEM_DONATE fills.
+ * @mem_donate.consumed_paddr:	SLOT_POOL_MEM_DONATE. Physical address of the
+ *				run the device took out of the donation.
+ * @mem_donate.consumed_size:	SLOT_POOL_MEM_DONATE. Size, in bytes, of the
+ *				run the device took, zero where it took none.
+ * @limits:			Fields SLOT_POOL_GET_LIMITS fills.
+ * @limits.meta_limit:		SLOT_POOL_GET_LIMITS. Slots the descriptors
+ *				donated so far have room for.
+ * @limits.mem_limit:		SLOT_POOL_GET_LIMITS. Slots the memory donated
+ *				so far has room to run.
+ * @limits.slots_per_donation:	SLOT_POOL_GET_LIMITS. Slots one donation of
+ *				memory adds to the memory the device runs a
+ *				slot out of.
+ * @limits.reserved14:		SLOT_POOL_GET_LIMITS. Bytes the device writes as
+ *				zero.
  */
 struct ne_pci_dev_cmd_reply {
 	s32	rc;
@@ -402,6 +433,16 @@ struct ne_pci_dev_cmd_reply {
 			u16	state;
 			u8	padding1[6];
 		} info;
+		struct {
+			u64	consumed_paddr;
+			u64	consumed_size;
+		} mem_donate;
+		struct {
+			u16	meta_limit;
+			u16	mem_limit;
+			u16	slots_per_donation;
+			u16	reserved14;
+		} limits;
 	};
 };
 
@@ -417,6 +458,15 @@ static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.mem_size) == 40);
 static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.nr_vcpus) == 48);
 static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.flags) == 56);
 static_assert(offsetof(struct ne_pci_dev_cmd_reply, info.state) == 64);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply,
+		       mem_donate.consumed_paddr) == 8);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply,
+		       mem_donate.consumed_size) == 16);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, limits.meta_limit) == 8);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, limits.mem_limit) == 10);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply,
+		       limits.slots_per_donation) == 12);
+static_assert(offsetof(struct ne_pci_dev_cmd_reply, limits.reserved14) == 14);
 
 /**
  * struct ne_pci_dev - Nitro Enclaves (NE) PCI device.
-- 
2.47.1

[PATCH 20/40] nitro_enclaves: Add a helper to grow the slot pool

[Alexander Graf]

In addition to the early boot pool allocations, users want to be able to
grow their pool at runtime. To the user, we only present a single maximum
slot count. Under the hood, there are 2 slot counters and the MIN() of
defines how many slots there really are. When growing, grow both of the
limits to satisfy the desired customer set maximum.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_pci_dev.c | 466 +++++++++++++++++++++--
 drivers/virt/nitro_enclaves/ne_pci_dev.h |  35 +-
 2 files changed, 474 insertions(+), 27 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index e42be1742a39c..9e304ffc2bd13 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -18,8 +18,10 @@
 #include <linux/module.h>
 #include <linux/mutex.h>
 #include <linux/nitro_enclaves.h>
+#include <linux/nodemask.h>
 #include <linux/pci.h>
 #include <linux/sizes.h>
+#include <linux/slab.h>
 #include <linux/sysfs.h>
 #include <linux/types.h>
 #include <linux/wait.h>
@@ -692,6 +694,61 @@ static void ne_slot_pool_read_geometry(struct pci_dev *pdev)
 	ne_pci_dev->slot_pool_gran = size;
 }
 
+/**
+ * ne_slot_pool_alloc_run() - Allocate one run for the slot pool.
+ * @nr_pages:	Number of pages in the run.
+ * @pref_nid:	NUMA node to try first.
+ * @out_cma:	Set to the CMA area the run came from.
+ *
+ * ne_alloc_contig() confines a run to the node it is given, which is what an
+ * enclave needs. A run in the slot pool is read and written by the device
+ * alone, so it can sit on any node that has memory, and a pool that stops
+ * growing because one node ran out is worth less than a pool spread over
+ * several. So walk the nodes that have memory, starting at @pref_nid, and take
+ * the first run any of them serves.
+ *
+ * A caller runs on a node that need not have memory of its own, and there is no
+ * run to be had from such a node, so a @pref_nid outside the set of nodes that
+ * have memory starts the walk at the first node that does. The walk then covers
+ * every node that can serve it.
+ *
+ * Context: Process context.
+ * Return: the first page of the run, NULL when no node served one, or the
+ *	   ERR_PTR() ne_alloc_contig() returned.
+ */
+static struct page *ne_slot_pool_alloc_run(unsigned long nr_pages,
+					   int pref_nid, struct cma **out_cma)
+{
+	int nid = pref_nid;
+	struct page *page;
+	int nr_nodes;
+	int i;
+
+	nr_nodes = num_node_state(N_MEMORY);
+	if (nr_nodes < 1)
+		return NULL;
+
+	if (pref_nid < 0 || pref_nid >= MAX_NUMNODES ||
+	    !node_state(pref_nid, N_MEMORY))
+		nid = first_node(node_states[N_MEMORY]);
+
+	for (i = 0; i < nr_nodes; i++) {
+		page = ne_alloc_contig(nr_pages, nid, out_cma);
+
+		/*
+		 * An error ends the walk along with a run: only a node that had
+		 * no run to give leaves NULL behind, and only then does the next
+		 * node get asked.
+		 */
+		if (page)
+			return page;
+
+		nid = next_node_in(nid, node_states[N_MEMORY]);
+	}
+
+	return NULL;
+}
+
 /**
  * ne_slot_pool_donate_base() - Donate the base run of the slot pool.
  * @pdev:	PCI device to donate to.
@@ -708,6 +765,7 @@ static void ne_slot_pool_read_geometry(struct pci_dev *pdev)
  *	    mutex held.
  * Return:
  * * 0 on success or if the base run has already been donated.
+ * * -EINTR if a fatal signal arrived while a run was being zeroed.
  * * Negative return value on failure.
  * * -ETIMEDOUT, or any other code ne_request_unanswered() names, if the
  *   request went unanswered, in which case the run stays recorded as donated
@@ -737,13 +795,13 @@ static int ne_slot_pool_donate_base(struct pci_dev *pdev)
 
 	nr_pages = size >> PAGE_SHIFT;
 
-	page = ne_alloc_contig(nr_pages, numa_node_id(), &cma);
+	page = ne_slot_pool_alloc_run(nr_pages, numa_node_id(), &cma);
 	if (IS_ERR_OR_NULL(page)) {
 		rc = page ? PTR_ERR(page) : -ENOMEM;
 
 		dev_warn(&pdev->dev,
-			 "Slot pool base size=%u not allocated on node %d [rc=%d]\n",
-			 size, numa_node_id(), rc);
+			 "Slot pool base size=%u not allocated [rc=%d]\n",
+			 size, rc);
 
 		return rc;
 	}
@@ -798,6 +856,7 @@ static int ne_slot_pool_donate_base(struct pci_dev *pdev)
  * Return:
  * * 0 on success.
  * * -ENOSPC if the pool has reached the ceiling the device reports.
+ * * -EINTR if a fatal signal arrived while a run was being zeroed.
  * * Negative return value on failure.
  * * -ETIMEDOUT, or any other code ne_request_unanswered() names, if the
  *   request went unanswered, in which case the run stays recorded as donated
@@ -809,14 +868,12 @@ static int ne_slot_pool_donate(struct pci_dev *pdev)
 	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
 	struct ne_pci_dev_cmd_reply cmd_reply = {};
 	struct slot_pool_donate_req donate_req = {};
+	struct ne_slot_pool_run *run;
 	struct cma *cma = NULL;
 	unsigned long nr_pages;
 	struct page *page;
 	int rc;
 
-	if (ne_pci_dev->slot_pool_page)
-		return 0;
-
 	if (!ne_pci_dev->slot_pool_gran)
 		return -EOPNOTSUPP;
 
@@ -826,17 +883,27 @@ static int ne_slot_pool_donate(struct pci_dev *pdev)
 
 	nr_pages = ne_pci_dev->slot_pool_gran >> PAGE_SHIFT;
 
-	page = ne_alloc_contig(nr_pages, numa_node_id(), &cma);
+	run = kzalloc_obj(*run);
+	if (!run)
+		return -ENOMEM;
+
+	page = ne_slot_pool_alloc_run(nr_pages, numa_node_id(), &cma);
 	if (IS_ERR_OR_NULL(page)) {
 		rc = page ? PTR_ERR(page) : -ENOMEM;
 
 		dev_warn(&pdev->dev,
-			 "Slot pool granularity size=%u not allocated on node %d [rc=%d]\n",
-			 ne_pci_dev->slot_pool_gran, numa_node_id(), rc);
+			 "Slot pool granularity size=%u not allocated [rc=%d]\n",
+			 ne_pci_dev->slot_pool_gran, rc);
+
+		kfree(run);
 
 		return rc;
 	}
 
+	run->page = page;
+	run->cma = cma;
+	run->nr_pages = nr_pages;
+
 	donate_req.paddr = page_to_phys(page);
 
 	rc = ne_do_request(pdev, SLOT_POOL_DONATE,
@@ -851,21 +918,362 @@ static int ne_slot_pool_donate(struct pci_dev *pdev)
 		 * does and for the same reason.
 		 */
 		if (ne_request_unanswered(rc)) {
-			ne_pci_dev->slot_pool_page = page;
-			ne_pci_dev->slot_pool_cma = cma;
+			list_add(&run->node, &ne_pci_dev->slot_pool_meta_runs);
+			ne_pci_dev->slot_pool_meta_nr++;
 		} else {
 			ne_contig_free_run(cma, page, nr_pages);
+			kfree(run);
 		}
 
 		return rc;
 	}
 
-	ne_pci_dev->slot_pool_page = page;
-	ne_pci_dev->slot_pool_cma = cma;
+	list_add(&run->node, &ne_pci_dev->slot_pool_meta_runs);
+	ne_pci_dev->slot_pool_meta_nr++;
 
 	return 0;
 }
 
+/*
+ * Size of one run of memory the slot pool takes, and of the span the driver
+ * offers so that the device can place such a run inside it. A run begins on a
+ * boundary of its own size in the device's address space, and an address the
+ * parent aligns does not stay aligned once the device has mapped it, so the
+ * parent cannot place the run itself. A span of twice the size holds a run the
+ * device can place wherever the span begins.
+ */
+#define NE_SLOT_POOL_MEM_SIZE	SZ_128M
+#define NE_SLOT_POOL_MEM_SPAN	(2 * NE_SLOT_POOL_MEM_SIZE)
+
+/**
+ * ne_slot_pool_donate_nid() - Pick the node the next memory donation comes from.
+ * @nr_donated:	Number of runs of memory the slot pool already holds.
+ *
+ * The device turns a run of donated memory into memory it gives to enclaves, and
+ * the host grows the zone it takes that memory from on the node the run sits on.
+ * Every donation coming from one node grows that node's zone alone, so the
+ * enclaves the host places on any other node run on memory out of a zone that
+ * never grew, and the host runs that zone out and takes every enclave on the
+ * machine down with it. So spread the donations instead: step through the nodes
+ * that have memory in time with the runs already held, and successive donations
+ * come from successive nodes.
+ *
+ * A node that has no run to give costs a donation nothing, because
+ * ne_slot_pool_alloc_run() carries on from the node named here to the rest.
+ *
+ * Context: Process context.
+ * Return: the node the donation is to come from first.
+ */
+static int ne_slot_pool_donate_nid(unsigned int nr_donated)
+{
+	int nr_nodes = num_node_state(N_MEMORY);
+	unsigned int steps;
+	unsigned int i;
+	int nid;
+
+	if (nr_nodes < 1)
+		return numa_node_id();
+
+	steps = nr_donated % (unsigned int)nr_nodes;
+	nid = first_node(node_states[N_MEMORY]);
+
+	for (i = 0; i < steps; i++)
+		nid = next_node_in(nid, node_states[N_MEMORY]);
+
+	return nid;
+}
+
+/**
+ * ne_slot_pool_mem_donate() - Donate a run of memory to the slot pool.
+ * @pdev:	PCI device to donate to.
+ *
+ * The device counts the slots the pool's records support and the slots its
+ * memory supports separately, and hands out the smaller of the two counts.
+ * This donation raises the second count.
+ *
+ * The driver offers a span and the device answers with the run it took out of
+ * it, so the pages before and after that run go back to the allocator and the
+ * run alone stays recorded. A device that answers without naming a run inside
+ * the span leaves the whole span recorded, because which of the span's pages
+ * the device is using is then not established.
+ *
+ * Context: Process context. This function is called with the enclaves list
+ *	    mutex held.
+ * Return:
+ * * 0 on success.
+ * * -EINTR if a fatal signal arrived while the span was being zeroed.
+ * * Negative return value on failure.
+ * * -ETIMEDOUT, or any other code ne_request_unanswered() names, if the
+ *   request went unanswered, in which case the span stays recorded as donated
+ *   because the driver cannot establish that the device refused it. A device
+ *   that answers such a code itself lands here too and costs one held span.
+ */
+static int ne_slot_pool_mem_donate(struct pci_dev *pdev)
+{
+	const unsigned long span_pages = NE_SLOT_POOL_MEM_SPAN >> PAGE_SHIFT;
+	const unsigned long run_pages = NE_SLOT_POOL_MEM_SIZE >> PAGE_SHIFT;
+	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
+	struct slot_pool_mem_donate_req donate_req = {};
+	struct ne_pci_dev_cmd_reply cmd_reply = {};
+	u64 span_paddr, taken_paddr, taken_size;
+	unsigned long head_pages, tail_pages;
+	struct ne_slot_pool_run *run;
+	struct cma *cma = NULL;
+	struct page *page;
+	int rc;
+
+	run = kzalloc_obj(*run);
+	if (!run)
+		return -ENOMEM;
+
+	page = ne_slot_pool_alloc_run(span_pages,
+				      ne_slot_pool_donate_nid(ne_pci_dev->slot_pool_mem_nr),
+				      &cma);
+	if (IS_ERR_OR_NULL(page)) {
+		rc = page ? PTR_ERR(page) : -ENOMEM;
+
+		dev_warn(&pdev->dev,
+			 "Slot pool memory span size=%u not allocated [rc=%d]\n",
+			 NE_SLOT_POOL_MEM_SPAN, rc);
+
+		kfree(run);
+
+		return rc;
+	}
+
+	span_paddr = page_to_phys(page);
+
+	run->page = page;
+	run->cma = cma;
+	run->nr_pages = span_pages;
+
+	donate_req.paddr = span_paddr;
+	donate_req.size = NE_SLOT_POOL_MEM_SPAN;
+
+	rc = ne_do_request(pdev, SLOT_POOL_MEM_DONATE,
+			   &donate_req, sizeof(donate_req),
+			   &cmd_reply, sizeof(cmd_reply));
+	if (rc < 0) {
+		dev_err_ratelimited(&pdev->dev,
+				    "Error in slot pool memory donate [rc=%d]\n",
+				    rc);
+
+		/*
+		 * Record a span whose donation went unanswered as donated
+		 * instead of giving it back, as ne_slot_pool_donate_base() does
+		 * and for the same reason. The whole span is recorded, because
+		 * the run the device took out of it is what the answer that
+		 * never came would have named.
+		 */
+		if (ne_request_unanswered(rc)) {
+			list_add(&run->node, &ne_pci_dev->slot_pool_mem_runs);
+			ne_pci_dev->slot_pool_mem_nr++;
+		} else {
+			ne_contig_free_run(cma, page, span_pages);
+			kfree(run);
+		}
+
+		return rc;
+	}
+
+	taken_paddr = cmd_reply.mem_donate.consumed_paddr;
+	taken_size = cmd_reply.mem_donate.consumed_size;
+
+	/*
+	 * The run the device names has to be one run's worth, inside the span
+	 * it was offered, and to begin on a page boundary of the span, since
+	 * the pages given back are counted from where it begins. The device
+	 * places the run on a boundary of its own address space, which the
+	 * span's physical address says nothing about, so where in the span it
+	 * begins is only bounded, not fixed.
+	 */
+	if (taken_size == NE_SLOT_POOL_MEM_SIZE &&
+	    taken_paddr >= span_paddr &&
+	    taken_paddr - span_paddr <=
+		    NE_SLOT_POOL_MEM_SPAN - NE_SLOT_POOL_MEM_SIZE &&
+	    IS_ALIGNED(taken_paddr - span_paddr, PAGE_SIZE)) {
+		head_pages = (taken_paddr - span_paddr) >> PAGE_SHIFT;
+		tail_pages = span_pages - head_pages - run_pages;
+
+		run->page = page + head_pages;
+		run->nr_pages = run_pages;
+
+		if (head_pages)
+			ne_contig_free_run(cma, page, head_pages);
+
+		if (tail_pages)
+			ne_contig_free_run(cma, run->page + run_pages,
+					   tail_pages);
+	} else {
+		dev_warn(&pdev->dev,
+			 "Slot pool memory span kept, the device took paddr=0x%llx size=0x%llx\n",
+			 taken_paddr, taken_size);
+	}
+
+	list_add(&run->node, &ne_pci_dev->slot_pool_mem_runs);
+	ne_pci_dev->slot_pool_mem_nr++;
+
+	return 0;
+}
+
+/**
+ * ne_slot_pool_limits() - Read how many slots each half of the slot pool has
+ *			   room for.
+ * @pdev:	PCI device to read the slot pool limits of.
+ * @cmd_reply:	Reply the limits are read out of, valid on success.
+ *
+ * The reply is a snapshot. A caller that donates on the strength of it and then
+ * reads it again holds the enclaves list mutex across both.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+static int ne_slot_pool_limits(struct pci_dev *pdev,
+			       struct ne_pci_dev_cmd_reply *cmd_reply)
+{
+	u8 unused = 0;
+
+	return ne_do_request(pdev, SLOT_POOL_GET_LIMITS,
+			     &unused, sizeof(unused),
+			     cmd_reply, sizeof(*cmd_reply));
+}
+
+/**
+ * ne_slot_pool_effective_limit() - Read how many slots the parent can use out
+ *				    of a reply the device filled.
+ * @cmd_reply:	Reply a SLOT_POOL_GET_LIMITS filled.
+ *
+ * A slot needs room in both halves of the pool, so the number of slots the
+ * parent can use is the smaller of the two the device reports.
+ *
+ * Context: Process context.
+ * Return: the number of slots the parent can use.
+ */
+static u16 ne_slot_pool_effective_limit(const struct ne_pci_dev_cmd_reply *cmd_reply)
+{
+	return min(cmd_reply->limits.meta_limit, cmd_reply->limits.mem_limit);
+}
+
+/**
+ * ne_slot_pool_grow() - Grow the slot pool until it holds @target slots.
+ * @pdev:	PCI device to donate to.
+ * @target:	Number of slots the pool is to hold.
+ *
+ * The device hands out the smaller of the slots its records support and the
+ * slots its memory supports, so each round donates to whichever of the two is
+ * behind and then asks the device for both counts again. Neither is behind
+ * when the two are level, and a round donates to both then, because a
+ * donation to one half of a level pair leaves the smaller count where it was.
+ *
+ * A round that leaves both counts where they were ends the growth: the device
+ * took every donation the round offered and got no further with any of them,
+ * and another round would allocate again for the same answer.
+ *
+ * The pool only grows. A run the device accepted cannot be taken back on its
+ * own, so growth that stops short leaves the pool at the size it reached.
+ *
+ * Context: Process context. This function is called with the enclaves list
+ *	    mutex held.
+ * Return:
+ * * 0 once the pool holds @target slots or more.
+ * * -EOPNOTSUPP on a device that takes no donation.
+ * * -EINVAL where @target is above the number of slots the device allows.
+ * * -ENOSPC where a round raised neither count and the pool is below @target.
+ * * -EINTR if a fatal signal arrived while a run was being zeroed.
+ * * Negative return value on failure.
+ */
+int ne_slot_pool_grow(struct pci_dev *pdev, u16 target)
+{
+	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
+	struct ne_pci_dev_cmd_reply cmd_reply = {};
+	u16 effective;
+	int rc;
+
+	if (!ne_pci_dev->slot_pool_max)
+		return -EOPNOTSUPP;
+
+	if (target > ne_pci_dev->slot_pool_max)
+		return -EINVAL;
+
+	rc = ne_slot_pool_donate_base(pdev);
+	if (rc < 0)
+		return rc;
+
+	rc = ne_slot_pool_limits(pdev, &cmd_reply);
+	if (rc < 0)
+		return rc;
+
+	effective = ne_slot_pool_effective_limit(&cmd_reply);
+
+	while (effective < target) {
+		u16 meta_before = cmd_reply.limits.meta_limit;
+		u16 mem_before = cmd_reply.limits.mem_limit;
+
+		if (cmd_reply.limits.mem_limit <= cmd_reply.limits.meta_limit) {
+			rc = ne_slot_pool_mem_donate(pdev);
+			if (rc < 0)
+				return rc;
+		}
+
+		if (cmd_reply.limits.meta_limit <= cmd_reply.limits.mem_limit) {
+			rc = ne_slot_pool_donate(pdev);
+			if (rc < 0)
+				return rc;
+		}
+
+		rc = ne_slot_pool_limits(pdev, &cmd_reply);
+		if (rc < 0)
+			return rc;
+
+		if (cmd_reply.limits.meta_limit <= meta_before &&
+		    cmd_reply.limits.mem_limit <= mem_before)
+			return -ENOSPC;
+
+		effective = ne_slot_pool_effective_limit(&cmd_reply);
+	}
+
+	return 0;
+}
+
+/**
+ * enum ne_slot_pool_runs_owner - Who keeps the runs of a drained list.
+ * @NE_RUNS_TO_ALLOCATOR:	The device handed the runs back, so they go back
+ *				to the allocator.
+ * @NE_RUNS_STAY_WITH_DEVICE:	The device may still use the runs, so only the
+ *				records of them go away.
+ */
+enum ne_slot_pool_runs_owner {
+	NE_RUNS_TO_ALLOCATOR,
+	NE_RUNS_STAY_WITH_DEVICE,
+};
+
+/**
+ * ne_slot_pool_runs_drain() - Empty a list of donated runs.
+ * @runs:	List of runs to empty.
+ * @nr:		Count of runs on @runs, set to zero.
+ * @owner:	Who keeps the runs.
+ *
+ * Context: Process context. This function is called with the enclaves list
+ *	    mutex held.
+ */
+static void ne_slot_pool_runs_drain(struct list_head *runs, unsigned int *nr,
+				    enum ne_slot_pool_runs_owner owner)
+{
+	struct ne_slot_pool_run *run, *next;
+
+	list_for_each_entry_safe(run, next, runs, node) {
+		if (owner == NE_RUNS_TO_ALLOCATOR)
+			ne_contig_free_run(run->cma, run->page, run->nr_pages);
+
+		list_del(&run->node);
+		kfree(run);
+	}
+
+	*nr = 0;
+}
+
 /**
  * ne_slot_pool_release() - Take the whole slot pool back from the device.
  * @pdev:	PCI device to take the pool back from.
@@ -904,17 +1312,15 @@ static void ne_slot_pool_release(struct pci_dev *pdev)
 	if (rc < 0) {
 		dev_err(&pdev->dev, "Error in slot pool shutdown [rc=%d]\n", rc);
 
-		goto unlock_mutex;
+		goto forget_runs;
 	}
 
-	if (ne_pci_dev->slot_pool_page) {
-		ne_contig_free_run(ne_pci_dev->slot_pool_cma,
-				   ne_pci_dev->slot_pool_page,
-				   ne_pci_dev->slot_pool_gran >> PAGE_SHIFT);
-
-		ne_pci_dev->slot_pool_page = NULL;
-		ne_pci_dev->slot_pool_cma = NULL;
-	}
+	ne_slot_pool_runs_drain(&ne_pci_dev->slot_pool_meta_runs,
+				&ne_pci_dev->slot_pool_meta_nr,
+				NE_RUNS_TO_ALLOCATOR);
+	ne_slot_pool_runs_drain(&ne_pci_dev->slot_pool_mem_runs,
+				&ne_pci_dev->slot_pool_mem_nr,
+				NE_RUNS_TO_ALLOCATOR);
 
 	ne_contig_free_run(ne_pci_dev->slot_pool_base_cma,
 			   ne_pci_dev->slot_pool_base_page,
@@ -923,6 +1329,20 @@ static void ne_slot_pool_release(struct pci_dev *pdev)
 	ne_pci_dev->slot_pool_base_page = NULL;
 	ne_pci_dev->slot_pool_base_cma = NULL;
 
+	goto unlock_mutex;
+
+forget_runs:
+	/*
+	 * The device keeps these runs, so only the records of them go, and the
+	 * pages they name stay the device's for as long as it is bound.
+	 */
+	ne_slot_pool_runs_drain(&ne_pci_dev->slot_pool_meta_runs,
+				&ne_pci_dev->slot_pool_meta_nr,
+				NE_RUNS_STAY_WITH_DEVICE);
+	ne_slot_pool_runs_drain(&ne_pci_dev->slot_pool_mem_runs,
+				&ne_pci_dev->slot_pool_mem_nr,
+				NE_RUNS_STAY_WITH_DEVICE);
+
 unlock_mutex:
 	mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
 }
@@ -976,6 +1396,8 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 	atomic_set(&ne_pci_dev->cmd_reply_avail, 0);
 	init_waitqueue_head(&ne_pci_dev->cmd_reply_wait_q);
 	INIT_LIST_HEAD(&ne_pci_dev->enclaves_list);
+	INIT_LIST_HEAD(&ne_pci_dev->slot_pool_meta_runs);
+	INIT_LIST_HEAD(&ne_pci_dev->slot_pool_mem_runs);
 	mutex_init(&ne_pci_dev->enclaves_list_mutex);
 	mutex_init(&ne_pci_dev->pci_dev_mutex);
 	ne_pci_dev->pdev = pdev;
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index 3718f893e4f65..cab5434fc3e0d 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -468,6 +468,24 @@ static_assert(offsetof(struct ne_pci_dev_cmd_reply,
 		       limits.slots_per_donation) == 12);
 static_assert(offsetof(struct ne_pci_dev_cmd_reply, limits.reserved14) == 14);
 
+/**
+ * struct ne_slot_pool_run - One run of pages the driver gave the slot pool.
+ * @node:	Entry on the list of runs donated to one of the pools.
+ * @page:	First page of the run.
+ * @cma:	CMA area the run came from.
+ * @nr_pages:	Number of pages in the run.
+ *
+ * A record is allocated before the run it describes, and the run is recorded
+ * before the request carrying its address goes out, so a run the device
+ * accepted is a run ne_slot_pool_release() finds.
+ */
+struct ne_slot_pool_run {
+	struct list_head	node;
+	struct page		*page;
+	struct cma		*cma;
+	unsigned long		nr_pages;
+};
+
 /**
  * struct ne_pci_dev - Nitro Enclaves (NE) PCI device.
  * @cmd_reply_avail:		Variable set if a reply has been sent by the
@@ -501,9 +519,11 @@ static_assert(offsetof(struct ne_pci_dev_cmd_reply, limits.reserved14) == 14);
  *				has been made.
  * @slot_pool_base_cma:		CMA area the base donation came from.
  * @slot_pool_base_nr_pages:	Number of pages in the base donation.
- * @slot_pool_page:		First page of the donation to the slot pool, NULL
- *				until it has been made.
- * @slot_pool_cma:		CMA area the donation came from.
+ * @slot_pool_meta_runs:	Runs donated to hold the pool's slot descriptors.
+ * @slot_pool_meta_nr:		Number of runs on @slot_pool_meta_runs.
+ * @slot_pool_mem_runs:		Runs donated for the enclaves in the pool's
+ *				slots to run in.
+ * @slot_pool_mem_nr:		Number of runs on @slot_pool_mem_runs.
  */
 struct ne_pci_dev {
 	atomic_t		cmd_reply_avail;
@@ -522,13 +542,18 @@ struct ne_pci_dev {
 	struct page		*slot_pool_base_page;
 	struct cma		*slot_pool_base_cma;
 	unsigned long		slot_pool_base_nr_pages;
-	struct page		*slot_pool_page;
-	struct cma		*slot_pool_cma;
+	struct list_head	slot_pool_meta_runs;
+	unsigned int		slot_pool_meta_nr;
+	struct list_head	slot_pool_mem_runs;
+	unsigned int		slot_pool_mem_nr;
 };
 
 /* Whether a failure left a submitted request in flight. */
 bool ne_request_unanswered(int rc);
 
+/* Grow the slot pool until it holds a given number of slots. */
+int ne_slot_pool_grow(struct pci_dev *pdev, u16 target);
+
 /**
  * ne_do_request() - Submit command request to the PCI device based on the command
  *		     type and retrieve the associated reply.
-- 
2.47.1

[PATCH 21/40] nitro_enclaves: Document the slot pool

[Alexander Graf]

An admin sizing a CMA reservation for Nitro Sandboxes needs to allow for
the memory the driver donates to the slot pool as well as for the
enclaves' own allocations. Document both halves of the pool: which
device asks for them, what probe donates, how they grow, where the
memory comes from and when it goes back.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 Documentation/virt/ne_overview.rst | 31 ++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/Documentation/virt/ne_overview.rst b/Documentation/virt/ne_overview.rst
index 3cfc50f788433..16d6b4a2771c7 100644
--- a/Documentation/virt/ne_overview.rst
+++ b/Documentation/virt/ne_overview.rst
@@ -215,3 +215,34 @@ An allocation lives until the last of the things holding it is gone: its
 descriptor, its mappings, the regions given from it and the enclave it is the
 info page of. So neither closing the descriptor nor unmapping takes memory away
 from a running enclave.
+
+Slot pool
+---------
+
+The device keeps a record for every enclave slot it hands out. On a device that
+holds those records in parent memory the driver donates the memory for them
+before the device can allocate a slot at all, and that device draws the memory
+an enclave runs on from parent memory too. It counts both donations in slots: a
+record pool that holds so many records, and a memory pool that backs so many
+enclaves. Neither number on its own says how many enclaves can exist, because
+each one needs a record and the memory behind it, so the answer is whichever
+pool carries fewer. Such a device reports the highest number of slots it will
+grow to, the size and alignment of one donation to the record pool, the size of
+the base donation that pool hangs off, and how many slots one donation to the
+memory pool is worth. A device that holds the records itself reports no slots
+and takes no donation.
+
+The driver reads those numbers when it probes and donates the base run the
+device asked for and one run of the granularity it reported, which is the
+record pool an enclave needs to be created at all. Both runs are asked of the
+node the probe runs on first, through the allocator an enclave allocation uses,
+so a per-node CMA reservation there covers the pool as well as the enclaves. A
+run that node cannot serve comes from any node that has memory, because a run
+in the pool is read by the device alone. Probing carries on when a donation
+fails, with a warning naming the error. The memory pool starts at the number of
+slots the device backs without a donation.
+
+The device holds every donated page until it is told to let go, so the driver
+takes the whole pool back when a probe fails after donating, when it is unbound
+and when the machine shuts down, and the pages reach the allocator only after
+that has returned. Pages the device will not give back stay with it.
-- 
2.47.1