[PATCH bpf-next 3/3] selftests/bpf: Add tests with u8/s16 kfunc return types

[Yonghong Song <yhs@fb.com>]

Add two program tests with kfunc return types u8/s16.
With previous patch, xlated codes looks like below:
  ...
  ; return bpf_kfunc_call_test4((struct sock *)sk, (1 << 16) + 0xff00, (1 << 16) + 0xff);
     5: (bf) r1 = r0
     6: (b4) w2 = 130816
     7: (b4) w3 = 65791
     8: (85) call bpf_kfunc_call_test4#8931696
     9: (67) r0 <<= 48
    10: (c7) r0 s>>= 48
    11: (bc) w6 = w0
  ; }
    12: (bc) w0 = w6
    13: (95) exit
  ...
  ; return bpf_kfunc_call_test5((struct sock *)sk, (1 << 8) + 1, (1 << 8) + 2);
     5: (bf) r1 = r0
     6: (b4) w2 = 257
     7: (b4) w3 = 258
     8: (85) call bpf_kfunc_call_test5#8931712
     9: (67) r0 <<= 56
    10: (77) r0 >>= 56
    11: (bc) w6 = w0
  ; }
    12: (bc) w0 = w6
    13: (95) exit

For return type s16, proper sign extension for the return value is done
for kfunc bpf_kfunc_call_test4(). For return type s8, proper zero
extension for the return value is done for bpf_kfunc_call_test5().

Without the previous patch, the test kfunc_call will fail with
  ...
  test_main:FAIL:test4-retval unexpected test4-retval: actual 196607 != expected 4294967295
  ...
  test_main:FAIL:test5-retval unexpected test5-retval: actual 515 != expected 3

Signed-off-by: Yonghong Song <yhs@fb.com>
---
 net/bpf/test_run.c                            | 12 +++++++
 .../selftests/bpf/prog_tests/kfunc_call.c     | 10 ++++++
 .../selftests/bpf/progs/kfunc_call_test.c     | 32 +++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/net/bpf/test_run.c b/net/bpf/test_run.c
index cbc9cd5058cb..3a17ab4107f5 100644
--- a/net/bpf/test_run.c
+++ b/net/bpf/test_run.c
@@ -551,6 +551,16 @@ struct sock * noinline bpf_kfunc_call_test3(struct sock *sk)
 	return sk;
 }
 
+s16 noinline bpf_kfunc_call_test4(struct sock *sk, u32 a, u32 b)
+{
+	return a + b;
+}
+
+u8 noinline bpf_kfunc_call_test5(struct sock *sk, u32 a, u32 b)
+{
+	return a + b;
+}
+
 struct prog_test_member1 {
 	int a;
 };
@@ -703,6 +713,8 @@ BTF_SET8_START(test_sk_check_kfunc_ids)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test1)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test2)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test3)
+BTF_ID_FLAGS(func, bpf_kfunc_call_test4)
+BTF_ID_FLAGS(func, bpf_kfunc_call_test5)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_acquire, KF_ACQUIRE | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_kfunc_call_memb_acquire, KF_ACQUIRE | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_release, KF_RELEASE)
diff --git a/tools/testing/selftests/bpf/prog_tests/kfunc_call.c b/tools/testing/selftests/bpf/prog_tests/kfunc_call.c
index c00eb974eb85..a355c98080f2 100644
--- a/tools/testing/selftests/bpf/prog_tests/kfunc_call.c
+++ b/tools/testing/selftests/bpf/prog_tests/kfunc_call.c
@@ -35,6 +35,16 @@ static void test_main(void)
 	ASSERT_OK(err, "bpf_prog_test_run(test_ref_btf_id)");
 	ASSERT_EQ(topts.retval, 0, "test_ref_btf_id-retval");
 
+	prog_fd = skel->progs.kfunc_call_test4.prog_fd;
+	err = bpf_prog_test_run_opts(prog_fd, &topts);
+	ASSERT_OK(err, "bpf_prog_test_run(test4)");
+	ASSERT_EQ(topts.retval, 0xffffffff, "test4-retval");
+
+	prog_fd = skel->progs.kfunc_call_test5.prog_fd;
+	err = bpf_prog_test_run_opts(prog_fd, &topts);
+	ASSERT_OK(err, "bpf_prog_test_run(test5)");
+	ASSERT_EQ(topts.retval, 3, "test5-retval");
+
 	kfunc_call_test_lskel__destroy(skel);
 }
 
diff --git a/tools/testing/selftests/bpf/progs/kfunc_call_test.c b/tools/testing/selftests/bpf/progs/kfunc_call_test.c
index 5aecbb9fdc68..0636cb13e574 100644
--- a/tools/testing/selftests/bpf/progs/kfunc_call_test.c
+++ b/tools/testing/selftests/bpf/progs/kfunc_call_test.c
@@ -6,6 +6,8 @@
 extern int bpf_kfunc_call_test2(struct sock *sk, __u32 a, __u32 b) __ksym;
 extern __u64 bpf_kfunc_call_test1(struct sock *sk, __u32 a, __u64 b,
 				  __u32 c, __u64 d) __ksym;
+extern __s16 bpf_kfunc_call_test4(struct sock *sk, __u32 a, __u32 b) __ksym;
+extern __u8 bpf_kfunc_call_test5(struct sock *sk, __u32 a, __u32 b) __ksym;
 
 extern struct prog_test_ref_kfunc *bpf_kfunc_call_test_acquire(unsigned long *sp) __ksym;
 extern void bpf_kfunc_call_test_release(struct prog_test_ref_kfunc *p) __ksym;
@@ -30,6 +32,36 @@ int kfunc_call_test2(struct __sk_buff *skb)
 	return bpf_kfunc_call_test2((struct sock *)sk, 1, 2);
 }
 
+SEC("tc")
+int kfunc_call_test4(struct __sk_buff *skb)
+{
+	struct bpf_sock *sk = skb->sk;
+
+	if (!sk)
+		return -1;
+
+	sk = bpf_sk_fullsock(sk);
+	if (!sk)
+		return -1;
+
+	return bpf_kfunc_call_test4((struct sock *)sk, (1 << 16) + 0xff00, (1 << 16) + 0xff);
+}
+
+SEC("tc")
+int kfunc_call_test5(struct __sk_buff *skb)
+{
+	struct bpf_sock *sk = skb->sk;
+
+	if (!sk)
+		return -1;
+
+	sk = bpf_sk_fullsock(sk);
+	if (!sk)
+		return -1;
+
+	return bpf_kfunc_call_test5((struct sock *)sk, (1 << 8) + 1, (1 << 8) + 2);
+}
+
 SEC("tc")
 int kfunc_call_test1(struct __sk_buff *skb)
 {
-- 
2.30.2