[PATCH bpf v1 2/2] selftests/bpf: Add test for SOCK_OPS_GET_SK with same src/dst register

[Jiayuan Chen <jiayuan.chen@linux.dev>]

Add a selftest that verifies SOCK_OPS_GET_SK() correctly returns NULL
when dst_reg == src_reg and is_fullsock == 0.

The BPF program reads ctx->is_fullsock, then loads ctx->sk using the
same register for source and destination. When is_fullsock == 0, sk
must be NULL. The test triggers this path via a TCP handshake (which
causes TCP_NEW_SYN_RECV state) and checks:
 - null_seen == 1: the is_fullsock==0 path was hit with correct NULL sk
 - bug_detected == 0: sk was never non-NULL when is_fullsock==0

Signed-off-by: Jiayuan Chen <jiayuan.chen@linux.dev>
---
 .../bpf/prog_tests/sock_ops_get_sk.c          | 62 +++++++++++++++++++
 .../selftests/bpf/progs/sock_ops_get_sk.c     | 44 +++++++++++++
 2 files changed, 106 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/sock_ops_get_sk.c
 create mode 100644 tools/testing/selftests/bpf/progs/sock_ops_get_sk.c

diff --git a/tools/testing/selftests/bpf/prog_tests/sock_ops_get_sk.c b/tools/testing/selftests/bpf/prog_tests/sock_ops_get_sk.c
new file mode 100644
index 0000000000000..9eaf97786c1d9
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/sock_ops_get_sk.c
@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <test_progs.h>
+#include "cgroup_helpers.h"
+#include "network_helpers.h"
+#include "sock_ops_get_sk.skel.h"
+
+/*
+ * Test that reading ctx->sk with dst_reg == src_reg in a sock_ops program
+ * correctly returns NULL when is_fullsock == 0 (TCP_NEW_SYN_RECV state).
+ *
+ * The bug was in SOCK_OPS_GET_SK() macro which failed to zero the destination
+ * register when it was the same as the source register and the socket was not
+ * a full socket. This left the ctx pointer in the register, which then passed
+ * the NULL check and could be used as a socket pointer, causing OOB access.
+ */
+void test_sock_ops_get_sk(void)
+{
+	struct sock_ops_get_sk *skel;
+	int cgroup_fd, server_fd, client_fd;
+	int prog_fd, err;
+
+	cgroup_fd = test__join_cgroup("/sock_ops_get_sk");
+	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup"))
+		return;
+
+	skel = sock_ops_get_sk__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "skel_open_load"))
+		goto close_cgroup;
+
+	prog_fd = bpf_program__fd(skel->progs.sock_ops_get_sk_same_reg);
+	err = bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_SOCK_OPS, 0);
+	if (!ASSERT_OK(err, "prog_attach"))
+		goto destroy_skel;
+
+	server_fd = start_server(AF_INET, SOCK_STREAM, NULL, 0, 0);
+	if (!ASSERT_GE(server_fd, 0, "start_server"))
+		goto detach;
+
+	/* Trigger TCP handshake which causes TCP_NEW_SYN_RECV state where
+	 * is_fullsock == 0. With the bug, ctx->sk loaded with same src/dst
+	 * register would return a stale ctx pointer instead of NULL.
+	 */
+	client_fd = connect_to_fd(server_fd, 0);
+	if (!ASSERT_GE(client_fd, 0, "connect_to_fd"))
+		goto close_server;
+
+	close(client_fd);
+
+	/* Verify that the is_fullsock == 0 path was hit and sk was NULL */
+	ASSERT_EQ(skel->bss->null_seen, 1, "null_seen");
+	ASSERT_EQ(skel->bss->bug_detected, 0, "bug_not_detected");
+
+close_server:
+	close(server_fd);
+detach:
+	bpf_prog_detach(cgroup_fd, BPF_CGROUP_SOCK_OPS);
+destroy_skel:
+	sock_ops_get_sk__destroy(skel);
+close_cgroup:
+	close(cgroup_fd);
+}
diff --git a/tools/testing/selftests/bpf/progs/sock_ops_get_sk.c b/tools/testing/selftests/bpf/progs/sock_ops_get_sk.c
new file mode 100644
index 0000000000000..4a75614b21eb5
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/sock_ops_get_sk.c
@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include "bpf_misc.h"
+
+/*
+ * Test that SOCK_OPS_GET_SK() correctly returns NULL when dst_reg == src_reg
+ * and is_fullsock == 0 (e.g., during TCP_NEW_SYN_RECV). Before the fix, the
+ * macro failed to zero the destination register, leaving a stale ctx pointer
+ * that bypassed the NULL check.
+ */
+
+int bug_detected;
+int null_seen;
+
+SEC("sockops")
+__naked void sock_ops_get_sk_same_reg(void)
+{
+	asm volatile (
+		"r7 = *(u32 *)(r1 + %[is_fullsock_off]);"
+		"r1 = *(u64 *)(r1 + %[sk_off]);"
+		"if r7 != 0 goto 2f;"
+		"if r1 == 0 goto 1f;"
+		"r1 = %[bug_detected] ll;"
+		"r2 = 1;"
+		"*(u32 *)(r1 + 0) = r2;"
+		"goto 2f;"
+	"1:"
+		"r1 = %[null_seen] ll;"
+		"r2 = 1;"
+		"*(u32 *)(r1 + 0) = r2;"
+	"2:"
+		"r0 = 1;"
+		"exit;"
+		:
+		: __imm_const(is_fullsock_off, offsetof(struct bpf_sock_ops, is_fullsock)),
+		  __imm_const(sk_off, offsetof(struct bpf_sock_ops, sk)),
+		  __imm_addr(bug_detected),
+		  __imm_addr(null_seen)
+		: __clobber_all);
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.43.0