162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0
262306a36Sopenharmony_ci#include <vmlinux.h>
362306a36Sopenharmony_ci#include <bpf/bpf_tracing.h>
462306a36Sopenharmony_ci#include <bpf/bpf_helpers.h>
562306a36Sopenharmony_ci#include <bpf/bpf_core_read.h>
662306a36Sopenharmony_ci#include "bpf_experimental.h"
762306a36Sopenharmony_ci#include "bpf_misc.h"
862306a36Sopenharmony_ci
962306a36Sopenharmony_cistruct node_data {
1062306a36Sopenharmony_ci	long key;
1162306a36Sopenharmony_ci	long data;
1262306a36Sopenharmony_ci	struct bpf_rb_node node;
1362306a36Sopenharmony_ci};
1462306a36Sopenharmony_ci
1562306a36Sopenharmony_ci#define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
1662306a36Sopenharmony_ciprivate(A) struct bpf_spin_lock glock;
1762306a36Sopenharmony_ciprivate(A) struct bpf_rb_root groot __contains(node_data, node);
1862306a36Sopenharmony_ciprivate(A) struct bpf_rb_root groot2 __contains(node_data, node);
1962306a36Sopenharmony_ci
2062306a36Sopenharmony_cistatic bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
2162306a36Sopenharmony_ci{
2262306a36Sopenharmony_ci	struct node_data *node_a;
2362306a36Sopenharmony_ci	struct node_data *node_b;
2462306a36Sopenharmony_ci
2562306a36Sopenharmony_ci	node_a = container_of(a, struct node_data, node);
2662306a36Sopenharmony_ci	node_b = container_of(b, struct node_data, node);
2762306a36Sopenharmony_ci
2862306a36Sopenharmony_ci	return node_a->key < node_b->key;
2962306a36Sopenharmony_ci}
3062306a36Sopenharmony_ci
3162306a36Sopenharmony_ciSEC("?tc")
3262306a36Sopenharmony_ci__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
3362306a36Sopenharmony_cilong rbtree_api_nolock_add(void *ctx)
3462306a36Sopenharmony_ci{
3562306a36Sopenharmony_ci	struct node_data *n;
3662306a36Sopenharmony_ci
3762306a36Sopenharmony_ci	n = bpf_obj_new(typeof(*n));
3862306a36Sopenharmony_ci	if (!n)
3962306a36Sopenharmony_ci		return 1;
4062306a36Sopenharmony_ci
4162306a36Sopenharmony_ci	bpf_rbtree_add(&groot, &n->node, less);
4262306a36Sopenharmony_ci	return 0;
4362306a36Sopenharmony_ci}
4462306a36Sopenharmony_ci
4562306a36Sopenharmony_ciSEC("?tc")
4662306a36Sopenharmony_ci__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
4762306a36Sopenharmony_cilong rbtree_api_nolock_remove(void *ctx)
4862306a36Sopenharmony_ci{
4962306a36Sopenharmony_ci	struct node_data *n;
5062306a36Sopenharmony_ci
5162306a36Sopenharmony_ci	n = bpf_obj_new(typeof(*n));
5262306a36Sopenharmony_ci	if (!n)
5362306a36Sopenharmony_ci		return 1;
5462306a36Sopenharmony_ci
5562306a36Sopenharmony_ci	bpf_spin_lock(&glock);
5662306a36Sopenharmony_ci	bpf_rbtree_add(&groot, &n->node, less);
5762306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
5862306a36Sopenharmony_ci
5962306a36Sopenharmony_ci	bpf_rbtree_remove(&groot, &n->node);
6062306a36Sopenharmony_ci	return 0;
6162306a36Sopenharmony_ci}
6262306a36Sopenharmony_ci
6362306a36Sopenharmony_ciSEC("?tc")
6462306a36Sopenharmony_ci__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
6562306a36Sopenharmony_cilong rbtree_api_nolock_first(void *ctx)
6662306a36Sopenharmony_ci{
6762306a36Sopenharmony_ci	bpf_rbtree_first(&groot);
6862306a36Sopenharmony_ci	return 0;
6962306a36Sopenharmony_ci}
7062306a36Sopenharmony_ci
7162306a36Sopenharmony_ciSEC("?tc")
7262306a36Sopenharmony_ci__failure __msg("rbtree_remove node input must be non-owning ref")
7362306a36Sopenharmony_cilong rbtree_api_remove_unadded_node(void *ctx)
7462306a36Sopenharmony_ci{
7562306a36Sopenharmony_ci	struct node_data *n, *m;
7662306a36Sopenharmony_ci	struct bpf_rb_node *res;
7762306a36Sopenharmony_ci
7862306a36Sopenharmony_ci	n = bpf_obj_new(typeof(*n));
7962306a36Sopenharmony_ci	if (!n)
8062306a36Sopenharmony_ci		return 1;
8162306a36Sopenharmony_ci
8262306a36Sopenharmony_ci	m = bpf_obj_new(typeof(*m));
8362306a36Sopenharmony_ci	if (!m) {
8462306a36Sopenharmony_ci		bpf_obj_drop(n);
8562306a36Sopenharmony_ci		return 1;
8662306a36Sopenharmony_ci	}
8762306a36Sopenharmony_ci
8862306a36Sopenharmony_ci	bpf_spin_lock(&glock);
8962306a36Sopenharmony_ci	bpf_rbtree_add(&groot, &n->node, less);
9062306a36Sopenharmony_ci
9162306a36Sopenharmony_ci	/* This remove should pass verifier */
9262306a36Sopenharmony_ci	res = bpf_rbtree_remove(&groot, &n->node);
9362306a36Sopenharmony_ci	n = container_of(res, struct node_data, node);
9462306a36Sopenharmony_ci
9562306a36Sopenharmony_ci	/* This remove shouldn't, m isn't in an rbtree */
9662306a36Sopenharmony_ci	res = bpf_rbtree_remove(&groot, &m->node);
9762306a36Sopenharmony_ci	m = container_of(res, struct node_data, node);
9862306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
9962306a36Sopenharmony_ci
10062306a36Sopenharmony_ci	if (n)
10162306a36Sopenharmony_ci		bpf_obj_drop(n);
10262306a36Sopenharmony_ci	if (m)
10362306a36Sopenharmony_ci		bpf_obj_drop(m);
10462306a36Sopenharmony_ci	return 0;
10562306a36Sopenharmony_ci}
10662306a36Sopenharmony_ci
10762306a36Sopenharmony_ciSEC("?tc")
10862306a36Sopenharmony_ci__failure __msg("Unreleased reference id=3 alloc_insn=10")
10962306a36Sopenharmony_cilong rbtree_api_remove_no_drop(void *ctx)
11062306a36Sopenharmony_ci{
11162306a36Sopenharmony_ci	struct bpf_rb_node *res;
11262306a36Sopenharmony_ci	struct node_data *n;
11362306a36Sopenharmony_ci
11462306a36Sopenharmony_ci	bpf_spin_lock(&glock);
11562306a36Sopenharmony_ci	res = bpf_rbtree_first(&groot);
11662306a36Sopenharmony_ci	if (!res)
11762306a36Sopenharmony_ci		goto unlock_err;
11862306a36Sopenharmony_ci
11962306a36Sopenharmony_ci	res = bpf_rbtree_remove(&groot, res);
12062306a36Sopenharmony_ci
12162306a36Sopenharmony_ci	if (res) {
12262306a36Sopenharmony_ci		n = container_of(res, struct node_data, node);
12362306a36Sopenharmony_ci		__sink(n);
12462306a36Sopenharmony_ci	}
12562306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
12662306a36Sopenharmony_ci
12762306a36Sopenharmony_ci	/* if (res) { bpf_obj_drop(n); } is missing here */
12862306a36Sopenharmony_ci	return 0;
12962306a36Sopenharmony_ci
13062306a36Sopenharmony_ciunlock_err:
13162306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
13262306a36Sopenharmony_ci	return 1;
13362306a36Sopenharmony_ci}
13462306a36Sopenharmony_ci
13562306a36Sopenharmony_ciSEC("?tc")
13662306a36Sopenharmony_ci__failure __msg("arg#1 expected pointer to allocated object")
13762306a36Sopenharmony_cilong rbtree_api_add_to_multiple_trees(void *ctx)
13862306a36Sopenharmony_ci{
13962306a36Sopenharmony_ci	struct node_data *n;
14062306a36Sopenharmony_ci
14162306a36Sopenharmony_ci	n = bpf_obj_new(typeof(*n));
14262306a36Sopenharmony_ci	if (!n)
14362306a36Sopenharmony_ci		return 1;
14462306a36Sopenharmony_ci
14562306a36Sopenharmony_ci	bpf_spin_lock(&glock);
14662306a36Sopenharmony_ci	bpf_rbtree_add(&groot, &n->node, less);
14762306a36Sopenharmony_ci
14862306a36Sopenharmony_ci	/* This add should fail since n already in groot's tree */
14962306a36Sopenharmony_ci	bpf_rbtree_add(&groot2, &n->node, less);
15062306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
15162306a36Sopenharmony_ci	return 0;
15262306a36Sopenharmony_ci}
15362306a36Sopenharmony_ci
15462306a36Sopenharmony_ciSEC("?tc")
15562306a36Sopenharmony_ci__failure __msg("dereference of modified ptr_or_null_ ptr R2 off=16 disallowed")
15662306a36Sopenharmony_cilong rbtree_api_use_unchecked_remove_retval(void *ctx)
15762306a36Sopenharmony_ci{
15862306a36Sopenharmony_ci	struct bpf_rb_node *res;
15962306a36Sopenharmony_ci
16062306a36Sopenharmony_ci	bpf_spin_lock(&glock);
16162306a36Sopenharmony_ci
16262306a36Sopenharmony_ci	res = bpf_rbtree_first(&groot);
16362306a36Sopenharmony_ci	if (!res)
16462306a36Sopenharmony_ci		goto err_out;
16562306a36Sopenharmony_ci	res = bpf_rbtree_remove(&groot, res);
16662306a36Sopenharmony_ci
16762306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
16862306a36Sopenharmony_ci
16962306a36Sopenharmony_ci	bpf_spin_lock(&glock);
17062306a36Sopenharmony_ci	/* Must check res for NULL before using in rbtree_add below */
17162306a36Sopenharmony_ci	bpf_rbtree_add(&groot, res, less);
17262306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
17362306a36Sopenharmony_ci	return 0;
17462306a36Sopenharmony_ci
17562306a36Sopenharmony_cierr_out:
17662306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
17762306a36Sopenharmony_ci	return 1;
17862306a36Sopenharmony_ci}
17962306a36Sopenharmony_ci
18062306a36Sopenharmony_ciSEC("?tc")
18162306a36Sopenharmony_ci__failure __msg("rbtree_remove node input must be non-owning ref")
18262306a36Sopenharmony_cilong rbtree_api_add_release_unlock_escape(void *ctx)
18362306a36Sopenharmony_ci{
18462306a36Sopenharmony_ci	struct node_data *n;
18562306a36Sopenharmony_ci
18662306a36Sopenharmony_ci	n = bpf_obj_new(typeof(*n));
18762306a36Sopenharmony_ci	if (!n)
18862306a36Sopenharmony_ci		return 1;
18962306a36Sopenharmony_ci
19062306a36Sopenharmony_ci	bpf_spin_lock(&glock);
19162306a36Sopenharmony_ci	bpf_rbtree_add(&groot, &n->node, less);
19262306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
19362306a36Sopenharmony_ci
19462306a36Sopenharmony_ci	bpf_spin_lock(&glock);
19562306a36Sopenharmony_ci	/* After add() in previous critical section, n should be
19662306a36Sopenharmony_ci	 * release_on_unlock and released after previous spin_unlock,
19762306a36Sopenharmony_ci	 * so should not be possible to use it here
19862306a36Sopenharmony_ci	 */
19962306a36Sopenharmony_ci	bpf_rbtree_remove(&groot, &n->node);
20062306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
20162306a36Sopenharmony_ci	return 0;
20262306a36Sopenharmony_ci}
20362306a36Sopenharmony_ci
20462306a36Sopenharmony_ciSEC("?tc")
20562306a36Sopenharmony_ci__failure __msg("rbtree_remove node input must be non-owning ref")
20662306a36Sopenharmony_cilong rbtree_api_first_release_unlock_escape(void *ctx)
20762306a36Sopenharmony_ci{
20862306a36Sopenharmony_ci	struct bpf_rb_node *res;
20962306a36Sopenharmony_ci	struct node_data *n;
21062306a36Sopenharmony_ci
21162306a36Sopenharmony_ci	bpf_spin_lock(&glock);
21262306a36Sopenharmony_ci	res = bpf_rbtree_first(&groot);
21362306a36Sopenharmony_ci	if (!res) {
21462306a36Sopenharmony_ci		bpf_spin_unlock(&glock);
21562306a36Sopenharmony_ci		return 1;
21662306a36Sopenharmony_ci	}
21762306a36Sopenharmony_ci	n = container_of(res, struct node_data, node);
21862306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
21962306a36Sopenharmony_ci
22062306a36Sopenharmony_ci	bpf_spin_lock(&glock);
22162306a36Sopenharmony_ci	/* After first() in previous critical section, n should be
22262306a36Sopenharmony_ci	 * release_on_unlock and released after previous spin_unlock,
22362306a36Sopenharmony_ci	 * so should not be possible to use it here
22462306a36Sopenharmony_ci	 */
22562306a36Sopenharmony_ci	bpf_rbtree_remove(&groot, &n->node);
22662306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
22762306a36Sopenharmony_ci	return 0;
22862306a36Sopenharmony_ci}
22962306a36Sopenharmony_ci
23062306a36Sopenharmony_cistatic bool less__bad_fn_call_add(struct bpf_rb_node *a, const struct bpf_rb_node *b)
23162306a36Sopenharmony_ci{
23262306a36Sopenharmony_ci	struct node_data *node_a;
23362306a36Sopenharmony_ci	struct node_data *node_b;
23462306a36Sopenharmony_ci
23562306a36Sopenharmony_ci	node_a = container_of(a, struct node_data, node);
23662306a36Sopenharmony_ci	node_b = container_of(b, struct node_data, node);
23762306a36Sopenharmony_ci	bpf_rbtree_add(&groot, &node_a->node, less);
23862306a36Sopenharmony_ci
23962306a36Sopenharmony_ci	return node_a->key < node_b->key;
24062306a36Sopenharmony_ci}
24162306a36Sopenharmony_ci
24262306a36Sopenharmony_cistatic bool less__bad_fn_call_remove(struct bpf_rb_node *a, const struct bpf_rb_node *b)
24362306a36Sopenharmony_ci{
24462306a36Sopenharmony_ci	struct node_data *node_a;
24562306a36Sopenharmony_ci	struct node_data *node_b;
24662306a36Sopenharmony_ci
24762306a36Sopenharmony_ci	node_a = container_of(a, struct node_data, node);
24862306a36Sopenharmony_ci	node_b = container_of(b, struct node_data, node);
24962306a36Sopenharmony_ci	bpf_rbtree_remove(&groot, &node_a->node);
25062306a36Sopenharmony_ci
25162306a36Sopenharmony_ci	return node_a->key < node_b->key;
25262306a36Sopenharmony_ci}
25362306a36Sopenharmony_ci
25462306a36Sopenharmony_cistatic bool less__bad_fn_call_first_unlock_after(struct bpf_rb_node *a, const struct bpf_rb_node *b)
25562306a36Sopenharmony_ci{
25662306a36Sopenharmony_ci	struct node_data *node_a;
25762306a36Sopenharmony_ci	struct node_data *node_b;
25862306a36Sopenharmony_ci
25962306a36Sopenharmony_ci	node_a = container_of(a, struct node_data, node);
26062306a36Sopenharmony_ci	node_b = container_of(b, struct node_data, node);
26162306a36Sopenharmony_ci	bpf_rbtree_first(&groot);
26262306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
26362306a36Sopenharmony_ci
26462306a36Sopenharmony_ci	return node_a->key < node_b->key;
26562306a36Sopenharmony_ci}
26662306a36Sopenharmony_ci
26762306a36Sopenharmony_cistatic __always_inline
26862306a36Sopenharmony_cilong add_with_cb(bool (cb)(struct bpf_rb_node *a, const struct bpf_rb_node *b))
26962306a36Sopenharmony_ci{
27062306a36Sopenharmony_ci	struct node_data *n;
27162306a36Sopenharmony_ci
27262306a36Sopenharmony_ci	n = bpf_obj_new(typeof(*n));
27362306a36Sopenharmony_ci	if (!n)
27462306a36Sopenharmony_ci		return 1;
27562306a36Sopenharmony_ci
27662306a36Sopenharmony_ci	bpf_spin_lock(&glock);
27762306a36Sopenharmony_ci	bpf_rbtree_add(&groot, &n->node, cb);
27862306a36Sopenharmony_ci	bpf_spin_unlock(&glock);
27962306a36Sopenharmony_ci	return 0;
28062306a36Sopenharmony_ci}
28162306a36Sopenharmony_ci
28262306a36Sopenharmony_ciSEC("?tc")
28362306a36Sopenharmony_ci__failure __msg("arg#1 expected pointer to allocated object")
28462306a36Sopenharmony_cilong rbtree_api_add_bad_cb_bad_fn_call_add(void *ctx)
28562306a36Sopenharmony_ci{
28662306a36Sopenharmony_ci	return add_with_cb(less__bad_fn_call_add);
28762306a36Sopenharmony_ci}
28862306a36Sopenharmony_ci
28962306a36Sopenharmony_ciSEC("?tc")
29062306a36Sopenharmony_ci__failure __msg("rbtree_remove not allowed in rbtree cb")
29162306a36Sopenharmony_cilong rbtree_api_add_bad_cb_bad_fn_call_remove(void *ctx)
29262306a36Sopenharmony_ci{
29362306a36Sopenharmony_ci	return add_with_cb(less__bad_fn_call_remove);
29462306a36Sopenharmony_ci}
29562306a36Sopenharmony_ci
29662306a36Sopenharmony_ciSEC("?tc")
29762306a36Sopenharmony_ci__failure __msg("can't spin_{lock,unlock} in rbtree cb")
29862306a36Sopenharmony_cilong rbtree_api_add_bad_cb_bad_fn_call_first_unlock_after(void *ctx)
29962306a36Sopenharmony_ci{
30062306a36Sopenharmony_ci	return add_with_cb(less__bad_fn_call_first_unlock_after);
30162306a36Sopenharmony_ci}
30262306a36Sopenharmony_ci
30362306a36Sopenharmony_cichar _license[] SEC("license") = "GPL";
304