[PATCH bpf-next 2/2] bpf: Add tests for bpf_sk_storage to bpf_tcp_ca

[Martin KaFai Lau <kafai@fb.com>]

This patch adds test to exercise the bpf_sk_storage_get()
and bpf_sk_storage_delete() helper from the bpf_dctcp.c.

Signed-off-by: Martin KaFai Lau <kafai@fb.com>
---
 .../selftests/bpf/prog_tests/bpf_tcp_ca.c     | 28 +++++++++++++++++--
 tools/testing/selftests/bpf/progs/bpf_dctcp.c | 16 +++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c b/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c
index 8482bbc67eec..9aaecce0bc3c 100644
--- a/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c
+++ b/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c
@@ -11,6 +11,7 @@
 static const unsigned int total_bytes = 10 * 1024 * 1024;
 static const struct timeval timeo_sec = { .tv_sec = 10 };
 static const size_t timeo_optlen = sizeof(timeo_sec);
+static int expected_stg = 0xeB9F;
 static int stop, duration;
 
 static int settimeo(int fd)
@@ -88,7 +89,7 @@ static void *server(void *arg)
 	return NULL;
 }
 
-static void do_test(const char *tcp_ca)
+static void do_test(const char *tcp_ca, const struct bpf_map *sk_stg_map)
 {
 	struct sockaddr_in6 sa6 = {};
 	ssize_t nr_recv = 0, bytes = 0;
@@ -110,6 +111,14 @@ static void do_test(const char *tcp_ca)
 		return;
 	}
 
+	if (sk_stg_map) {
+		err = bpf_map_update_elem(bpf_map__fd(sk_stg_map), &fd,
+					  &expected_stg, BPF_NOEXIST);
+		if (CHECK(err, "bpf_map_update_elem(sk_stg_map)",
+			  "err:%d errno:%d\n", err, errno))
+			goto done;
+	}
+
 	if (settcpca(lfd, tcp_ca) || settcpca(fd, tcp_ca) ||
 	    settimeo(lfd) || settimeo(fd))
 		goto done;
@@ -149,6 +158,16 @@ static void do_test(const char *tcp_ca)
 	CHECK(bytes != total_bytes, "recv", "%zd != %u nr_recv:%zd errno:%d\n",
 	      bytes, total_bytes, nr_recv, errno);
 
+	if (sk_stg_map) {
+		int tmp_stg;
+
+		err = bpf_map_lookup_elem(bpf_map__fd(sk_stg_map), &fd,
+					  &tmp_stg);
+		CHECK(!err || errno != ENOENT,
+		      "bpf_map_lookup_elem(sk_stg_map)",
+		      "err:%d errno:%d\n", err, errno);
+	}
+
 wait_thread:
 	WRITE_ONCE(stop, 1);
 	pthread_join(srv_thread, &thread_ret);
@@ -175,7 +194,7 @@ static void test_cubic(void)
 		return;
 	}
 
-	do_test("bpf_cubic");
+	do_test("bpf_cubic", NULL);
 
 	bpf_link__destroy(link);
 	bpf_cubic__destroy(cubic_skel);
@@ -197,7 +216,10 @@ static void test_dctcp(void)
 		return;
 	}
 
-	do_test("bpf_dctcp");
+	do_test("bpf_dctcp", dctcp_skel->maps.sk_stg_map);
+	CHECK(dctcp_skel->bss->stg_result != expected_stg,
+	      "Unexpected stg_result", "stg_result (%x) != expected_stg (%x)\n",
+	      dctcp_skel->bss->stg_result, expected_stg);
 
 	bpf_link__destroy(link);
 	bpf_dctcp__destroy(dctcp_skel);
diff --git a/tools/testing/selftests/bpf/progs/bpf_dctcp.c b/tools/testing/selftests/bpf/progs/bpf_dctcp.c
index 127ea762a062..5c1fc584f3ae 100644
--- a/tools/testing/selftests/bpf/progs/bpf_dctcp.c
+++ b/tools/testing/selftests/bpf/progs/bpf_dctcp.c
@@ -6,6 +6,7 @@
  * the kernel BPF logic.
  */
 
+#include <stddef.h>
 #include <linux/bpf.h>
 #include <linux/types.h>
 #include <bpf/bpf_helpers.h>
@@ -14,6 +15,15 @@
 
 char _license[] SEC("license") = "GPL";
 
+static volatile int stg_result;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, int);
+	__type(value, int);
+} sk_stg_map SEC(".maps");
+
 #define DCTCP_MAX_ALPHA	1024U
 
 struct dctcp {
@@ -43,12 +53,18 @@ void BPF_PROG(dctcp_init, struct sock *sk)
 {
 	const struct tcp_sock *tp = tcp_sk(sk);
 	struct dctcp *ca = inet_csk_ca(sk);
+	int *stg;
 
 	ca->prior_rcv_nxt = tp->rcv_nxt;
 	ca->dctcp_alpha = min(dctcp_alpha_on_init, DCTCP_MAX_ALPHA);
 	ca->loss_cwnd = 0;
 	ca->ce_state = 0;
 
+	stg = bpf_sk_storage_get(&sk_stg_map, (void *)tp, NULL, 0);
+	if (stg) {
+		stg_result = *stg;
+		bpf_sk_storage_delete(&sk_stg_map, (void *)tp);
+	}
 	dctcp_reset(tp, ca);
 }
 
-- 
2.17.1