vhost_iotlb: bound map allocation in add_range

vhost_iotlb_add_range_ctx() only retires an old entry when the table
has a non-zero limit, has exactly reached that limit and has
VHOST_IOTLB_FLAG_RETIRE set. Non-retiring tables can keep allocating
entries after reaching their configured limit.

Existing vhost devices allocate their IOTLB with max_iotlb_entries from
vhost.c, which defaults to 2048 and is tunable by module parameter. Use
the caller-provided limit at the allocation point instead of adding a
separate default in the common IOTLB helper, and reject non-positive
values in vhost paths that can report an error.

Other vhost IOTLB users should not create zero-limit tables when entries
can be populated from userspace or guest-controlled requests. Add
caller-side max_iotlb_entries parameters for mlx5 vDPA, VDUSE and
vhost-vDPA. Reject non-positive VDUSE and vhost-vDPA values, and require
at least two entries for vdpa_sim and mlx5 vDPA paths that install
full-range mappings, since those mappings are split into two IOTLB
entries.

Handle full-range mappings in the common helper by checking that the
IOTLB can hold both split entries before inserting the first half. This
avoids returning an error after leaving a half mapping behind.

When the table is full, keep the existing retire behavior for retiring
tables and return -ENOSPC for non-retiring tables. Reuse the retired map
node instead of freeing it and allocating a replacement, so a stream of
IOTLB updates cannot keep forcing GFP_ATOMIC allocations after the table
has reached its limit. If a zero-limit IOTLB still reaches the common
helper, treat it as a configuration error and return -EINVAL.

I found this bug myself, though the patch was written with AI assistance.

Fixes: 0bbe30668d89 ("vhost: factor out IOTLB")
Assisted-by: OpenAI-Codex:GPT-5
Signed-off-by: Linfeng Sun <linfeng.sun.dev@gamil.com>
Message-ID: <AMYAtgAiKmgYcSQT5ukl-4qq.3.1781960405943.Hmail.241270009@hdu.edu.cn>
Signed-off-by: Michael S. Tsirkin <mst@redhat.com>

3 files changed, 49 insertions, 15 deletions

diff --git a/drivers/vhost/iotlb.c b/drivers/vhost/iotlb.c
index e1414c774c34..a1d4376a5b87 100644
--- a/drivers/vhost/iotlb.c
+++ b/drivers/vhost/iotlb.c
@@ -20,6 +20,14 @@ INTERVAL_TREE_DEFINE(struct vhost_iotlb_map,
 		     rb, __u64, __subtree_last,
 		     START, LAST, static inline, vhost_iotlb_itree);
 
