[PATCH bpf-next v1 0/4] Follow ups for kptr series

[PATCH bpf-next v1 0/4] Follow ups for kptr series

[Kumar Kartikeya Dwivedi]

Fix a build time warning, and address comments from Alexei on the merged version
[0].

  [0]: https://lore.kernel.org/bpf/20220424214901.2743946-1-memxor@gmail.com

Kumar Kartikeya Dwivedi (4):
  bpf: Fix sparse warning for bpf_kptr_xchg_proto
  bpf: Prepare prog_test_struct kfuncs for runtime tests
  selftests/bpf: Add negative C tests for kptrs
  selftests/bpf: Add tests for kptr_ref refcounting

 include/linux/bpf.h                           |   1 +
 net/bpf/test_run.c                            |  32 +-
 .../selftests/bpf/prog_tests/map_kptr.c       | 109 ++++-
 tools/testing/selftests/bpf/progs/map_kptr.c  | 109 ++++-
 .../selftests/bpf/progs/map_kptr_fail.c       | 419 ++++++++++++++++++
 .../testing/selftests/bpf/verifier/map_kptr.c |   4 +-
 6 files changed, 660 insertions(+), 14 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/progs/map_kptr_fail.c

-- 
2.35.1

[PATCH bpf-next v1 1/4] bpf: Fix sparse warning for bpf_kptr_xchg_proto

[Kumar Kartikeya Dwivedi]

Kernel Test Robot complained about missing static storage class
annotation for bpf_kptr_xchg_proto variable.

sparse: symbol 'bpf_kptr_xchg_proto' was not declared. Should it be static?

This caused by missing extern definition in the header. Add it to
suppress the sparse warning.

Reported-by: kernel test robot <lkp@intel.com>
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 include/linux/bpf.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 551b7198ae8a..d1871eacf728 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -2224,6 +2224,7 @@ extern const struct bpf_func_proto bpf_find_vma_proto;
 extern const struct bpf_func_proto bpf_loop_proto;
 extern const struct bpf_func_proto bpf_strncmp_proto;
 extern const struct bpf_func_proto bpf_copy_from_user_task_proto;
+extern const struct bpf_func_proto bpf_kptr_xchg_proto;
 
 const struct bpf_func_proto *tracing_prog_func_proto(
   enum bpf_func_id func_id, const struct bpf_prog *prog);
-- 
2.35.1

[PATCH bpf-next v1 2/4] bpf: Prepare prog_test_struct kfuncs for runtime tests

[Kumar Kartikeya Dwivedi]

In an effort to actually test the refcounting logic at runtime, add a
refcount_t member to prog_test_ref_kfunc and use it in selftests to
verify and test the whole logic more exhaustively.

To ensure reading the count to verify it remains stable, make
prog_test_ref_kfunc a per-CPU variable, so that inside a BPF program the
count can be read reliably based on number of acquisitions made. Then,
pairing them with releases and reading from the global per-CPU variable
will allow verifying whether release operations put the refcount.

We don't actually rely on preemption being disabled, but migration must
be disabled, as BPF program is the only context where these acquire and
release functions are called. As such, when an object is acquired on one
CPU, only that CPU should manipulate its refcount. Likewise, for
bpf_this_cpu_ptr, the returned pointer and acquired pointer must match,
so that refcount can actually be read and verified.

The kfunc calls for prog_test_member do not require runtime refcounting,
as they are only used for verifier selftests, not during runtime
execution. Hence, their implementation now has a WARN_ON_ONCE as it is
not meant to be reachable code at runtime. It is strictly used in tests
triggering failure cases in the verifier. bpf_kfunc_call_memb_release is
called from map free path, since prog_test_member is embedded in map
value for some verifier tests, so we skip WARN_ON_ONCE for it.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 net/bpf/test_run.c                            | 32 +++++++++++++------
 .../testing/selftests/bpf/verifier/map_kptr.c |  4 +--
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/net/bpf/test_run.c b/net/bpf/test_run.c
index 7a1579c91432..16d489b03000 100644
--- a/net/bpf/test_run.c
+++ b/net/bpf/test_run.c
@@ -564,31 +564,39 @@ struct prog_test_ref_kfunc {
 	int b;
 	struct prog_test_member memb;
 	struct prog_test_ref_kfunc *next;
+	refcount_t cnt;
 };
 
-static struct prog_test_ref_kfunc prog_test_struct = {
+DEFINE_PER_CPU(struct prog_test_ref_kfunc, prog_test_struct) = {
 	.a = 42,
 	.b = 108,
-	.next = &prog_test_struct,
+	.cnt = REFCOUNT_INIT(1),
 };
 
 noinline struct prog_test_ref_kfunc *
 bpf_kfunc_call_test_acquire(unsigned long *scalar_ptr)
 {
-	/* randomly return NULL */
-	if (get_jiffies_64() % 2)
-		return NULL;
-	return &prog_test_struct;
+	struct prog_test_ref_kfunc *p;
+
+	p = this_cpu_ptr(&prog_test_struct);
+	p->next = p;
+	refcount_inc(&p->cnt);
+	return p;
 }
 
 noinline struct prog_test_member *
 bpf_kfunc_call_memb_acquire(void)
 {
-	return &prog_test_struct.memb;
+	WARN_ON_ONCE(1);
+	return NULL;
 }
 
 noinline void bpf_kfunc_call_test_release(struct prog_test_ref_kfunc *p)
 {
+	if (!p)
+		return;
+
+	refcount_dec(&p->cnt);
 }
 
 noinline void bpf_kfunc_call_memb_release(struct prog_test_member *p)
@@ -597,12 +605,18 @@ noinline void bpf_kfunc_call_memb_release(struct prog_test_member *p)
 
 noinline void bpf_kfunc_call_memb1_release(struct prog_test_member1 *p)
 {
+	WARN_ON_ONCE(1);
 }
 
 noinline struct prog_test_ref_kfunc *
-bpf_kfunc_call_test_kptr_get(struct prog_test_ref_kfunc **p, int a, int b)
+bpf_kfunc_call_test_kptr_get(struct prog_test_ref_kfunc **pp, int a, int b)
 {
-	return &prog_test_struct;
+	struct prog_test_ref_kfunc *p = READ_ONCE(*pp);
+
+	if (!p)
+		return NULL;
+	refcount_inc(&p->cnt);
+	return p;
 }
 
 struct prog_test_pass1 {
diff --git a/tools/testing/selftests/bpf/verifier/map_kptr.c b/tools/testing/selftests/bpf/verifier/map_kptr.c
index 9113834640e6..6914904344c0 100644
--- a/tools/testing/selftests/bpf/verifier/map_kptr.c
+++ b/tools/testing/selftests/bpf/verifier/map_kptr.c
@@ -212,13 +212,13 @@
 	BPF_LDX_MEM(BPF_DW, BPF_REG_0, BPF_REG_0, 0),
 	BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 1),
 	BPF_EXIT_INSN(),
-	BPF_LDX_MEM(BPF_DW, BPF_REG_0, BPF_REG_0, 24),
+	BPF_LDX_MEM(BPF_DW, BPF_REG_0, BPF_REG_0, 32),
 	BPF_EXIT_INSN(),
 	},
 	.prog_type = BPF_PROG_TYPE_SCHED_CLS,
 	.fixup_map_kptr = { 1 },
 	.result = REJECT,
