bpf: Add batch and iter ops for resizable hashtab

Signed-off-by: Mykyta Yatsenko <[email protected]>

Author: mykyta5

kernel/bpf/hashtab.c

@@ -2982,66 +2982,278 @@ static long bpf_for_each_rhash_elem(struct bpf_map *map,
 
 static u64 rhtab_map_mem_usage(const struct bpf_map *map)
 {
-	return 1;
+	struct bpf_rhtab *rhtab = container_of(map, struct bpf_rhtab, map);
+	u64 num_entries;
+	u64 usage = sizeof(struct bpf_rhtab);
+
+	num_entries = atomic_read(&rhtab->ht.nelems);
+	usage += rhtab->elem_size * num_entries;
+
+	/* Add rhashtable internal overhead estimate */
+	usage += rhtab->ht.tbl ? sizeof(struct bucket_table) : 0;
+
+	return usage;
+}
+
+static int __rhtab_map_lookup_and_delete_batch(struct bpf_map *map,
+					       const union bpf_attr *attr,
+					       union bpf_attr __user *uattr,
+					       bool do_delete)
+{
+	struct bpf_rhtab *rhtab = container_of(map, struct bpf_rhtab, map);
+	void __user *uvalues = u64_to_user_ptr(attr->batch.values);
+	void __user *ukeys = u64_to_user_ptr(attr->batch.keys);
+	void __user *ubatch = u64_to_user_ptr(attr->batch.in_batch);
+	u32 batch, max_count, total, key_size, value_size;
+	void *keys = NULL, *values = NULL, *dst_key, *dst_val;
+	struct rhtab_elem **elems_to_delete = NULL;
+	struct rhashtable_iter iter;
+	struct rhtab_elem *l;
+	u64 elem_map_flags, map_flags;
+	u32 buf_size;
+	int ret = 0;
+
+	elem_map_flags = attr->batch.elem_flags;
+	if ((elem_map_flags & ~BPF_F_LOCK) ||
+	    ((elem_map_flags & BPF_F_LOCK) &&
+	     !btf_record_has_field(map->record, BPF_SPIN_LOCK)))
+		return -EINVAL;
+
+	map_flags = attr->batch.flags;
+	if (map_flags)
+		return -EINVAL;
+
+	max_count = attr->batch.count;
+	if (!max_count)
+		return 0;
+
+	if (put_user(0, &uattr->batch.count))
+		return -EFAULT;
+
+	batch = 0;
+	if (ubatch && copy_from_user(&batch, ubatch, sizeof(batch)))
+		return -EFAULT;
+
+	key_size = map->key_size;
+	value_size = map->value_size;
+
+	/* Allocate buffers for keys and values */
+	buf_size = max_count;
+	keys = kvmalloc_array(key_size, buf_size, GFP_USER | __GFP_NOWARN);
+	values = kvmalloc_array(value_size, buf_size, GFP_USER | __GFP_NOWARN);
+	if (!keys || !values) {
+		ret = -ENOMEM;
+		goto out;
+	}
+
+	if (do_delete) {
+		elems_to_delete = kvmalloc_array(sizeof(*elems_to_delete),
+						 buf_size,
+						 GFP_USER | __GFP_NOWARN);
+		if (!elems_to_delete) {
+			ret = -ENOMEM;
+			goto out;
+		}
+	}
+
+	dst_key = keys;
+	dst_val = values;
+	total = 0;
+
+	rcu_read_lock();
+	rhashtable_walk_enter(&rhtab->ht, &iter);
+	rhashtable_walk_start(&iter);
+
+	/* Skip elements up to batch position */
+	while (batch > 0 && (l = rhashtable_walk_next(&iter))) {
+		if (IS_ERR(l)) {
+			if (PTR_ERR(l) == -EAGAIN)
+				continue;
+			break;
+		}
+		batch--;
+	}
+
+	/* Collect elements */
+	while (total < max_count && (l = rhashtable_walk_next(&iter))) {
+		if (IS_ERR(l)) {
+			if (PTR_ERR(l) == -EAGAIN)
+				continue;
+			break;
+		}
+
+		memcpy(dst_key, l->data, key_size);
+		memcpy(dst_val, rhtab_elem_value(l), value_size);
+		check_and_init_map_value(map, dst_val);
+
+		if (do_delete)
+			elems_to_delete[total] = l;
+
+		dst_key += key_size;
+		dst_val += value_size;
+		total++;
+	}
+
+	/* Delete collected elements while still in RCU critical section */
+	if (do_delete && total > 0) {
+		u32 i;
+
+		for (i = 0; i < total; i++) {
+			l = elems_to_delete[i];
+			rhashtable_remove_fast(&rhtab->ht, &l->node,
+					       rhtab->params);
+			bpf_mem_cache_free_rcu(&rhtab->ma, l);
+		}
+	}
+
+	rhashtable_walk_stop(&iter);
+	rhashtable_walk_exit(&iter);
+	rcu_read_unlock();
+
+	if (total == 0) {
+		ret = -ENOENT;
+		goto out;
+	}
+
+	/* Copy results to userspace */
+	if (copy_to_user(ukeys, keys, total * key_size) ||
+	    copy_to_user(uvalues, values, total * value_size)) {
+		ret = -EFAULT;
+		goto out;
+	}
+
+	if (put_user(total, &uattr->batch.count))
+		ret = -EFAULT;
+
+	/* Update batch cursor for next iteration */
+	if (ubatch) {
+		u32 next_batch;
+
+		if (copy_from_user(&next_batch, ubatch, sizeof(next_batch))) {
+			ret = -EFAULT;
+			goto out;
+		}
+		next_batch += total;
+		if (copy_to_user(ubatch, &next_batch, sizeof(next_batch)))
+			ret = -EFAULT;
+	}
+
+out:
+	kvfree(keys);
+	kvfree(values);
+	kvfree(elems_to_delete);
+	return ret;
 }
 
 static int rhtab_map_lookup_batch(struct bpf_map *map,
 				  const union bpf_attr *attr,
 				  union bpf_attr __user *uattr)
 {
-	return 0;
+	return __rhtab_map_lookup_and_delete_batch(map, attr, uattr, false);
 }
 
 static int rhtab_map_lookup_and_delete_batch(struct bpf_map *map,
 					     const union bpf_attr *attr,
 					     union bpf_attr __user *uattr)
 {
-	return 0;
+	return __rhtab_map_lookup_and_delete_batch(map, attr, uattr, true);
 }
 
 struct bpf_iter_seq_rhash_map_info {
 	struct bpf_map *map;
 	struct bpf_rhtab *rhtab;
-	void *percpu_value_buf; // non-zero means percpu hash
-	u32 bucket_id;
-	u32 skip_elems;
+	struct rhashtable_iter iter;
+	bool iter_active;
 };
 
-static struct htab_elem *
-bpf_rhash_map_seq_find_next(struct bpf_iter_seq_hash_map_info *info,
-			    struct htab_elem *prev_elem)
+static void *bpf_rhash_map_seq_start(struct seq_file *seq, loff_t *pos)
 {
-	return NULL;
+	struct bpf_iter_seq_rhash_map_info *info = seq->private;
+	struct rhtab_elem *elem;
+
+	rhashtable_walk_enter(&info->rhtab->ht, &info->iter);
+	rhashtable_walk_start(&info->iter);
+	info->iter_active = true;
+
+	elem = rhtab_iter_next(&info->iter);
+	if (!elem)
+		return NULL;
+
+	if (*pos == 0)
+		++*pos;
+	return elem;
 }
 
-static void *bpf_rhash_map_seq_start(struct seq_file *seq, loff_t *pos)
+static void *bpf_rhash_map_seq_next(struct seq_file *seq, void *v, loff_t *pos)
 {
-	return NULL;
+	struct bpf_iter_seq_rhash_map_info *info = seq->private;
+
+	++*pos;
+
+	return rhtab_iter_next(&info->iter);
 }
 
-static void *bpf_rhash_map_seq_next(struct seq_file *seq, void *v,
-				    loff_t *pos) /*  */
+static int __bpf_rhash_map_seq_show(struct seq_file *seq,
+				    struct rhtab_elem *elem)
 {
-	return NULL;
+	struct bpf_iter_seq_rhash_map_info *info = seq->private;
+	struct bpf_iter__bpf_map_elem ctx = {};
+	struct bpf_iter_meta meta;
+	struct bpf_prog *prog;
+	int ret = 0;
+
+	meta.seq = seq;
+	prog = bpf_iter_get_info(&meta, elem == NULL);
+	if (prog) {
+		ctx.meta = &meta;
+		ctx.map = info->map;
+		if (elem) {
+			ctx.key = elem->data;
+			ctx.value = rhtab_elem_value(elem);
+		}
+		ret = bpf_iter_run_prog(prog, &ctx);
+	}
+
+	return ret;
 }
 
 static int bpf_rhash_map_seq_show(struct seq_file *seq, void *v)
 {
-	return 0;
+	return __bpf_rhash_map_seq_show(seq, v);
 }
 
 static void bpf_rhash_map_seq_stop(struct seq_file *seq, void *v)
 {
+	struct bpf_iter_seq_rhash_map_info *info = seq->private;
+
+	if (!v)
+		(void)__bpf_rhash_map_seq_show(seq, NULL);
+
+	if (info->iter_active) {
+		rhashtable_walk_stop(&info->iter);
+		rhashtable_walk_exit(&info->iter);
+		info->iter_active = false;
+	}
 }
 
 static int bpf_iter_init_rhash_map(void *priv_data,
 				   struct bpf_iter_aux_info *aux)
 {
+	struct bpf_iter_seq_rhash_map_info *info = priv_data;
+	struct bpf_map *map = aux->map;
+
+	bpf_map_inc_with_uref(map);
+	info->map = map;
+	info->rhtab = container_of(map, struct bpf_rhtab, map);
+	info->iter_active = false;
 	return 0;
 }
 
 static void bpf_iter_fini_rhash_map(void *priv_data)
 {
+	struct bpf_iter_seq_rhash_map_info *info = priv_data;
+
+	bpf_map_put_with_uref(info->map);
 }
 
 static const struct seq_operations bpf_rhash_map_seq_ops = {