[PATCH v1 1/2] iommu/virtio: Avoid use-after-put in viommu_get_by_fwnode

From: Xiong Weimin <[email protected]>

bus_find_device() returns a device reference that must be released with
put_device(). viommu_get_by_fwnode() currently drops that reference
before dereferencing the device to fetch the virtio-IOMMU private data.

Fetch the private data while the reference is still held, then release
the device reference before returning.

Signed-off-by: Xiong Weimin <[email protected]>
---
 drivers/iommu/virtio-iommu.c | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/drivers/iommu/virtio-iommu.c b/drivers/iommu/virtio-iommu.c
index 587fc1319..342785c76 100644
--- a/drivers/iommu/virtio-iommu.c
+++ b/drivers/iommu/virtio-iommu.c
@@ -1009,12 +1009,16 @@ static int viommu_match_node(struct device *dev, const void *data)
 
 static struct viommu_dev *viommu_get_by_fwnode(struct fwnode_handle *fwnode)
 {
+	struct viommu_dev *viommu = NULL;
 	struct device *dev = bus_find_device(virtio_bus_type, NULL, fwnode,
 					     viommu_match_node);
 
-	put_device(dev);
+	if (dev) {
+		viommu = dev_to_virtio(dev)->priv;
+		put_device(dev);
+	}
 
-	return dev ? dev_to_virtio(dev)->priv : NULL;
+	return viommu;
 }
 
 static struct iommu_device *viommu_probe_device(struct device *dev)
-- 
2.43.0