15d8d6634SDave Marchevsky // SPDX-License-Identifier: GPL-2.0
25d8d6634SDave Marchevsky /* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
35d8d6634SDave Marchevsky 
45d8d6634SDave Marchevsky #include <vmlinux.h>
55d8d6634SDave Marchevsky #include <bpf/bpf_tracing.h>
65d8d6634SDave Marchevsky #include <bpf/bpf_helpers.h>
75d8d6634SDave Marchevsky #include <bpf/bpf_core_read.h>
88e9af821SJiri Olsa #include "../bpf_experimental.h"
98e9af821SJiri Olsa #include "../bpf_testmod/bpf_testmod_kfunc.h"
105d8d6634SDave Marchevsky 
115d8d6634SDave Marchevsky struct node_data {
125d8d6634SDave Marchevsky 	long key;
135d8d6634SDave Marchevsky 	long data;
145d8d6634SDave Marchevsky 	struct bpf_rb_node node;
155d8d6634SDave Marchevsky };
165d8d6634SDave Marchevsky 
17*001fedacSYonghong Song struct plain_local {
18*001fedacSYonghong Song 	long key;
19*001fedacSYonghong Song 	long data;
20*001fedacSYonghong Song };
21*001fedacSYonghong Song 
225d8d6634SDave Marchevsky struct map_value {
235d8d6634SDave Marchevsky 	struct prog_test_ref_kfunc *not_kptr;
245d8d6634SDave Marchevsky 	struct prog_test_ref_kfunc __kptr *val;
255d8d6634SDave Marchevsky 	struct node_data __kptr *node;
26*001fedacSYonghong Song 	struct plain_local __kptr *plain;
275d8d6634SDave Marchevsky };
285d8d6634SDave Marchevsky 
295d8d6634SDave Marchevsky /* This is necessary so that LLVM generates BTF for node_data struct
305d8d6634SDave Marchevsky  * If it's not included, a fwd reference for node_data will be generated but
315d8d6634SDave Marchevsky  * no struct. Example BTF of "node" field in map_value when not included:
325d8d6634SDave Marchevsky  *
335d8d6634SDave Marchevsky  * [10] PTR '(anon)' type_id=35
345d8d6634SDave Marchevsky  * [34] FWD 'node_data' fwd_kind=struct
355d8d6634SDave Marchevsky  * [35] TYPE_TAG 'kptr_ref' type_id=34
365d8d6634SDave Marchevsky  *
375d8d6634SDave Marchevsky  * (with no node_data struct defined)
385d8d6634SDave Marchevsky  * Had to do the same w/ bpf_kfunc_call_test_release below
395d8d6634SDave Marchevsky  */
405d8d6634SDave Marchevsky struct node_data *just_here_because_btf_bug;
415d8d6634SDave Marchevsky 
425d8d6634SDave Marchevsky struct {
435d8d6634SDave Marchevsky 	__uint(type, BPF_MAP_TYPE_ARRAY);
445d8d6634SDave Marchevsky 	__type(key, int);
455d8d6634SDave Marchevsky 	__type(value, struct map_value);
465d8d6634SDave Marchevsky 	__uint(max_entries, 2);
475d8d6634SDave Marchevsky } some_nodes SEC(".maps");
485d8d6634SDave Marchevsky 
create_and_stash(int idx,int val)495d8d6634SDave Marchevsky static int create_and_stash(int idx, int val)
505d8d6634SDave Marchevsky {
515d8d6634SDave Marchevsky 	struct map_value *mapval;
525d8d6634SDave Marchevsky 	struct node_data *res;
535d8d6634SDave Marchevsky 
545d8d6634SDave Marchevsky 	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
555d8d6634SDave Marchevsky 	if (!mapval)
565d8d6634SDave Marchevsky 		return 1;
575d8d6634SDave Marchevsky 
585d8d6634SDave Marchevsky 	res = bpf_obj_new(typeof(*res));
595d8d6634SDave Marchevsky 	if (!res)
605d8d6634SDave Marchevsky 		return 1;
615d8d6634SDave Marchevsky 	res->key = val;
625d8d6634SDave Marchevsky 
635d8d6634SDave Marchevsky 	res = bpf_kptr_xchg(&mapval->node, res);
645d8d6634SDave Marchevsky 	if (res)
655d8d6634SDave Marchevsky 		bpf_obj_drop(res);
665d8d6634SDave Marchevsky 	return 0;
675d8d6634SDave Marchevsky }
685d8d6634SDave Marchevsky 
695d8d6634SDave Marchevsky SEC("tc")
stash_rb_nodes(void * ctx)705d8d6634SDave Marchevsky long stash_rb_nodes(void *ctx)
715d8d6634SDave Marchevsky {
725d8d6634SDave Marchevsky 	return create_and_stash(0, 41) ?: create_and_stash(1, 42);
735d8d6634SDave Marchevsky }
745d8d6634SDave Marchevsky 
755d8d6634SDave Marchevsky SEC("tc")
stash_plain(void * ctx)76*001fedacSYonghong Song long stash_plain(void *ctx)
77*001fedacSYonghong Song {
78*001fedacSYonghong Song 	struct map_value *mapval;
79*001fedacSYonghong Song 	struct plain_local *res;
80*001fedacSYonghong Song 	int idx = 0;
81*001fedacSYonghong Song 
82*001fedacSYonghong Song 	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
83*001fedacSYonghong Song 	if (!mapval)
84*001fedacSYonghong Song 		return 1;
85*001fedacSYonghong Song 
86*001fedacSYonghong Song 	res = bpf_obj_new(typeof(*res));
87*001fedacSYonghong Song 	if (!res)
88*001fedacSYonghong Song 		return 1;
89*001fedacSYonghong Song 	res->key = 41;
90*001fedacSYonghong Song 
91*001fedacSYonghong Song 	res = bpf_kptr_xchg(&mapval->plain, res);
92*001fedacSYonghong Song 	if (res)
93*001fedacSYonghong Song 		bpf_obj_drop(res);
94*001fedacSYonghong Song 	return 0;
95*001fedacSYonghong Song }
96*001fedacSYonghong Song 
97*001fedacSYonghong Song SEC("tc")
unstash_rb_node(void * ctx)985d8d6634SDave Marchevsky long unstash_rb_node(void *ctx)
995d8d6634SDave Marchevsky {
1005d8d6634SDave Marchevsky 	struct map_value *mapval;
1015d8d6634SDave Marchevsky 	struct node_data *res;
1025d8d6634SDave Marchevsky 	long retval;
1035d8d6634SDave Marchevsky 	int key = 1;
1045d8d6634SDave Marchevsky 
1055d8d6634SDave Marchevsky 	mapval = bpf_map_lookup_elem(&some_nodes, &key);
1065d8d6634SDave Marchevsky 	if (!mapval)
1075d8d6634SDave Marchevsky 		return 1;
1085d8d6634SDave Marchevsky 
1095d8d6634SDave Marchevsky 	res = bpf_kptr_xchg(&mapval->node, NULL);
1105d8d6634SDave Marchevsky 	if (res) {
1115d8d6634SDave Marchevsky 		retval = res->key;
1125d8d6634SDave Marchevsky 		bpf_obj_drop(res);
1135d8d6634SDave Marchevsky 		return retval;
1145d8d6634SDave Marchevsky 	}
1155d8d6634SDave Marchevsky 	return 1;
1165d8d6634SDave Marchevsky }
1175d8d6634SDave Marchevsky 
1185d8d6634SDave Marchevsky SEC("tc")
stash_test_ref_kfunc(void * ctx)1195d8d6634SDave Marchevsky long stash_test_ref_kfunc(void *ctx)
1205d8d6634SDave Marchevsky {
1215d8d6634SDave Marchevsky 	struct prog_test_ref_kfunc *res;
1225d8d6634SDave Marchevsky 	struct map_value *mapval;
1235d8d6634SDave Marchevsky 	int key = 0;
1245d8d6634SDave Marchevsky 
1255d8d6634SDave Marchevsky 	mapval = bpf_map_lookup_elem(&some_nodes, &key);
1265d8d6634SDave Marchevsky 	if (!mapval)
1275d8d6634SDave Marchevsky 		return 1;
1285d8d6634SDave Marchevsky 
1295d8d6634SDave Marchevsky 	res = bpf_kptr_xchg(&mapval->val, NULL);
1305d8d6634SDave Marchevsky 	if (res)
1315d8d6634SDave Marchevsky 		bpf_kfunc_call_test_release(res);
1325d8d6634SDave Marchevsky 	return 0;
1335d8d6634SDave Marchevsky }
1345d8d6634SDave Marchevsky 
1355d8d6634SDave Marchevsky char _license[] SEC("license") = "GPL";
136