[PATCH 5/8] nitro_enclaves: Add NE_SET_ALLOC_NUMA_NODE

[Alexander Graf <graf@amazon.com>]

A CPU pool that spans NUMA nodes, which pool setup now permits, still
hands every enclave the same allocation target: the node that owns the
pool's first core. Allocation from that target does not spill to another
node, so once that node's contribution is exhausted NE_ADD_VCPU fails
with NE_ERR_NO_CPUS_AVAIL_IN_POOL (272) while cores on the pool's other
nodes sit unused. An enclave wider than one node's share of the pool
cannot be built, and nothing a caller passes to NE_ADD_VCPU changes
which node it draws from.

Give the enclave fd a sticky allocation target instead. NE_ADD_VCPU with
vcpu_id 0 draws its core from that node and returns the id it picked, so
a caller placing an enclave across two nodes sets the target once per
node and drains the count. NE_ALLOC_NUMA_NODE_ANY drops the constraint
and takes any free pool core. That constant is spelled (-1) in the uapi
header rather than reusing NUMA_NO_NODE, which is defined in
linux/nodemask_types.h and reachable from no uapi header at all.

I can think of two ways to let a caller name a node: a target that lives
on the fd, or a second NE_ADD_VCPU carrying the node on every call. I
picked the target. NE_ADD_VCPU already reports the id it chose, so a
caller spreading an enclave over several nodes needs no new call at all,
only the target and the existing count. The variant would make that same
caller learn a new ioctl to get behaviour it already has, and would
price every future input to a placement decision the same way: one more
ioctl number each.

Assisted-by: Kiro:claude-opus-5
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 38 ++++++++++++++++++++
 include/uapi/linux/nitro_enclaves.h       | 42 +++++++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index e019bb1f6594..eb0091de1182 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -22,6 +22,7 @@
 #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>
@@ -1477,6 +1478,43 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 		return 0;
 	}
 
+	case NE_SET_ALLOC_NUMA_NODE: {
+		struct ne_alloc_numa_node n;
+		int nid;
+
+		if (copy_from_user(&n, (void __user *)arg, sizeof(n)))
+			return -EFAULT;
+
+		if (n.flags)
+			return -EINVAL;
+
+		nid = n.numa_node;
+		if (nid != NUMA_NO_NODE &&
+		    (nid < 0 || nid >= nr_node_ids || !node_state(nid, N_POSSIBLE))) {
+			dev_err_ratelimited(ne_misc_dev.this_device,
+					    "Invalid NUMA node %d\n", nid);
+
+			return -EINVAL;
+		}
+
+		mutex_lock(&ne_enclave->enclave_info_mutex);
+
+		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;
+		}
+
+		ne_enclave->alloc_nid = nid;
+
+		mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+		return 0;
+	}
+
 	default:
 		return -ENOTTY;
 	}
diff --git a/include/uapi/linux/nitro_enclaves.h b/include/uapi/linux/nitro_enclaves.h
index 7c6ec8dfe451..8ddc4b150b7e 100644
--- a/include/uapi/linux/nitro_enclaves.h
+++ b/include/uapi/linux/nitro_enclaves.h
@@ -184,6 +184,31 @@
  */
 #define NE_START_ENCLAVE		_IOWR(0xAE, 0x24, struct ne_enclave_start_info)
 
+/**
+ * NE_SET_ALLOC_NUMA_NODE - Set the NUMA node of the primary VM used for
+ *			    subsequent kernel-side allocations on this enclave
+ *			    fd: NE_ADD_VCPU auto-pick (vcpu_id == 0) draws a
+ *			    core from this node. The setting is sticky
+ *			    until changed or the fd is closed. Pass
+ *			    %NE_ALLOC_NUMA_NODE_ANY for node-agnostic
+ *			    allocation (any pool core).
+ *
+ *			    Without this ioctl the target is the first node
+ *			    that owns a core in the CPU pool, and allocation
+ *			    never spills to another node, so a caller that
+ *			    cares which node it lands on names it here.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0				- On success.
+ * * -1				- On failure, errno is set to:
+ * *   EFAULT			- copy_from_user() failed.
+ * *   EINVAL			- flags is non-zero, or node id is not a
+ *				  possible node.
+ * *   NE_ERR_NOT_IN_INIT_STATE	- Enclave is not in init state.
+ */
+#define NE_SET_ALLOC_NUMA_NODE		_IOW(0xAE, 0x25, struct ne_alloc_numa_node)
+
 /**
  * DOC: NE specific error codes
  */
@@ -297,6 +322,23 @@
 
 #define NE_IMAGE_LOAD_MAX_FLAG_VAL	(0x02)
 
+/**
+ * NE_ALLOC_NUMA_NODE_ANY - Node id for a node-agnostic allocation target.
+ */
+#define NE_ALLOC_NUMA_NODE_ANY		(-1)
+
+/**
+ * struct ne_alloc_numa_node - Argument for %NE_SET_ALLOC_NUMA_NODE.
+ * @numa_node:	NUMA node of the primary VM for subsequent kernel-side
+ *		allocations (NE_ADD_VCPU auto-pick), or
+ *		%NE_ALLOC_NUMA_NODE_ANY for node-agnostic allocation.
+ * @flags:	Must be 0.
+ */
+struct ne_alloc_numa_node {
+	__s32	numa_node;
+	__u32	flags;
+};
+
 /**
  * struct ne_image_load_info - Info necessary for in-memory enclave image
  *			       loading (in / out).
-- 
2.47.1