[PATCH 00/19] nitro_enclaves: Fix enclave release, command retries and the CPU pool

[PATCH 00/19] nitro_enclaves: Fix enclave release, command retries and the CPU pool

[Alexander Graf]

While working with AI on the nitro_enclaves driver, it stumbled over a few
issues in the code. None of them are out of the ordinary, but they all fix
legit corner cases that are better to have fixed upstream.

While at it, I also realized that the MAINTAINERS file is outdated.
Alexandru Ciobotaru no longer works on Nitro Enclaves. I'm taking over
the maintainer role instead.

Alex

Alexander Graf (19):
  nitro_enclaves: Release the reference on a rejected region page
  nitro_enclaves: Split freeing the slot from freeing the enclave
  nitro_enclaves: Ignore the reply of a failed slot info command
  nitro_enclaves: Retry a command the device asks us to retry
  nitro_enclaves: Stop sending commands to a broken device
  nitro_enclaves: Set up the device state before requesting its
    interrupts
  nitro_enclaves: Wait for file operations before unbinding the PCI
    device
  nitro_enclaves: Let the region tracking arrays fall back to vmalloc
  nitro_enclaves: Do not return enclave threads to a replaced CPU pool
  nitro_enclaves: Do not report a torn down CPU pool as set
  nitro_enclaves: Keep the CPU pool when the new CPU list does not parse
  nitro_enclaves: Put the CPU pool back when the new CPU list is refused
  nitro_enclaves: Initialise the CPU pool mutex statically
  nitro_enclaves: Build the boot-time CPU pool once the CPUs are online
  nitro_enclaves: Set up the CPU pool under the CPU hotplug read lock
  nitro_enclaves: Let 32-bit processes use the ioctl interface
  nitro_enclaves: Index CPU pool cores by sibling mask
  nitro_enclaves: Time the wait for the device to turn off by the clock
  MAINTAINERS: Take over the Nitro Enclaves driver

 MAINTAINERS                               |   2 +-
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 658 +++++++++++++++-------
 drivers/virt/nitro_enclaves/ne_misc_dev.h |  24 +-
 drivers/virt/nitro_enclaves/ne_pci_dev.c  | 236 +++++++-
 drivers/virt/nitro_enclaves/ne_pci_dev.h  |   4 +
 5 files changed, 688 insertions(+), 236 deletions(-)


base-commit: 6628e7d4c6b25c8d18f1345e3a20d21829aae7e4
-- 
2.47.1

[PATCH 01/19] nitro_enclaves: Release the reference on a rejected region page

[Alexander Graf]

A process that passes ordinary anonymous memory to
NE_SET_USER_MEMORY_REGION can leak one page reference each time the driver
rejects it with NE_ERR_MEM_NOT_HUGE_PAGE. Any process with permission to open
/dev/nitro_enclaves and create an enclave can repeat that, which may lead
to situations where the system runs out of memory.

The ioctl takes a reference on one page at a time and only counts it in
ne_mem_region->nr_pages once that page has passed both of the checks
that follow, ne_sanity_check_user_mem_region_page() and
ne_merge_phys_contig_memory_regions(). The put_pages unwind walks
[0, nr_pages), so it misses the page whose check just failed.

Count the reference as soon as get_user_pages_unlocked() reports it, so
that the unwind covers it.

