[PATCH bpf-next v3 11/11] selftests/bpf: Test that a kfunc gets its __int128 stack argument aligned

[Yonghong Song <yonghong.song@linux.dev>]

Add bpf_kfunc_call_test_i128_arg_align(), which takes an __int128 after
seven u64s, in the stack slot after a hole, and returns the low four
bits of its address. OPTIMIZER_HIDE_VAR() keeps the compiler from taking
the alignment the ABI promises for granted and folding them to 0.

Call it from two subprograms whose frames are 8 bytes apart, so that one
of the calls starts with RSP off by 8 whatever alignment the program
was given, and check that both find the __int128 16-byte aligned.
Without the previous patch, the x86-64 JIT fails it with

	run_subtest:FAIL:1552 Unexpected retval: 2 != 0
	aggregate_arg/aggregate_arg_kfunc_int128_align:FAIL

Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
---
 .../selftests/bpf/progs/aggregate_arg_kfunc.c | 42 +++++++++++++++++++
 .../selftests/bpf/test_kmods/bpf_testmod.c    | 15 +++++++
 .../bpf/test_kmods/bpf_testmod_kfunc.h        |  2 +
 3 files changed, 59 insertions(+)

diff --git a/tools/testing/selftests/bpf/progs/aggregate_arg_kfunc.c b/tools/testing/selftests/bpf/progs/aggregate_arg_kfunc.c
index c87e886423f4..3c3bcb3a02d5 100644
--- a/tools/testing/selftests/bpf/progs/aggregate_arg_kfunc.c
+++ b/tools/testing/selftests/bpf/progs/aggregate_arg_kfunc.c
@@ -222,6 +222,48 @@ int aggregate_arg_kfunc_int128_pad(struct __sk_buff *skb)
 	return 0;
 }
 
+/*
+ * The kfunc returns the low four bits of the address of v, which both
+ * conventions pass in a 16-byte aligned stack slot. A program's frame is
+ * only 8-byte aligned on x86-64, so call it from two frames 8 bytes apart:
+ * one of them starts the call off by 8, whatever the program was given.
+ */
+static __noinline __u64 int128_align_frame1(__u64 a, __u64 b)
+{
+	volatile __u64 pad[1];
+	u128 v = ((u128)a << 64) | b;
+
+	pad[0] = a;
+	return bpf_kfunc_call_test_i128_arg_align(1, 2, 3, 4, 5, 6, pad[0], v);
+}
+
+static __noinline __u64 int128_align_frame2(__u64 a, __u64 b)
+{
+	volatile __u64 pad[2];
+	u128 v = ((u128)a << 64) | b;
+
+	pad[0] = a;
+	pad[1] = b;
+	return bpf_kfunc_call_test_i128_arg_align(1, 2, 3, 4, 5, pad[1], pad[0], v);
+}
+
+SEC("tc")
+__arch_x86_64 __arch_arm64
+__load_if_JITed()
+__success __retval(0)
+int aggregate_arg_kfunc_int128_align(struct __sk_buff *skb)
+{
+	__u64 a = skb->len ^ MIX_A;
+	__u64 b = skb->len ^ MIX_B;
+
+	if (int128_align_frame1(a, b))
+		return 1;
+	if (int128_align_frame2(a, b))
+		return 2;
+
+	return 0;
+}
+
 #endif /* __SIZEOF_INT128__ */
 
 #endif
diff --git a/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c b/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c
index 38edb4addc5b..8550897fdfa6 100644
--- a/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c
+++ b/tools/testing/selftests/bpf/test_kmods/bpf_testmod.c
@@ -1120,6 +1120,20 @@ __bpf_kfunc u64 bpf_kfunc_call_test_i128_arg_pad(u64 a, u64 b, u64 c, u64 d, u64
 	       (u64)v * 8 + (u64)((unsigned __int128)v >> 64) * 9;
 }
 
+/*
+ * The low four bits of the address of v, which follows a stack slot hole: 0 if
+ * the caller passed it in a 16-byte aligned slot, as the x86-64 ABI asks.
+ */
+__bpf_kfunc u64 bpf_kfunc_call_test_i128_arg_align(u64 a, u64 b, u64 c, u64 d, u64 e,
+						   u64 f, u64 g, __int128 v)
+{
+	unsigned long addr = (unsigned long)&v;
+
+	/* or the compiler takes the alignment the ABI promises for granted */
+	OPTIMIZER_HIDE_VAR(addr);
+	return addr & 15;
+}
+
 __bpf_kfunc u64 bpf_kfunc_call_test_pair_arg_nofit(u64 a, u64 b, u64 c, u64 d,
 						   struct prog_test_pair_arg s)
 {
@@ -1840,6 +1854,7 @@ BTF_ID_FLAGS(func, bpf_kfunc_call_test_i128_arg_odd)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_i128_arg_odd_many)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_i128_arg_stack_hole)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_i128_arg_pad)
+BTF_ID_FLAGS(func, bpf_kfunc_call_test_i128_arg_align)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_pair_arg_nofit)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_pair_arg_tail)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test_pair_arg_split8)
diff --git a/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h b/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h
index a6a93c901416..da0c9b0c38a4 100644
--- a/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h
+++ b/tools/testing/selftests/bpf/test_kmods/bpf_testmod_kfunc.h
@@ -249,6 +249,8 @@ __u64 bpf_kfunc_call_test_i128_arg_stack_hole(__u64 a, __u64 b, __u64 c, __u64 d
 					      struct prog_test_i128_arg w) __ksym;
 __u64 bpf_kfunc_call_test_i128_arg_pad(__u64 a, __u64 b, __u64 c, __u64 d, __u64 e,
 				       __u64 f, __u64 g, __int128 v) __ksym;
+__u64 bpf_kfunc_call_test_i128_arg_align(__u64 a, __u64 b, __u64 c, __u64 d, __u64 e,
+					 __u64 f, __u64 g, __int128 v) __ksym;
 #endif
 __u64 bpf_kfunc_call_test_pair_arg_nofit(__u64 a, __u64 b, __u64 c, __u64 d,
 					 struct prog_test_pair_arg s) __ksym;
-- 
2.53.0-Meta