From: Alexei Starovoitov <ast@kernel.org>

The register of an implicit argument is not live before the call:

  kfunc                                 live
  ------------------------------------  -----
  bpf_stream_vprintk()                  r1-r4
  bpf_kfunc_implicit_arg_legacy()       r1-r2
  bpf_kfunc_implicit_arg_legacy_impl()  r1-r3

The last kfunc has no KF_IMPLICIT_ARGS: the prog passes all three
arguments.

kfunc_implicit_arg_in_a_loop calls bpf_stream_print_stack() 4 times. r2 is
its implicit argument. The test expects that the loop is counted and that
each live register has a formula.

Two more progs run. One jumps to the call, one has may_goto in front of the
call. The kfunc has to get aux in both.

Signed-off-by: Alexei Starovoitov <ast@kernel.org>
---
 .../bpf/progs/compute_live_registers.c        | 13 +++++
 .../selftests/bpf/progs/kfunc_implicit_args.c | 55 +++++++++++++++++++
 .../selftests/bpf/progs/verifier_scev.c       | 33 +++++++++++
 3 files changed, 101 insertions(+)

diff --git a/tools/testing/selftests/bpf/progs/compute_live_registers.c b/tools/testing/selftests/bpf/progs/compute_live_registers.c
index 1ff24bbc7beb..de0496ce7d5e 100644
--- a/tools/testing/selftests/bpf/progs/compute_live_registers.c
+++ b/tools/testing/selftests/bpf/progs/compute_live_registers.c
@@ -538,6 +538,19 @@ void gotox(void)
 
 #endif /* gotox: x86, arm64, riscv64 */
 
+extern int bpf_stream_vprintk(int stream_id, const char *fmt__str, const void *args,
+			      __u32 len__sz) __weak __ksym;
+
+/* The verifier fills in the implicit argument of the kfunc. The prog does not pass r5. */
+SEC("syscall")
+__log_level(2)
+__msg(": .1234..... (85) call bpf_stream_vprintk")
+int kfunc_implicit_arg(void *ctx)
+{
+	bpf_stream_printk(1 /* BPF_STDOUT */, "x");
+	return 0;
+}
+
 /* to retain debug info for BTF generation */
 void kfunc_root(void)
 {
diff --git a/tools/testing/selftests/bpf/progs/kfunc_implicit_args.c b/tools/testing/selftests/bpf/progs/kfunc_implicit_args.c
index 89b6a47e22dd..250fc0c7b355 100644
--- a/tools/testing/selftests/bpf/progs/kfunc_implicit_args.c
+++ b/tools/testing/selftests/bpf/progs/kfunc_implicit_args.c
@@ -27,15 +27,70 @@ int test_kfunc_implicit_arg_impl_illegal(void *ctx)
 }
 
 SEC("syscall")
+__log_level(2)
+__msg(": .12....... (85) call bpf_kfunc_implicit_arg_legacy#")
 __retval(7)
 int test_kfunc_implicit_arg_legacy(void *ctx)
 {
 	return bpf_kfunc_implicit_arg_legacy(3, 4);
 }
 
+/* The prog passes all three arguments of this kfunc. */
 SEC("syscall")
+__log_level(2)
+__msg(": .123...... (85) call bpf_kfunc_implicit_arg_legacy_impl#")
 __retval(11)
 int test_kfunc_implicit_arg_legacy_impl(void *ctx)
 {
 	return bpf_kfunc_implicit_arg_legacy_impl(5, 6, NULL);
 }
+
+/*
+ * The verifier adds 'r2 = aux' in front of the call. A jump to the call has to
+ * execute the new insn: the kfunc returns -EINVAL when aux is NULL.
+ */
+SEC("syscall")
+__retval(5)
+__naked void implicit_arg_jump_to_call(void)
+{
+	asm volatile (
+	"call %[bpf_get_prandom_u32];"
+	"r1 = 5;"
+	"r2 = 0;"
+	"if r0 != 0 goto 1f;"
+	"r2 = 0;"
+"1:"
+	"call %[bpf_kfunc_implicit_arg];"
+	"exit;"
+	:
+	: __imm(bpf_get_prandom_u32),
+	  __imm(bpf_kfunc_implicit_arg)
+	: __clobber_all);
+}
+
+/*
+ * r0 and r1 are live at may_goto, r2 is the first register that is not. The
+ * insns that the verifier adds for may_goto are free to use r2: 'r2 = aux'
+ * comes behind them.
+ */
+SEC("syscall")
+__retval(5)
+__naked void implicit_arg_after_may_goto(void)
+{
+	asm volatile (
+	"r6 = 1;"
+"1:"
+	"r0 = r6;"
+	"r1 = r6;"
+	"may_goto 2f;"
+	"call %[bpf_kfunc_implicit_arg];"
+	"if r0 != r6 goto 2f;"
+	"r6 += 1;"
+	"if r6 < 5 goto 1b;"
+	"r0 = r6;"
+"2:"
+	"exit;"
+	:
+	: __imm(bpf_kfunc_implicit_arg)
+	: __clobber_all);
+}
diff --git a/tools/testing/selftests/bpf/progs/verifier_scev.c b/tools/testing/selftests/bpf/progs/verifier_scev.c
index fa01a616af9d..bede3f48fe38 100644
--- a/tools/testing/selftests/bpf/progs/verifier_scev.c
+++ b/tools/testing/selftests/bpf/progs/verifier_scev.c
@@ -1411,6 +1411,39 @@ __naked void widened_reg_is_not_linked(void)
 	: __clobber_all);
 }
 
+extern int bpf_stream_print_stack(int stream_id) __weak __ksym;
+
+/*
+ * The verifier fills in r2, the implicit argument of the kfunc. The call clobbers r2 and nothing
+ * sets it: if r2 is live at the head of the loop it has no formula.
+ */
+SEC("syscall")
+__success __retval(0)
+__log_level(2)
+__msg("loop header at 1, header_count is 4 ")
+__not_msg("can't widen")
+__naked void kfunc_implicit_arg_in_a_loop(void)
+{
+	asm volatile (
+	"r6 = 0;"
+"1:"
+	"r1 = 2;"
+	"call %[bpf_stream_print_stack];"
+	"r6 += 1;"
+	"if r6 < 4 goto 1b;"
+	"r0 = 0;"
+	"exit;"
+	:
+	: __imm(bpf_stream_print_stack)
+	: __clobber_all);
+}
+
+/* to retain debug info for BTF generation */
+void kfunc_root(void)
+{
+	bpf_stream_print_stack(0);
+}
+
 /* The unknown SCEV of an unbound STACK_MISC slot must not prevent widening. */
 SEC("tracepoint")
 __success
-- 
2.55.0