[PATCH bpf v3 3/3] selftests/bpf: Test bpf_sock_destroy() on TIME_WAIT and listener socks

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

Add two subtests.

tcp_timewait: the client shuts down first and the server closes after
it, so the client sock ends up in TIME_WAIT. A tcp iterator then finds
the timewait sock by the cookie it inherited from the client sock and
destroys it. Iterate once more to make sure it is gone. Without the
first fix bpf_sock_destroy() reads past the timewait sock and KASAN
complains.

tcp_listen_pending: connect to a listener but never accept, so the child
sits in the accept queue, then destroy the listener. Without the second
fix the cond_resched() in inet_csk_listen_stop() trips the might_sleep
check under rcu_read_lock().

  ./test_progs -a sock_destroy

  #444/1   sock_destroy/tcp_client:OK
  #444/2   sock_destroy/tcp_server:OK
  #444/3   sock_destroy/tcp_listen_pending:OK
  #444/4   sock_destroy/tcp_timewait:OK
  #444/5   sock_destroy/udp_client:OK
  #444/6   sock_destroy/udp_server:OK
  #444/7   sock_destroy/trace_tcp_destroy_sock:OK
  #444     sock_destroy:OK
  Summary: 1/7 PASSED, 0 SKIPPED, 0/0 FAILED

Signed-off-by: Jiayuan Chen <jiayuan.chen@linux.dev>
---
 .../selftests/bpf/prog_tests/sock_destroy.c   | 121 ++++++++++++++++++
 .../selftests/bpf/progs/sock_destroy_prog.c   |  30 +++++
 2 files changed, 151 insertions(+)

diff --git a/tools/testing/selftests/bpf/prog_tests/sock_destroy.c b/tools/testing/selftests/bpf/prog_tests/sock_destroy.c
index 9c11938fe597d..78d642a02bdb6 100644
--- a/tools/testing/selftests/bpf/prog_tests/sock_destroy.c
+++ b/tools/testing/selftests/bpf/prog_tests/sock_destroy.c
@@ -1,4 +1,5 @@
 // SPDX-License-Identifier: GPL-2.0
+#include <poll.h>
 #include <test_progs.h>
 #include <bpf/bpf_endian.h>
 
@@ -110,6 +111,122 @@ static void test_tcp_server(struct sock_destroy_prog *skel)
 		close(serv);
 }
 
