[PATCH bpf-next 2/2] selftests/bpf: Cover refcount acquire node offsets

[Yiyang Chen]

Add regression coverage for bpf_refcount_acquire() on graph-node-derived
pointers.

The accepted case pops a list node and normalizes it with container_of()
before acquiring a refcount. The rejected cases pass popped list and rbtree
node pointers directly to bpf_refcount_acquire(), which must fail because
those pointers carry non-zero fixed offsets.

Signed-off-by: Yiyang Chen <chenyy23@xxxxxxxxxxxxxxxxxxxxx>
---
 .../selftests/bpf/progs/refcounted_kptr.c     | 33 ++++++++
 .../bpf/progs/refcounted_kptr_fail.c          | 84 +++++++++++++++++++
 2 files changed, 117 insertions(+)

diff --git a/tools/testing/selftests/bpf/progs/refcounted_kptr.c b/tools/testing/selftests/bpf/progs/refcounted_kptr.c
index 61906f480..7955a9973 100644
--- a/tools/testing/selftests/bpf/progs/refcounted_kptr.c
+++ b/tools/testing/selftests/bpf/progs/refcounted_kptr.c
@@ -741,6 +741,39 @@ int list_push_back_uninit_head(void *ctx)
 	return ret;
 }
 
+SEC("tc")
+__description("refcount_acquire_list_pop_container: acquire normalized list pop")
+__success __retval(0)
+int refcount_acquire_list_pop_container(void *ctx)
+{
+	struct node_data *node, *base, *ref;
+	struct bpf_list_node *list_node;
+	long err;
+
+	node = bpf_obj_new(typeof(*node));
+	if (!node)
+		return -1;
+
+	bpf_spin_lock(&lock);
+	err = bpf_list_push_front(&head, &node->l);
+	if (err) {
+		bpf_spin_unlock(&lock);
+		bpf_obj_drop(node);
+		return -2;
+	}
+
+	list_node = bpf_list_pop_front(&head);
+	bpf_spin_unlock(&lock);
+	if (!list_node)
+		return -3;
+
+	base = container_of(list_node, struct node_data, l);
+	ref = bpf_refcount_acquire(base);
+	bpf_obj_drop(ref);
+	bpf_obj_drop(base);
+	return 0;
+}
+
 SEC("?tc")
 __failure __msg("bpf_spin_lock at off=32 must be held for bpf_list_head")
 long list_del_without_lock_fail(void *ctx)
diff --git a/tools/testing/selftests/bpf/progs/refcounted_kptr_fail.c b/tools/testing/selftests/bpf/progs/refcounted_kptr_fail.c
index 7247a20c0..f6cb084af 100644
--- a/tools/testing/selftests/bpf/progs/refcounted_kptr_fail.c
+++ b/tools/testing/selftests/bpf/progs/refcounted_kptr_fail.c
@@ -13,12 +13,22 @@ struct node_acquire {
 	struct bpf_refcount refcount;
 };
 
+struct node_refcounted {
+	long key;
+	struct bpf_rb_node rb;
+	struct bpf_list_node list;
+	struct bpf_refcount refcount;
+};
+
 extern void bpf_rcu_read_lock(void) __ksym;
 extern void bpf_rcu_read_unlock(void) __ksym;
 
 #define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
 private(A) struct bpf_spin_lock glock;
 private(A) struct bpf_rb_root groot __contains(node_acquire, node);
+private(B) struct bpf_spin_lock lock;
+private(B) struct bpf_rb_root root __contains(node_refcounted, rb);
+private(B) struct bpf_list_head head __contains(node_refcounted, list);
 
 static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
 {
@@ -31,6 +41,17 @@ static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
 	return node_a->key < node_b->key;
 }
 
+static bool less_refcounted(struct bpf_rb_node *a, const struct bpf_rb_node *b)
+{
+	struct node_refcounted *node_a;
+	struct node_refcounted *node_b;
+
+	node_a = container_of(a, struct node_refcounted, rb);
+	node_b = container_of(b, struct node_refcounted, rb);
+
+	return node_a->key < node_b->key;
+}
+
 SEC("?tc")
 __failure __msg("Unreleased reference id=4 alloc_insn={{[0-9]+}}")
 long rbtree_refcounted_node_ref_escapes(void *ctx)
@@ -93,6 +114,69 @@ long rbtree_refcounted_node_ref_escapes_owning_input(void *ctx)
 	return 0;
 }
 
+SEC("?tc")
+__failure __msg("R1 must have zero offset when passed to bpf_refcount_acquire")
+long refcount_acquire_list_node_offset(void *ctx)
+{
+	struct node_refcounted *node, *base, *ref;
+	struct bpf_list_node *list_node;
+	long err;
+
+	node = bpf_obj_new(typeof(*node));
+	if (!node)
+		return 1;
+
+	bpf_spin_lock(&lock);
+	err = bpf_list_push_front(&head, &node->list);
+	if (err) {
+		bpf_spin_unlock(&lock);
+		bpf_obj_drop(node);
+		return 2;
+	}
+
+	list_node = bpf_list_pop_front(&head);
+	bpf_spin_unlock(&lock);
+	if (!list_node)
+		return 3;
+
+	base = container_of(list_node, struct node_refcounted, list);
+	ref = bpf_refcount_acquire(list_node);
+	if (ref)
+		bpf_obj_drop(ref);
+	bpf_obj_drop(base);
+	return 0;
+}
+
+SEC("?tc")
+__failure __msg("R1 must have zero offset when passed to bpf_refcount_acquire")
+long refcount_acquire_rbtree_node_offset(void *ctx)
+{
+	struct node_refcounted *node, *base, *ref;
+	struct bpf_rb_node *rb_node;
+
+	node = bpf_obj_new(typeof(*node));
+	if (!node)
+		return 1;
+
+	node->key = 1;
+
+	bpf_spin_lock(&lock);
+	bpf_rbtree_add(&root, &node->rb, less_refcounted);
+	rb_node = bpf_rbtree_first(&root);
+	if (rb_node)
+		rb_node = bpf_rbtree_remove(&root, rb_node);
+	bpf_spin_unlock(&lock);
+	if (!rb_node)
+		return 2;
+
+	base = container_of(rb_node, struct node_refcounted, rb);
+	ref = bpf_refcount_acquire(rb_node);
+	if (ref)
+		bpf_obj_drop(ref);
+	bpf_obj_drop(base);
+	return 0;
+}
+
 SEC("?fentry.s/" SYS_PREFIX "sys_getpgid")
 __failure __msg("function calls are not allowed while holding a lock")
 int BPF_PROG(rbtree_fail_sleepable_lock_across_rcu,
-- 
2.34.1