+static void vhost_iotlb_map_unlink(struct vhost_iotlb *iotlb,
+				   struct vhost_iotlb_map *map)
+{
+	vhost_iotlb_itree_remove(map, &iotlb->root);
+	list_del(&map->link);
+	iotlb->nmaps--;
+}
+
 /**
  * vhost_iotlb_map_free - remove a map node and free it
  * @iotlb: the IOTLB
@@ -28,10 +36,8 @@ INTERVAL_TREE_DEFINE(struct vhost_iotlb_map,
 void vhost_iotlb_map_free(struct vhost_iotlb *iotlb,
 			  struct vhost_iotlb_map *map)
 {
-	vhost_iotlb_itree_remove(map, &iotlb->root);
-	list_del(&map->link);
+	vhost_iotlb_map_unlink(iotlb, map);
 	kfree(map);
-	iotlb->nmaps--;
 }
 EXPORT_SYMBOL_GPL(vhost_iotlb_map_free);
 
@@ -57,14 +63,25 @@ int vhost_iotlb_add_range_ctx(struct vhost_iotlb *iotlb,
 	if (last < start)
 		return -EFAULT;
 
+	if (!iotlb->limit)
+		return -EINVAL;
+
 	/* If the range being mapped is [0, ULONG_MAX], split it into two entries
 	 * otherwise its size would overflow u64.
 	 */
 	if (start == 0 && last == ULONG_MAX) {
 		u64 mid = last / 2;
-		int err = vhost_iotlb_add_range_ctx(iotlb, start, mid, addr,
-				perm, opaque);
+		int err;
+
+		if (iotlb->limit < 2)
+			return -ENOSPC;
 
+		if (!(iotlb->flags & VHOST_IOTLB_FLAG_RETIRE) &&
+		    iotlb->nmaps > iotlb->limit - 2)
+			return -ENOSPC;
+
+		err = vhost_iotlb_add_range_ctx(iotlb, start, mid, addr,
+						perm, opaque);
 		if (err)
 			return err;
 
@@ -72,17 +89,19 @@ int vhost_iotlb_add_range_ctx(struct vhost_iotlb *iotlb,
 		start = mid + 1;
 	}
 
-	if (iotlb->limit &&
-	    iotlb->nmaps == iotlb->limit &&
-	    iotlb->flags & VHOST_IOTLB_FLAG_RETIRE) {
-		map = list_first_entry(&iotlb->list, typeof(*map), link);
-		vhost_iotlb_map_free(iotlb, map);
+	if (iotlb->nmaps >= iotlb->limit) {
+		if (iotlb->flags & VHOST_IOTLB_FLAG_RETIRE) {
+			map = list_first_entry(&iotlb->list, typeof(*map), link);
+			vhost_iotlb_map_unlink(iotlb, map);
+		} else {
+			return -ENOSPC;
+		}
+	} else {
+		map = kmalloc_obj(*map, GFP_ATOMIC);
+		if (!map)
+			return -ENOMEM;
 	}
 
-	map = kmalloc_obj(*map, GFP_ATOMIC);
-	if (!map)
-		return -ENOMEM;
-
 	map->start = start;
 	map->size = last - start + 1;
 	map->last = last;
diff --git a/drivers/vhost/vdpa.c b/drivers/vhost/vdpa.c
index ac55275fa0d0..ef642bc9f97e 100644
--- a/drivers/vhost/vdpa.c
+++ b/drivers/vhost/vdpa.c
@@ -34,6 +34,11 @@ enum {
 
 #define VHOST_VDPA_DEV_MAX (1U << MINORBITS)
 
+static int max_iotlb_entries = 2048;
+module_param(max_iotlb_entries, int, 0444);
+MODULE_PARM_DESC(max_iotlb_entries,
+		 "Maximum number of iotlb entries. (default: 2048)");
+
 #define VHOST_VDPA_IOTLB_BUCKETS 16
 
 struct vhost_vdpa_as {
@@ -109,12 +114,14 @@ static struct vhost_vdpa_as *vhost_vdpa_alloc_as(struct vhost_vdpa *v, u32 asid)
 
 	if (asid >= v->vdpa->nas)
 		return NULL;
+	if (max_iotlb_entries <= 0)
+		return NULL;
 
 	as = kmalloc_obj(*as);
 	if (!as)
 		return NULL;
 
-	vhost_iotlb_init(&as->iotlb, 0, 0);
+	vhost_iotlb_init(&as->iotlb, max_iotlb_entries, 0);
 	as->id = asid;
 	hlist_add_head(&as->hash_link, head);
 
diff --git a/drivers/vhost/vhost.c b/drivers/vhost/vhost.c
index 4c525b3e16ea..ae29a7ef7baa 100644
--- a/drivers/vhost/vhost.c
+++ b/drivers/vhost/vhost.c
@@ -1137,6 +1137,9 @@ EXPORT_SYMBOL_GPL(vhost_dev_set_owner);
 
 static struct vhost_iotlb *iotlb_alloc(void)
 {
+	if (max_iotlb_entries <= 0)
+		return NULL;
+
 	return vhost_iotlb_alloc(max_iotlb_entries,
 				 VHOST_IOTLB_FLAG_RETIRE);
 }
@@ -1981,6 +1984,8 @@ static long vhost_set_memory(struct vhost_dev *d, struct vhost_memory __user *m)
 		return -EOPNOTSUPP;
 	if (mem.nregions > max_mem_regions)
 		return -E2BIG;
+	if (max_iotlb_entries <= 0)
+		return -EINVAL;
 	newmem = kvzalloc_flex(*newmem, regions, mem.nregions);
 	if (!newmem)
 		return -ENOMEM;
@@ -2275,6 +2280,9 @@ int vhost_init_device_iotlb(struct vhost_dev *d)
 	struct vhost_iotlb *niotlb, *oiotlb;
 	int i;
 
+	if (max_iotlb_entries <= 0)
+		return -EINVAL;
+
 	niotlb = iotlb_alloc();
 	if (!niotlb)
 		return -ENOMEM;