[virtio_magma] Support magma_get_buffer_handle2

Includes proper lifetime management of the buffer handle.
This foundation allows for upcoming magma_query_returns_buffer2.

Bug:83610

Change-Id: I3cd3b96a8e12848e2483214ecb77adb5f9f2536b
diff --git a/drivers/virtio/virtio_magma.c b/drivers/virtio/virtio_magma.c
index 14def0e..98b4344 100644
--- a/drivers/virtio/virtio_magma.c
+++ b/drivers/virtio/virtio_magma.c
@@ -112,10 +112,10 @@
 };
 
 struct virtmagma_buffer_fd_priv {
-	struct virtmagma_instance *instance;
+	struct virtmagma_info *vi;
 	struct mutex mutex_lock;
-	u64 connection_id;
-	u64 buffer_id;
+	u32 buffer_handle;
+	u64 buffer_size;
 	u64 phys_addr;
 };
 
@@ -493,31 +493,26 @@
 	return 0;
 }
 
-static int virtmagma_buffer_fd_release(struct inode *inodep, struct file *filp)
+static int virtmagma_buffer_fd_unmap(struct virtmagma_buffer_fd_priv *priv)
 {
-	struct virtmagma_buffer_fd_priv *priv = filp->private_data;
-	struct virtmagma_instance *instance = priv->instance;
-	struct virtmagma_connection *connection = NULL;
-	struct virtmagma_connection_object *object = NULL;
-	struct virtio_magma_internal_unmap_ctrl *request = NULL;
-	struct virtio_magma_internal_unmap_resp *response = NULL;
+	struct virtio_magma_internal_unmap2_ctrl *request = NULL;
+	struct virtio_magma_internal_unmap2_resp *response = NULL;
 	int ret;
 
-	if (instance) {
-		connection = get_connection(instance, priv->connection_id);
+	if (!priv->phys_addr)
+		return -EINVAL;
+
+	request = kzalloc(sizeof(*request), GFP_KERNEL);
+	if (!request)
+		return -ENOMEM;
+
+	response = kzalloc(sizeof(*response), GFP_KERNEL);
+	if (!response) {
+		kfree(request);
+		return -ENOMEM;
 	}
 
-	if (connection) {
-		object = get_connection_object(connection, priv->buffer_id,
-					       MAGMA_BUFFER);
-	}
-
-	if (connection && object && priv->phys_addr) {
-		request = kmem_cache_alloc(instance->msg_cache, GFP_KERNEL);
-		response = kmem_cache_alloc(instance->msg_cache, GFP_KERNEL);
-	}
-
-	if (request && response) {
+	{
 		struct virtmagma_virtio_command command = {
 			.request_ptr = request,
 			.request_size = sizeof(*request),
@@ -525,24 +520,72 @@
 			.response_size = sizeof(*response)
 		};
 
-		request->hdr.type = VIRTIO_MAGMA_CMD_INTERNAL_UNMAP;
+		request->hdr.type = VIRTIO_MAGMA_CMD_INTERNAL_UNMAP2;
 		request->hdr.flags = 0;
-		request->connection = priv->connection_id;
-		request->buffer = priv->buffer_id;
+		request->buffer = priv->buffer_handle;
 		request->address = priv->phys_addr;
 
-		ret = vq_out_send_sync(instance->vi, &command);
+		ret = vq_out_send_sync(priv->vi, &command);
 		if (ret == 0) {
 			virtmagma_check_expected_response_type(request,
 							       response);
 		}
 	}
 
-	if (request)
-		kmem_cache_free(instance->msg_cache, request);
+	kfree(request);
+	kfree(response);
 
-	if (response)
-		kmem_cache_free(instance->msg_cache, response);
+	return 0;
+}
+
+static int virtmagma_release_handle(struct virtmagma_info *vi,
+	uint32_t handle)
+{
+	struct virtio_magma_internal_release_handle_ctrl *request;
+	struct virtio_magma_internal_release_handle_resp *response;
+	int ret;
+
+	request = kzalloc(sizeof(*request), GFP_KERNEL);
+	if (!request)
+		return -ENOMEM;
+
+	response = kzalloc(sizeof(*response), GFP_KERNEL);
+	if (!response) {
+		kfree(request);
+		return -ENOMEM;
+	}
+
+	request->hdr.type = VIRTIO_MAGMA_CMD_INTERNAL_RELEASE_HANDLE;
+	request->hdr.flags = 0;
+	request->handle = handle;
+
+	{
+		struct virtmagma_virtio_command command = {
+			.request_ptr = request,
+			.request_size = sizeof(*request),
+			.response_ptr = response,
+			.response_size = sizeof(*response)
+		};
+
+		ret = vq_out_send_sync(vi, &command);
+	}
+
+	if (ret == 0) {
+		ret = virtmagma_check_expected_response_type(request, response);
+	}
+
+	kfree(request);
+	kfree(response);
+
+	return ret;
+}
+
+static int virtmagma_buffer_fd_release(struct inode *inodep, struct file *filp)
+{
+	struct virtmagma_buffer_fd_priv *priv = filp->private_data;
+
+	virtmagma_buffer_fd_unmap(priv);
+	virtmagma_release_handle(priv->vi, priv->buffer_handle);
 
 	mutex_destroy(&priv->mutex_lock);
 	kfree(priv);
@@ -557,42 +600,26 @@
 				    struct vm_area_struct *vma)
 {
 	struct virtmagma_buffer_fd_priv *priv = filp->private_data;
-	struct virtmagma_instance *instance = priv->instance;
-	struct virtmagma_connection *connection;
-	struct virtmagma_connection_object *object;
-	struct virtio_magma_internal_map_ctrl *request;
-	struct virtio_magma_internal_map_resp *response;
+	struct virtio_magma_internal_map2_ctrl *request;
+	struct virtio_magma_internal_map2_resp *response;
 	unsigned long vm_size = vma->vm_end - vma->vm_start;
 	size_t max_map_size;
 	int ret = 0;
 
-	if (!instance)
-		return -ENODEV;
-
-	connection = get_connection(instance, priv->connection_id);
-	if (!connection)
-		return -EINVAL;
-
-	object = get_connection_object(connection, priv->buffer_id,
-				       MAGMA_BUFFER);
-	if (!object)
-		return -EINVAL;
-
-	request = kmem_cache_alloc(instance->msg_cache, GFP_KERNEL);
+	request = kzalloc(sizeof(*request), GFP_KERNEL);
 	if (!request)
 		return -ENOMEM;
 
-	response = kmem_cache_alloc(instance->msg_cache, GFP_KERNEL);
+	response = kzalloc(sizeof(*response), GFP_KERNEL);
 	if (!response) {
-		kmem_cache_free(instance->msg_cache, request);
+		kfree(request);
 		return -ENOMEM;
 	}
 
-	request->hdr.type = VIRTIO_MAGMA_CMD_INTERNAL_MAP;
+	request->hdr.type = VIRTIO_MAGMA_CMD_INTERNAL_MAP2;
 	request->hdr.flags = 0;
-	request->connection = priv->connection_id;
-	request->buffer = priv->buffer_id;
-	request->length = object->buffer.size_allocated;
+	request->buffer = priv->buffer_handle;
+	request->length = priv->buffer_size;
 
 	mutex_lock(&priv->mutex_lock);
 	if (!priv->phys_addr) {
@@ -603,7 +630,7 @@
 			.response_size = sizeof(*response)
 		};
 
-		ret = vq_out_send_sync(instance->vi, &command);
+		ret = vq_out_send_sync(priv->vi, &command);
 		if (ret == 0) {
 			ret = virtmagma_check_expected_response_type(request,
 								     response);
@@ -615,13 +642,13 @@
 	}
 	mutex_unlock(&priv->mutex_lock);
 
-	kmem_cache_free(instance->msg_cache, request);
-	kmem_cache_free(instance->msg_cache, response);
+	kfree(request);
+	kfree(response);
 
 	if (ret)
 		return ret;
 
-	max_map_size = PAGE_ALIGN(object->buffer.size_allocated);
+	max_map_size = PAGE_ALIGN(priv->buffer_size);
 
 	if (vma->vm_pgoff * PAGE_SIZE + vm_size > max_map_size) {
 		pr_warn("virtmagma: user tried to mmap with offset (%ld) and size (%ld) exceededing the buffer's size (%ld)",
@@ -1578,54 +1605,133 @@
 	return 0;
 }
 
-static int virtmagma_command_magma_get_buffer_handle(
+// Assumes that command->response_ptr is large enough for the response struct +
+// uint64_t for the buffer size.
+static int virtmagma_command_magma_get_buffer_handle2_core(
 	struct virtmagma_instance *instance,
 	struct virtmagma_virtio_command *command)
 {
-	struct virtmagma_connection *connection;
-	struct virtmagma_connection_object *object;
 	struct virtmagma_buffer_fd_priv *priv;
-	struct virtio_magma_get_buffer_handle_ctrl *request =
+	struct virtio_magma_get_buffer_handle2_ctrl *request =
 		command->request_ptr;
-	struct virtio_magma_get_buffer_handle_resp *response =
+	struct virtio_magma_get_buffer_handle2_resp *response =
 		command->response_ptr;
+	uint32_t handle;
 	int ret;
 
-	connection = get_connection(instance, request->connection);
-	if (!connection)
+	if (!COMMAND_OK(command, request, response))
 		return -EINVAL;
 
-	object = get_connection_object(connection, request->buffer,
-				       MAGMA_BUFFER);
-	if (!object)
-		return -EINVAL;
+	ret = vq_out_send_sync(instance->vi, command);
+	if (ret)
+		return ret;
+
+	ret = virtmagma_check_expected_response_type(request, response);
+	if (ret)
+		return ret;
+
+	if (response->result_return) {
+		pr_warn("virtmagma: magma_get_buffer_handle2 returned %d",
+			(int32_t)response->result_return);
+		return 0; /* the ioctl is still successful */
+	}
+
+	handle = response->handle_out;
 
 	priv = kzalloc(sizeof(struct virtmagma_buffer_fd_priv), GFP_KERNEL);
-	if (!priv)
+	if (!priv) {
+		virtmagma_release_handle(instance->vi, handle);
 		return -ENOMEM;
+	}
 
 	mutex_init(&priv->mutex_lock);
 
-	// Simple reference; the client must ensure the fd is used only while
-	// the buffer remains valid.
-	priv->instance = instance;
-	priv->connection_id = request->connection;
-	priv->buffer_id = request->buffer;
+	priv->vi = instance->vi;
+	priv->buffer_handle = handle;
+	priv->buffer_size = *(uint64_t*)(response + 1);
 
 	ret = anon_inode_getfd("[virtmagma]", &virtmagma_buffer_fd_fops, priv,
 			       O_RDWR);
 	if (ret < 0) {
 		pr_err("virtmagma: failed to create fd: %d", ret);
+		virtmagma_release_handle(instance->vi, handle);
+		kfree(priv);
 		return ret;
 	}
 
-	response->hdr.type = VIRTIO_MAGMA_RESP_GET_BUFFER_HANDLE;
 	response->handle_out = ret;
-	response->result_return = 0;
 
 	return 0;
 }
 
+static int virtmagma_command_magma_get_buffer_handle2(
+	struct virtmagma_instance *instance,
+	struct virtmagma_virtio_command *command)
+{
+	struct virtio_magma_get_buffer_handle2_resp *response;
+
+	/* Reallocate response buffer with additional space for buffer size;
+	   note that response_size is not modified, as we only want the response struct
+	   itself to be copied back to the user by our caller.
+	   Memory will be freed by caller.
+	*/
+	response = kzalloc(sizeof(*response) + sizeof(uint64_t), GFP_KERNEL);
+	if (!response)
+		return -ENOMEM;
+
+	command->response_ptr = response;
+
+	return virtmagma_command_magma_get_buffer_handle2_core(instance, command);
+}
+
+static int virtmagma_command_magma_get_buffer_handle(
+	struct virtmagma_instance *instance,
+	struct virtmagma_virtio_command *command)
+{
+	struct virtio_magma_get_buffer_handle_ctrl *request = command->request_ptr;
+	struct virtio_magma_get_buffer_handle_resp *response = command->response_ptr;
+	struct virtio_magma_get_buffer_handle2_ctrl *request2;
+	struct virtio_magma_get_buffer_handle2_resp *response2;
+	int ret;
+
+	BUILD_BUG_ON(sizeof(*request2) > MESSAGE_CACHE_OBJECT_SIZE);
+	BUILD_BUG_ON(sizeof(*response2) + sizeof(uint64_t) > MESSAGE_CACHE_OBJECT_SIZE);
+
+	request2 = kmem_cache_alloc(instance->msg_cache, GFP_KERNEL);
+	if (!request2)
+		return -ENOMEM;
+
+	response2 = kmem_cache_alloc(instance->msg_cache, GFP_KERNEL);
+	if (!response2) {
+		kmem_cache_free(instance->msg_cache, request2);
+		return -ENOMEM;
+	}
+
+	request2->hdr.type = VIRTIO_MAGMA_CMD_GET_BUFFER_HANDLE2;
+	request2->hdr.flags = 0;
+	request2->buffer = request->buffer;
+
+	{
+		struct virtmagma_virtio_command command2 = {
+			.request_ptr = request2,
+			.request_size = sizeof(*request2),
+			.response_ptr = response2,
+			.response_size = sizeof(*response2)
+		};
+
+		ret = virtmagma_command_magma_get_buffer_handle2_core(instance, &command2);
+	}
+
+	response->hdr.type = VIRTIO_MAGMA_RESP_GET_BUFFER_HANDLE;
+	response->handle_out = response2->handle_out;
+	response->result_return = response2->result_return;
+
+	kmem_cache_free(instance->msg_cache, request2);
+	kmem_cache_free(instance->msg_cache, response2);
+
+	return ret;
+}
+
 static int virtmagma_ioctl_handshake(struct file *filp, void __user *ptr)
 {
 	struct virtmagma_ioctl_args_handshake ioctl_args;
@@ -1760,6 +1866,10 @@
 		ret = virtmagma_command_magma_get_buffer_handle(instance,
 								&command);
 		break;
+	case VIRTIO_MAGMA_CMD_GET_BUFFER_HANDLE2:
+		ret = virtmagma_command_magma_get_buffer_handle2(instance,
+								&command);
+		break;
 	case VIRTIO_MAGMA_CMD_VIRT_CREATE_IMAGE:
 		ret = virtmagma_command_magma_virt_create_image(instance,
 								&command);
diff --git a/include/uapi/linux/virtio_magma.h b/include/uapi/linux/virtio_magma.h
index fea2e0e..55be49e 100644
--- a/include/uapi/linux/virtio_magma.h
+++ b/include/uapi/linux/virtio_magma.h
@@ -71,8 +71,12 @@
 	VIRTIO_MAGMA_CMD_VIRT_GET_IMAGE_INFO = 0x1048,
 	VIRTIO_MAGMA_CMD_SYNC = 0x1049,
 	VIRTIO_MAGMA_CMD_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2 = 0x104A,
+	VIRTIO_MAGMA_CMD_GET_BUFFER_HANDLE2 = 0x104B,
 	VIRTIO_MAGMA_CMD_INTERNAL_MAP = 0x1044,
 	VIRTIO_MAGMA_CMD_INTERNAL_UNMAP = 0x1045,
+	VIRTIO_MAGMA_CMD_INTERNAL_RELEASE_HANDLE = 0x104C,
+	VIRTIO_MAGMA_CMD_INTERNAL_MAP2 = 0x104D,
+	VIRTIO_MAGMA_CMD_INTERNAL_UNMAP2 = 0x104E,
 	/* magma success responses
  */
 	VIRTIO_MAGMA_RESP_RELEASE_CONNECTION = 0x2004,
@@ -125,8 +129,12 @@
 	VIRTIO_MAGMA_RESP_VIRT_GET_IMAGE_INFO = 0x2048,
 	VIRTIO_MAGMA_RESP_SYNC = 0x2049,
 	VIRTIO_MAGMA_RESP_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2 = 0x204A,
+	VIRTIO_MAGMA_RESP_GET_BUFFER_HANDLE2 = 0x204B,
 	VIRTIO_MAGMA_RESP_INTERNAL_MAP = 0x2044,
 	VIRTIO_MAGMA_RESP_INTERNAL_UNMAP = 0x2045,
+	VIRTIO_MAGMA_RESP_INTERNAL_RELEASE_HANDLE = 0x204C,
+	VIRTIO_MAGMA_RESP_INTERNAL_MAP2 = 0x204D,
+	VIRTIO_MAGMA_RESP_INTERNAL_UNMAP2 = 0x204E,
 	/* magma error responses
  */
 	VIRTIO_MAGMA_RESP_ERR_UNIMPLEMENTED = 0x3001,
@@ -239,10 +247,18 @@
 		case VIRTIO_MAGMA_RESP_SYNC: return "VIRTIO_MAGMA_RESP_SYNC";
 		case VIRTIO_MAGMA_CMD_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2: return "VIRTIO_MAGMA_CMD_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2";
 		case VIRTIO_MAGMA_RESP_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2: return "VIRTIO_MAGMA_RESP_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2";
+		case VIRTIO_MAGMA_CMD_GET_BUFFER_HANDLE2: return "VIRTIO_MAGMA_CMD_GET_BUFFER_HANDLE2";
+		case VIRTIO_MAGMA_RESP_GET_BUFFER_HANDLE2: return "VIRTIO_MAGMA_RESP_GET_BUFFER_HANDLE2";
 		case VIRTIO_MAGMA_CMD_INTERNAL_MAP: return "VIRTIO_MAGMA_CMD_INTERNAL_MAP";
 		case VIRTIO_MAGMA_RESP_INTERNAL_MAP: return "VIRTIO_MAGMA_RESP_INTERNAL_MAP";
 		case VIRTIO_MAGMA_CMD_INTERNAL_UNMAP: return "VIRTIO_MAGMA_CMD_INTERNAL_UNMAP";
 		case VIRTIO_MAGMA_RESP_INTERNAL_UNMAP: return "VIRTIO_MAGMA_RESP_INTERNAL_UNMAP";
+		case VIRTIO_MAGMA_CMD_INTERNAL_RELEASE_HANDLE: return "VIRTIO_MAGMA_CMD_INTERNAL_RELEASE_HANDLE";
+		case VIRTIO_MAGMA_RESP_INTERNAL_RELEASE_HANDLE: return "VIRTIO_MAGMA_RESP_INTERNAL_RELEASE_HANDLE";
+		case VIRTIO_MAGMA_CMD_INTERNAL_MAP2: return "VIRTIO_MAGMA_CMD_INTERNAL_MAP2";
+		case VIRTIO_MAGMA_RESP_INTERNAL_MAP2: return "VIRTIO_MAGMA_RESP_INTERNAL_MAP2";
+		case VIRTIO_MAGMA_CMD_INTERNAL_UNMAP2: return "VIRTIO_MAGMA_CMD_INTERNAL_UNMAP2";
+		case VIRTIO_MAGMA_RESP_INTERNAL_UNMAP2: return "VIRTIO_MAGMA_RESP_INTERNAL_UNMAP2";
 		case VIRTIO_MAGMA_RESP_ERR_UNIMPLEMENTED: return "VIRTIO_MAGMA_RESP_ERR_UNIMPLEMENTED";
 		case VIRTIO_MAGMA_RESP_ERR_INTERNAL: return "VIRTIO_MAGMA_RESP_ERR_INTERNAL";
 		case VIRTIO_MAGMA_RESP_ERR_HOST_DISCONNECTED: return "VIRTIO_MAGMA_RESP_ERR_HOST_DISCONNECTED";
@@ -305,8 +321,12 @@
 		case VIRTIO_MAGMA_CMD_VIRT_GET_IMAGE_INFO: return VIRTIO_MAGMA_RESP_VIRT_GET_IMAGE_INFO;
 		case VIRTIO_MAGMA_CMD_SYNC: return VIRTIO_MAGMA_RESP_SYNC;
 		case VIRTIO_MAGMA_CMD_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2: return VIRTIO_MAGMA_RESP_EXECUTE_COMMAND_BUFFER_WITH_RESOURCES2;
+		case VIRTIO_MAGMA_CMD_GET_BUFFER_HANDLE2: return VIRTIO_MAGMA_RESP_GET_BUFFER_HANDLE2;
 		case VIRTIO_MAGMA_CMD_INTERNAL_MAP: return VIRTIO_MAGMA_RESP_INTERNAL_MAP;
 		case VIRTIO_MAGMA_CMD_INTERNAL_UNMAP: return VIRTIO_MAGMA_RESP_INTERNAL_UNMAP;
+		case VIRTIO_MAGMA_CMD_INTERNAL_RELEASE_HANDLE: return VIRTIO_MAGMA_RESP_INTERNAL_RELEASE_HANDLE;
+		case VIRTIO_MAGMA_CMD_INTERNAL_MAP2: return VIRTIO_MAGMA_RESP_INTERNAL_MAP2;
+		case VIRTIO_MAGMA_CMD_INTERNAL_UNMAP2: return VIRTIO_MAGMA_RESP_INTERNAL_UNMAP2;
 		default: return VIRTIO_MAGMA_RESP_ERR_INVALID_COMMAND;
 	}
 }
@@ -890,6 +910,17 @@
 	__le64 result_return;
 } __attribute((packed));
 
+struct virtio_magma_get_buffer_handle2_ctrl {
+	struct virtio_magma_ctrl_hdr hdr;
+	__le64 buffer;
+} __attribute((packed));
+
+struct virtio_magma_get_buffer_handle2_resp {
+	struct virtio_magma_ctrl_hdr hdr;
+	uintptr_t handle_out;
+	__le64 result_return;
+} __attribute((packed));
+
 struct virtio_magma_internal_map_ctrl {
 	struct virtio_magma_ctrl_hdr hdr;
 	__le64 connection;
@@ -915,5 +946,38 @@
 	__le64 result_return;
 } __attribute((packed));
 
+struct virtio_magma_internal_release_handle_ctrl {
+	struct virtio_magma_ctrl_hdr hdr;
+	__le32 handle;
+} __attribute((packed));
+
+struct virtio_magma_internal_release_handle_resp {
+	struct virtio_magma_ctrl_hdr hdr;
+	__le64 result_return;
+} __attribute((packed));
+
+struct virtio_magma_internal_map2_ctrl {
+	struct virtio_magma_ctrl_hdr hdr;
+	__le32 buffer;
+	__le64 length;
+} __attribute((packed));
+
+struct virtio_magma_internal_map2_resp {
+	struct virtio_magma_ctrl_hdr hdr;
+	uintptr_t address_out;
+	__le64 result_return;
+} __attribute((packed));
+
+struct virtio_magma_internal_unmap2_ctrl {
+	struct virtio_magma_ctrl_hdr hdr;
+	__le32 buffer;
+	uintptr_t address;
+} __attribute((packed));
+
+struct virtio_magma_internal_unmap2_resp {
+	struct virtio_magma_ctrl_hdr hdr;
+	__le64 result_return;
+} __attribute((packed));
+
 #endif /* _LINUX_VIRTIO_MAGMA_H
  */