-	.errstr = "access beyond struct prog_test_ref_kfunc at off 24 size 8",
+	.errstr = "access beyond struct prog_test_ref_kfunc at off 32 size 8",
 },
 {
 	"map_kptr: unref: inherit PTR_UNTRUSTED on struct walk",
-- 
2.35.1

Re: [PATCH bpf-next v1 2/4] bpf: Prepare prog_test_struct kfuncs for runtime tests

[Alexei Starovoitov]

On Tue, May 10, 2022 at 2:17 PM Kumar Kartikeya Dwivedi
<memxor@gmail.com> wrote:
>
> In an effort to actually test the refcounting logic at runtime, add a
> refcount_t member to prog_test_ref_kfunc and use it in selftests to
> verify and test the whole logic more exhaustively.
>
> To ensure reading the count to verify it remains stable, make
> prog_test_ref_kfunc a per-CPU variable, so that inside a BPF program the
> count can be read reliably based on number of acquisitions made. Then,
> pairing them with releases and reading from the global per-CPU variable
> will allow verifying whether release operations put the refcount.

The patches look good, but the per-cpu part is a puzzle.
The test is not parallel. Everything looks sequential
and there are no races.
It seems to me if it was
static struct prog_test_ref_kfunc prog_test_struct = {..};
and none of [bpf_]this_cpu_ptr()
the test would work the same way.
What am I missing?

Re: [PATCH bpf-next v1 2/4] bpf: Prepare prog_test_struct kfuncs for runtime tests

[Kumar Kartikeya Dwivedi]

On Wed, May 11, 2022 at 10:07:35AM IST, Alexei Starovoitov wrote:
> On Tue, May 10, 2022 at 2:17 PM Kumar Kartikeya Dwivedi
> <memxor@gmail.com> wrote:
> >
> > In an effort to actually test the refcounting logic at runtime, add a
> > refcount_t member to prog_test_ref_kfunc and use it in selftests to
> > verify and test the whole logic more exhaustively.
> >
> > To ensure reading the count to verify it remains stable, make
> > prog_test_ref_kfunc a per-CPU variable, so that inside a BPF program the
> > count can be read reliably based on number of acquisitions made. Then,
> > pairing them with releases and reading from the global per-CPU variable
> > will allow verifying whether release operations put the refcount.
>
> The patches look good, but the per-cpu part is a puzzle.
> The test is not parallel. Everything looks sequential
> and there are no races.
> It seems to me if it was
> static struct prog_test_ref_kfunc prog_test_struct = {..};
> and none of [bpf_]this_cpu_ptr()
> the test would work the same way.
> What am I missing?

You are not missing anything. It would work the same. I just made it per-CPU for
the off chance that someone runs ./test_progs -t map_kptr in parallel on the
same machine. Then one or both might fail, since count won't just be inc/dec by
us, and reading it would produce something other than what we expect.

One other benefit is getting non-ref PTR_TO_BTF_ID to prog_test_struct to
inspect cnt after releasing acquired pointer (using bpf_this_cpu_ptr), but that
can also be done by non-ref kfunc returning a pointer to it.

If you feel it's not worth it, I will drop it in the next version.

--
Kartikeya

Re: [PATCH bpf-next v1 2/4] bpf: Prepare prog_test_struct kfuncs for runtime tests

[Alexei Starovoitov]

On Tue, May 10, 2022 at 11:01 PM Kumar Kartikeya Dwivedi
<memxor@gmail.com> wrote:
>
> On Wed, May 11, 2022 at 10:07:35AM IST, Alexei Starovoitov wrote:
> > On Tue, May 10, 2022 at 2:17 PM Kumar Kartikeya Dwivedi
> > <memxor@gmail.com> wrote:
> > >
> > > In an effort to actually test the refcounting logic at runtime, add a
> > > refcount_t member to prog_test_ref_kfunc and use it in selftests to
> > > verify and test the whole logic more exhaustively.
> > >
> > > To ensure reading the count to verify it remains stable, make
> > > prog_test_ref_kfunc a per-CPU variable, so that inside a BPF program the
> > > count can be read reliably based on number of acquisitions made. Then,
> > > pairing them with releases and reading from the global per-CPU variable
> > > will allow verifying whether release operations put the refcount.
> >
> > The patches look good, but the per-cpu part is a puzzle.
> > The test is not parallel. Everything looks sequential
> > and there are no races.
> > It seems to me if it was
> > static struct prog_test_ref_kfunc prog_test_struct = {..};
> > and none of [bpf_]this_cpu_ptr()
> > the test would work the same way.
> > What am I missing?
>
> You are not missing anything. It would work the same. I just made it per-CPU for
> the off chance that someone runs ./test_progs -t map_kptr in parallel on the
> same machine. Then one or both might fail, since count won't just be inc/dec by
> us, and reading it would produce something other than what we expect.

I see. You should have mentioned that in the commit log.
But how per-cpu helps in this case?
prog_run is executed with cpu=0, so both test_progs -t map_kptr
will collide on the same cpu.
At the end it's the same. one or both might fail?

In general all serial_ tests in test_progs will fail in
parallel run.
Even non-serial tests might fail.
The non-serial tests are ok for test_progs -j.
They're parallel between themselves, but there are no guarantees
that every individual test can be run parallel with itself.
Majority will probably be fine, but not all.

> One other benefit is getting non-ref PTR_TO_BTF_ID to prog_test_struct to
> inspect cnt after releasing acquired pointer (using bpf_this_cpu_ptr), but that
> can also be done by non-ref kfunc returning a pointer to it.

Not following. non-ref == ptr_untrusted. That doesn't preclude
bpf prog from reading refcnt directly, but disallows passing
into helpers.
So with non-percpu the following hack
 bpf_kfunc_call_test_release(p);
 if (p_cpu->cnt.refs.counter ...)
wouldn't be necessary.
The prog could release(p) and read p->cnt.refs.counter right after.
While with per-cpu approach you had to do
p_cpu = bpf_this_cpu_ptr(&prog_test_struct);
hack and rely on intimate knowledge of the kernel side.

> If you feel it's not worth it, I will drop it in the next version.

So far I see the downsides.
Also check the CI. test_progs-no_alu32 fails:
test_map_kptr_fail_prog:PASS:map_kptr__load must fail 0 nsec
test_map_kptr_fail_prog:FAIL:expected error message unexpected error: -22
Expected: Unreleased reference id=5 alloc_insn=18
Verifier: 0: R1=ctx(off=0,imm=0) R10=fp0

Re: [PATCH bpf-next v1 2/4] bpf: Prepare prog_test_struct kfuncs for runtime tests

[Kumar Kartikeya Dwivedi]

On Wed, May 11, 2022 at 11:23:59PM IST, Alexei Starovoitov wrote:
> On Tue, May 10, 2022 at 11:01 PM Kumar Kartikeya Dwivedi
> <memxor@gmail.com> wrote:
> >
> > On Wed, May 11, 2022 at 10:07:35AM IST, Alexei Starovoitov wrote:
> > > On Tue, May 10, 2022 at 2:17 PM Kumar Kartikeya Dwivedi
> > > <memxor@gmail.com> wrote:
> > > >
> > > > In an effort to actually test the refcounting logic at runtime, add a
> > > > refcount_t member to prog_test_ref_kfunc and use it in selftests to
> > > > verify and test the whole logic more exhaustively.
> > > >
> > > > To ensure reading the count to verify it remains stable, make
> > > > prog_test_ref_kfunc a per-CPU variable, so that inside a BPF program the
> > > > count can be read reliably based on number of acquisitions made. Then,
> > > > pairing them with releases and reading from the global per-CPU variable
> > > > will allow verifying whether release operations put the refcount.
> > >
> > > The patches look good, but the per-cpu part is a puzzle.
> > > The test is not parallel. Everything looks sequential
> > > and there are no races.
> > > It seems to me if it was
> > > static struct prog_test_ref_kfunc prog_test_struct = {..};
> > > and none of [bpf_]this_cpu_ptr()
> > > the test would work the same way.
> > > What am I missing?
> >
> > You are not missing anything. It would work the same. I just made it per-CPU for
> > the off chance that someone runs ./test_progs -t map_kptr in parallel on the
> > same machine. Then one or both might fail, since count won't just be inc/dec by
> > us, and reading it would produce something other than what we expect.
>
> I see. You should have mentioned that in the commit log.
> But how per-cpu helps in this case?
> prog_run is executed with cpu=0, so both test_progs -t map_kptr
> will collide on the same cpu.

Right, I was thinking bpf_prog_run disabled preemption, so that would prevent
collisions, but it seems my knowledge is now outdated (only migration is
disabled). Also, just realising, we rely on observing a specific count across
test_run invocations, which won't be protected against for parallel runs
anyway.

> At the end it's the same. one or both might fail?
>
> In general all serial_ tests in test_progs will fail in
> parallel run.
> Even non-serial tests might fail.
> The non-serial tests are ok for test_progs -j.
> They're parallel between themselves, but there are no guarantees
> that every individual test can be run parallel with itself.
> Majority will probably be fine, but not all.
>

I'll drop it and go with a global struct.

> > One other benefit is getting non-ref PTR_TO_BTF_ID to prog_test_struct to
> > inspect cnt after releasing acquired pointer (using bpf_this_cpu_ptr), but that
> > can also be done by non-ref kfunc returning a pointer to it.
>
> Not following. non-ref == ptr_untrusted. That doesn't preclude

By non-ref PTR_TO_BTF_ID I meant normal (not untrusted) PTR_TO_BTF_ID with
ref_obj_id = 0.

bpf_this_cpu_ptr returns a normal PTR_TO_BTF_ID, not an untrusted one.

> bpf prog from reading refcnt directly, but disallows passing
> into helpers.
> So with non-percpu the following hack
>  bpf_kfunc_call_test_release(p);
>  if (p_cpu->cnt.refs.counter ...)
> wouldn't be necessary.
> The prog could release(p) and read p->cnt.refs.counter right after.

release(p) will kill p, so that won't work. I have a better idea, since
p->next points to itself, just loading that will give me a pointer I can
read after release(p).

As an aside, do you think we should change the behaviour of killing released
registers and skip it for refcounted PTR_TO_BTF_ID (perhaps mark it as untrusted
pointer instead, with ref_obj_id reset to zero)? So loads are allowed into it,
but passing into the kernel isn't, wdyt?

p = acq();	  // p.type = PTR_TO_BTF_ID, ref_obj_id=X
foo(p);		  // works
bar(p->a + p->b); // works
rel(p);		  // p.type = PTR_TO_BTF_ID | PTR_UNTRUSTED, ref_obj_id=0
		  // Instead of mark_reg_unknown(p)

There is still the case where you can do:
p2 = p->next;
rel(p);
p3 = p->next;

Now p2 is trusted PTR_TO_BTF_ID, while p3 is untrusted, but this is a separate
problem which requires a more general fix, and needs more discussion.

A bit of a digression, but I would like to know what you and other BPF
developers think.

So far my thinking (culminating towards an RFC) is this:

For a refcounted PTR_TO_BTF_ID, it is marked as trusted.

When loading from it, by default all loads yield untrusted pointers, except
those which are specifically marked with some annotation ("bpf_ptr_trust") which
indicates that parent holds reference to member pointer. This is a loose
description to mean that for the lifetime of trusted parent pointer, member
pointer may also be trusted. If lifetime can end (due to release), trusted
member pointers will become untrusted. If it cannot (e.g. function arguments),
it remains valid.

This will use BTF tags.
Known cases in the kernel which are useful and safe can be whitelisted.

Such loads yield trusted pointers linked to refcounted PTR_TO_BTF_ID. Linked
means the source refcounted PTR_TO_BTF_ID owns it.

When releasing PTR_TO_BTF_ID, all registers with same ref_obj_id, and all linked
PTR_TO_BTF_ID are marked as untrusted.

As an example:

struct foo {
	struct bar __ref *br;
	struct baz *bz;
};

struct foo *f = acq(); // f.type = PTR_TO_BTF_ID, ref_obj_id=X
br = f->br;	       // br.type = PTR_TO_BTF_ID, linked_to=X
bz = f->bz;	       // bz.type = PTR_TO_BTF_ID | PTR_UNTRUSTED
rel(f);		       // f.type = PTR_TO_BTF_ID | PTR_UNTRUSTED
		       // and since br.linked_to == f.ref_obj_id,
		       // br.type = PTR_TO_BTF_ID | PTR_UNTRUSTED

For trusted loads from br, linked_to will be same as X, so they will also be
marked as untrusted, and so on.

For tp_btf/LSM programs, pointer arguments will be non-refcounted trusted
PTR_TO_BTF_ID. All rules as above apply, but since it cannot be released,
trusted pointers obtained from them remain valid till BPF_EXIT.

I have no idea how much backwards compat this will break, or how much of it can
be tolerated.

> While with per-cpu approach you had to do
> p_cpu = bpf_this_cpu_ptr(&prog_test_struct);
> hack and rely on intimate knowledge of the kernel side.
>
> > If you feel it's not worth it, I will drop it in the next version.
>
> So far I see the downsides.
> Also check the CI. test_progs-no_alu32 fails:
> test_map_kptr_fail_prog:PASS:map_kptr__load must fail 0 nsec
> test_map_kptr_fail_prog:FAIL:expected error message unexpected error: -22
> Expected: Unreleased reference id=5 alloc_insn=18
> Verifier: 0: R1=ctx(off=0,imm=0) R10=fp0

Yes, noticed that. It is because alloc_insn= is different for no_alu32. I'll
drop the matching on specific insn idx number.

Thanks for your feedback.

--
Kartikeya

Re: [PATCH bpf-next v1 2/4] bpf: Prepare prog_test_struct kfuncs for runtime tests

[Alexei Starovoitov]

On Thu, May 12, 2022 at 12:37:24AM +0530, Kumar Kartikeya Dwivedi wrote:
> On Wed, May 11, 2022 at 11:23:59PM IST, Alexei Starovoitov wrote:
> > On Tue, May 10, 2022 at 11:01 PM Kumar Kartikeya Dwivedi
> > <memxor@gmail.com> wrote:
> > >
> > > On Wed, May 11, 2022 at 10:07:35AM IST, Alexei Starovoitov wrote:
> > > > On Tue, May 10, 2022 at 2:17 PM Kumar Kartikeya Dwivedi
> > > > <memxor@gmail.com> wrote:
> > > > >
> > > > > In an effort to actually test the refcounting logic at runtime, add a
> > > > > refcount_t member to prog_test_ref_kfunc and use it in selftests to
> > > > > verify and test the whole logic more exhaustively.
> > > > >
> > > > > To ensure reading the count to verify it remains stable, make
> > > > > prog_test_ref_kfunc a per-CPU variable, so that inside a BPF program the
> > > > > count can be read reliably based on number of acquisitions made. Then,
> > > > > pairing them with releases and reading from the global per-CPU variable
> > > > > will allow verifying whether release operations put the refcount.
> > > >
> > > > The patches look good, but the per-cpu part is a puzzle.
> > > > The test is not parallel. Everything looks sequential
> > > > and there are no races.
> > > > It seems to me if it was
> > > > static struct prog_test_ref_kfunc prog_test_struct = {..};
> > > > and none of [bpf_]this_cpu_ptr()
> > > > the test would work the same way.
> > > > What am I missing?
> > >
> > > You are not missing anything. It would work the same. I just made it per-CPU for
> > > the off chance that someone runs ./test_progs -t map_kptr in parallel on the
> > > same machine. Then one or both might fail, since count won't just be inc/dec by
> > > us, and reading it would produce something other than what we expect.
> >
> > I see. You should have mentioned that in the commit log.
> > But how per-cpu helps in this case?
> > prog_run is executed with cpu=0, so both test_progs -t map_kptr
> > will collide on the same cpu.
> 
> Right, I was thinking bpf_prog_run disabled preemption, so that would prevent
> collisions, but it seems my knowledge is now outdated (only migration is
> disabled). Also, just realising, we rely on observing a specific count across
> test_run invocations, which won't be protected against for parallel runs
> anyway.
> 
> > At the end it's the same. one or both might fail?
> >
> > In general all serial_ tests in test_progs will fail in
> > parallel run.
> > Even non-serial tests might fail.
> > The non-serial tests are ok for test_progs -j.
> > They're parallel between themselves, but there are no guarantees
> > that every individual test can be run parallel with itself.
> > Majority will probably be fine, but not all.
> >
> 
> I'll drop it and go with a global struct.
> 
> > > One other benefit is getting non-ref PTR_TO_BTF_ID to prog_test_struct to
> > > inspect cnt after releasing acquired pointer (using bpf_this_cpu_ptr), but that
> > > can also be done by non-ref kfunc returning a pointer to it.
> >
> > Not following. non-ref == ptr_untrusted. That doesn't preclude
> 
> By non-ref PTR_TO_BTF_ID I meant normal (not untrusted) PTR_TO_BTF_ID with
> ref_obj_id = 0.
> 
> bpf_this_cpu_ptr returns a normal PTR_TO_BTF_ID, not an untrusted one.
> 
> > bpf prog from reading refcnt directly, but disallows passing
> > into helpers.
> > So with non-percpu the following hack
> >  bpf_kfunc_call_test_release(p);
> >  if (p_cpu->cnt.refs.counter ...)
> > wouldn't be necessary.
> > The prog could release(p) and read p->cnt.refs.counter right after.
> 
> release(p) will kill p, so that won't work. I have a better idea, since
> p->next points to itself, just loading that will give me a pointer I can
> read after release(p).
> 
> As an aside, do you think we should change the behaviour of killing released
> registers and skip it for refcounted PTR_TO_BTF_ID (perhaps mark it as untrusted
> pointer instead, with ref_obj_id reset to zero)? So loads are allowed into it,
> but passing into the kernel isn't, wdyt?
> 
> p = acq();	  // p.type = PTR_TO_BTF_ID, ref_obj_id=X
> foo(p);		  // works
> bar(p->a + p->b); // works
> rel(p);		  // p.type = PTR_TO_BTF_ID | PTR_UNTRUSTED, ref_obj_id=0
> 		  // Instead of mark_reg_unknown(p)
> 
> There is still the case where you can do:
> p2 = p->next;
> rel(p);
> p3 = p->next;

It's probably better to keep existing behavior since acquire/release is mainly
used in the networking context today.
Probe reading a socket after release is technically safe, but right now
such usage will be rejected by the verifier and the user will have to
fix such bug. If we relax it such bugs might be much harder to spot.

> Now p2 is trusted PTR_TO_BTF_ID, while p3 is untrusted, but this is a separate
> problem which requires a more general fix, and needs more discussion.

It might not be a bug. p3 might still be trusted and valid pointer.
rel(p) releases p only.
The verifier doesn't know semantics.
It's a general link list walking issue.
It needs separate discussion.

> A bit of a digression, but I would like to know what you and other BPF
> developers think.
> 
> So far my thinking (culminating towards an RFC) is this:
> 
> For a refcounted PTR_TO_BTF_ID, it is marked as trusted.
> 
> When loading from it, by default all loads yield untrusted pointers, except
> those which are specifically marked with some annotation ("bpf_ptr_trust") which
> indicates that parent holds reference to member pointer. This is a loose
> description to mean that for the lifetime of trusted parent pointer, member
> pointer may also be trusted. If lifetime can end (due to release), trusted
> member pointers will become untrusted. If it cannot (e.g. function arguments),
> it remains valid.
> 
> This will use BTF tags.
> Known cases in the kernel which are useful and safe can be whitelisted.
> 
> Such loads yield trusted pointers linked to refcounted PTR_TO_BTF_ID. Linked
> means the source refcounted PTR_TO_BTF_ID owns it.
> 
> When releasing PTR_TO_BTF_ID, all registers with same ref_obj_id, and all linked
> PTR_TO_BTF_ID are marked as untrusted.
> 
> As an example:
> 
> struct foo {
> 	struct bar __ref *br;
> 	struct baz *bz;
> };
> 
> struct foo *f = acq(); // f.type = PTR_TO_BTF_ID, ref_obj_id=X
> br = f->br;	       // br.type = PTR_TO_BTF_ID, linked_to=X
> bz = f->bz;	       // bz.type = PTR_TO_BTF_ID | PTR_UNTRUSTED
> rel(f);		       // f.type = PTR_TO_BTF_ID | PTR_UNTRUSTED
> 		       // and since br.linked_to == f.ref_obj_id,
> 		       // br.type = PTR_TO_BTF_ID | PTR_UNTRUSTED
> 
> For trusted loads from br, linked_to will be same as X, so they will also be
> marked as untrusted, and so on.
> 
> For tp_btf/LSM programs, pointer arguments will be non-refcounted trusted
> PTR_TO_BTF_ID. All rules as above apply, but since it cannot be released,
> trusted pointers obtained from them remain valid till BPF_EXIT.
> 
> I have no idea how much backwards compat this will break, or how much of it can
> be tolerated.

Exactly. It's not practical to mark all such fields in the kernel with __ref tags.
bpf_lsm is already in the wild and is using multi level pointer walks
and then passes them into helpers. We cannot break them.
In other words if you try really hard you can crash the kernel.
bpf tracing is only 99.99% safe.

I'll start a separate thread about link lists in bpf.

[PATCH bpf-next v1 3/4] selftests/bpf: Add negative C tests for kptrs

[Kumar Kartikeya Dwivedi]

This uses the newly added SEC("?foo") naming to disable autoload of
programs, and then loads them one by one for the object and verifies
that loading fails and matches the returned error string from verifier.
This is similar to already existing verifier tests but provides coverage
for BPF C.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/prog_tests/map_kptr.c       |  87 +++-
 .../selftests/bpf/progs/map_kptr_fail.c       | 419 ++++++++++++++++++
 2 files changed, 505 insertions(+), 1 deletion(-)
 create mode 100644 tools/testing/selftests/bpf/progs/map_kptr_fail.c

diff --git a/tools/testing/selftests/bpf/prog_tests/map_kptr.c b/tools/testing/selftests/bpf/prog_tests/map_kptr.c
index 9e2fbda64a65..ffef3a319bac 100644
--- a/tools/testing/selftests/bpf/prog_tests/map_kptr.c
+++ b/tools/testing/selftests/bpf/prog_tests/map_kptr.c
@@ -2,8 +2,86 @@
 #include <test_progs.h>
 
 #include "map_kptr.skel.h"
+#include "map_kptr_fail.skel.h"
 
-void test_map_kptr(void)
+static char log_buf[1024 * 1024];
+
+struct {
+	const char *prog_name;
+	const char *err_msg;
+} map_kptr_fail_tests[] = {
+	{ "size_not_bpf_dw", "kptr access size must be BPF_DW" },
+	{ "non_const_var_off", "kptr access cannot have variable offset" },
+	{ "non_const_var_off_kptr_xchg", "R1 doesn't have constant offset. kptr has to be" },
+	{ "misaligned_access_write", "kptr access misaligned expected=8 off=7" },
+	{ "misaligned_access_read", "kptr access misaligned expected=8 off=1" },
+	{ "reject_var_off_store", "variable untrusted_ptr_ access var_off=(0x0; 0x1e0)" },
+	{ "reject_bad_type_match", "invalid kptr access, R1 type=untrusted_ptr_prog_test_ref_kfunc" },
+	{ "marked_as_untrusted_or_null", "R1 type=untrusted_ptr_or_null_ expected=percpu_ptr_" },
+	{ "correct_btf_id_check_size", "access beyond struct prog_test_ref_kfunc at off 32 size 4" },
+	{ "inherit_untrusted_on_walk", "R1 type=untrusted_ptr_ expected=percpu_ptr_" },
+	{ "reject_kptr_xchg_on_unref", "off=8 kptr isn't referenced kptr" },
+	{ "reject_kptr_get_no_map_val", "arg#0 expected pointer to map value" },
+	{ "reject_kptr_get_no_null_map_val", "arg#0 expected pointer to map value" },
+	{ "reject_kptr_get_no_kptr", "arg#0 no referenced kptr at map value offset=0" },
+	{ "reject_kptr_get_on_unref", "arg#0 no referenced kptr at map value offset=8" },
+	{ "reject_kptr_get_bad_type_match", "kernel function bpf_kfunc_call_test_kptr_get args#0" },
+	{ "mark_ref_as_untrusted_or_null", "R1 type=untrusted_ptr_or_null_ expected=percpu_ptr_" },
+	{ "reject_untrusted_store_to_ref", "store to referenced kptr disallowed" },
+	{ "reject_bad_type_xchg", "invalid kptr access, R2 type=ptr_prog_test_ref_kfunc expected=ptr_prog_test_member" },
+	{ "reject_untrusted_xchg", "R2 type=untrusted_ptr_ expected=ptr_" },
+	{ "reject_member_of_ref_xchg", "invalid kptr access, R2 type=ptr_prog_test_ref_kfunc" },
+	{ "reject_indirect_helper_access", "kptr cannot be accessed indirectly by helper" },
+	{ "reject_indirect_global_func_access", "kptr cannot be accessed indirectly by helper" },
+	{ "kptr_xchg_ref_state", "Unreleased reference id=5 alloc_insn=18" },
+	{ "kptr_get_ref_state", "Unreleased reference id=3 alloc_insn=12" },
+};
+
+static void test_map_kptr_fail_prog(const char *prog_name, const char *err_msg)
+{
+	LIBBPF_OPTS(bpf_object_open_opts, opts, .kernel_log_buf = log_buf,
+						.kernel_log_size = sizeof(log_buf),
+						.kernel_log_level = 1);
+	struct map_kptr_fail *skel;
+	struct bpf_program *prog;
+	int ret;
+
+	skel = map_kptr_fail__open_opts(&opts);
+	if (!ASSERT_OK_PTR(skel, "map_kptr_fail__open_opts"))
+		return;
+
+	prog = bpf_object__find_program_by_name(skel->obj, prog_name);
+	if (!ASSERT_OK_PTR(prog, "bpf_object__find_program_by_name"))
+		goto end;
+
+	bpf_program__set_autoload(prog, true);
+
+	ret = map_kptr_fail__load(skel);
+	if (!ASSERT_ERR(ret, "map_kptr__load must fail"))
+		goto end;
+
+	if (!ASSERT_OK_PTR(strstr(log_buf, err_msg), "expected error message")) {
+		fprintf(stderr, "Expected: %s\n", err_msg);
+		fprintf(stderr, "Verifier: %s\n", log_buf);
+	}
+
+end:
+	map_kptr_fail__destroy(skel);
+}
+
+static void test_map_kptr_fail(void)
+{
+	int i;
+
+	for (i = 0; i < ARRAY_SIZE(map_kptr_fail_tests); i++) {
+		if (!test__start_subtest(map_kptr_fail_tests[i].prog_name))
+			continue;
+		test_map_kptr_fail_prog(map_kptr_fail_tests[i].prog_name,
+					map_kptr_fail_tests[i].err_msg);
+	}
+}
+
+static void test_map_kptr_success(void)
 {
 	struct map_kptr *skel;
 	int key = 0, ret;
@@ -35,3 +113,10 @@ void test_map_kptr(void)
 
 	map_kptr__destroy(skel);
 }
+
+void test_map_kptr(void)
+{
+	if (test__start_subtest("success"))
+		test_map_kptr_success();
+	test_map_kptr_fail();
+}
diff --git a/tools/testing/selftests/bpf/progs/map_kptr_fail.c b/tools/testing/selftests/bpf/progs/map_kptr_fail.c
new file mode 100644
index 000000000000..c4ef48224362
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/map_kptr_fail.c
@@ -0,0 +1,419 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <vmlinux.h>
+#include <bpf/bpf_tracing.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_core_read.h>
+
+struct map_value {
+	char buf[8];
+	struct prog_test_ref_kfunc __kptr *unref_ptr;
+	struct prog_test_ref_kfunc __kptr_ref *ref_ptr;
+	struct prog_test_member __kptr_ref *ref_memb_ptr;
+};
+
+struct array_map {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__type(key, int);
+	__type(value, struct map_value);
+	__uint(max_entries, 1);
+} array_map SEC(".maps");
+
+extern struct prog_test_ref_kfunc *bpf_kfunc_call_test_acquire(unsigned long *sp) __ksym;
+extern struct prog_test_ref_kfunc *
+bpf_kfunc_call_test_kptr_get(struct prog_test_ref_kfunc **p, int a, int b) __ksym;
+
+SEC("?tc")
+int size_not_bpf_dw(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	*(u32 *)&v->unref_ptr = 0;
+	return 0;
+}
+
+SEC("?tc")
+int non_const_var_off(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0, id;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	id = ctx->protocol;
+	if (id < 4 || id > 12)
+		return 0;
+	*(u64 *)((void *)v + id) = 0;
+
+	return 0;
+}
+
+SEC("?tc")
+int non_const_var_off_kptr_xchg(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0, id;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	id = ctx->protocol;
+	if (id < 4 || id > 12)
+		return 0;
+	bpf_kptr_xchg((void *)v + id, NULL);
+
+	return 0;
+}
+
+SEC("?tc")
+int misaligned_access_write(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	*(void **)((void *)v + 7) = NULL;
+
+	return 0;
+}
+
+SEC("?tc")
+int misaligned_access_read(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	return *(u64 *)((void *)v + 1);
+}
+
+SEC("?tc")
+int reject_var_off_store(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *unref_ptr;
+	struct map_value *v;
+	int key = 0, id;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	unref_ptr = v->unref_ptr;
+	if (!unref_ptr)
+		return 0;
+	id = ctx->protocol;
+	if (id < 4 || id > 12)
+		return 0;
+	unref_ptr += id;
+	v->unref_ptr = unref_ptr;
+
+	return 0;
+}
+
+SEC("?tc")
+int reject_bad_type_match(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *unref_ptr;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	unref_ptr = v->unref_ptr;
+	if (!unref_ptr)
+		return 0;
+	unref_ptr = (void *)unref_ptr + 4;
+	v->unref_ptr = unref_ptr;
+
+	return 0;
+}
+
+SEC("?tc")
+int marked_as_untrusted_or_null(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_this_cpu_ptr(v->unref_ptr);
+	return 0;
+}
+
+SEC("?tc")
+int correct_btf_id_check_size(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *p;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	p = v->unref_ptr;
+	if (!p)
+		return 0;
+	return *(int *)((void *)p + bpf_core_type_size(struct prog_test_ref_kfunc));
+}
+
+SEC("?tc")
+int inherit_untrusted_on_walk(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *unref_ptr;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	unref_ptr = v->unref_ptr;
+	if (!unref_ptr)
+		return 0;
+	unref_ptr = unref_ptr->next;
+	bpf_this_cpu_ptr(unref_ptr);
+	return 0;
+}
+
+SEC("?tc")
+int reject_kptr_xchg_on_unref(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_kptr_xchg(&v->unref_ptr, NULL);
+	return 0;
+}
+
+SEC("?tc")
+int reject_kptr_get_no_map_val(struct __sk_buff *ctx)
+{
+	bpf_kfunc_call_test_kptr_get((void *)&ctx, 0, 0);
+	return 0;
+}
+
+SEC("?tc")
+int reject_kptr_get_no_null_map_val(struct __sk_buff *ctx)
+{
+	bpf_kfunc_call_test_kptr_get(bpf_map_lookup_elem(&array_map, &(int){0}), 0, 0);
+	return 0;
+}
+
+SEC("?tc")
+int reject_kptr_get_no_kptr(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_kfunc_call_test_kptr_get((void *)v, 0, 0);
+	return 0;
+}
+
+SEC("?tc")
+int reject_kptr_get_on_unref(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_kfunc_call_test_kptr_get(&v->unref_ptr, 0, 0);
+	return 0;
+}
+
+SEC("?tc")
+int reject_kptr_get_bad_type_match(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_kfunc_call_test_kptr_get((void *)&v->ref_memb_ptr, 0, 0);
+	return 0;
+}
+
+SEC("?tc")
+int mark_ref_as_untrusted_or_null(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_this_cpu_ptr(v->ref_ptr);
+	return 0;
+}
+
+SEC("?tc")
+int reject_untrusted_store_to_ref(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *p;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	p = v->ref_ptr;
+	if (!p)
+		return 0;
+	/* Checkmate, clang */
+	*(struct prog_test_ref_kfunc * volatile *)&v->ref_ptr = p;
+	return 0;
+}
+
+SEC("?tc")
+int reject_untrusted_xchg(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *p;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	p = v->ref_ptr;
+	if (!p)
+		return 0;
+	bpf_kptr_xchg(&v->ref_ptr, p);
+	return 0;
+}
+
+SEC("?tc")
+int reject_bad_type_xchg(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *ref_ptr;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	ref_ptr = bpf_kfunc_call_test_acquire(&(unsigned long){0});
+	if (!ref_ptr)
+		return 0;
+	bpf_kptr_xchg(&v->ref_memb_ptr, ref_ptr);
+	return 0;
+}
+
+SEC("?tc")
+int reject_member_of_ref_xchg(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *ref_ptr;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	ref_ptr = bpf_kfunc_call_test_acquire(&(unsigned long){0});
+	if (!ref_ptr)
+		return 0;
+	bpf_kptr_xchg(&v->ref_memb_ptr, &ref_ptr->memb);
+	return 0;
+}
+
+SEC("?syscall")
+int reject_indirect_helper_access(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_get_current_comm(v, sizeof(v->buf) + 1);
+	return 0;
+}
+
+__noinline
+int write_func(int *p)
+{
+
+	return p ? *p = 42 : 0;
+}
+
+SEC("?tc")
+int reject_indirect_global_func_access(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	return write_func((void *)v + 5);
+}
+
+SEC("?tc")
+int kptr_xchg_ref_state(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *p;
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
+	if (!p)
+		return 0;
+	bpf_kptr_xchg(&v->ref_ptr, p);
+	return 0;
+}
+
+SEC("?tc")
+int kptr_get_ref_state(struct __sk_buff *ctx)
+{
+	struct map_value *v;
+	int key = 0;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 0;
+
+	bpf_kfunc_call_test_kptr_get(&v->ref_ptr, 0, 0);
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.35.1

[PATCH bpf-next v1 4/4] selftests/bpf: Add tests for kptr_ref refcounting

[Kumar Kartikeya Dwivedi]

Check at runtime how various operations for kptr_ref affect its refcount
and verify against the actual count. We use the per-CPU prog_test_struct
to get an isolated object to inspect the refcount of.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../selftests/bpf/prog_tests/map_kptr.c       |  28 ++++-
 tools/testing/selftests/bpf/progs/map_kptr.c  | 109 +++++++++++++++++-
 2 files changed, 132 insertions(+), 5 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/map_kptr.c b/tools/testing/selftests/bpf/prog_tests/map_kptr.c
index ffef3a319bac..bcee3e54e3ed 100644
--- a/tools/testing/selftests/bpf/prog_tests/map_kptr.c
+++ b/tools/testing/selftests/bpf/prog_tests/map_kptr.c
@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0
 #include <test_progs.h>
+#include <network_helpers.h>
 
 #include "map_kptr.skel.h"
 #include "map_kptr_fail.skel.h"
@@ -81,8 +82,14 @@ static void test_map_kptr_fail(void)
 	}
 }
 
