[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