[PATCH 4/6] iommufd/device: Add mock_device support in iommufd_device_get_info()

[Yi Liu <yi.l.liu@xxxxxxxxx>]

This prepares for adding selftest for IOMMU_DEVICE_GET_INFO. Selftest
uses mock device, while physical device uses iommufd_device, so add
a helper iommufd_obj_dev() to get struct device from the iommufd_object
for a given dev_id.

Signed-off-by: Nicolin Chen <nicolinc@xxxxxxxxxx>
Signed-off-by: Yi Liu <yi.l.liu@xxxxxxxxx>
---
 drivers/iommu/iommufd/device.c          | 22 ++++++++++++++++++++--
 drivers/iommu/iommufd/iommufd_private.h |  1 +
 drivers/iommu/iommufd/selftest.c        |  7 +++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/drivers/iommu/iommufd/device.c b/drivers/iommu/iommufd/device.c
index 3b64aef24807..470838e6902d 100644
--- a/drivers/iommu/iommufd/device.c
+++ b/drivers/iommu/iommufd/device.c
@@ -145,6 +145,20 @@ static int iommufd_zero_fill_user(u64 ptr, int bytes)
 	return 0;
 }
 
+static struct device *
+iommufd_obj_dev(struct iommufd_object *obj)
+{
+	struct device *dev = NULL;
+
+	if (obj->type == IOMMUFD_OBJ_DEVICE)
+		dev = container_of(obj, struct iommufd_device, obj)->dev;
+#ifdef CONFIG_IOMMUFD_TEST
+	else if (obj->type == IOMMUFD_OBJ_SELFTEST)
+		dev = iommufd_selftest_obj_to_dev(obj);
+#endif
+	return dev;
+}
+
 int iommufd_device_get_info(struct iommufd_ucmd *ucmd)
 {
 	struct iommu_device_info *cmd = ucmd->cmd;
@@ -159,11 +173,15 @@ int iommufd_device_get_info(struct iommufd_ucmd *ucmd)
 		return -EOPNOTSUPP;
 
 	dev_obj = iommufd_get_object(ucmd->ictx, cmd->dev_id,
-				     IOMMUFD_OBJ_DEVICE);
+				     IOMMUFD_OBJ_ANY);
 	if (IS_ERR(dev_obj))
 		return PTR_ERR(dev_obj);
 
-	dev = container_of(dev_obj, struct iommufd_device, obj)->dev;
+	dev = iommufd_obj_dev(dev_obj);
+	if (!dev) {
+		rc = -EINVAL;
+		goto out_put;
+	}
 
 	ops = dev_iommu_ops(dev);
 	if (!ops || !ops->hw_info) {
diff --git a/drivers/iommu/iommufd/iommufd_private.h b/drivers/iommu/iommufd/iommufd_private.h
index 4a0a1a7fdae1..7748117e36f9 100644
--- a/drivers/iommu/iommufd/iommufd_private.h
+++ b/drivers/iommu/iommufd/iommufd_private.h
@@ -284,6 +284,7 @@ iommufd_device_selftest_attach(struct iommufd_ctx *ictx,
 			       struct device *mock_dev);
 void iommufd_device_selftest_detach(struct iommufd_ctx *ictx,
 				    struct iommufd_hw_pagetable *hwpt);
+struct device *iommufd_selftest_obj_to_dev(struct iommufd_object *obj);
 int iommufd_test(struct iommufd_ucmd *ucmd);
 void iommufd_selftest_destroy(struct iommufd_object *obj);
 extern size_t iommufd_test_memory_limit;
diff --git a/drivers/iommu/iommufd/selftest.c b/drivers/iommu/iommufd/selftest.c
index b94870f93138..2ecde22a60f4 100644
--- a/drivers/iommu/iommufd/selftest.c
+++ b/drivers/iommu/iommufd/selftest.c
@@ -802,6 +802,13 @@ void iommufd_selftest_destroy(struct iommufd_object *obj)
 	}
 }
 
+struct device *iommufd_selftest_obj_to_dev(struct iommufd_object *obj)
+{
+	struct selftest_obj *sobj = container_of(obj, struct selftest_obj, obj);
+
+	return &sobj->idev.mock_dev;
+}
+
 int iommufd_test(struct iommufd_ucmd *ucmd)
 {
 	struct iommu_test_cmd *cmd = ucmd->cmd;
-- 
2.34.1