[PATCH bpf-next v1 1/4] bpf: Save struct_ops instance pointer in bpf_prog_aux

[Amery Hung <ameryhung@gmail.com>]

Allows struct_ops implementors to infer the calling struct_ops instance
inside a kfunc through prog->aux->this_st_ops. A new field, flags, is
added to bpf_struct_ops. If BPF_STRUCT_OPS_F_THIS_PTR is set in flags,
a pointer to the struct_ops structure registered to the kernel (i.e.,
kvalue->data) will be saved to prog->aux->this_st_ops. To access it in
a kfunc, use BPF_STRUCT_OPS_F_THIS_PTR with __prog argument [0]. The
verifier will fixup the argument with a pointer to prog->aux. this_st_ops
is protected by rcu and is valid until a struct_ops map is unregistered
updated.

For a struct_ops map with BPF_STRUCT_OPS_F_THIS_PTR, to make sure all
programs in it have the same this_st_ops, cmpxchg is used. Only if a
program is not already used in another struct_ops map also with
BPF_STRUCT_OPS_F_THIS_PTR can it be assigned to the current struct_ops
map.

[0]
commit bc049387b41f ("bpf: Add support for __prog argument suffix to
pass in prog->aux")
https://lore.kernel.org/r/20250513142812.1021591-1-memxor@gmail.com

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 include/linux/bpf.h         | 10 ++++++++++
 kernel/bpf/bpf_struct_ops.c | 38 +++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 83c56f40842b..25c3488ab926 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1640,6 +1640,11 @@ struct bpf_prog_aux {
 		struct work_struct work;
 		struct rcu_head	rcu;
 	};
+	/* pointer to struct_ops struct registered to the kernel.
+	 * Only valid when BPF_STRUCT_OPS_F_THIS_PTR is set in
+	 * bpf_struct_ops.flags
+	 */
+	void __rcu *this_st_ops;
 };
 
 struct bpf_prog {
@@ -1788,6 +1793,10 @@ struct bpf_token {
 struct bpf_struct_ops_value;
 struct btf_member;
 
+#define BPF_STRUCT_OPS_F_THIS_PTR BIT(0)
+
+#define BPF_STRUCT_OPS_FLAG_MASK BPF_STRUCT_OPS_F_THIS_PTR
+
 #define BPF_STRUCT_OPS_MAX_NR_MEMBERS 64
 /**
  * struct bpf_struct_ops - A structure of callbacks allowing a subsystem to
@@ -1853,6 +1862,7 @@ struct bpf_struct_ops {
 	struct module *owner;
 	const char *name;
 	struct btf_func_model func_models[BPF_STRUCT_OPS_MAX_NR_MEMBERS];
+	u32 flags;
 };
 
 /* Every member of a struct_ops type has an instance even a member is not
diff --git a/kernel/bpf/bpf_struct_ops.c b/kernel/bpf/bpf_struct_ops.c
index 96113633e391..717a2cec0b0f 100644
--- a/kernel/bpf/bpf_struct_ops.c
+++ b/kernel/bpf/bpf_struct_ops.c
@@ -533,6 +533,17 @@ static void bpf_struct_ops_map_put_progs(struct bpf_struct_ops_map *st_map)
 	}
 }
 
+static void bpf_struct_ops_map_clear_this_ptr(struct bpf_struct_ops_map *st_map)
+{
+	u32 i;
+
+	for (i = 0; i < st_map->funcs_cnt; i++) {
+		if (!st_map->links[i])
+			break;
+		RCU_INIT_POINTER(st_map->links[i]->prog->aux->this_st_ops, NULL);
+	}
+}
+
 static void bpf_struct_ops_map_free_image(struct bpf_struct_ops_map *st_map)
 {
 	int i;
@@ -695,6 +706,9 @@ static long bpf_struct_ops_map_update_elem(struct bpf_map *map, void *key,
 	if (flags)
 		return -EINVAL;
 
+	if (st_ops->flags & ~BPF_STRUCT_OPS_FLAG_MASK)
+		return -EINVAL;
+
 	if (*(u32 *)key != 0)
 		return -E2BIG;
 
@@ -801,6 +815,19 @@ static long bpf_struct_ops_map_update_elem(struct bpf_map *map, void *key,
 			goto reset_unlock;
 		}
 
+		if (st_ops->flags & BPF_STRUCT_OPS_F_THIS_PTR) {
+			/* Make sure a struct_ops map will not have programs with
+			 * different this_st_ops. Once a program is associated with
+			 * a struct_ops map, it cannot be used in another struct_ops
+			 * map also with BPF_STRUCT_OPS_F_THIS_PTR
+			 */
+			if (cmpxchg(&prog->aux->this_st_ops, NULL, kdata)) {
+				bpf_prog_put(prog);
+				err = -EINVAL;
+				goto reset_unlock;
+			}
+		}
+
 		link = kzalloc(sizeof(*link), GFP_USER);
 		if (!link) {
 			bpf_prog_put(prog);
@@ -894,6 +921,8 @@ static long bpf_struct_ops_map_update_elem(struct bpf_map *map, void *key,
 	bpf_struct_ops_map_free_ksyms(st_map);
 	bpf_struct_ops_map_free_image(st_map);
 	bpf_struct_ops_map_put_progs(st_map);
+	if (st_ops->flags & BPF_STRUCT_OPS_F_THIS_PTR)
+		bpf_struct_ops_map_clear_this_ptr(st_map);
 	memset(uvalue, 0, map->value_size);
 	memset(kvalue, 0, map->value_size);
 unlock:
@@ -919,6 +948,8 @@ static long bpf_struct_ops_map_delete_elem(struct bpf_map *map, void *key)
 	switch (prev_state) {
 	case BPF_STRUCT_OPS_STATE_INUSE:
 		st_map->st_ops_desc->st_ops->unreg(&st_map->kvalue.data, NULL);
+		if (st_map->st_ops_desc->st_ops->flags & BPF_STRUCT_OPS_F_THIS_PTR)
+			bpf_struct_ops_map_clear_this_ptr(st_map);
 		bpf_map_put(map);
 		return 0;
 	case BPF_STRUCT_OPS_STATE_TOBEFREE:
@@ -1194,6 +1225,8 @@ static void bpf_struct_ops_map_link_dealloc(struct bpf_link *link)
 		rcu_dereference_protected(st_link->map, true);
 	if (st_map) {
 		st_map->st_ops_desc->st_ops->unreg(&st_map->kvalue.data, link);
+		if (st_map->st_ops_desc->st_ops->flags & BPF_STRUCT_OPS_F_THIS_PTR)
+			bpf_struct_ops_map_clear_this_ptr(st_map);
 		bpf_map_put(&st_map->map);
 	}
 	kfree(st_link);
@@ -1268,6 +1301,9 @@ static int bpf_struct_ops_map_link_update(struct bpf_link *link, struct bpf_map
 	if (err)
 		goto err_out;
 
+	if (st_map->st_ops_desc->st_ops->flags & BPF_STRUCT_OPS_F_THIS_PTR)
+		bpf_struct_ops_map_clear_this_ptr(st_map);
+
 	bpf_map_inc(new_map);
 	rcu_assign_pointer(st_link->map, new_map);
 	bpf_map_put(old_map);
@@ -1294,6 +1330,8 @@ static int bpf_struct_ops_map_link_detach(struct bpf_link *link)
 	st_map = container_of(map, struct bpf_struct_ops_map, map);
 
 	st_map->st_ops_desc->st_ops->unreg(&st_map->kvalue.data, link);
+	if (st_map->st_ops_desc->st_ops->flags & BPF_STRUCT_OPS_F_THIS_PTR)
+		bpf_struct_ops_map_clear_this_ptr(st_map);
 
 	RCU_INIT_POINTER(st_link->map, NULL);
 	/* Pair with bpf_map_get() in bpf_struct_ops_link_create() or
-- 
2.47.1