[PATCH] vdpa/mlx5: Bound the KLM count of the user MR

[PATCH] vdpa/mlx5: Bound the KLM count of the user MR

[Vladimir Petrzhikovskii]

create_user_mr() builds one indirect MKey that covers the whole vhost
IOTLB. It fills each gap between two mappings with null KLMs, one per
GiB, and never checks the total number of KLMs against the device
limit, MLX5_CAP_GEN(log_max_klm_list_size).

With a virtio-IOMMU the guest chooses the IOVAs: the VMM (QEMU)
forwards each map the guest makes to the vhost-vDPA device. If the
guest maps one page at IOVA 0 and one at IOVA 2^58, the driver needs
more than 2^28 KLMs. create_indirect_key() computes the buffer size in
an int, which overflows: the driver allocates 336 bytes, and
fill_indir() then writes 16 bytes per KLM, 4 GiB in total, far past
the end of the buffer.

nnuls is an int, so a gap of 2^61 bytes or more is miscounted on every
architecture, and on 32-bit the unsigned long num_klms wraps. The gap
is rounded up to whole GiB as (size + 2^30 - 1) >> 30. For a gap
larger than 2^64 - 2^30, for example below a page mapped in the top
GiB, that addition wraps and the driver counts no null KLMs for the
gap, while fill_indir() still writes one for every GiB of it.

Count KLMs in u64 and round the gap up without the overflowing
addition. Before creating any direct key, return -E2BIG with a
rate-limited warning if the MR needs more KLMs than the device limit,
or more than fit in a CREATE_MKEY command whose length is an int.
Allocate the command with kvzalloc(), because a list within the device
limit can be too large for kzalloc(). A normal guest is unaffected:
memory at 0-3 GiB and 4-68 GiB needs 68 KLMs.

Found by an LLM code review agent auditing vhost-vDPA backends. We have
no mlx5 hardware, so the fix is untested on it; it builds with W=1 on
x86_64 and i386.

Fixes: 94abbccdf291 ("vdpa/mlx5: Add shared memory registration code")
Cc: stable@vger.kernel.org
Co-developed-by: Nikolai Kalinin <nikolai@kalinidze.me>
Signed-off-by: Nikolai Kalinin <nikolai@kalinidze.me>
Signed-off-by: Vladimir Petrzhikovskii <vladimir@dfb.sh>
---
Dragos, one question on the limit: the patch takes
MLX5_CAP_GEN(log_max_klm_list_size) as the largest KLM list an indirect
CREATE_MKEY accepts, as mlx5_ib does for fast-reg lists. If firmware
applies a different limit to CREATE_MKEY, please say so and I will
change the bound.

The guest kernel change used for the reachability check is available
on request.
---
 drivers/vdpa/mlx5/core/mlx5_vdpa.h |  2 +-
 drivers/vdpa/mlx5/core/mr.c        | 30 +++++++++++++++++++++++++-----
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/drivers/vdpa/mlx5/core/mlx5_vdpa.h b/drivers/vdpa/mlx5/core/mlx5_vdpa.h
index 42f2f44b383c..1d84824ca429 100644
--- a/drivers/vdpa/mlx5/core/mlx5_vdpa.h
+++ b/drivers/vdpa/mlx5/core/mlx5_vdpa.h
@@ -32,7 +32,7 @@ struct mlx5_vdpa_mr {
 	/* list of direct MRs descendants of this indirect mr */
 	struct list_head head;
 	unsigned long num_directs;
-	unsigned long num_klms;
+	u64 num_klms;
 
 	struct vhost_iotlb *iotlb;
 
diff --git a/drivers/vdpa/mlx5/core/mr.c b/drivers/vdpa/mlx5/core/mr.c
index b0c5ff23d022..03db8345c1bb 100644
--- a/drivers/vdpa/mlx5/core/mr.c
+++ b/drivers/vdpa/mlx5/core/mr.c
@@ -206,6 +206,16 @@ static int klm_byte_size(int nklms)
 	return 16 * ALIGN(nklms, 4);
 }
 
+/* Largest KLM count whose CREATE_MKEY input length still fits in an int. */
+#define MLX5_VDPA_MAX_KLMS_INLEN \
+	ALIGN_DOWN((INT_MAX - MLX5_ST_SZ_BYTES(create_mkey_in)) / 16, 4)
+
+static u64 max_klms(struct mlx5_vdpa_dev *mvdev)
+{
+	return min_t(u64, BIT_ULL(MLX5_CAP_GEN(mvdev->mdev, log_max_klm_list_size)),
+		     MLX5_VDPA_MAX_KLMS_INLEN);
+}
+
 #define MLX5_VDPA_MTT_ALIGN 16
 
 static int create_direct_keys(struct mlx5_vdpa_dev *mvdev, struct mlx5_vdpa_mr *mr)
