diff --git a/drivers/iommu/iommufd/iommufd_test.h b/drivers/iommu/iommufd/iommufd_test.h
index 1d9ba7454605..3f3644375bf1 100644
--- a/drivers/iommu/iommufd/iommufd_test.h
+++ b/drivers/iommu/iommufd/iommufd_test.h
@@ -109,18 +109,4 @@ struct iommu_test_hw_info {
 	__u32 test_reg;
 };
 
-/* Should not be equal to any defined value in enum iommu_hwpt_type */
-#define IOMMU_HWPT_TYPE_SELFTEST		0xdead
-
-/**
- * struct iommu_hwpt_selftest
- *
- * @flags: page table entry attributes, must be 0
- * @test_type: can be either IOMMU_HW_INFO_TYPE_NONE or IOMMU_HWPT_TYPE_SELFTEST
- */
-struct iommu_hwpt_selftest {
-	__u64 flags;
-	__u32 test_type;
-};
-
 #endif
diff --git a/drivers/iommu/iommufd/selftest.c b/drivers/iommu/iommufd/selftest.c
index 3cd1b17638bc..0bd7534be257 100644
--- a/drivers/iommu/iommufd/selftest.c
+++ b/drivers/iommu/iommufd/selftest.c
@@ -88,8 +88,6 @@ void iommufd_test_syz_conv_iova_id(struct iommufd_ucmd *ucmd,
 struct mock_iommu_domain {
 	struct iommu_domain domain;
 	struct xarray pfns;
-	/* mock domain test data */
-	enum iommu_hwpt_type hwpt_type;
 };
 
 enum selftest_obj_type {
@@ -150,10 +148,17 @@ static void *mock_domain_hw_info(struct device *dev, u32 *length, u32 *type)
 
 static const struct iommu_ops mock_ops;
 
+union mock_domain_alloc_data {
+	struct iommu_hwpt_default default_data;
+};
+
 static struct iommu_domain *
-__mock_domain_alloc(unsigned int iommu_domain_type,
-		    const struct iommu_hwpt_selftest *user_cfg)
+__mock_domain_alloc_default(unsigned int iommu_domain_type,
+			    const union mock_domain_alloc_data *data)
 {
+	const struct iommu_hwpt_default *user_cfg =
+		(const struct iommu_hwpt_default *)data;
+	dma_addr_t aperture_end = MOCK_APERTURE_LAST;
 	struct mock_iommu_domain *mock;
 
 	if (iommu_domain_type == IOMMU_DOMAIN_BLOCKED)
@@ -162,16 +167,21 @@ __mock_domain_alloc(unsigned int iommu_domain_type,
 	if (iommu_domain_type != IOMMU_DOMAIN_UNMANAGED)
 		return NULL;
 
+	if (user_cfg) {
+		if (user_cfg->max_addr > MOCK_APERTURE_LAST ||
+		    user_cfg->max_addr <= MOCK_APERTURE_START)
+			return ERR_PTR(-EINVAL);
+		aperture_end = user_cfg->max_addr;
+	}
+
 	mock = kzalloc(sizeof(*mock), GFP_KERNEL);
 	if (!mock)
 		return ERR_PTR(-ENOMEM);
 	mock->domain.type = iommu_domain_type;
 	mock->domain.geometry.aperture_start = MOCK_APERTURE_START;
-	mock->domain.geometry.aperture_end = MOCK_APERTURE_LAST;
+	mock->domain.geometry.aperture_end = aperture_end;
 	mock->domain.ops = mock_ops.default_domain_ops;
 	mock->domain.pgsize_bitmap = MOCK_IO_PAGE_SIZE;
-	if (user_cfg)
-		mock->hwpt_type = user_cfg->test_type;
 	xa_init(&mock->pfns);
 	return &mock->domain;
 }
@@ -183,7 +193,7 @@ static struct iommu_domain *mock_domain_alloc(unsigned int iommu_domain_type)
 	if (iommu_domain_type != IOMMU_DOMAIN_BLOCKED &&
 	    iommu_domain_type != IOMMU_DOMAIN_UNMANAGED)
 		return NULL;
-	domain = __mock_domain_alloc(iommu_domain_type, NULL);
+	domain = __mock_domain_alloc_default(iommu_domain_type, NULL);
 	if (IS_ERR(domain))
 		domain = NULL;
 	return domain;
@@ -193,27 +203,31 @@ static struct iommu_domain *mock_domain_alloc_user(struct device *dev,
 						   enum iommu_hwpt_type hwpt_type,
 						   const struct iommu_user_data *user_data)
 {
-	const size_t min_len =
-		offsetofend(struct iommu_hwpt_selftest, test_type);
+	struct iommu_domain *(*alloc_fn)(unsigned int iommu_domain_type,
+					 const union mock_domain_alloc_data *data);
 	unsigned int iommu_domain_type = IOMMU_DOMAIN_UNMANAGED;
-	struct iommu_hwpt_selftest data, *user_cfg = NULL;
-
-	if (hwpt_type != IOMMU_HWPT_TYPE_DEFAULT &&
-	    hwpt_type != IOMMU_HWPT_TYPE_SELFTEST)
+	union mock_domain_alloc_data data, *user_cfg = NULL;
+	size_t data_len, min_len;
+
+	switch (hwpt_type) {
+	case IOMMU_HWPT_TYPE_DEFAULT:
+		alloc_fn = __mock_domain_alloc_default;
+		data_len = sizeof(struct iommu_hwpt_default);
+		min_len = offsetofend(struct iommu_hwpt_default, max_addr);
+		break;
+	default:
 		return ERR_PTR(-EINVAL);
+	}
 
 	if (user_data) {
 		int rc = iommu_copy_user_data(&data, user_data,
-					      sizeof(data), min_len);
+					      data_len, min_len);
 		if (rc)
 			return ERR_PTR(rc);
-		/* Expecting test program to pass hwpt_type in test data too */
-		if (data.test_type != hwpt_type)
-			return ERR_PTR(rc);
 		user_cfg = &data;
 	}
 
-	return __mock_domain_alloc(iommu_domain_type, user_cfg);
+	return alloc_fn(iommu_domain_type, user_cfg);
 }
 
 static void mock_domain_free(struct iommu_domain *domain)
diff --git a/tools/testing/selftests/iommu/iommufd.c b/tools/testing/selftests/iommu/iommufd.c
index 7eb7b88459bd..2cbe47c80c09 100644
--- a/tools/testing/selftests/iommu/iommufd.c
+++ b/tools/testing/selftests/iommu/iommufd.c
@@ -274,20 +274,31 @@ TEST_F(iommufd_ioas, ioas_destroy)
 TEST_F(iommufd_ioas, hwpt_alloc)
 {
 	const uint32_t min_data_len =
-		offsetofend(struct iommu_hwpt_selftest, test_type);
-	struct iommu_hwpt_selftest data = {
+		offsetofend(struct iommu_hwpt_default, max_addr);
+	struct iommu_hwpt_default data = {
 		.flags = 0,
-		.test_type = IOMMU_HWPT_TYPE_SELFTEST,
 	};
 	uint32_t new_hwpt_id[2] = {};
 
 	if (self->stdev_id && self->device_id) {
 		test_cmd_hwpt_alloc(self->device_id,
 				    self->ioas_id, &new_hwpt_id[0]);
+		/* Try alloc_user w/o data */
+		test_err_cmd_hwpt_alloc_user(EINVAL,
+					     self->device_id, self->ioas_id,
+					     &new_hwpt_id[1],
+					     IOMMU_HWPT_TYPE_DEFAULT,
+					     0, &data);
+		test_err_cmd_hwpt_alloc_user(EINVAL,
+					     self->device_id, self->ioas_id,
+					     &new_hwpt_id[1],
+					     IOMMU_HWPT_TYPE_DEFAULT,
+					     min_data_len, NULL);
 		test_cmd_hwpt_alloc_user(self->device_id, self->ioas_id,
 					 &new_hwpt_id[1],
-					 IOMMU_HWPT_TYPE_SELFTEST,
-					 min_data_len, &data);
+					 IOMMU_HWPT_TYPE_DEFAULT, 0, NULL);
+
+		/* Replace the auto domain with new_hwpt_id[0] */
 		test_cmd_mock_domain_replace(self->stdev_id, new_hwpt_id[0]);
 		/* hw_pagetable cannot be freed if a device is attached to it */
 		EXPECT_ERRNO(EBUSY,
@@ -302,6 +313,25 @@ TEST_F(iommufd_ioas, hwpt_alloc)
 		/* Detach from the new hw_pagetable[1] and try again */
 		test_cmd_mock_domain_replace(self->stdev_id, self->ioas_id);
 		test_ioctl_destroy(new_hwpt_id[1]);
+
+		/* Try alloc_user with data */
+		data.max_addr = MOCK_APERTURE_START - 1;
+		test_err_cmd_hwpt_alloc_user(EINVAL,
+					     self->device_id, self->ioas_id,
+					     &new_hwpt_id[1],
+					     IOMMU_HWPT_TYPE_DEFAULT,
+					     min_data_len, &data);
+		data.max_addr = (uint64_t)MOCK_APERTURE_LAST + 1;
+		test_err_cmd_hwpt_alloc_user(EINVAL,
+					     self->device_id, self->ioas_id,
+					     &new_hwpt_id[1],
+					     IOMMU_HWPT_TYPE_DEFAULT,
+					     min_data_len, &data);
+		data.max_addr = MOCK_APERTURE_LAST;
+		test_cmd_hwpt_alloc_user(self->device_id, self->ioas_id,
+					 &new_hwpt_id[1],
+					 IOMMU_HWPT_TYPE_DEFAULT,
+					 min_data_len, &data);
 	} else {
 		test_err_cmd_hwpt_alloc(ENOENT,
 					self->device_id, self->ioas_id,
@@ -309,7 +339,7 @@ TEST_F(iommufd_ioas, hwpt_alloc)
 		test_err_cmd_hwpt_alloc_user(ENOENT,
 					     self->device_id, self->ioas_id,
 					     &new_hwpt_id[1],
-					     IOMMU_HWPT_TYPE_SELFTEST,
+					     IOMMU_HWPT_TYPE_DEFAULT,
 					     min_data_len, &data);
 		test_err_mock_domain_replace(ENOENT,
 					     self->stdev_id, new_hwpt_id[0]);