tools/testing/selftests/bpf/progs/local_kptr_stash.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
   3
   4 #include <vmlinux.h>
   5 #include <bpf/bpf_tracing.h>
   6 #include <bpf/bpf_helpers.h>
   7 #include <bpf/bpf_core_read.h>
   8 #include "../bpf_experimental.h"
   9 #include "../bpf_testmod/bpf_testmod_kfunc.h"
  10
  11 struct node_data {
  12         long key;
  13         long data;
  14         struct bpf_rb_node node;
  15 };
  16
  17 struct plain_local {
  18         long key;
  19         long data;
  20 };
  21
  22 struct map_value {
  23         struct prog_test_ref_kfunc *not_kptr;
  24         struct prog_test_ref_kfunc __kptr *val;
  25         struct node_data __kptr *node;
  26         struct plain_local __kptr *plain;
  27 };
  28
  29 /* This is necessary so that LLVM generates BTF for node_data struct
  30  * If it's not included, a fwd reference for node_data will be generated but
  31  * no struct. Example BTF of "node" field in map_value when not included:
  32  *
  33  * [10] PTR '(anon)' type_id=35
  34  * [34] FWD 'node_data' fwd_kind=struct
  35  * [35] TYPE_TAG 'kptr_ref' type_id=34
  36  *
  37  * (with no node_data struct defined)
  38  * Had to do the same w/ bpf_kfunc_call_test_release below
  39  */
  40 struct node_data *just_here_because_btf_bug;
  41
  42 struct {
  43         __uint(type, BPF_MAP_TYPE_ARRAY);
  44         __type(key, int);
  45         __type(value, struct map_value);
  46         __uint(max_entries, 2);
  47 } some_nodes SEC(".maps");
  48
  49 static int create_and_stash(int idx, int val)
  50 {
  51         struct map_value *mapval;
  52         struct node_data *res;
  53
  54         mapval = bpf_map_lookup_elem(&some_nodes, &idx);
  55         if (!mapval)
  56                 return 1;
  57
  58         res = bpf_obj_new(typeof(*res));
  59         if (!res)
  60                 return 1;
  61         res->key = val;
  62
  63         res = bpf_kptr_xchg(&mapval->node, res);
  64         if (res)
  65                 bpf_obj_drop(res);
  66         return 0;
  67 }
  68
  69 SEC("tc")
  70 long stash_rb_nodes(void *ctx)
  71 {
  72         return create_and_stash(0, 41) ?: create_and_stash(1, 42);
  73 }
  74
  75 SEC("tc")
  76 long stash_plain(void *ctx)
  77 {
  78         struct map_value *mapval;
  79         struct plain_local *res;
  80         int idx = 0;
  81
  82         mapval = bpf_map_lookup_elem(&some_nodes, &idx);
  83         if (!mapval)
  84                 return 1;
  85
  86         res = bpf_obj_new(typeof(*res));
  87         if (!res)
  88                 return 1;
  89         res->key = 41;
  90
  91         res = bpf_kptr_xchg(&mapval->plain, res);
  92         if (res)
  93                 bpf_obj_drop(res);
  94         return 0;
  95 }
  96
  97 SEC("tc")
  98 long unstash_rb_node(void *ctx)
  99 {
 100         struct map_value *mapval;
 101         struct node_data *res;
 102         long retval;
 103         int key = 1;
 104
 105         mapval = bpf_map_lookup_elem(&some_nodes, &key);
 106         if (!mapval)
 107                 return 1;
 108
 109         res = bpf_kptr_xchg(&mapval->node, NULL);
 110         if (res) {
 111                 retval = res->key;
 112                 bpf_obj_drop(res);
 113                 return retval;
 114         }
 115         return 1;
 116 }
 117
 118 SEC("tc")
 119 long stash_test_ref_kfunc(void *ctx)
 120 {
 121         struct prog_test_ref_kfunc *res;
 122         struct map_value *mapval;
 123         int key = 0;
 124
 125         mapval = bpf_map_lookup_elem(&some_nodes, &key);
 126         if (!mapval)
 127                 return 1;
 128
 129         res = bpf_kptr_xchg(&mapval->val, NULL);
 130         if (res)
 131                 bpf_kfunc_call_test_release(res);
 132         return 0;
 133 }
 134
 135 char _license[] SEC("license") = "GPL";