@@ -338,7 +348,7 @@ static int create_indirect_key(struct mlx5_vdpa_dev *mvdev, struct mlx5_vdpa_mr
 		return -EINVAL;
 
 	inlen = MLX5_ST_SZ_BYTES(create_mkey_in) + klm_byte_size(mr->num_klms);
-	in = kzalloc(inlen, GFP_KERNEL);
+	in = kvzalloc(inlen, GFP_KERNEL);
 	if (!in)
 		return -ENOMEM;
 
@@ -355,7 +365,7 @@ static int create_indirect_key(struct mlx5_vdpa_dev *mvdev, struct mlx5_vdpa_mr
 	MLX5_SET(create_mkey_in, in, translations_octword_actual_size, mr->num_klms);
 	fill_indir(mvdev, mr, in);
 	err = mlx5_vdpa_create_mkey(mvdev, &mr->mkey, in, inlen);
-	kfree(in);
+	kvfree(in);
 	return err;
 }
 
@@ -501,6 +511,7 @@ static int create_user_mr(struct mlx5_vdpa_dev *mvdev,
 {
 	struct mlx5_vdpa_direct_mr *dmr;
 	struct mlx5_vdpa_direct_mr *n;
+	u64 klm_limit = max_klms(mvdev);
 	struct vhost_iotlb_map *map;
 	u32 pperm = U16_MAX;
 	u64 last = U64_MAX;
@@ -508,7 +519,7 @@ static int create_user_mr(struct mlx5_vdpa_dev *mvdev,
 	u64 pe = U64_MAX;
 	u64 start = 0;
 	int err = 0;
-	int nnuls;
+	u64 nnuls;
 
 	INIT_LIST_HEAD(&mr->head);
 	for (map = vhost_iotlb_itree_first(iotlb, start, last); map;
@@ -521,9 +532,10 @@ static int create_user_mr(struct mlx5_vdpa_dev *mvdev,
 				if (pe < map->start) {
 					/* We have a hole in the map. Check how
 					 * many null keys are required to fill it.
+					 * The hole can be nearly 2^64 bytes, so
+					 * round up without adding to its size.
 					 */
-					nnuls = MLX5_DIV_ROUND_UP_POW2(map->start - pe,
-								       LOG_MAX_KLM_SIZE);
+					nnuls = ((map->start - pe - 1) >> LOG_MAX_KLM_SIZE) + 1;
 					mr->num_klms += nnuls;
 				}
 				err = add_direct_chain(mvdev, mr, ps, pe - ps, pperm, iotlb);
@@ -539,6 +551,14 @@ static int create_user_mr(struct mlx5_vdpa_dev *mvdev,
 	if (err)
 		goto err_chain;
 
+	if (mr->num_klms > klm_limit) {
+		dev_warn_ratelimited(mvdev->mdev->device,
+				     "map needs %llu KLMs, limit is %llu\n",
+				     mr->num_klms, klm_limit);
+		err = -E2BIG;
+		goto err_chain;
+	}
+
 	err = create_direct_keys(mvdev, mr);
 	if (err)
 		goto err_chain;

---
base-commit: d48723711bae9b8cdc56719c448ba42fef28cbee
change-id: 20261109-b4-mlx5-klm-798a398a159e

Best regards,
--  
Vladimir Petrzhikovskii <vladimir@dfb.sh>

Re: [PATCH] vdpa/mlx5: Bound the KLM count of the user MR

[sashiko-bot]

> create_user_mr() builds one indirect MKey that covers the whole vhost
> IOTLB. It fills each gap between two mappings with null KLMs, one per
> GiB, and never checks the total number of KLMs against the device
> limit, MLX5_CAP_GEN(log_max_klm_list_size).
> 
> With a virtio-IOMMU the guest chooses the IOVAs: the VMM (QEMU)
> forwards each map the guest makes to the vhost-vDPA device. If the
> guest maps one page at IOVA 0 and one at IOVA 2^58, the driver needs
> more than 2^28 KLMs. create_indirect_key() computes the buffer size in
> an int, which overflows: the driver allocates 336 bytes, and
> fill_indir() then writes 16 bytes per KLM, 4 GiB in total, far past
> the end of the buffer.
> 
> nnuls is an int, so a gap of 2^61 bytes or more is miscounted on every
> architecture, and on 32-bit the unsigned long num_klms wraps. The gap
> [ ... ]
> Fixes: 94abbccdf291 ("vdpa/mlx5: Add shared memory registration code")
> Cc: stable@vger.kernel.org
> Co-developed-by: Nikolai Kalinin <nikolai@kalinidze.me>
> Signed-off-by: Nikolai Kalinin <nikolai@kalinidze.me>
> Signed-off-by: Vladimir Petrzhikovskii <vladimir@dfb.sh>

Sashiko has reviewed this patch and found no issues.

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261109-b4-mlx5-klm-v1-1-787b45145719@dfb.sh?part=1