1 files changed, 55 insertions, 0 deletions
diff --git a/tools/testing/selftests/bpf/progs/rbtree_fail.c b/tools/testing/selftests/bpf/progs/rbtree_fail.c
index 803419a47c62..4504608196ab 100644
--- a/tools/testing/selftests/bpf/progs/rbtree_fail.c
+++ b/tools/testing/selftests/bpf/progs/rbtree_fail.c
@@ -272,6 +272,47 @@ static bool less__bad_res_spin_unlock(struct bpf_rb_node *a, const struct bpf_rb
 	return false;
 }
 
+static __noinline void rbtree_cb_unlock_relock(void)
+{
+	bpf_spin_unlock(&glock);
+	bpf_spin_lock(&glock);
+}
+
+static __noinline void rbtree_cb_nested_unlock(void)
+{
+	rbtree_cb_unlock_relock();
+	asm volatile ("");
+}
+
+static bool less__bad_subprog_unlock(struct bpf_rb_node *a, const struct bpf_rb_node *b)
+{
+	struct node_data *node_a;
+	struct node_data *node_b;
+
+	node_a = container_of(a, struct node_data, node);
+	node_b = container_of(b, struct node_data, node);
+	rbtree_cb_nested_unlock();
+
+	return node_a->key < node_b->key;
+}
+
+static __noinline void rbtree_cb_noop(void)
+{
+	asm volatile ("");
+}
+
+static bool less__subprog_allowed(struct bpf_rb_node *a, const struct bpf_rb_node *b)
+{
+	struct node_data *node_a;
+	struct node_data *node_b;
+
+	node_a = container_of(a, struct node_data, node);
+	node_b = container_of(b, struct node_data, node);
+	rbtree_cb_noop();
+
+	return node_a->key < node_b->key;
+}
+
 static __always_inline
 long add_with_cb(bool (cb)(struct bpf_rb_node *a, const struct bpf_rb_node *b))
 {
@@ -330,4 +371,18 @@ long rbtree_api_add_bad_cb_res_spin_unlock(void *ctx)
 	return 0;
 }
 
+SEC("?tc")
+__failure __msg("can't spin_{lock,unlock} in rbtree cb")
+long rbtree_api_add_bad_cb_subprog_unlock(void *ctx)
+{
+	return add_with_cb(less__bad_subprog_unlock);
+}
+
+SEC("?tc")
+__success
+long rbtree_api_add_cb_subprog_allowed(void *ctx)
+{
+	return add_with_cb(less__subprog_allowed);
+}
+
 char _license[] SEC("license") = "GPL";