+static void test_tcp_listen_pending(struct sock_destroy_prog *skel)
+{
+	int serv = -1, clien = -1, accept_serv = -1, n, serv_port;
+	struct pollfd pfd = { .events = POLLIN };
+	char buf[1];
+
+	serv = start_server(AF_INET6, SOCK_STREAM, NULL, 0, 0);
+	if (!ASSERT_GE(serv, 0, "start_server"))
+		goto cleanup;
+	serv_port = get_socket_local_port(serv);
+	if (!ASSERT_GE(serv_port, 0, "get_sock_local_port"))
+		goto cleanup;
+	skel->bss->serv_port = (__be16)serv_port;
+
+	/*
+	 * Connect but never accept, so the child sits in the accept queue
+	 * of the listener. Wait until it's actually there.
+	 */
+	clien = connect_to_fd(serv, 0);
+	if (!ASSERT_GE(clien, 0, "connect_to_fd"))
+		goto cleanup;
+	pfd.fd = serv;
+	if (!ASSERT_EQ(poll(&pfd, 1, -1), 1, "poll listener"))
+		goto cleanup;
+
+	/* Run iterator program that destroys server sockets. */
+	start_iter_sockets(skel->progs.iter_tcp6_server);
+
+	accept_serv = accept(serv, NULL, NULL);
+	if (!ASSERT_LT(accept_serv, 0, "accept on destroyed listener"))
+		goto cleanup;
+	ASSERT_EQ(errno, EINVAL, "error code on destroyed listener");
+
+	/* The unaccepted child was reset along with the listener. */
+	n = recv(clien, buf, sizeof(buf), 0);
+	if (!ASSERT_LT(n, 0, "client recv on reset child"))
+		goto cleanup;
+	ASSERT_EQ(errno, ECONNRESET, "error code on reset child");
+
+cleanup:
+	if (clien != -1)
+		close(clien);
+	if (accept_serv != -1)
+		close(accept_serv);
+	if (serv != -1)
+		close(serv);
+}
+
+static void test_tcp_timewait(struct sock_destroy_prog *skel)
+{
+	int serv = -1, clien = -1, accept_serv = -1, n;
+	struct timeval tv = {};
+	char buf[1];
+
+	serv = start_server(AF_INET6, SOCK_STREAM, NULL, 0, 0);
+	if (!ASSERT_GE(serv, 0, "start_server"))
+		goto cleanup;
+
+	clien = connect_to_fd(serv, 0);
+	if (!ASSERT_GE(clien, 0, "connect_to_fd"))
+		goto cleanup;
+
+	accept_serv = accept(serv, NULL, NULL);
+	if (!ASSERT_GE(accept_serv, 0, "serv accept"))
+		goto cleanup;
+
+	/*
+	 * Active close from the client, then close the server side. Once
+	 * recv() sees EOF the server FIN has been processed and the client
+	 * sock is in TIME_WAIT. Block without timeout so a loaded CI box
+	 * can't race us.
+	 */
+	if (!ASSERT_OK(setsockopt(clien, SOL_SOCKET, SO_RCVTIMEO, &tv,
+				  sizeof(tv)), "clear rcvtimeo"))
+		goto cleanup;
+	if (!ASSERT_OK(shutdown(clien, SHUT_WR), "client shutdown"))
+		goto cleanup;
+
+	/*
+	 * Make sure the server has seen the client FIN before it closes,
+	 * so the two FINs never cross.
+	 */
+	n = recv(accept_serv, buf, sizeof(buf), 0);
+	if (!ASSERT_EQ(n, 0, "server recv EOF"))
+		goto cleanup;
+
+	close(accept_serv);
+	accept_serv = -1;
+
+	/* block until return EOF */
+	n = recv(clien, buf, sizeof(buf), 0);
+	if (!ASSERT_EQ(n, 0, "client recv EOF"))
+		goto cleanup;
+
+	/* Run iterator program that destroys the timewait client sock. */
+	skel->bss->tw_found = 0;
+	start_iter_sockets(skel->progs.iter_tcp6_timewait);
+	if (!ASSERT_EQ(skel->bss->tw_found, 1, "timewait sock found"))
+		goto cleanup;
+
+	ASSERT_OK(skel->bss->tw_destroy_err, "destroy timewait sock");
+
+	/* The destroyed timewait sock must be gone. */
+	skel->bss->tw_found = 0;
+	start_iter_sockets(skel->progs.iter_tcp6_timewait);
+	ASSERT_EQ(skel->bss->tw_found, 0, "timewait sock destroyed");
+
+cleanup:
+	if (clien != -1)
+		close(clien);
+	if (accept_serv != -1)
+		close(accept_serv);
+	if (serv != -1)
+		close(serv);
+}
+
 static void test_udp_client(struct sock_destroy_prog *skel)
 {
 	int serv = -1, clien = -1, n = 0;
@@ -204,6 +321,10 @@ void test_sock_destroy(void)
 		test_tcp_client(skel);
 	if (test__start_subtest("tcp_server"))
 		test_tcp_server(skel);
+	if (test__start_subtest("tcp_listen_pending"))
+		test_tcp_listen_pending(skel);
+	if (test__start_subtest("tcp_timewait"))
+		test_tcp_timewait(skel);
 	if (test__start_subtest("udp_client"))
 		test_udp_client(skel);
 	if (test__start_subtest("udp_server"))
diff --git a/tools/testing/selftests/bpf/progs/sock_destroy_prog.c b/tools/testing/selftests/bpf/progs/sock_destroy_prog.c
index 9e0bf7a54cec9..0a8887543218a 100644
--- a/tools/testing/selftests/bpf/progs/sock_destroy_prog.c
+++ b/tools/testing/selftests/bpf/progs/sock_destroy_prog.c
@@ -7,6 +7,8 @@
 #include "bpf_tracing_net.h"
 
 __be16 serv_port = 0;
+int tw_found = 0;
+int tw_destroy_err = 0;
 
 int bpf_sock_destroy(struct sock_common *sk) __ksym;
 
@@ -100,6 +102,34 @@ int iter_tcp6_server(struct bpf_iter__tcp *ctx)
 	return 0;
 }
 
+SEC("iter/tcp")
+int iter_tcp6_timewait(struct bpf_iter__tcp *ctx)
+{
+	struct sock_common *sk_common = ctx->sk_common;
+	__u64 *val;
+	int key = 0;
+
+	if (!sk_common)
+		return 0;
+
+	if (sk_common->skc_family != AF_INET6)
+		return 0;
+
+	if (!bpf_skc_to_tcp_timewait_sock(sk_common))
+		return 0;
+
+	val = bpf_map_lookup_elem(&tcp_conn_sockets, &key);
+	if (!val)
+		return 0;
+	/* The timewait sock inherits the cookie of the closed client sock. */
+	if (bpf_get_socket_cookie(sk_common) != *val)
+		return 0;
+
+	tw_found++;
+	tw_destroy_err = bpf_sock_destroy(sk_common);
+
+	return 0;
+}
 
 SEC("iter/udp")
 int iter_udp6_client(struct bpf_iter__udp *ctx)
-- 
2.43.0