Fixes: 7dc9d4309fdb ("nitro_enclaves: Add logic for setting an enclave memory region")
Cc: stable@vger.kernel.org
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 0d7bdad924b73..a50dc4f737b55 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -963,8 +963,8 @@ static int ne_set_user_memory_region_ioctl(struct ne_enclave *ne_enclave,
 		gup_rc = get_user_pages_unlocked(mem_region.userspace_addr + memory_size, 1,
 						 ne_mem_region->pages + i, FOLL_GET);
 
-		if (gup_rc < 0) {
-			rc = gup_rc;
+		if (gup_rc != 1) {
+			rc = gup_rc < 0 ? gup_rc : -EFAULT;
 
 			dev_err_ratelimited(ne_misc_dev.this_device,
 					    "Error in get user pages [rc=%d]\n", rc);
@@ -972,6 +972,12 @@ static int ne_set_user_memory_region_ioctl(struct ne_enclave *ne_enclave,
 			goto put_pages;
 		}
 
+		/*
+		 * Count the reference before the checks below, so that the
+		 * put_pages unwind releases this page as well.
+		 */
+		ne_mem_region->nr_pages++;
+
 		rc = ne_sanity_check_user_mem_region_page(ne_enclave, ne_mem_region->pages[i]);
 		if (rc < 0)
 			goto put_pages;
@@ -983,8 +989,6 @@ static int ne_set_user_memory_region_ioctl(struct ne_enclave *ne_enclave,
 			goto put_pages;
 
 		memory_size += page_size(ne_mem_region->pages[i]);
-
-		ne_mem_region->nr_pages++;
 	} while (memory_size < mem_region.memory_size);
 
 	if ((ne_enclave->nr_mem_regions + phys_contig_mem_regions.num) >
-- 
2.47.1

[PATCH 02/19] nitro_enclaves: Split freeing the slot from freeing the enclave

[Alexander Graf]

Under rare circumstances, the nitro_enclaves device may turn down an
ENCLAVE_STOP or SLOT_FREE command. In preparation to handle that case,
split ne_enclave_release() into ne_enclave_free_slot() and
ne_enclave_free().

ne_enclave_free_slot() asks the device to stop the enclave and to free
its slot. ne_enclave_free() drops the references on the memory region
pages, puts the enclave's threads back into the NE CPU pool and frees
the enclave. It only runs once the device has confirmed that it no
longer owns what the enclave was given, which is the default behavior
today.

With the two apart, all donated Enclave resources like memory and CPUs
can stay tracked as in-use when the device turns a command down, for as
long as the device may still use them.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 111 ++++++++++++++--------
 1 file changed, 74 insertions(+), 37 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index a50dc4f737b55..bb5fdc377b051 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -1386,7 +1386,8 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
  *						from the enclave data structure.
  * @ne_enclave :	Private data associated with the current enclave.
  *
- * Context: Process context. This function is called with the ne_enclave mutex held.
+ * 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_mem_region_entries(struct ne_enclave *ne_enclave)
 {
@@ -1413,7 +1414,8 @@ static void ne_enclave_remove_all_mem_region_entries(struct ne_enclave *ne_encla
  *					     the enclave data structure.
  * @ne_enclave :	Private data associated with the current enclave.
  *
- * Context: Process context. This function is called with the ne_enclave mutex held.
+ * 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_vcpu_id_entries(struct ne_enclave *ne_enclave)
 {
@@ -1464,41 +1466,24 @@ static void ne_pci_dev_remove_enclave_entry(struct ne_enclave *ne_enclave,
 }
 
 /**
- * ne_enclave_release() - Release function provided by the enclave file.
- * @inode:	Inode associated with this file release function.
- * @file:	File associated with this release function.
+ * ne_enclave_free_slot() - Stop an enclave and free its slot in the device.
+ * @ne_pci_dev :	Private data associated with the PCI device.
+ * @ne_enclave :	Private data associated with the current enclave.
  *
- * Context: Process context.
+ * Context: Process context. This function is called with the ne_pci_dev enclave
+ *	    mutex and the ne_enclave mutex held.
  * Return:
- * * 0 on success.
+ * * 0 once the device has freed the slot.
  * * Negative return value on failure.
  */
-static int ne_enclave_release(struct inode *inode, struct file *file)
+static 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 = {};
-	struct ne_enclave *ne_enclave = file->private_data;
-	struct ne_pci_dev *ne_pci_dev = ne_devs.ne_pci_dev;
-	struct pci_dev *pdev = ne_pci_dev->pdev;
-	int rc = -EINVAL;
 	struct slot_free_req slot_free_req = {};
-
-	if (!ne_enclave)
-		return 0;
-
-	/*
-	 * Early exit in case there is an error in the enclave creation logic
-	 * and fput() is called on the cleanup path.
-	 */
-	if (!ne_enclave->slot_uid)
-		return 0;
-
-	/*
-	 * Acquire the enclave list mutex before the enclave mutex
-	 * in order to avoid deadlocks with @ref ne_event_work_handler.
-	 */
-	mutex_lock(&ne_pci_dev->enclaves_list_mutex);
-	mutex_lock(&ne_enclave->enclave_info_mutex);
+	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;
@@ -1510,7 +1495,7 @@ static int ne_enclave_release(struct inode *inode, struct file *file)
 			dev_err_ratelimited(ne_misc_dev.this_device,
 					    "Error in enclave stop [rc=%d]\n", rc);
 
-			goto unlock_mutex;
+			return rc;
 		}
 
 		memset(&cmd_reply, 0, sizeof(cmd_reply));
@@ -1525,22 +1510,74 @@ static int ne_enclave_release(struct inode *inode, struct file *file)
 		dev_err_ratelimited(ne_misc_dev.this_device,
 				    "Error in slot free [rc=%d]\n", rc);
 
-		goto unlock_mutex;
+		return rc;
 	}
 
-	ne_pci_dev_remove_enclave_entry(ne_enclave, ne_pci_dev);
+	return 0;
+}
+
+/**
+ * ne_enclave_free() - Give back what an enclave was given and free it.
+ * @ne_enclave :	Private data associated with the current enclave.
+ *
+ * Context: Process context. This function is called once the device has freed
+ *	    the slot of the enclave, on an enclave that is on no list and has no
+ *	    file left.
+ */
+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);
 
-	mutex_unlock(&ne_enclave->enclave_info_mutex);
-	mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
-
 	kfree(ne_enclave);
+}
 
-	return 0;
+/**
+ * ne_enclave_release() - Release function provided by the enclave file.
+ * @inode:	Inode associated with this file release function.
+ * @file:	File associated with this release function.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+static int ne_enclave_release(struct inode *inode, struct file *file)
+{
+	struct ne_enclave *ne_enclave = file->private_data;
+	struct ne_pci_dev *ne_pci_dev = ne_devs.ne_pci_dev;
+	int rc;
+
+	if (!ne_enclave)
+		return 0;
+
+	/*
+	 * Early exit in case there is an error in the enclave creation logic
+	 * and fput() is called on the cleanup path.
+	 */
+	if (!ne_enclave->slot_uid)
+		return 0;
+
+	/*
+	 * Acquire the enclave list mutex before the enclave mutex
+	 * in order to avoid deadlocks with @ref ne_event_work_handler.
+	 */
+	mutex_lock(&ne_pci_dev->enclaves_list_mutex);
+	mutex_lock(&ne_enclave->enclave_info_mutex);
+
+	/*
+	 * A slot the device did not free may still run on the memory and the
+	 * CPUs of the enclave, so the enclave keeps them.
+	 */
+	rc = ne_enclave_free_slot(ne_pci_dev, ne_enclave);
+	if (!rc)
+		ne_pci_dev_remove_enclave_entry(ne_enclave, ne_pci_dev);
 
-unlock_mutex:
 	mutex_unlock(&ne_enclave->enclave_info_mutex);
+
+	if (!rc)
+		ne_enclave_free(ne_enclave);
+
 	mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
 
 	return rc;
-- 
2.47.1

[PATCH 03/19] nitro_enclaves: Ignore the reply of a failed slot info command

[Alexander Graf]

When the nitro_enclaves device fails a SLOT_INFO command,
ne_event_work_handler() logs the failure and still compares the
enclave's state against the reply, which then holds zero or the reply
for the previous enclave. Zero is NE_STATE_INIT, so a running enclave is
recorded as not started: its owner is woken for a state change that did
not happen, and NE_ADD_VCPU, NE_SET_USER_MEMORY_REGION and
NE_START_ENCLAVE become callable on it again.

Skip the enclave when SLOT_INFO fails, so that it stays in
NE_STATE_RUNNING and a later event can examine it.

Fixes: e5d616d8bb50 ("nitro_enclaves: Handle out-of-band PCI device events")
Cc: stable@vger.kernel.org
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_pci_dev.c | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index 3e1455d505269..c37777bb95b49 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -243,9 +243,12 @@ static void ne_event_work_handler(struct work_struct *work)
 		rc = ne_do_request(pdev, SLOT_INFO,
 				   &slot_info_req, sizeof(slot_info_req),
 				   &cmd_reply, sizeof(cmd_reply));
-		if (rc < 0)
+		if (rc < 0) {
 			dev_err(&pdev->dev, "Error in slot info [rc=%d]\n", rc);
 
+			goto unlock;
+		}
+
 		/* Notify enclave process that the enclave state changed. */
 		if (ne_enclave->state != cmd_reply.state) {
 			ne_enclave->state = cmd_reply.state;
-- 
2.47.1

[PATCH 04/19] nitro_enclaves: Retry a command the device asks us to retry

[Alexander Graf]

The nitro_enclaves device may return -EAGAIN under rare circumstances,
which asks the system to perform the command again. Handle it as gracefully
as possible inside the kernel driver.

Most commands finish almost instantly; give them a 40s timeout. But commands
like free may indicate through -EAGAIN that they are still actively working
on freeling a slot which can take significantly longer. Use 10 minutes for
those.

To ensure that user space still reacts to events during such long timeout
periods, rework some of the locking ordering and use killable locks.

With this patch applied, a launch or a release that runs into such a
window waits for it to pass. When the device still asks for a retry once
the time is up, the caller gets -EAGAIN.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 28 +++++---
 drivers/virt/nitro_enclaves/ne_pci_dev.c  | 83 ++++++++++++++++++++++-
 2 files changed, 101 insertions(+), 10 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index bb5fdc377b051..0e9d6157dcc64 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -1155,7 +1155,9 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 		if (copy_from_user(&vcpu_id, (void __user *)arg, sizeof(vcpu_id)))
 			return -EFAULT;
 
-		mutex_lock(&ne_enclave->enclave_info_mutex);
+		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,
@@ -1219,11 +1221,14 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 
 	case NE_GET_IMAGE_LOAD_INFO: {
 		struct ne_image_load_info image_load_info = {};
+		int rc;
 
 		if (copy_from_user(&image_load_info, (void __user *)arg, sizeof(image_load_info)))
 			return -EFAULT;
 
-		mutex_lock(&ne_enclave->enclave_info_mutex);
+		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,
@@ -1267,7 +1272,9 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 			return -NE_ERR_INVALID_FLAG_VALUE;
 		}
 
-		mutex_lock(&ne_enclave->enclave_info_mutex);
+		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,
@@ -1347,7 +1354,9 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 			return -NE_ERR_INVALID_ENCLAVE_CID;
 		}
 
-		mutex_lock(&ne_enclave->enclave_info_mutex);
+		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,
@@ -1777,15 +1786,18 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 {
 	switch (cmd) {
 	case NE_CREATE_VM: {
-		int enclave_fd = -1;
 		struct ne_pci_dev *ne_pci_dev = ne_devs.ne_pci_dev;
 		u64 __user *slot_uid = (void __user *)arg;
+		int rc;
 
-		mutex_lock(&ne_pci_dev->enclaves_list_mutex);
-		enclave_fd = ne_create_vm_ioctl(ne_pci_dev, slot_uid);
+		rc = mutex_lock_killable(&ne_pci_dev->enclaves_list_mutex);
+		if (rc)
+			return rc;
+
+		rc = ne_create_vm_ioctl(ne_pci_dev, slot_uid);
 		mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
 
-		return enclave_fd;
+		return rc;
 	}
 
 	default:
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index c37777bb95b49..199c3c7f8c42e 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -26,6 +26,32 @@
  */
 #define NE_DEFAULT_TIMEOUT_MSECS	(120000) /* 120 sec */
 
+/**
+ * NE_CMD_RETRY_INITIAL_MSECS - Delay before the first retry of a command the
+ *				NE PCI device asked to be sent again.
+ */
+#define NE_CMD_RETRY_INITIAL_MSECS	(10)
+
+/**
+ * NE_CMD_RETRY_MAX_MSECS - Upper bound on the delay between two retries of a
+ *			    command the NE PCI device asked to be sent again.
+ */
+#define NE_CMD_RETRY_MAX_MSECS		(160)
+
+/**
+ * NE_CMD_RETRY_TOTAL_MSECS - How long after its first refusal a command the
+ *			      NE PCI device asks to be sent again may still go
+ *			      out, before -EAGAIN is reported to the caller.
+ */
+#define NE_CMD_RETRY_TOTAL_MSECS	(40000) /* 40 sec */
+
+/**
+ * NE_FREE_TIMEOUT_MSECS - How long after a first -EAGAIN ENCLAVE_STOP and
+ *			   SLOT_FREE may still go out. Scrubbing a large enclave
+ *			   can take minutes.
+ */
+#define NE_FREE_TIMEOUT_MSECS		(600000) /* 10 min */
+
 static const struct pci_device_id ne_pci_ids[] = {
 	{ PCI_DEVICE(PCI_VENDOR_ID_AMAZON, PCI_DEVICE_ID_NE) },
 	{ 0, }
@@ -96,11 +122,24 @@ static int ne_wait_for_reply(struct pci_dev *pdev)
 	return 0;
 }
 
+/**
+ * ne_cmd_frees_slot() - Whether a command is one of the two that free a slot.
+ * @cmd_type:	Command type of the request.
+ *
+ * Return: true for ENCLAVE_STOP and SLOT_FREE.
+ */
+static bool ne_cmd_frees_slot(enum ne_pci_dev_cmd_type cmd_type)
+{
+	return cmd_type == ENCLAVE_STOP || cmd_type == SLOT_FREE;
+}
+
 int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 		  void *cmd_request, size_t cmd_request_size,
 		  struct ne_pci_dev_cmd_reply *cmd_reply, size_t cmd_reply_size)
 {
+	unsigned int delay_msecs = NE_CMD_RETRY_INITIAL_MSECS;
 	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
+	unsigned long deadline = 0;
 	int rc = -EINVAL;
 
 	if (cmd_type <= INVALID_CMD || cmd_type >= MAX_CMD) {
@@ -137,12 +176,23 @@ int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 		return -EINVAL;
 	}
 
+retry:
 	/*
 	 * Use this mutex so that the PCI device handles one command request at
 	 * a time.
 	 */
 	mutex_lock(&ne_pci_dev->pci_dev_mutex);
 
+	if (deadline && time_after_eq(jiffies, deadline)) {
+		dev_err_ratelimited(&pdev->dev,
+				    "Device asked to retry cmd type=%u, giving up\n",
+				    cmd_type);
+
+		rc = -EAGAIN;
+
+		goto unlock_mutex;
+	}
+
 	atomic_set(&ne_pci_dev->cmd_reply_avail, 0);
 
 	ne_submit_request(pdev, cmd_type, cmd_request, cmd_request_size);
@@ -159,6 +209,27 @@ int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 
 	atomic_set(&ne_pci_dev->cmd_reply_avail, 0);
 
+	if (cmd_reply->rc == -EAGAIN) {
+		/*
+		 * The device answered, so nothing is in flight and another
+		 * command may go out while this one waits to be sent again.
+		 */
+		mutex_unlock(&ne_pci_dev->pci_dev_mutex);
+
+		if (!deadline) {
+			unsigned int budget_msecs = ne_cmd_frees_slot(cmd_type) ?
+						    NE_FREE_TIMEOUT_MSECS :
+						    NE_CMD_RETRY_TOTAL_MSECS;
+
+			deadline = jiffies + msecs_to_jiffies(budget_msecs);
+		}
+
+		msleep(delay_msecs);
+		delay_msecs = min(delay_msecs * 2, NE_CMD_RETRY_MAX_MSECS);
+
+		goto retry;
+	}
+
 	if (cmd_reply->rc < 0) {
 		rc = cmd_reply->rc;
 
@@ -221,7 +292,14 @@ static void ne_event_work_handler(struct work_struct *work)
 	int rc = -EINVAL;
 	struct slot_info_req slot_info_req = {};
 
-	mutex_lock(&ne_pci_dev->enclaves_list_mutex);
+	/*
+	 * A release holds these locks for as long as the device takes to free
+	 * the slot, which can be minutes, so wait killable: that is not
+	 * reported as a hung task. The worker gets no signal, so the locks do
+	 * not fail.
+	 */
+	if (mutex_lock_killable(&ne_pci_dev->enclaves_list_mutex))
+		return;
 
 	/*
 	 * Iterate over all enclaves registered for the Nitro Enclaves
@@ -229,7 +307,8 @@ static void ne_event_work_handler(struct work_struct *work)
 	 * is corresponding to.
 	 */
 	list_for_each_entry(ne_enclave, &ne_pci_dev->enclaves_list, enclave_list_entry) {
-		mutex_lock(&ne_enclave->enclave_info_mutex);
+		if (mutex_lock_killable(&ne_enclave->enclave_info_mutex))
+			break;
 
 		/*
 		 * Enclaves that were never started cannot receive out-of-band
-- 
2.47.1

[PATCH 05/19] nitro_enclaves: Stop sending commands to a broken device

[Alexander Graf]

The Nitro Enclaves device always answers quickly to its commands. Consider
it broken if it does not reply to a command within 120s. Except for
STOP/FREE commands which legitimately may take longer because the
hypervisor performs cleanup work that may take some time. For those,
use a 10 minute timeout.

Rework the request handler logic to enforce these limits. But still keep
processes killable during that time. Except for stop commands, which
are used as part of the slot shutdown sequence and have to block.

If a STOP/FREE command failed to respond in time (10 minutes), the
device is clearly broken. Switch to a mode then where it no longer
accepts any new commands.

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

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 0e9d6157dcc64..964cb21b27c77 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -1500,29 +1500,17 @@ static int ne_enclave_free_slot(struct ne_pci_dev *ne_pci_dev,
 		rc = ne_do_request(pdev, ENCLAVE_STOP,
 				   &enclave_stop_request, sizeof(enclave_stop_request),
 				   &cmd_reply, sizeof(cmd_reply));
-		if (rc < 0) {
-			dev_err_ratelimited(ne_misc_dev.this_device,
-					    "Error in enclave stop [rc=%d]\n", rc);
-
+		if (rc < 0)
 			return rc;
-		}
 
 		memset(&cmd_reply, 0, sizeof(cmd_reply));
 	}
 
 	slot_free_req.slot_uid = ne_enclave->slot_uid;
 
-	rc = ne_do_request(pdev, SLOT_FREE,
-			   &slot_free_req, sizeof(slot_free_req),
-			   &cmd_reply, sizeof(cmd_reply));
-	if (rc < 0) {
-		dev_err_ratelimited(ne_misc_dev.this_device,
-				    "Error in slot free [rc=%d]\n", rc);
-
-		return rc;
-	}
-
-	return 0;
+	return ne_do_request(pdev, SLOT_FREE,
+			     &slot_free_req, sizeof(slot_free_req),
+			     &cmd_reply, sizeof(cmd_reply));
 }
 
 /**
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index 199c3c7f8c42e..70fda4d29891f 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -46,9 +46,10 @@
 #define NE_CMD_RETRY_TOTAL_MSECS	(40000) /* 40 sec */
 
 /**
- * NE_FREE_TIMEOUT_MSECS - How long after a first -EAGAIN ENCLAVE_STOP and
- *			   SLOT_FREE may still go out. Scrubbing a large enclave
- *			   can take minutes.
+ * NE_FREE_TIMEOUT_MSECS - How long the driver waits for the reply to
+ *			   ENCLAVE_STOP or SLOT_FREE, and how long after a first
+ *			   -EAGAIN it sends either again. Scrubbing a large
+ *			   enclave can take minutes.
  */
 #define NE_FREE_TIMEOUT_MSECS		(600000) /* 10 min */
 
@@ -98,13 +99,14 @@ static void ne_retrieve_reply(struct pci_dev *pdev, struct ne_pci_dev_cmd_reply
 /**
  * ne_wait_for_reply() - Wait for a reply of a PCI device command.
  * @pdev:	PCI device for which a reply is waited.
+ * @timeout_msecs:	How long to wait for the reply.
  *
  * Context: Process context. This function is called with the ne_pci_dev mutex held.
  * Return:
  * * 0 on success.
  * * Negative return value on failure.
  */
-static int ne_wait_for_reply(struct pci_dev *pdev)
+static int ne_wait_for_reply(struct pci_dev *pdev, unsigned int timeout_msecs)
 {
 	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
 	int rc = -EINVAL;
@@ -112,10 +114,13 @@ static int ne_wait_for_reply(struct pci_dev *pdev)
 	/*
 	 * TODO: Update to _interruptible and handle interrupted wait event
 	 * e.g. -ERESTARTSYS, incoming signals + update timeout, if needed.
+	 *
+	 * A slot that is being scrubbed holds the reply for minutes, which is
+	 * no hung task, so wait idle.
 	 */
-	rc = wait_event_timeout(ne_pci_dev->cmd_reply_wait_q,
-				atomic_read(&ne_pci_dev->cmd_reply_avail) != 0,
-				msecs_to_jiffies(NE_DEFAULT_TIMEOUT_MSECS));
+	rc = wait_event_idle_timeout(ne_pci_dev->cmd_reply_wait_q,
+				     atomic_read(&ne_pci_dev->cmd_reply_avail) != 0,
+				     msecs_to_jiffies(timeout_msecs));
 	if (!rc)
 		return -ETIMEDOUT;
 
@@ -179,9 +184,23 @@ int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 retry:
 	/*
 	 * Use this mutex so that the PCI device handles one command request at
-	 * a time.
+	 * 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.
 	 */
-	mutex_lock(&ne_pci_dev->pci_dev_mutex);
+	if (ne_cmd_frees_slot(cmd_type)) {
+		mutex_lock(&ne_pci_dev->pci_dev_mutex);
+	} else {
+		rc = mutex_lock_killable(&ne_pci_dev->pci_dev_mutex);
+		if (rc)
+			return rc;
+	}
+
+	if (ne_pci_dev->failed) {
+		rc = -ENODEV;
+
+		goto unlock_mutex;
+	}
 
 	if (deadline && time_after_eq(jiffies, deadline)) {
 		dev_err_ratelimited(&pdev->dev,
@@ -197,10 +216,18 @@ int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 
 	ne_submit_request(pdev, cmd_type, cmd_request, cmd_request_size);
 
-	rc = ne_wait_for_reply(pdev);
+	rc = ne_wait_for_reply(pdev, ne_cmd_frees_slot(cmd_type) ?
+				     NE_FREE_TIMEOUT_MSECS :
+				     NE_DEFAULT_TIMEOUT_MSECS);
 	if (rc < 0) {
-		dev_err_ratelimited(&pdev->dev, "Error in wait for reply for cmd type=%u [rc=%d]\n",
-				    cmd_type, rc);
+		/*
+		 * A reply does not name the command it answers, so a reply that
+		 * comes after this would be read as the answer to the next one.
+		 */
+		ne_pci_dev->failed = true;
+
+		dev_err(&pdev->dev, "No reply for cmd type=%u, not sending more cmds\n",
+			cmd_type);
 
 		goto unlock_mutex;
 	}
@@ -242,6 +269,18 @@ int ne_do_request(struct pci_dev *pdev, enum ne_pci_dev_cmd_type cmd_type,
 	rc = 0;
 
 unlock_mutex:
+	/*
+	 * A device that does not stop or free a slot leaves the driver unable
+	 * to tell what the slot still uses, so treat it like one that did not
+	 * answer.
+	 */
+	if (rc < 0 && !ne_pci_dev->failed && ne_cmd_frees_slot(cmd_type)) {
+		ne_pci_dev->failed = true;
+
+		dev_err(&pdev->dev, "Slot not freed, cmd type=%u [rc=%d], not sending more cmds\n",
+			cmd_type, rc);
+	}
+
 	mutex_unlock(&ne_pci_dev->pci_dev_mutex);
 
 	return rc;
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.h b/drivers/virt/nitro_enclaves/ne_pci_dev.h
index 6e9f28971a4e0..19e129d14a3a0 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.h
@@ -279,6 +279,9 @@ struct ne_pci_dev_cmd_reply {
  * struct ne_pci_dev - Nitro Enclaves (NE) PCI device.
  * @cmd_reply_avail:		Variable set if a reply has been sent by the
  *				PCI device.
+ * @failed:			A command got no reply in time, or the device did
+ *				not stop or free a slot. No command is sent to
+ *				the device after that.
  * @cmd_reply_wait_q:		Wait queue for handling command reply from the
  *				PCI device.
  * @enclaves_list:		List of the enclaves managed by the PCI device.
@@ -294,6 +297,7 @@ struct ne_pci_dev_cmd_reply {
  */
 struct ne_pci_dev {
 	atomic_t		cmd_reply_avail;
+	bool			failed;
 	wait_queue_head_t	cmd_reply_wait_q;
 	struct list_head	enclaves_list;
 	struct mutex		enclaves_list_mutex;
-- 
2.47.1

[PATCH 06/19] nitro_enclaves: Set up the device state before requesting its interrupts

[Alexander Graf]

A nitro_enclaves device left enabled by a previous kernel, for example
before a kdump capture kernel binds it, may raise an interrupt as soon
as the driver requests its MSI-X vectors. ne_pci_probe() requests them
before it initialises the reply wait queue, the enclave list and the
mutexes the handlers use, so such an interrupt oopses the kernel while
it binds the device.

Initialise those fields ahead of ne_setup_msix(). ne_devs.ne_pci_dev
stays beside misc_register(), because setting it publishes the device to
the ioctl paths.

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

diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index 70fda4d29891f..d2c430d10758f 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -620,6 +620,13 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 
 	pci_set_drvdata(pdev, ne_pci_dev);
 
+	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);
+	mutex_init(&ne_pci_dev->enclaves_list_mutex);
+	mutex_init(&ne_pci_dev->pci_dev_mutex);
+	ne_pci_dev->pdev = pdev;
+
 	rc = ne_setup_msix(pdev);
 	if (rc < 0) {
 		dev_err(&pdev->dev, "Error in pci dev msix setup [rc=%d]\n", rc);
@@ -636,13 +643,6 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 		goto teardown_msix;
 	}
 
-	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);
-	mutex_init(&ne_pci_dev->enclaves_list_mutex);
-	mutex_init(&ne_pci_dev->pci_dev_mutex);
-	ne_pci_dev->pdev = pdev;
-
 	ne_devs.ne_pci_dev = ne_pci_dev;
 
 	rc = misc_register(ne_devs.ne_misc_dev);
-- 
2.47.1

[PATCH 07/19] nitro_enclaves: Wait for file operations before unbinding the PCI device

[Alexander Graf]

User space may hold enclave fds and /dev/nitro_enclaves fds across a
sysfs unbind or a hot-unplug of the NE PCI device. Their ioctls and
their release then use ne_devs.ne_pci_dev and ne_misc_dev.this_device
after ne_pci_remove() freed both, or oops on the NULL pointer it
leaves behind.

Add ne_devs.lock, a rw_semaphore that probe and remove hold for writing
and the file operations hold for reading, so that remove waits for the
ones in flight. A release or a remove can hold it for minutes while the
device frees a slot, so the ioctls take it with down_read_killable().
Remove then marks every enclave detached, which wakes its poll like an
out-of-band event, and an ioctl on a detached enclave, or on the misc
device while no device is bound, fails with -ENODEV. A detached enclave
keeps its memory and CPUs on release, because the device that could
confirm that they are free is gone.

A second NE PCI function would overwrite ne_devs.ne_pci_dev at probe, so
refuse it with -EBUSY.

Kiro, an AI coding assistant, found this while checking device removal
with enclave fds still open.

Fixes: 38907e124088 ("nitro_enclaves: Add logic for creating an enclave VM")
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 93 +++++++++++++++++++----
 drivers/virt/nitro_enclaves/ne_misc_dev.h |  8 ++
 drivers/virt/nitro_enclaves/ne_pci_dev.c  | 62 +++++++++++++++
 3 files changed, 148 insertions(+), 15 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 964cb21b27c77..821624c93795b 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -78,6 +78,7 @@ static struct miscdevice ne_misc_dev = {
 
 struct ne_devs ne_devs = {
 	.ne_misc_dev	= &ne_misc_dev,
+	.lock		= __RWSEM_INITIALIZER(ne_devs.lock),
 };
 
 /*
@@ -147,18 +148,22 @@ struct ne_phys_contig_mem_regions {
  */
 static bool ne_check_enclaves_created(void)
 {
-	struct ne_pci_dev *ne_pci_dev = ne_devs.ne_pci_dev;
+	struct ne_pci_dev *ne_pci_dev = NULL;
 	bool ret = false;
 
-	if (!ne_pci_dev)
-		return ret;
+	down_read(&ne_devs.lock);
 
-	mutex_lock(&ne_pci_dev->enclaves_list_mutex);
+	ne_pci_dev = ne_devs.ne_pci_dev;
+	if (ne_pci_dev) {
+		mutex_lock(&ne_pci_dev->enclaves_list_mutex);
 
-	if (!list_empty(&ne_pci_dev->enclaves_list))
-		ret = true;
+		if (!list_empty(&ne_pci_dev->enclaves_list))
+			ret = true;
 
-	mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+		mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+	}
+
+	up_read(&ne_devs.lock);
 
 	return ret;
 }
@@ -1133,17 +1138,18 @@ static int ne_start_enclave_ioctl(struct ne_enclave *ne_enclave,
 }
 
 /**
- * ne_enclave_ioctl() - Ioctl function provided by the enclave file.
+ * ne_enclave_do_ioctl() - Ioctl function provided by the enclave file.
  * @file:	File associated with this ioctl function.
  * @cmd:	The command that is set for the ioctl call.
  * @arg:	The argument that is provided for the ioctl call.
  *
- * Context: Process context.
+ * 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.
  * * Negative return value on failure.
  */
-static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
+static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 {
 	struct ne_enclave *ne_enclave = file->private_data;
 
@@ -1390,6 +1396,36 @@ static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long
 	return 0;
 }
 
+/**
+ * ne_enclave_ioctl() - Ioctl function provided by the enclave file.
+ * @file:	File associated with this ioctl function.
+ * @cmd:	The command that is set for the ioctl call.
+ * @arg:	The argument that is provided for the ioctl call.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+static long ne_enclave_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
+{
+	struct ne_enclave *ne_enclave = file->private_data;
+	long rc;
+
+	rc = down_read_killable(&ne_devs.lock);
+	if (rc)
+		return rc;
+
+	if (ne_enclave->detached)
+		rc = -ENODEV;
+	else
+		rc = ne_enclave_do_ioctl(file, cmd, arg);
+
+	up_read(&ne_devs.lock);
+
+	return rc;
+}
+
 /**
  * ne_enclave_remove_all_mem_region_entries() - Remove all memory region entries
  *						from the enclave data structure.
@@ -1542,7 +1578,7 @@ static void ne_enclave_free(struct ne_enclave *ne_enclave)
 static int ne_enclave_release(struct inode *inode, struct file *file)
 {
 	struct ne_enclave *ne_enclave = file->private_data;
-	struct ne_pci_dev *ne_pci_dev = ne_devs.ne_pci_dev;
+	struct ne_pci_dev *ne_pci_dev = NULL;
 	int rc;
 
 	if (!ne_enclave)
@@ -1555,6 +1591,21 @@ static int ne_enclave_release(struct inode *inode, struct file *file)
 	if (!ne_enclave->slot_uid)
 		return 0;
 
+	down_read(&ne_devs.lock);
+
+	/*
+	 * 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.
+	 */
+	if (ne_enclave->detached) {
+		rc = -ENODEV;
+
+		goto unlock_devs;
+	}
+
+	ne_pci_dev = ne_devs.ne_pci_dev;
+
 	/*
 	 * Acquire the enclave list mutex before the enclave mutex
 	 * in order to avoid deadlocks with @ref ne_event_work_handler.
@@ -1577,6 +1628,9 @@ static int ne_enclave_release(struct inode *inode, struct file *file)
 
 	mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
 
+unlock_devs:
+	up_read(&ne_devs.lock);
+
 	return rc;
 }
 
@@ -1774,16 +1828,25 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 {
 	switch (cmd) {
 	case NE_CREATE_VM: {
-		struct ne_pci_dev *ne_pci_dev = ne_devs.ne_pci_dev;
+		struct ne_pci_dev *ne_pci_dev = NULL;
 		u64 __user *slot_uid = (void __user *)arg;
 		int rc;
 
-		rc = mutex_lock_killable(&ne_pci_dev->enclaves_list_mutex);
+		rc = down_read_killable(&ne_devs.lock);
 		if (rc)
 			return rc;
 
-		rc = ne_create_vm_ioctl(ne_pci_dev, slot_uid);
-		mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+		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, slot_uid);
+				mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+			}
+		}
+
+		up_read(&ne_devs.lock);
 
 		return rc;
 	}
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 2a4d2224babac..4031d5d23eb00 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -12,6 +12,7 @@
 #include <linux/mm.h>
 #include <linux/mutex.h>
 #include <linux/pci.h>
+#include <linux/rwsem.h>
 #include <linux/wait.h>
 
 #include "ne_pci_dev.h"
@@ -35,6 +36,8 @@ struct ne_mem_region {
 /**
  * struct ne_enclave - Per-enclave data used for enclave lifetime management.
  * @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.
  * @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
@@ -63,6 +66,7 @@ struct ne_mem_region {
  */
 struct ne_enclave {
 	struct mutex		enclave_info_mutex;
+	bool			detached;
 	struct list_head	enclave_list_entry;
 	wait_queue_head_t	eventq;
 	bool			has_event;
@@ -97,10 +101,14 @@ enum ne_state {
  * struct ne_devs - Data structure to keep refs to the NE misc and PCI devices.
  * @ne_misc_dev:	Nitro Enclaves misc device.
  * @ne_pci_dev :	Nitro Enclaves PCI device.
+ * @lock:	Held for writing while the NE PCI device binds or unbinds, and for
+ *		reading by the file operations that reach it or log through the
+ *		misc device, which goes away with it.
  */
 struct ne_devs {
 	struct miscdevice	*ne_misc_dev;
 	struct ne_pci_dev	*ne_pci_dev;
+	struct rw_semaphore	lock;
 };
 
 /* Nitro Enclaves (NE) data structure for keeping refs to the NE misc and PCI devices. */
diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index d2c430d10758f..cd8e934e968c3 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -574,6 +574,38 @@ static void ne_pci_dev_disable(struct pci_dev *pdev)
 		dev_err(&pdev->dev, "Error in pci dev disable cmd\n");
 }
 
+/**
+ * ne_pci_dev_detach_enclaves() - Cut the enclaves off the NE PCI device that
+ *				  goes away.
+ * @ne_pci_dev:	Private data associated with the PCI device.
+ *
+ * Context: Process context. This function is called with ne_devs.lock held
+ *	    for writing.
+ */
+static void ne_pci_dev_detach_enclaves(struct ne_pci_dev *ne_pci_dev)
+{
+	struct ne_enclave *ne_enclave = NULL;
+	struct ne_enclave *ne_enclave_tmp = NULL;
+
+	mutex_lock(&ne_pci_dev->enclaves_list_mutex);
+
+	list_for_each_entry_safe(ne_enclave, ne_enclave_tmp,
+				 &ne_pci_dev->enclaves_list, enclave_list_entry) {
+		list_del(&ne_enclave->enclave_list_entry);
+
+		mutex_lock(&ne_enclave->enclave_info_mutex);
+
+		ne_enclave->detached = true;
+		ne_enclave->has_event = true;
+
+		wake_up_interruptible(&ne_enclave->eventq);
+
+		mutex_unlock(&ne_enclave->enclave_info_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.
@@ -643,6 +675,20 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 		goto teardown_msix;
 	}
 
+	down_write(&ne_devs.lock);
+
+	/*
+	 * The driver keeps a single device in ne_devs, and the enclaves and the
+	 * misc device belong to it.
+	 */
+	if (ne_devs.ne_pci_dev) {
+		dev_err(&pdev->dev, "Another NE PCI device is bound already\n");
+
+		rc = -EBUSY;
+
+		goto unlock_ne_devs;
+	}
+
 	ne_devs.ne_pci_dev = ne_pci_dev;
 
 	rc = misc_register(ne_devs.ne_misc_dev);
@@ -652,10 +698,14 @@ static int ne_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 		goto disable_ne_pci_dev;
 	}
 
+	up_write(&ne_devs.lock);
+
 	return 0;
 
 disable_ne_pci_dev:
 	ne_devs.ne_pci_dev = NULL;
+unlock_ne_devs:
+	up_write(&ne_devs.lock);
 	ne_pci_dev_disable(pdev);
 teardown_msix:
 	ne_teardown_msix(pdev);
@@ -682,6 +732,10 @@ static void ne_pci_remove(struct pci_dev *pdev)
 {
 	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
 
+	down_write(&ne_devs.lock);
+
+	ne_pci_dev_detach_enclaves(ne_pci_dev);
+
 	misc_deregister(ne_devs.ne_misc_dev);
 
 	ne_devs.ne_pci_dev = NULL;
@@ -699,6 +753,8 @@ static void ne_pci_remove(struct pci_dev *pdev)
 	pci_disable_device(pdev);
 
 	kfree(ne_pci_dev);
+
+	up_write(&ne_devs.lock);
 }
 
 /**
@@ -714,6 +770,10 @@ static void ne_pci_shutdown(struct pci_dev *pdev)
 	if (!ne_pci_dev)
 		return;
 
+	down_write(&ne_devs.lock);
+
+	ne_pci_dev_detach_enclaves(ne_pci_dev);
+
 	misc_deregister(ne_devs.ne_misc_dev);
 
 	ne_devs.ne_pci_dev = NULL;
@@ -731,6 +791,8 @@ static void ne_pci_shutdown(struct pci_dev *pdev)
 	pci_disable_device(pdev);
 
 	kfree(ne_pci_dev);
+
+	up_write(&ne_devs.lock);
 }
 
 /*
-- 
2.47.1

[PATCH 08/19] nitro_enclaves: Let the region tracking arrays fall back to vmalloc

[Alexander Graf]

NE_SET_USER_MEMORY_REGION allocates two tracking arrays that grow with
the region, and for a large region kzalloc_objs() fails with a WARN.
Neither needs physically contiguous memory, so allocate both with
kvzalloc_objs(), with __GFP_NOWARN since a region can push even
kvmalloc() past its limit.

Fixes: 7dc9d4309fdb ("nitro_enclaves: Add logic for setting an enclave memory region")
Cc: stable@vger.kernel.org
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 821624c93795b..ac33fbcd41f35 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -938,15 +938,17 @@ static int ne_set_user_memory_region_ioctl(struct ne_enclave *ne_enclave,
 
 	max_nr_pages = mem_region.memory_size / NE_MIN_MEM_REGION_SIZE;
 
-	ne_mem_region->pages = kzalloc_objs(*ne_mem_region->pages, max_nr_pages);
+	ne_mem_region->pages = kvzalloc_objs(*ne_mem_region->pages, max_nr_pages,
+					     GFP_KERNEL | __GFP_NOWARN);
 	if (!ne_mem_region->pages) {
 		rc = -ENOMEM;
 
 		goto free_mem_region;
 	}
 
-	phys_contig_mem_regions.regions = kzalloc_objs(*phys_contig_mem_regions.regions,
-						       max_nr_pages);
+	phys_contig_mem_regions.regions = kvzalloc_objs(*phys_contig_mem_regions.regions,
+							max_nr_pages,
+							GFP_KERNEL | __GFP_NOWARN);
 	if (!phys_contig_mem_regions.regions) {
 		rc = -ENOMEM;
 
@@ -1036,7 +1038,7 @@ static int ne_set_user_memory_region_ioctl(struct ne_enclave *ne_enclave,
 			dev_err_ratelimited(ne_misc_dev.this_device,
 					    "Error in slot add mem [rc=%d]\n", rc);
 
-			kfree(phys_contig_mem_regions.regions);
+			kvfree(phys_contig_mem_regions.regions);
 
 			/*
 			 * Exit here without put pages as memory regions may
@@ -1049,7 +1051,7 @@ static int ne_set_user_memory_region_ioctl(struct ne_enclave *ne_enclave,
 		ne_enclave->nr_mem_regions++;
 	}
 
-	kfree(phys_contig_mem_regions.regions);
+	kvfree(phys_contig_mem_regions.regions);
 
 	return 0;
 
@@ -1057,8 +1059,8 @@ static int ne_set_user_memory_region_ioctl(struct ne_enclave *ne_enclave,
 	for (i = 0; i < ne_mem_region->nr_pages; i++)
 		put_page(ne_mem_region->pages[i]);
 free_mem_region:
-	kfree(phys_contig_mem_regions.regions);
-	kfree(ne_mem_region->pages);
+	kvfree(phys_contig_mem_regions.regions);
+	kvfree(ne_mem_region->pages);
 	kfree(ne_mem_region);
 
 	return rc;
@@ -1448,7 +1450,7 @@ static void ne_enclave_remove_all_mem_region_entries(struct ne_enclave *ne_encla
 		for (i = 0; i < ne_mem_region->nr_pages; i++)
 			put_page(ne_mem_region->pages[i]);
 
-		kfree(ne_mem_region->pages);
+		kvfree(ne_mem_region->pages);
 
 		kfree(ne_mem_region);
 	}
-- 
2.47.1

[PATCH 09/19] nitro_enclaves: Do not return enclave threads to a replaced CPU pool

[Alexander Graf]

A write to ne_cpus checks that no enclave exists, drops
enclaves_list_mutex and then replaces the CPU pool, so an NE_CREATE_VM
in between gets an enclave whose release writes into the freed pool.
Hold the mutex across the check, the teardown and the setup, and take it
and ne_devs.lock killable, as the ioctls do.

Fixes: 9c8eb50fe9e2 ("nitro_enclaves: Add logic for terminating an enclave")
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 76 +++++++++++------------
 1 file changed, 37 insertions(+), 39 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index ac33fbcd41f35..0a1080443fc72 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -137,37 +137,6 @@ struct ne_phys_contig_mem_regions {
 	struct range  *regions;
 };
 
-/**
- * ne_check_enclaves_created() - Verify if at least one enclave has been created.
- * @void:	No parameters provided.
- *
- * Context: Process context.
- * Return:
- * * True if at least one enclave is created.
- * * False otherwise.
- */
-static bool ne_check_enclaves_created(void)
-{
-	struct ne_pci_dev *ne_pci_dev = NULL;
-	bool ret = false;
-
-	down_read(&ne_devs.lock);
-
-	ne_pci_dev = ne_devs.ne_pci_dev;
-	if (ne_pci_dev) {
-		mutex_lock(&ne_pci_dev->enclaves_list_mutex);
-
-		if (!list_empty(&ne_pci_dev->enclaves_list))
-			ret = true;
-
-		mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
-	}
-
-	up_read(&ne_devs.lock);
-
-	return ret;
-}
-
 /**
  * ne_setup_cpu_pool() - Set the NE CPU pool after handling sanity checks such
  *			 as not sharing CPU cores with the primary / parent VM
@@ -176,7 +145,8 @@ static bool ne_check_enclaves_created(void)
  *			 pool after the checks passed.
  * @ne_cpu_list:	The CPU list used for setting NE CPU pool.
  *
- * Context: Process context.
+ * Context: Process context. This function is called with the ne_pci_dev enclave
+ *	    mutex held, if there is a NE PCI device.
  * Return:
  * * 0 on success.
  * * Negative return value on failure.
@@ -385,7 +355,8 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
  *			    CPU pool.
  * @void:	No parameters provided.
  *
- * Context: Process context.
+ * Context: Process context. This function is called with the ne_pci_dev enclave
+ *	    mutex held, if there is a NE PCI device.
  */
 static void ne_teardown_cpu_pool(void)
 {
@@ -435,15 +406,37 @@ static void ne_teardown_cpu_pool(void)
 static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 {
 	char error_val[] = "";
+	struct ne_pci_dev *ne_pci_dev = NULL;
 	int rc = -EINVAL;
 
 	if (!capable(CAP_SYS_ADMIN))
 		return -EPERM;
 
-	if (ne_check_enclaves_created()) {
+	rc = down_read_killable(&ne_devs.lock);
+	if (rc)
+		return rc;
+
+	ne_pci_dev = ne_devs.ne_pci_dev;
+
+	/*
+	 * Keep the enclave list from gaining an entry while the pool is being
+	 * replaced, so that no enclave ends up holding a core count that
+	 * describes a pool that is already gone. The list lives in the PCI
+	 * device, which is not bound yet when this runs for a value given on
+	 * the kernel command line; nothing can create an enclave then either.
+	 */
+	if (ne_pci_dev) {
+		rc = mutex_lock_killable(&ne_pci_dev->enclaves_list_mutex);
+		if (rc)
+			goto unlock_devs;
+	}
+
+	if (ne_pci_dev && !list_empty(&ne_pci_dev->enclaves_list)) {
 		pr_err("%s: The CPU pool is used by enclave(s)\n", ne_misc_dev.name);
 
-		return -EPERM;
+		rc = -EPERM;
+
+		goto unlock;
 	}
 
 	ne_teardown_cpu_pool();
@@ -454,7 +447,7 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 
 		param_set_copystring(error_val, kp);
 
-		return rc;
+		goto unlock;
 	}
 
 	rc = param_set_copystring(val, kp);
@@ -464,11 +457,16 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 		ne_teardown_cpu_pool();
 
 		param_set_copystring(error_val, kp);
-
-		return rc;
 	}
 
-	return 0;
+unlock:
+	if (ne_pci_dev)
+		mutex_unlock(&ne_pci_dev->enclaves_list_mutex);
+
+unlock_devs:
+	up_read(&ne_devs.lock);
+
+	return rc;
 }
 
 /**
-- 
2.47.1

[PATCH 10/19] nitro_enclaves: Do not report a torn down CPU pool as set

[Alexander Graf]

An admin who writes a CPU list to ne_cpus can be told that the write
succeeded while the driver is left with no CPU pool: remove_cpu()
returns 1 for a CPU that went offline after the list was checked, and
ne_set_kernel_param() only tests for a negative value. Return -EINVAL,
as for a list that names an offline CPU.

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

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 0a1080443fc72..1c2a4994257d1 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -318,6 +318,9 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 			pr_err("%s: CPU %d is not offlined [rc=%d]\n",
 			       ne_misc_dev.name, cpu, rc);
 
+			if (rc > 0)
+				rc = -EINVAL;
+
 			goto online_cpus;
 		}
 	}
-- 
2.47.1

[PATCH 11/19] nitro_enclaves: Keep the CPU pool when the new CPU list does not parse

[Alexander Graf]

An admin who mistypes a CPU list into ne_cpus loses the pool that was
there, because ne_set_kernel_param() tears the pool down before
ne_setup_cpu_pool() parses the new value. The write fails, but the
pool's CPUs come back online and ne_cpus reads empty.

Parse the list before the teardown and pass the mask down. The other
checks stay after the teardown, because they want every CPU of the list
online, and a list that extends the current pool names offline CPUs.

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

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 1c2a4994257d1..de3bf93a4ef0f 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -143,7 +143,7 @@ struct ne_phys_contig_mem_regions {
  *			 or not using CPU 0, which should remain available for
  *			 the primary / parent VM. Offline the CPUs from the
  *			 pool after the checks passed.
- * @ne_cpu_list:	The CPU list used for setting NE CPU pool.
+ * @cpu_pool:	The CPU mask used for setting NE CPU pool.
  *
  * Context: Process context. This function is called with the ne_pci_dev enclave
  *	    mutex held, if there is a NE PCI device.
@@ -151,35 +151,24 @@ struct ne_phys_contig_mem_regions {
  * * 0 on success.
  * * Negative return value on failure.
  */
-static int ne_setup_cpu_pool(const char *ne_cpu_list)
+static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 {
 	int core_id = -1;
 	unsigned int cpu = 0;
-	cpumask_var_t cpu_pool;
 	unsigned int cpu_sibling = 0;
 	unsigned int i = 0;
 	int numa_node = -1;
 	int rc = -EINVAL;
 
-	if (!zalloc_cpumask_var(&cpu_pool, GFP_KERNEL))
-		return -ENOMEM;
-
 	mutex_lock(&ne_cpu_pool.mutex);
 
-	rc = cpulist_parse(ne_cpu_list, cpu_pool);
-	if (rc < 0) {
-		pr_err("%s: Error in cpulist parse [rc=%d]\n", ne_misc_dev.name, rc);
-
-		goto free_pool_cpumask;
-	}
-
 	cpu = cpumask_any(cpu_pool);
 	if (cpu >= nr_cpu_ids) {
 		pr_err("%s: No CPUs available in CPU pool\n", ne_misc_dev.name);
 
 		rc = -EINVAL;
 
-		goto free_pool_cpumask;
+		goto reset_cpu_pool;
 	}
 
 	/*
@@ -193,7 +182,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 			rc = -EINVAL;
 
-			goto free_pool_cpumask;
+			goto reset_cpu_pool;
 		}
 
 	/*
@@ -208,7 +197,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 				rc = -EINVAL;
 
-				goto free_pool_cpumask;
+				goto reset_cpu_pool;
 			}
 		} else {
 			if (numa_node != cpu_to_node(cpu)) {
@@ -217,7 +206,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 				rc = -EINVAL;
 
-				goto free_pool_cpumask;
+				goto reset_cpu_pool;
 			}
 		}
 
@@ -230,7 +219,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 		rc = -EINVAL;
 
-		goto free_pool_cpumask;
+		goto reset_cpu_pool;
 	}
 
 	for_each_cpu(cpu_sibling, topology_sibling_cpumask(0)) {
@@ -240,7 +229,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 			rc = -EINVAL;
 
-			goto free_pool_cpumask;
+			goto reset_cpu_pool;
 		}
 	}
 
@@ -257,7 +246,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 
 				rc = -EINVAL;
 
-				goto free_pool_cpumask;
+				goto reset_cpu_pool;
 			}
 		}
 	}
@@ -274,7 +263,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 	if (!ne_cpu_pool.avail_threads_per_core) {
 		rc = -ENOMEM;
 
-		goto free_pool_cpumask;
+		goto reset_cpu_pool;
 	}
 
 	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
@@ -325,8 +314,6 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 		}
 	}
 
-	free_cpumask_var(cpu_pool);
-
 	ne_cpu_pool.numa_node = numa_node;
 
 	mutex_unlock(&ne_cpu_pool.mutex);
@@ -343,8 +330,7 @@ static int ne_setup_cpu_pool(const char *ne_cpu_list)
 	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);
-free_pool_cpumask:
-	free_cpumask_var(cpu_pool);
+reset_cpu_pool:
 	ne_cpu_pool.nr_parent_vm_cores = 0;
 	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = -1;
@@ -409,15 +395,19 @@ static void ne_teardown_cpu_pool(void)
 static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 {
 	char error_val[] = "";
+	cpumask_var_t cpu_pool;
 	struct ne_pci_dev *ne_pci_dev = NULL;
 	int rc = -EINVAL;
 
 	if (!capable(CAP_SYS_ADMIN))
 		return -EPERM;
 
+	if (!zalloc_cpumask_var(&cpu_pool, GFP_KERNEL))
+		return -ENOMEM;
+
 	rc = down_read_killable(&ne_devs.lock);
 	if (rc)
-		return rc;
+		goto free_cpumasks;
 
 	ne_pci_dev = ne_devs.ne_pci_dev;
 
@@ -442,9 +432,20 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 		goto unlock;
 	}
 
+	/*
+	 * Parse the list before the teardown below drops the pool the driver
+	 * already has, so that a value that does not parse leaves it in place.
+	 */
+	rc = cpulist_parse(val, cpu_pool);
+	if (rc < 0) {
+		pr_err("%s: Error in cpulist parse [rc=%d]\n", ne_misc_dev.name, rc);
+
+		goto unlock;
+	}
+
 	ne_teardown_cpu_pool();
 
-	rc = ne_setup_cpu_pool(val);
+	rc = ne_setup_cpu_pool(cpu_pool);
 	if (rc < 0) {
 		pr_err("%s: Error in setup CPU pool [rc=%d]\n", ne_misc_dev.name, rc);
 
@@ -469,6 +470,9 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 unlock_devs:
 	up_read(&ne_devs.lock);
 
+free_cpumasks:
+	free_cpumask_var(cpu_pool);
+
 	return rc;
 }
 
-- 
2.47.1

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

[Alexander Graf]

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

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

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

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

[PATCH 13/19] nitro_enclaves: Initialise the CPU pool mutex statically

[Alexander Graf]

The ne_cpus setter can run before ne_init(), when the parameter is given
at boot or to modprobe, and then takes ne_cpu_pool.mutex before
mutex_init() ran on it. Initialise the mutex in its definition instead.

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

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 70d740d64ee17..44a4432be8164 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -125,7 +125,9 @@ struct ne_cpu_pool {
 	int		numa_node;
 };
 
-static struct ne_cpu_pool ne_cpu_pool;
+static struct ne_cpu_pool ne_cpu_pool = {
+	.mutex = __MUTEX_INITIALIZER(ne_cpu_pool.mutex),
+};
 
 /**
  * struct ne_phys_contig_mem_regions - Contiguous physical memory regions.
@@ -1902,8 +1904,6 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 
 static int __init ne_init(void)
 {
-	mutex_init(&ne_cpu_pool.mutex);
-
 	return pci_register_driver(&ne_pci_driver);
 }
 
-- 
2.47.1

[PATCH 14/19] nitro_enclaves: Build the boot-time CPU pool once the CPUs are online

[Alexander Graf]

An admin who passes nitro_enclaves.ne_cpus= to a kernel with the driver
built in gets no CPU pool, and this not very informative error message:

  Booting kernel: `2-5' invalid for parameter `nitro_enclaves.ne_cpus'

The parameter setter runs from parse_args() in start_kernel(), before
the page allocator and the secondary CPUs are up. Building a pool needs
both.

Keep a value that arrives while system_state is SYSTEM_BOOTING in
ne_cpus, note its parameter in ne_cpus_kp, and build the pool from it in
ne_init(), which runs once both are up. The pool building moves into
ne_replace_cpu_pool() for that. That way a valid list on the command
line gets its pool, and an invalid one reports why.

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

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 44a4432be8164..a897036dee2dd 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -14,6 +14,7 @@
 #include <linux/device.h>
 #include <linux/file.h>
 #include <linux/hugetlb.h>
+#include <linux/kernel.h>
 #include <linux/limits.h>
 #include <linux/list.h>
 #include <linux/miscdevice.h>
@@ -102,6 +103,13 @@ module_param_cb(ne_cpus, &ne_cpu_pool_ops, &ne_cpus_arg, 0644);
 /* https://www.kernel.org/doc/html/latest/admin-guide/kernel-parameters.html#cpu-lists */
 MODULE_PARM_DESC(ne_cpus, "<cpu-list> - CPU pool used for Nitro Enclaves");
 
+/*
+ * The parameter a value from the kernel command line is kept in until ne_init()
+ * builds its pool. Nothing else writes ne_cpus, so it is set whenever ne_cpus
+ * holds a value, and no sysfs write can reach the setter before ne_init() runs.
+ */
+static const struct kernel_param *ne_cpus_kp;
+
 /**
  * struct ne_cpu_pool - CPU pool used for Nitro Enclaves.
  * @avail_threads_per_core:	Available full CPU cores to be dedicated to
@@ -403,7 +411,7 @@ static void ne_snapshot_cpu_pool(struct cpumask *cpu_pool)
 }
 
 /**
- * ne_set_kernel_param() - Set the NE CPU pool value via the NE kernel parameter.
+ * ne_replace_cpu_pool() - Replace the NE CPU pool with the one a CPU list names.
  * @val:	NE CPU pool string value.
  * @kp :	NE kernel parameter associated with the NE CPU pool.
  *
@@ -412,7 +420,7 @@ static void ne_snapshot_cpu_pool(struct cpumask *cpu_pool)
  * * 0 on success.
  * * Negative return value on failure.
  */
-static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
+static int ne_replace_cpu_pool(const char *val, const struct kernel_param *kp)
 {
 	char error_val[] = "";
 	cpumask_var_t cpu_pool;
@@ -420,9 +428,6 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 	struct ne_pci_dev *ne_pci_dev = NULL;
 	int rc = -EINVAL;
 
-	if (!capable(CAP_SYS_ADMIN))
-		return -EPERM;
-
 	if (!zalloc_cpumask_var(&cpu_pool, GFP_KERNEL))
 		return -ENOMEM;
 
@@ -509,6 +514,39 @@ static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
 	return rc;
 }
 
+/**
+ * ne_set_kernel_param() - Set the NE CPU pool value via the NE kernel parameter.
+ * @val:	NE CPU pool string value.
+ * @kp :	NE kernel parameter associated with the NE CPU pool.
+ *
+ * Context: Process context.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+static int ne_set_kernel_param(const char *val, const struct kernel_param *kp)
+{
+	if (!capable(CAP_SYS_ADMIN))
+		return -EPERM;
+
+	/*
+	 * A value on the kernel command line reaches this from the parse_args()
+	 * call in start_kernel(), which runs before mm_core_init() has brought
+	 * up the allocator a cpumask comes from and long before smp_init() has
+	 * brought online the CPUs a pool is built out of. Keep the value and
+	 * leave the pool to ne_init(), which runs from do_basic_setup() with
+	 * both of those behind it. A value from modprobe or from sysfs arrives
+	 * with the system already up and is applied here.
+	 */
+	if (system_state == SYSTEM_BOOTING) {
+		ne_cpus_kp = kp;
+
+		return param_set_copystring(val, kp);
+	}
+
+	return ne_replace_cpu_pool(val, kp);
+}
+
 /**
  * ne_donated_cpu() - Check if the provided CPU is already used by the enclave.
  * @ne_enclave :	Private data associated with the current enclave.
@@ -1904,6 +1942,21 @@ static long ne_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 
 static int __init ne_init(void)
 {
+	static char cpu_list[NE_CPUS_SIZE] __initdata;
+
+	/*
+	 * Build the pool of a CPU list that came from the kernel command line,
+	 * which the setter only kept. The value is taken out of the parameter
+	 * string first, because building the pool writes that string.
+	 */
+	if (ne_cpus_kp && ne_cpus[0]) {
+		memcpy(cpu_list, ne_cpus, sizeof(cpu_list));
+
+		if (ne_replace_cpu_pool(cpu_list, ne_cpus_kp) < 0)
+			pr_err("%s: No CPU pool from nitro_enclaves.ne_cpus=%s\n",
+			       ne_misc_dev.name, cpu_list);
+	}
+
 	return pci_register_driver(&ne_pci_driver);
 }
 
-- 
2.47.1

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

[Alexander Graf]

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

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

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

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

[PATCH 16/19] nitro_enclaves: Let 32-bit processes use the ioctl interface

[Alexander Graf]

A 32-bit process on a 64-bit kernel gets -ENOTTY for every NE ioctl,
because neither the misc device nor an enclave fd sets .compat_ioctl.
The ioctl arguments don't require translation, so just add the default
compat_ioctl pointer to them and hence enable 32bit user space to drive
the device as well.

Fixes: 38907e124088 ("nitro_enclaves: Add logic for creating an enclave VM")
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index f8b066c990ce1..355eb0a684ec2 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -68,6 +68,7 @@ static const struct file_operations ne_fops = {
 	.owner		= THIS_MODULE,
 	.llseek		= noop_llseek,
 	.unlocked_ioctl	= ne_ioctl,
+	.compat_ioctl	= compat_ptr_ioctl,
 };
 
 static struct miscdevice ne_misc_dev = {
@@ -1762,6 +1763,7 @@ static const struct file_operations ne_enclave_fops = {
 	.llseek		= noop_llseek,
 	.poll		= ne_enclave_poll,
 	.unlocked_ioctl	= ne_enclave_ioctl,
+	.compat_ioctl	= compat_ptr_ioctl,
 	.release	= ne_enclave_release,
 };
 
-- 
2.47.1

[PATCH 17/19] nitro_enclaves: Index CPU pool cores by sibling mask

[Alexander Graf]

An enclave on a host whose CPU pool spans two packages can get the
threads of a core on a different package. The pool indexes
avail_threads_per_core[] by topology_core_id(), which is package-local
on x86, so two cores in different packages share one entry.

Walk the sibling mask of each pool CPU instead, in
ne_build_core_slots(), which gives each core its own entry and counts
the cores of the pool. An enclave copies that count at NE_CREATE_VM.
NE_ADD_VCPU now bounds vcpu_id by nr_cpu_ids instead of the thread count
of the pool, so an id above the old bound that the pool does not hold
answers NE_ERR_VCPU_NOT_IN_CPU_POOL instead of NE_ERR_INVALID_VCPU.

Fixes: ff8a4d3e3a99 ("nitro_enclaves: Add logic for setting an enclave vCPU")
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_misc_dev.c | 168 ++++++++++++++--------
 drivers/virt/nitro_enclaves/ne_misc_dev.h |  16 +--
 2 files changed, 113 insertions(+), 71 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.c b/drivers/virt/nitro_enclaves/ne_misc_dev.c
index 355eb0a684ec2..cb1efc95a7f34 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.c
@@ -18,6 +18,7 @@
 #include <linux/limits.h>
 #include <linux/list.h>
 #include <linux/miscdevice.h>
+#include <linux/minmax.h>
 #include <linux/mm.h>
 #include <linux/mman.h>
 #include <linux/module.h>
@@ -114,23 +115,23 @@ static const struct kernel_param *ne_cpus_kp;
 /**
  * struct ne_cpu_pool - CPU pool used for Nitro Enclaves.
  * @avail_threads_per_core:	Available full CPU cores to be dedicated to
- *				enclave(s). The cpumasks from the array, indexed
- *				by core id, contain all the threads from the
- *				available cores, that are not set for created
- *				enclave(s). The full CPU cores are part of the
- *				NE CPU pool.
+ *				enclave(s). Each cpumask in the array holds the
+ *				threads of one core in the pool that are not set
+ *				for created enclave(s). Slots are handed out as
+ *				cores are discovered and carry no meaning
+ *				outside the driver.
  * @mutex:			Mutex for the access to the NE CPU pool.
- * @nr_parent_vm_cores :	The size of the available threads per core array.
- *				The total number of CPU cores available on the
- *				primary / parent VM.
- * @nr_threads_per_core:	The number of threads that a full CPU core has.
+ * @nr_parent_vm_cores :	The size of the available threads per core array,
+ *				one slot per full CPU core in the pool. Counted
+ *				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.
  */
 struct ne_cpu_pool {
 	cpumask_var_t	*avail_threads_per_core;
 	struct mutex	mutex;
 	unsigned int	nr_parent_vm_cores;
-	unsigned int	nr_threads_per_core;
 	int		numa_node;
 };
 
@@ -169,6 +170,88 @@ static void ne_free_core_slots(void)
 	ne_cpu_pool.avail_threads_per_core = NULL;
 }
 
+/**
+ * ne_build_core_slots() - Group the CPU pool into full cores and give each one
+ *			   a slot of the NE CPU pool arrays.
+ * @cpu_pool:	The CPU pool as parsed out of the ne_cpus parameter.
+ *
+ * The slots preserve the CPU topology of the pool once its CPUs are offlined.
+ * topology_core_id() is package-local on x86, so two cores in different
+ * packages report the same id and would share a slot. The per-CPU sibling
+ * cpumask identifies a core on the architectures this driver builds for, so
+ * walk that instead: once to count the cores the pool holds, and once to hand
+ * each of them the next free slot. Counting first sizes the arrays to the pool,
+ * so the walk has a slot for every core it finds.
+ *
+ * Context: Process context. This function is called with the ne_cpu_pool mutex
+ *	    and the CPU hotplug read lock held.
+ * Return:
+ * * 0 on success.
+ * * Negative return value on failure.
+ */
+static int ne_build_core_slots(const struct cpumask *cpu_pool)
+{
+	unsigned int cpu = 0;
+	unsigned int cpu_sibling = 0;
+	unsigned int i = 0;
+	unsigned int next_core_idx = 0;
+	unsigned int nr_cores = 0;
+	cpumask_var_t processed;
+	int rc = -ENOMEM;
+
+	if (!zalloc_cpumask_var(&processed, GFP_KERNEL))
+		return -ENOMEM;
+
+	for_each_cpu(cpu, cpu_pool) {
+		if (cpumask_test_cpu(cpu, processed))
+			continue;
+
+		cpumask_or(processed, processed, topology_sibling_cpumask(cpu));
+
+		nr_cores++;
+	}
+
+	ne_cpu_pool.avail_threads_per_core = kzalloc_objs(*ne_cpu_pool.avail_threads_per_core,
+							  nr_cores);
+	if (!ne_cpu_pool.avail_threads_per_core)
+		goto free_processed;
+
+	ne_cpu_pool.nr_parent_vm_cores = nr_cores;
+
+	for (i = 0; i < nr_cores; i++)
+		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL))
+			goto free_slots;
+
+	cpumask_clear(processed);
+
+	for_each_cpu(cpu, cpu_pool) {
+		if (cpumask_test_cpu(cpu, processed))
+			continue;
+
+		for_each_cpu(cpu_sibling, topology_sibling_cpumask(cpu)) {
+			if (!cpumask_test_cpu(cpu_sibling, cpu_pool))
+				continue;
+
+			cpumask_set_cpu(cpu_sibling,
+					ne_cpu_pool.avail_threads_per_core[next_core_idx]);
+			cpumask_set_cpu(cpu_sibling, processed);
+		}
+
+		next_core_idx++;
+	}
+
+	free_cpumask_var(processed);
+
+	return 0;
+
+free_slots:
+	ne_free_core_slots();
+free_processed:
+	free_cpumask_var(processed);
+
+	return rc;
+}
+
 /**
  * ne_setup_cpu_pool() - Set the NE CPU pool after handling sanity checks such
  *			 as not sharing CPU cores with the primary / parent VM
@@ -177,7 +260,7 @@ static void ne_free_core_slots(void)
  *			 pool after the checks passed.
  * @cpu_pool:	The CPU mask used for setting NE CPU pool.
  *
- * The topology of the pool - online state, sibling masks, core ids - is read
+ * The topology of the pool - online state, sibling masks, core slots - is read
  * under the CPU hotplug read lock. The lock is dropped before the pool CPUs are
  * offlined, because remove_cpu() takes the write side of the same percpu-rwsem.
  * Nothing reads topology after the drop.
@@ -190,10 +273,8 @@ static void ne_free_core_slots(void)
  */
 static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 {
-	int core_id = -1;
 	unsigned int cpu = 0;
 	unsigned int cpu_sibling = 0;
-	unsigned int i = 0;
 	int numa_node = -1;
 	int rc = -EINVAL;
 
@@ -211,7 +292,7 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 
 	/*
 	 * Check if the CPUs are online, to further get info about them
-	 * e.g. numa node, core id, siblings.
+	 * e.g. numa node, siblings.
 	 */
 	for_each_cpu(cpu, cpu_pool)
 		if (cpu_is_offline(cpu)) {
@@ -289,45 +370,9 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 		}
 	}
 
-	/* Calculate the number of threads from a full CPU core. */
-	cpu = cpumask_any(cpu_pool);
-	for_each_cpu(cpu_sibling, topology_sibling_cpumask(cpu))
-		ne_cpu_pool.nr_threads_per_core++;
-
-	ne_cpu_pool.nr_parent_vm_cores = nr_cpu_ids / ne_cpu_pool.nr_threads_per_core;
-
-	ne_cpu_pool.avail_threads_per_core = kzalloc_objs(*ne_cpu_pool.avail_threads_per_core,
-							  ne_cpu_pool.nr_parent_vm_cores);
-	if (!ne_cpu_pool.avail_threads_per_core) {
-		rc = -ENOMEM;
-
+	rc = ne_build_core_slots(cpu_pool);
+	if (rc < 0)
 		goto unlock_hotplug;
-	}
-
-	for (i = 0; i < ne_cpu_pool.nr_parent_vm_cores; i++)
-		if (!zalloc_cpumask_var(&ne_cpu_pool.avail_threads_per_core[i], GFP_KERNEL)) {
-			rc = -ENOMEM;
-
-			goto free_slots;
-		}
-
-	/*
-	 * Split the NE CPU pool in threads per core to keep the CPU topology
-	 * after offlining the CPUs.
-	 */
-	for_each_cpu(cpu, cpu_pool) {
-		core_id = topology_core_id(cpu);
-		if (core_id < 0 || core_id >= ne_cpu_pool.nr_parent_vm_cores) {
-			pr_err("%s: Invalid core id  %d for CPU %d\n",
-			       ne_misc_dev.name, core_id, cpu);
-
-			rc = -EINVAL;
-
-			goto free_slots;
-		}
-
-		cpumask_set_cpu(cpu, ne_cpu_pool.avail_threads_per_core[core_id]);
-	}
 
 	cpus_read_unlock();
 
@@ -365,13 +410,10 @@ static int ne_setup_cpu_pool(const struct cpumask *cpu_pool)
 		add_cpu(cpu);
 	ne_free_core_slots();
 	goto reset_pool;
-free_slots:
-	ne_free_core_slots();
 unlock_hotplug:
 	cpus_read_unlock();
 reset_pool:
 	ne_cpu_pool.nr_parent_vm_cores = 0;
-	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = -1;
 	mutex_unlock(&ne_cpu_pool.mutex);
 
@@ -410,7 +452,6 @@ static void ne_teardown_cpu_pool(void)
 
 	ne_free_core_slots();
 	ne_cpu_pool.nr_parent_vm_cores = 0;
-	ne_cpu_pool.nr_threads_per_core = 0;
 	ne_cpu_pool.numa_node = -1;
 
 	mutex_unlock(&ne_cpu_pool.mutex);
@@ -647,7 +688,8 @@ static int ne_set_enclave_threads_per_core(struct ne_enclave *ne_enclave,
 		return -NE_ERR_VCPU_NOT_IN_CPU_POOL;
 	}
 
-	if (core_id >= ne_enclave->nr_parent_vm_cores) {
+	if (core_id >= min(ne_enclave->nr_parent_vm_cores,
+			   ne_cpu_pool.nr_parent_vm_cores)) {
 		dev_err_ratelimited(ne_misc_dev.this_device,
 				    "Invalid core id %d - ne_enclave\n", core_id);
 
@@ -1276,8 +1318,7 @@ static long ne_enclave_do_ioctl(struct file *file, unsigned int cmd, unsigned lo
 			return -NE_ERR_NOT_IN_INIT_STATE;
 		}
 
-		if (vcpu_id >= (ne_enclave->nr_parent_vm_cores *
-		    ne_enclave->nr_threads_per_core)) {
+		if (vcpu_id >= nr_cpu_ids) {
 			dev_err_ratelimited(ne_misc_dev.this_device,
 					    "vCPU id higher than max CPU id\n");
 
@@ -1566,18 +1607,22 @@ static void ne_enclave_remove_all_mem_region_entries(struct ne_enclave *ne_encla
  */
 static void ne_enclave_remove_all_vcpu_id_entries(struct ne_enclave *ne_enclave)
 {
+	unsigned int nr_cores = 0;
 	unsigned int cpu = 0;
 	unsigned int i = 0;
 
 	mutex_lock(&ne_cpu_pool.mutex);
 
-	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++) {
+	nr_cores = min(ne_enclave->nr_parent_vm_cores,
+		       ne_cpu_pool.nr_parent_vm_cores);
+
+	for (i = 0; i < nr_cores; 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]);
 
+	for (i = 0; i < ne_enclave->nr_parent_vm_cores; i++)
 		free_cpumask_var(ne_enclave->threads_per_core[i]);
-	}
 
 	mutex_unlock(&ne_cpu_pool.mutex);
 
@@ -1816,7 +1861,6 @@ static int ne_create_vm_ioctl(struct ne_pci_dev *ne_pci_dev, u64 __user *slot_ui
 	mutex_lock(&ne_cpu_pool.mutex);
 
 	ne_enclave->nr_parent_vm_cores = ne_cpu_pool.nr_parent_vm_cores;
-	ne_enclave->nr_threads_per_core = ne_cpu_pool.nr_threads_per_core;
 	ne_enclave->numa_node = ne_cpu_pool.numa_node;
 
 	mutex_unlock(&ne_cpu_pool.mutex);
diff --git a/drivers/virt/nitro_enclaves/ne_misc_dev.h b/drivers/virt/nitro_enclaves/ne_misc_dev.h
index 4031d5d23eb00..283c17558f2a3 100644
--- a/drivers/virt/nitro_enclaves/ne_misc_dev.h
+++ b/drivers/virt/nitro_enclaves/ne_misc_dev.h
@@ -50,18 +50,17 @@ struct ne_mem_region {
  * @mem_size:			Enclave memory size.
  * @mm :			Enclave process abstraction mm data struct.
  * @nr_mem_regions:		Number of memory regions associated with the enclave.
- * @nr_parent_vm_cores :	The size of the threads per core array. The
- *				total number of CPU cores available on the
- *				parent / primary VM.
- * @nr_threads_per_core:	The number of threads that a full CPU core has.
+ * @nr_parent_vm_cores :	The size of the threads per core array, one slot
+ *				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.
  * @slot_uid:			Slot unique id mapped to the enclave.
  * @state:			Enclave state, updated during enclave lifetime.
- * @threads_per_core:		Enclave full CPU cores array, indexed by core id,
- *				consisting of cpumasks with all their threads.
- *				Full CPU cores are taken from the NE CPU pool
- *				and are available to the enclave.
+ * @threads_per_core:		Enclave full CPU cores array. Each cpumask in the
+ *				array holds the threads of one core taken from
+ *				the NE CPU pool and available to the enclave,
+ *				at the index of that core's slot in the pool.
  * @vcpu_ids:			Cpumask of the vCPUs that are set for the enclave.
  */
 struct ne_enclave {
@@ -76,7 +75,6 @@ struct ne_enclave {
 	struct mm_struct	*mm;
 	unsigned int		nr_mem_regions;
 	unsigned int		nr_parent_vm_cores;
-	unsigned int		nr_threads_per_core;
 	unsigned int		nr_vcpus;
 	int			numa_node;
 	u64			slot_uid;
-- 
2.47.1

[PATCH 18/19] nitro_enclaves: Time the wait for the device to turn off by the clock

[Alexander Graf]

Unbind, shutdown and the next probe give the nitro_enclaves device
NE_DEFAULT_TIMEOUT_MSECS, 120 s, to turn off. The device turns off only
once it has stopped every enclave and scrubbed its memory, which can
take as long as a single ENCLAVE_STOP. ne_pci_dev_disable() also counts
the wait in 10 ms steps of msleep_interruptible(). At HZ=100 each step
sleeps two jiffies, so the wait takes 240 s, and with a signal pending
it does not sleep at all.

Allow the device NE_FREE_TIMEOUT_MSECS, 10 minutes, as for ENCLAVE_STOP,
and poll the register with readx_poll_timeout(), which keeps a ktime
deadline.

Kiro, an AI coding assistant, found this while timing an unbind on
hardware.

Fixes: 89308c11ae3b ("nitro_enclaves: Init PCI device driver")
Signed-off-by: Alexander Graf <graf@amazon.com>
---
 drivers/virt/nitro_enclaves/ne_pci_dev.c | 21 ++++++++-------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/drivers/virt/nitro_enclaves/ne_pci_dev.c b/drivers/virt/nitro_enclaves/ne_pci_dev.c
index cd8e934e968c3..edc72d4f623e5 100644
--- a/drivers/virt/nitro_enclaves/ne_pci_dev.c
+++ b/drivers/virt/nitro_enclaves/ne_pci_dev.c
@@ -9,6 +9,7 @@
 
 #include <linux/delay.h>
 #include <linux/device.h>
+#include <linux/iopoll.h>
 #include <linux/list.h>
 #include <linux/module.h>
 #include <linux/mutex.h>
@@ -550,8 +551,7 @@ static void ne_pci_dev_disable(struct pci_dev *pdev)
 {
 	u8 dev_disable_reply = 0;
 	struct ne_pci_dev *ne_pci_dev = pci_get_drvdata(pdev);
-	const unsigned int sleep_time = 10; /* 10 ms */
-	unsigned int sleep_time_count = 0;
+	int rc;
 
 	iowrite8(NE_ENABLE_OFF, ne_pci_dev->iomem_base + NE_ENABLE);
 
@@ -560,17 +560,12 @@ static void ne_pci_dev_disable(struct pci_dev *pdev)
 	 * state is not immediately set to disabled and going through a
 	 * transitory state of disabling.
 	 */
-	while (sleep_time_count < NE_DEFAULT_TIMEOUT_MSECS) {
-		dev_disable_reply = ioread8(ne_pci_dev->iomem_base + NE_ENABLE);
-		if (dev_disable_reply == NE_ENABLE_OFF)
-			return;
-
-		msleep_interruptible(sleep_time);
-		sleep_time_count += sleep_time;
-	}
-
-	dev_disable_reply = ioread8(ne_pci_dev->iomem_base + NE_ENABLE);
-	if (dev_disable_reply != NE_ENABLE_OFF)
+	rc = readx_poll_timeout(ioread8, ne_pci_dev->iomem_base + NE_ENABLE,
+				dev_disable_reply,
+				dev_disable_reply == NE_ENABLE_OFF,
+				10 * USEC_PER_MSEC,
+				NE_FREE_TIMEOUT_MSECS * USEC_PER_MSEC);
+	if (rc < 0)
 		dev_err(&pdev->dev, "Error in pci dev disable cmd\n");
 }
 
-- 
2.47.1

[PATCH 19/19] MAINTAINERS: Take over the Nitro Enclaves driver

[Alexander Graf]

Alexandru Ciobotaru has moved on from Nitro Enclaves, so the driver needs
a new maintainer. I take it over.

Signed-off-by: Alexander Graf <graf@amazon.com>
---
 MAINTAINERS | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/MAINTAINERS b/MAINTAINERS
index 93abbb315065d..20e4a56a197ee 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -19834,7 +19834,7 @@ T:	git git://git.kernel.org/pub/scm/linux/kernel/git/dinguyen/linux.git
 F:	arch/nios2/
 
 NITRO ENCLAVES (NE)
-M:	Alexandru Ciobotaru <alcioa@amazon.com>
+M:	Alexander Graf <graf@amazon.com>
 R:	The AWS Nitro Enclaves Team <aws-nitro-enclaves-devel@amazon.com>
 L:	linux-kernel@vger.kernel.org
 S:	Supported
-- 
2.47.1