[PATCH] iommu/iommufd: avoid selftest dirty bitmap size wrap

[Samuel Moelius]

IOMMU_TEST_OP_DIRTY sizes its temporary dirty bitmap from
length / page_size.  Very large selftest ranges can make the
DIV_ROUND_UP() additions wrap before allocation, producing a
zero-length allocation while the later test_bit() loop still walks the
original number of bits.

Reject bitmap sizes whose byte or unsigned-long counts cannot be
computed without overflow, and use array_size() for the allocation.

Assisted-by: Codex:gpt-5.5-cyber-preview
Signed-off-by: Samuel Moelius <sam.moelius@xxxxxxxxxxxxxxx>
---
 drivers/iommu/iommufd/selftest.c | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/drivers/iommu/iommufd/selftest.c b/drivers/iommu/iommufd/selftest.c
index af07c642a526..08c36d9ba3e4 100644
--- a/drivers/iommu/iommufd/selftest.c
+++ b/drivers/iommu/iommufd/selftest.c
@@ -1685,8 +1685,9 @@ static int iommufd_test_dirty(struct iommufd_ucmd *ucmd, unsigned int mockpt_id,
 			      unsigned long page_size, void __user *uptr,
 			      u32 flags)
 {
-	unsigned long i, max;
 	struct iommu_test_cmd *cmd = ucmd->cmd;
+	unsigned long i, max, nr_bytes;
+	unsigned long nr_longs;
 	struct iommufd_hw_pagetable *hwpt;
 	struct mock_iommu_domain *mock;
 	int rc, count = 0;
@@ -1706,14 +1707,22 @@ static int iommufd_test_dirty(struct iommufd_ucmd *ucmd, unsigned int mockpt_id,
 	}
 
 	max = length / page_size;
-	tmp = kvzalloc(DIV_ROUND_UP(max, BITS_PER_LONG) * sizeof(unsigned long),
+	if (check_add_overflow(max, BITS_PER_LONG - 1, &nr_longs) ||
+	    check_add_overflow(max, BITS_PER_BYTE - 1, &nr_bytes)) {
+		rc = -EOVERFLOW;
+		goto out_put;
+	}
+	nr_longs /= BITS_PER_LONG;
+	nr_bytes /= BITS_PER_BYTE;
+
+	tmp = kvzalloc(array_size(nr_longs, sizeof(unsigned long)),
 		       GFP_KERNEL_ACCOUNT);
 	if (!tmp) {
 		rc = -ENOMEM;
 		goto out_put;
 	}
 
-	if (copy_from_user(tmp, uptr, DIV_ROUND_UP(max, BITS_PER_BYTE))) {
+	if (copy_from_user(tmp, uptr, nr_bytes)) {
 		rc = -EFAULT;
 		goto out_free;
 	}
-- 
2.43.0