[PATCH v2 1/2] iommu: allocate the fwspec before dev->iommu in iommu_fwspec_init()

[Tal Zussman]

iommu_fwspec_init() allocates dev->iommu and then the fwspec. If the
fwspec allocation fails it returns -ENOMEM with dev->iommu left
allocated, which only iommu_probe_device() knows to clean up.
iommu_mock_device_add() calls it outside of a probe and returns straight
away on error, so dev->iommu leaks.

kmemleak reports it when iommufd_fail_nth injects the allocation
failure:

  unreferenced object 0xffff88811c138b00 (size 192):
    comm "iommufd_fail_nt", pid 783
    backtrace:
      __kmalloc_cache_noprof
      dev_iommu_get
      iommu_fwspec_init
      iommu_mock_device_add
      iommufd_test
      iommufd_fops_ioctl

Allocate the fwspec first and free it if dev_iommu_get() fails, so an
error leaves nothing for the caller to unwind.

Fixes: 2a918911ed3d ("iommufd: Register iommufd mock devices with fwspec")
Signed-off-by: Tal Zussman <tz2294@xxxxxxxxxxxx>
---
 drivers/iommu/iommu.c | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index cd1bca7ede9a..8cdecdd119b6 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -3083,14 +3083,16 @@ int iommu_fwspec_init(struct device *dev, struct fwnode_handle *iommu_fwnode)
 	if (fwspec)
 		return iommu->ops == iommu_fwspec_ops(fwspec) ? 0 : -EINVAL;
 
-	if (!dev_iommu_get(dev))
-		return -ENOMEM;
-
 	/* Preallocate for the overwhelmingly common case of 1 ID */
 	fwspec = kzalloc_flex(*fwspec, ids, 1);
 	if (!fwspec)
 		return -ENOMEM;
 
+	if (!dev_iommu_get(dev)) {
+		kfree(fwspec);
+		return -ENOMEM;
+	}
+
 	fwnode_handle_get(iommu_fwnode);
 	fwspec->iommu_fwnode = iommu_fwnode;
 	dev_iommu_fwspec_set(dev, fwspec);

-- 
2.39.5