-static void test_map_kptr_success(void)
+static void test_map_kptr_success(bool test_run)
 {
+	LIBBPF_OPTS(bpf_test_run_opts, opts,
+		.data_in = &pkt_v4,
+		.data_size_in = sizeof(pkt_v4),
+		.repeat = 1,
+		.cpu = 0,
+	);
 	struct map_kptr *skel;
 	int key = 0, ret;
 	char buf[24];
@@ -91,6 +98,16 @@ static void test_map_kptr_success(void)
 	if (!ASSERT_OK_PTR(skel, "map_kptr__open_and_load"))
 		return;
 
+	ret = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.test_map_kptr_ref), &opts);
+	ASSERT_OK(ret, "test_map_kptr_ref refcount");
+	ASSERT_OK(opts.retval, "test_map_kptr_ref retval");
+	ret = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.test_map_kptr_ref2), &opts);
+	ASSERT_OK(ret, "test_map_kptr_ref2 refcount");
+	ASSERT_OK(opts.retval, "test_map_kptr_ref2 retval");
+
+	if (test_run)
+		return;
+
 	ret = bpf_map_update_elem(bpf_map__fd(skel->maps.array_map), &key, buf, 0);
 	ASSERT_OK(ret, "array_map update");
 	ret = bpf_map_update_elem(bpf_map__fd(skel->maps.array_map), &key, buf, 0);
