[PATCH RFC bpf-next 5/7] libbpf: Implement bpf_gen__struct_ops_update_elem() for light skeletons

[Siddharth Nayyar <sidnayyar@google.com>]

Add bpf_gen__struct_ops_update_elem() in gen_loader.c and call it from
bpf_object_prepare_struct_ops() when generating a light skeleton.

Once all programs are loaded, bpf_gen__struct_ops_update_elem():
1. Closes the attach_btf_obj_fd stored in the upper 32 bits of
   st_ops->type_id_off.
2. Copies each callback program's FD from the loader stack into
   st_ops->kern_vdata_off at the dynamic kernel member offset stored in
   st_ops->kern_func_off[i].
3. Emits BPF_MAP_UPDATE_ELEM to populate the struct_ops map with the
   prepared kern_vdata.

Signed-off-by: Siddharth Nayyar <sidnayyar@google.com>
---
 tools/lib/bpf/bpf_gen_internal.h |  3 +++
 tools/lib/bpf/gen_loader.c       | 58 ++++++++++++++++++++++++++++++++++++++++
 tools/lib/bpf/libbpf.c           | 16 ++++++++---
 tools/lib/bpf/libbpf_internal.h  |  2 ++
 4 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/tools/lib/bpf/bpf_gen_internal.h b/tools/lib/bpf/bpf_gen_internal.h
index b424410e8839..cd603f82f6ec 100644
--- a/tools/lib/bpf/bpf_gen_internal.h
+++ b/tools/lib/bpf/bpf_gen_internal.h
@@ -88,5 +88,8 @@ void bpf_gen__init_kern_struct_ops(struct bpf_gen *gen, struct btf *btf,
 				   struct bpf_struct_ops *st_ops,
 				   __u32 *vtype_id_off,
 				   unsigned int *val_sz_off);
+void bpf_gen__struct_ops_update_elem(struct bpf_gen *gen, int map_idx,
+				     struct btf *btf,
+				     struct bpf_struct_ops *st_ops);
 
 #endif
diff --git a/tools/lib/bpf/gen_loader.c b/tools/lib/bpf/gen_loader.c
index aaa491463521..6dcba83dcc09 100644
--- a/tools/lib/bpf/gen_loader.c
+++ b/tools/lib/bpf/gen_loader.c
@@ -1573,3 +1573,61 @@ void bpf_gen__init_kern_struct_ops(struct bpf_gen *gen, struct btf *btf,
 		}
 	}
 }
+
+void bpf_gen__struct_ops_update_elem(struct bpf_gen *gen, int map_idx,
+				     struct btf *btf,
+				     struct bpf_struct_ops *st_ops)
+{
+	int attr_size = offsetofend(union bpf_attr, flags);
+	const struct btf_type *type;
+	union bpf_attr attr;
+	int key, zero = 0;
+	__u32 i;
+
+	/* 1. Close attach_btf_obj_fd (upper 32 bits of st_ops->type_id_off) */
+	emit2(gen, BPF_LD_IMM64_RAW_FULL(BPF_REG_0, BPF_PSEUDO_MAP_IDX_VALUE,
+					 0, 0, 0, st_ops->type_id_off));
+	emit(gen, BPF_LDX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0));
+	emit(gen, BPF_ALU64_IMM(BPF_RSH, BPF_REG_1, 32));
+	__emit_sys_close(gen);
+
+	/* 2. Populate prog_fd into kern_vdata for all callback members */
+	type = btf__type_by_id(btf, st_ops->type_id);
+	for (i = 0; i < btf_vlen(type); i++) {
+		struct bpf_program *prog = st_ops->progs[i];
+		int prog_idx, off_slot;
+
+		if (!prog)
+			continue;
+
+		prog_idx = bpf_program_idx(prog);
+		off_slot = st_ops->kern_func_off[i];
+
+		/* R2 = dynamic kernel member offset from off_slot */
+		emit2(gen, BPF_LD_IMM64_RAW_FULL(BPF_REG_1, BPF_PSEUDO_MAP_IDX_VALUE,
+						 0, 0, 0, off_slot));
+		emit(gen, BPF_LDX_MEM(BPF_W, BPF_REG_2, BPF_REG_1, 0));
+		emit(gen, BPF_JMP_IMM(BPF_JLE, BPF_REG_2,
+				      KERN_VDATA_MAX_SIZE - sizeof(__u64), 2));
+		emit(gen, BPF_MOV64_IMM(BPF_REG_7, -E2BIG));
+		emit_check_err(gen);
+
+		/* R1 = kern_vdata_off + offset */
+		emit2(gen, BPF_LD_IMM64_RAW_FULL(BPF_REG_1, BPF_PSEUDO_MAP_IDX_VALUE,
+						 0, 0, 0, st_ops->kern_vdata_off));
+		emit(gen, BPF_ALU64_REG(BPF_ADD, BPF_REG_1, BPF_REG_2));
+
+		/* R0 = prog_fd from stack (stack_off(prog_fd[prog_idx])) */
+		emit(gen, BPF_LDX_MEM(BPF_W, BPF_REG_0, BPF_REG_10,
+				      stack_off(prog_fd[prog_idx])));
+
+		/* Store prog_fd (R0) into kern_vdata + offset (R1) */
+		emit(gen, BPF_STX_MEM(BPF_DW, BPF_REG_1, BPF_REG_0, 0));
+	}
+
+	/* 3. Emit BPF_MAP_UPDATE_ELEM command for struct_ops map */
+	memset(&attr, 0, attr_size);
+	key = add_data(gen, &zero, sizeof(zero));
+	emit_map_update_elem(gen, map_idx, &attr, attr_size, key,
+			     st_ops->kern_vdata_off, KERN_VDATA_MAX_SIZE);
+}
diff --git a/tools/lib/bpf/libbpf.c b/tools/lib/bpf/libbpf.c
index 7754d6640167..d00ff5f8b0b0 100644
--- a/tools/lib/bpf/libbpf.c
+++ b/tools/lib/bpf/libbpf.c
@@ -9772,9 +9772,14 @@ static int bpf_object_prepare_struct_ops(struct bpf_object *obj)
 		if (!map->autocreate)
 			continue;
 
-		err = bpf_map_prepare_vdata(map);
-		if (err)
-			return err;
+		if (obj->gen_loader) {
+			bpf_gen__struct_ops_update_elem(obj->gen_loader, i,
+							obj->btf, map->st_ops);
+		} else {
+			err = bpf_map_prepare_vdata(map);
+			if (err)
+				return err;
+		}
 	}
 
 	return 0;
@@ -10720,6 +10725,11 @@ int bpf_program__set_expected_attach_type(struct bpf_program *prog,
 	return 0;
 }
 
+int bpf_program_idx(const struct bpf_program *prog)
+{
+	return prog - prog->obj->programs;
+}
+
 __u32 bpf_program__flags(const struct bpf_program *prog)
 {
 	return prog->prog_flags;
diff --git a/tools/lib/bpf/libbpf_internal.h b/tools/lib/bpf/libbpf_internal.h
index 845d6557c922..94c3a4019b11 100644
--- a/tools/lib/bpf/libbpf_internal.h
+++ b/tools/lib/bpf/libbpf_internal.h
@@ -317,6 +317,8 @@ struct bpf_struct_ops {
 	__u32 type_id;
 };
 
+int bpf_program_idx(const struct bpf_program *prog);
+
 int parse_btf_map_def(const char *map_name, struct btf *btf,
 		      const struct btf_type *def_t, bool strict,
 		      struct btf_map_def *map_def, struct btf_map_def *inner_def);

-- 
2.56.0.455.g6c5a503b38-goog