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

[Alexander Graf <graf@amazon.com>]

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