@@ -116,7 +133,12 @@ static void test_map_kptr_success(void)
 
 void test_map_kptr(void)
 {
-	if (test__start_subtest("success"))
-		test_map_kptr_success();
+	if (test__start_subtest("success")) {
+		test_map_kptr_success(false);
+		/* Do test_run twice, so that we see refcount going back to 1
+		 * after we leave it in map from first iteration.
+		 */
+		test_map_kptr_success(true);
+	}
 	test_map_kptr_fail();
 }
diff --git a/tools/testing/selftests/bpf/progs/map_kptr.c b/tools/testing/selftests/bpf/progs/map_kptr.c
index 1b0e0409eaa5..569d7522bb9f 100644
--- a/tools/testing/selftests/bpf/progs/map_kptr.c
+++ b/tools/testing/selftests/bpf/progs/map_kptr.c
@@ -61,6 +61,7 @@ extern struct prog_test_ref_kfunc *bpf_kfunc_call_test_acquire(unsigned long *sp
 extern struct prog_test_ref_kfunc *
 bpf_kfunc_call_test_kptr_get(struct prog_test_ref_kfunc **p, int a, int b) __ksym;
 extern void bpf_kfunc_call_test_release(struct prog_test_ref_kfunc *p) __ksym;
+extern struct prog_test_ref_kfunc prog_test_struct __ksym;
 
 static void test_kptr_unref(struct map_value *v)
 {
@@ -141,7 +142,7 @@ SEC("tc")
 int test_map_kptr(struct __sk_buff *ctx)
 {
 	struct map_value *v;
-	int i, key = 0;
+	int key = 0;
 
 #define TEST(map)					\
 	v = bpf_map_lookup_elem(&map, &key);		\
@@ -162,7 +163,7 @@ SEC("tc")
 int test_map_in_map_kptr(struct __sk_buff *ctx)
 {
 	struct map_value *v;
-	int i, key = 0;
+	int key = 0;
 	void *map;
 
 #define TEST(map_in_map)                                \
@@ -187,4 +188,108 @@ int test_map_in_map_kptr(struct __sk_buff *ctx)
 	return 0;
 }
 
+SEC("tc")
+int test_map_kptr_ref(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *volatile p, *p_cpu;
+	unsigned long arg = 0;
+	struct map_value *v;
+	int key = 0, ret;
+
+	p_cpu = bpf_this_cpu_ptr(&prog_test_struct);
+	if (p_cpu->cnt.refs.counter != 1)
+		return 1;
+
+	p = bpf_kfunc_call_test_acquire(&arg);
+	if (!p)
+		return 2;
+	if (p != p_cpu || p_cpu->cnt.refs.counter != 2) {
+		ret = 3;
+		goto end;
+	}
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v) {
+		ret = 4;
+		goto end;
+	}
+
+	p = bpf_kptr_xchg(&v->ref_ptr, p);
+	if (p) {
+		ret = 5;
+		goto end;
+	}
+	if (p_cpu->cnt.refs.counter != 2)
+		return 6;
+
+	p = bpf_kfunc_call_test_kptr_get(&v->ref_ptr, 0, 0);
+	if (!p)
+		return 7;
+	if (p_cpu->cnt.refs.counter != 3) {
+		ret = 8;
+		goto end;
+	}
+	bpf_kfunc_call_test_release(p);
+	if (p_cpu->cnt.refs.counter != 2)
+		return 9;
+
+	p = bpf_kptr_xchg(&v->ref_ptr, NULL);
+	if (!p)
+		return 10;
+	bpf_kfunc_call_test_release(p);
+	if (p_cpu->cnt.refs.counter != 1)
+		return 11;
+
+	p = bpf_kfunc_call_test_acquire(&arg);
+	if (!p)
+		return 12;
+	p = bpf_kptr_xchg(&v->ref_ptr, p);
+	if (p) {
+		ret = 13;
+		goto end;
+	}
+	if (p_cpu->cnt.refs.counter != 2)
+		return 14;
+	/* Leave in map */
+
+	return 0;
+end:
+	bpf_kfunc_call_test_release(p);
+	return ret;
+}
+
+SEC("tc")
+int test_map_kptr_ref2(struct __sk_buff *ctx)
+{
+	struct prog_test_ref_kfunc *volatile p, *p_cpu;
+	struct map_value *v;
+	int key = 0;
+
+	p_cpu = bpf_this_cpu_ptr(&prog_test_struct);
+	if (p_cpu->cnt.refs.counter != 2)
+		return 1;
+
+	v = bpf_map_lookup_elem(&array_map, &key);
+	if (!v)
+		return 2;
+
+	p = bpf_kptr_xchg(&v->ref_ptr, NULL);
+	if (!p)
+		return 3;
+	if (p != p_cpu || p_cpu->cnt.refs.counter != 2) {
+		bpf_kfunc_call_test_release(p);
+		return 4;
+	}
+
+	p = bpf_kptr_xchg(&v->ref_ptr, p);
+	if (p) {
+		bpf_kfunc_call_test_release(p);
+		return 5;
+	}
+	if (p_cpu->cnt.refs.counter != 2)
+		return 6;
+
+	return 0;
+}
+
 char _license[] SEC("license") = "GPL";
-- 
2.35.1