[PATCH bpf-next 2/3] selftests/bpf: add tests for TCP_BPF_SOCK_OPS_CB_FLAGS

[Alan Maguire <alan.maguire@xxxxxxxxxx>]

Add tests to set/get TCP sockopt TCP_BPF_SOCK_OPS_CB_FLAGS via
bpf_setsockopt() and also add a cgroup/setsockopt program that
catches setsockopt() for this option and uses bpf_setsockopt()
to set it.  The latter allows us to support modifying sockops
cb flags on a per-socket basis via setsockopt() without adding
support into core setsockopt() itself.

Signed-off-by: Alan Maguire <alan.maguire@xxxxxxxxxx>
---
 .../selftests/bpf/prog_tests/setget_sockopt.c | 11 ++++++
 .../selftests/bpf/progs/setget_sockopt.c      | 37 +++++++++++++++++--
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/setget_sockopt.c b/tools/testing/selftests/bpf/prog_tests/setget_sockopt.c
index 7d4a9b3d3722..b9c54217a489 100644
--- a/tools/testing/selftests/bpf/prog_tests/setget_sockopt.c
+++ b/tools/testing/selftests/bpf/prog_tests/setget_sockopt.c
@@ -42,6 +42,7 @@ static int create_netns(void)
 static void test_tcp(int family)
 {
 	struct setget_sockopt__bss *bss = skel->bss;
+	int cb_flags = BPF_SOCK_OPS_STATE_CB_FLAG | BPF_SOCK_OPS_RTO_CB_FLAG;
 	int sfd, cfd;
 
 	memset(bss, 0, sizeof(*bss));
@@ -56,6 +57,9 @@ static void test_tcp(int family)
 		close(sfd);
 		return;
 	}
+	ASSERT_EQ(setsockopt(sfd, SOL_TCP, TCP_BPF_SOCK_OPS_CB_FLAGS,
+			     &cb_flags, sizeof(cb_flags)),
+		  0, "setsockopt cb_flags");
 	close(sfd);
 	close(cfd);
 
@@ -65,6 +69,8 @@ static void test_tcp(int family)
 	ASSERT_EQ(bss->nr_passive, 1, "nr_passive");
 	ASSERT_EQ(bss->nr_socket_post_create, 2, "nr_socket_post_create");
 	ASSERT_EQ(bss->nr_binddev, 2, "nr_bind");
+	ASSERT_GE(bss->nr_state, 1, "nr_state");
+	ASSERT_EQ(bss->nr_setsockopt, 1, "nr_setsockopt");
 }
 
 static void test_udp(int family)
@@ -185,6 +191,11 @@ void test_setget_sockopt(void)
 	if (!ASSERT_OK_PTR(skel->links.socket_post_create, "attach_cgroup"))
 		goto done;
 
+	skel->links.tcp_setsockopt =
+		bpf_program__attach_cgroup(skel->progs.tcp_setsockopt, cg_fd);
+	if (!ASSERT_OK_PTR(skel->links.tcp_setsockopt, "attach_setsockopt"))
+		goto done;
+
 	test_tcp(AF_INET6);
 	test_tcp(AF_INET);
 	test_udp(AF_INET6);
diff --git a/tools/testing/selftests/bpf/progs/setget_sockopt.c b/tools/testing/selftests/bpf/progs/setget_sockopt.c
index 60518aed1ffc..920af9e21e84 100644
--- a/tools/testing/selftests/bpf/progs/setget_sockopt.c
+++ b/tools/testing/selftests/bpf/progs/setget_sockopt.c
@@ -20,6 +20,8 @@ int nr_connect;
 int nr_binddev;
 int nr_socket_post_create;
 int nr_fin_wait1;
+int nr_state;
+int nr_setsockopt;
 
 struct sockopt_test {
 	int opt;
@@ -59,6 +61,8 @@ static const struct sockopt_test sol_tcp_tests[] = {
 	{ .opt = TCP_THIN_LINEAR_TIMEOUTS, .flip = 1, },
 	{ .opt = TCP_USER_TIMEOUT, .new = 123400, .expected = 123400, },
 	{ .opt = TCP_NOTSENT_LOWAT, .new = 1314, .expected = 1314, },
+	{ .opt = TCP_BPF_SOCK_OPS_CB_FLAGS, .new = BPF_SOCK_OPS_ALL_CB_FLAGS,
+	  .expected = BPF_SOCK_OPS_ALL_CB_FLAGS, .restore = BPF_SOCK_OPS_STATE_CB_FLAG, },
 	{ .opt = 0, },
 };
 
@@ -124,6 +128,7 @@ static int bpf_test_sockopt_int(void *ctx, struct sock *sk,
 
 	if (bpf_setsockopt(ctx, level, opt, &new, sizeof(new)))
 		return 1;
+
 	if (bpf_getsockopt(ctx, level, opt, &tmp, sizeof(tmp)) ||
 	    tmp != expected)
 		return 1;
@@ -384,11 +389,14 @@ int skops_sockopt(struct bpf_sock_ops *skops)
 		nr_passive += !(bpf_test_sockopt(skops, sk) ||
 				test_tcp_maxseg(skops, sk) ||
 				test_tcp_saved_syn(skops, sk));
-		bpf_sock_ops_cb_flags_set(skops,
-					  skops->bpf_sock_ops_cb_flags |
-					  BPF_SOCK_OPS_STATE_CB_FLAG);
+
+		/* no need to set sockops cb flags here as sockopt
+		 * tests and user-space originated setsockopt() will
+		 * set flags to include BPF_SOCK_OPS_STATE_CB.
+		 */
 		break;
 	case BPF_SOCK_OPS_STATE_CB:
+		nr_state++;
 		if (skops->args[1] == BPF_TCP_CLOSE_WAIT)
 			nr_fin_wait1 += !bpf_test_sockopt(skops, sk);
 		break;
@@ -397,4 +405,27 @@ int skops_sockopt(struct bpf_sock_ops *skops)
 	return 1;
 }
 
+SEC("cgroup/setsockopt")
+int tcp_setsockopt(struct bpf_sockopt *ctx)
+{
+	struct bpf_sock *sk = ctx->sk;
+	__u8 *optval_end = ctx->optval_end;
+	__u8 *optval = ctx->optval;
+	int val = 0;
+
+	if (!sk || ctx->level != SOL_TCP || ctx->optname != TCP_BPF_SOCK_OPS_CB_FLAGS)
+		return 1;
+	if (optval + sizeof(int) > optval_end)
+		return 0;
+	if (ctx->optlen != sizeof(int))
+		return 0;
+	val = *(int *)optval;
+	if (bpf_setsockopt(sk, ctx->level, ctx->optname, &val, sizeof(val)))
+		return 0;
+	nr_setsockopt++;
+	/* BPF has handled this no need to call "real" setsockopt() */
+	ctx->optlen = -1;
+	return 1;
+}
+
 char _license[] SEC("license") = "GPL";
-- 
2.31.1