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

[Alexander Graf <graf@amazon.com>]

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

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

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

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