selftests/bpf: Add tests for bucket resume logic in established sockets

Replicate the set of test cases used for UDP socket iterators to test
similar scenarios for TCP established sockets.

Signed-off-by: Jordan Rife <jordan@jrife.io>
Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
Acked-by: Stanislav Fomichev <sdf@fomichev.me>

diff --git a/tools/testing/selftests/bpf/prog_tests/sock_iter_batch.c b/tools/testing/selftests/bpf/prog_tests/sock_iter_batch.c
index 0ccc90d1d1ef00f83129dcdc1e40025871f6d7bc..27781df8f2fbdbef3a18b17981df3a0e1b54fa64 100644 (file)
--- a/tools/testing/selftests/bpf/prog_tests/sock_iter_batch.c
+++ b/tools/testing/selftests/bpf/prog_tests/sock_iter_batch.c
@@ -120,6 +120,45 @@ done:
        return nth_sock_idx;
 }
 
+static void destroy(int fd)
+{
+       struct sock_iter_batch *skel = NULL;
+       __u64 cookie = socket_cookie(fd);
+       struct bpf_link *link = NULL;
+       int iter_fd = -1;
+       int nread;
+       __u64 out;
+
+       skel = sock_iter_batch__open();
+       if (!ASSERT_OK_PTR(skel, "sock_iter_batch__open"))
+               goto done;
+
+       skel->rodata->destroy_cookie = cookie;
+
+       if (!ASSERT_OK(sock_iter_batch__load(skel), "sock_iter_batch__load"))
+               goto done;
+
+       link = bpf_program__attach_iter(skel->progs.iter_tcp_destroy, NULL);
+       if (!ASSERT_OK_PTR(link, "bpf_program__attach_iter"))
+               goto done;
+
+       iter_fd = bpf_iter_create(bpf_link__fd(link));
+       if (!ASSERT_OK_FD(iter_fd, "bpf_iter_create"))
+               goto done;
+
+       /* Delete matching socket. */
+       nread = read(iter_fd, &out, sizeof(out));
+       ASSERT_GE(nread, 0, "nread");
+       if (nread)
+               ASSERT_EQ(out, cookie, "cookie matches");
+done:
+       if (iter_fd >= 0)
+               close(iter_fd);
+       bpf_link__destroy(link);
+       sock_iter_batch__destroy(skel);
+       close(fd);
+}
+
 static int get_seen_count(int fd, struct sock_count counts[], int n)
 {
        __u64 cookie = socket_cookie(fd);
@@ -247,6 +286,43 @@ static void remove_seen(int family, int sock_type, const char *addr, __u16 port,
                               counts_len);
 }
 
+static void remove_seen_established(int family, int sock_type, const char *addr,
+                                   __u16 port, int *listen_socks,
+                                   int listen_socks_len, int *established_socks,
+                                   int established_socks_len,
+                                   struct sock_count *counts, int counts_len,
+                                   struct bpf_link *link, int iter_fd)
+{
+       int close_idx;
+
+       /* Iterate through all listening sockets. */
+       read_n(iter_fd, listen_socks_len, counts, counts_len);
+
+       /* Make sure we saw all listening sockets exactly once. */
+       check_n_were_seen_once(listen_socks, listen_socks_len, listen_socks_len,
+                              counts, counts_len);
+
+       /* Leave one established socket. */
+       read_n(iter_fd, established_socks_len - 1, counts, counts_len);
+
+       /* Close a socket we've already seen to remove it from the bucket. */
+       close_idx = get_nth_socket(established_socks, established_socks_len,
+                                  link, listen_socks_len + 1);
+       if (!ASSERT_GE(close_idx, 0, "close_idx"))
+               return;
+       destroy(established_socks[close_idx]);
+       established_socks[close_idx] = -1;
+
+       /* Iterate through the rest of the sockets. */
+       read_n(iter_fd, -1, counts, counts_len);
+
+       /* Make sure the last socket wasn't skipped and that there were no
+        * repeats.
+        */
+       check_n_were_seen_once(established_socks, established_socks_len,
+                              established_socks_len - 1, counts, counts_len);
+}
+
 static void remove_unseen(int family, int sock_type, const char *addr,
                          __u16 port, int *socks, int socks_len,
                          int *established_socks, int established_socks_len,
@@ -280,6 +356,51 @@ static void remove_unseen(int family, int sock_type, const char *addr,
                               counts_len);
 }
 
+static void remove_unseen_established(int family, int sock_type,
+                                     const char *addr, __u16 port,
+                                     int *listen_socks, int listen_socks_len,
+                                     int *established_socks,
+                                     int established_socks_len,
+                                     struct sock_count *counts, int counts_len,
+                                     struct bpf_link *link, int iter_fd)
+{
+       int close_idx;
+
+       /* Iterate through all listening sockets. */
+       read_n(iter_fd, listen_socks_len, counts, counts_len);
+
+       /* Make sure we saw all listening sockets exactly once. */
+       check_n_were_seen_once(listen_socks, listen_socks_len, listen_socks_len,
+                              counts, counts_len);
+
+       /* Iterate through the first established socket. */
+       read_n(iter_fd, 1, counts, counts_len);
+
+       /* Make sure we saw one established socks. */
+       check_n_were_seen_once(established_socks, established_socks_len, 1,
+                              counts, counts_len);
+
+       /* Close what would be the next socket in the bucket to exercise the
+        * condition where we need to skip past the first cookie we remembered.
+        */
+       close_idx = get_nth_socket(established_socks, established_socks_len,
+                                  link, listen_socks_len + 1);
+       if (!ASSERT_GE(close_idx, 0, "close_idx"))
+               return;
+
+       destroy(established_socks[close_idx]);
+       established_socks[close_idx] = -1;
+
+       /* Iterate through the rest of the sockets. */
+       read_n(iter_fd, -1, counts, counts_len);
+
+       /* Make sure the remaining sockets were seen exactly once and that we
+        * didn't repeat the socket that was already seen.
+        */
+       check_n_were_seen_once(established_socks, established_socks_len,
+                              established_socks_len - 1, counts, counts_len);
+}
+
 static void remove_all(int family, int sock_type, const char *addr,
                       __u16 port, int *socks, int socks_len,
                       int *established_socks, int established_socks_len,
@@ -309,6 +430,54 @@ static void remove_all(int family, int sock_type, const char *addr,
        ASSERT_EQ(read_n(iter_fd, -1, counts, counts_len), 0, "read_n");
 }
 
+static void remove_all_established(int family, int sock_type, const char *addr,
+                                  __u16 port, int *listen_socks,
+                                  int listen_socks_len, int *established_socks,
+                                  int established_socks_len,
+                                  struct sock_count *counts, int counts_len,
+                                  struct bpf_link *link, int iter_fd)
+{
+       int *close_idx = NULL;
+       int i;
+
+       /* Iterate through all listening sockets. */
+       read_n(iter_fd, listen_socks_len, counts, counts_len);
+
+       /* Make sure we saw all listening sockets exactly once. */
+       check_n_were_seen_once(listen_socks, listen_socks_len, listen_socks_len,
+                              counts, counts_len);
+
+       /* Iterate through the first established socket. */
+       read_n(iter_fd, 1, counts, counts_len);
+
+       /* Make sure we saw one established socks. */
+       check_n_were_seen_once(established_socks, established_socks_len, 1,
+                              counts, counts_len);
+
+       /* Close all remaining sockets to exhaust the list of saved cookies and
+        * exit without putting any sockets into the batch on the next read.
+        */
+       close_idx = malloc(sizeof(int) * (established_socks_len - 1));
+       if (!ASSERT_OK_PTR(close_idx, "close_idx malloc"))
+               return;
+       for (i = 0; i < established_socks_len - 1; i++) {
+               close_idx[i] = get_nth_socket(established_socks,
+                                             established_socks_len, link,
+                                             listen_socks_len + i);
+               if (!ASSERT_GE(close_idx[i], 0, "close_idx"))
+                       return;
+       }
+
+       for (i = 0; i < established_socks_len - 1; i++) {
+               destroy(established_socks[close_idx[i]]);
+               established_socks[close_idx[i]] = -1;
+       }
+
+       /* Make sure there are no more sockets returned */
+       ASSERT_EQ(read_n(iter_fd, -1, counts, counts_len), 0, "read_n");
+       free(close_idx);
+}
+
 static void add_some(int family, int sock_type, const char *addr, __u16 port,
                     int *socks, int socks_len, int *established_socks,
                     int established_socks_len, struct sock_count *counts,
@@ -339,6 +508,49 @@ done:
        free_fds(new_socks, socks_len);
 }
 
+static void add_some_established(int family, int sock_type, const char *addr,
+                                __u16 port, int *listen_socks,
+                                int listen_socks_len, int *established_socks,
+                                int established_socks_len,
+                                struct sock_count *counts,
+                                int counts_len, struct bpf_link *link,
+                                int iter_fd)
+{
+       int *new_socks = NULL;
+
+       /* Iterate through all listening sockets. */
+       read_n(iter_fd, listen_socks_len, counts, counts_len);
+
+       /* Make sure we saw all listening sockets exactly once. */
+       check_n_were_seen_once(listen_socks, listen_socks_len, listen_socks_len,
+                              counts, counts_len);
+
+       /* Iterate through the first established_socks_len - 1 sockets. */
+       read_n(iter_fd, established_socks_len - 1, counts, counts_len);
+
+       /* Make sure we saw established_socks_len - 1 sockets exactly once. */
+       check_n_were_seen_once(established_socks, established_socks_len,
+                              established_socks_len - 1, counts, counts_len);
+
+       /* Double the number of established sockets in the bucket. */
+       new_socks = connect_to_server(family, sock_type, addr, port,
+                                     established_socks_len / 2, listen_socks,
+                                     listen_socks_len);
+       if (!ASSERT_OK_PTR(new_socks, "connect_to_server"))
+               goto done;
+
+       /* Iterate through the rest of the sockets. */
+       read_n(iter_fd, -1, counts, counts_len);
+
+       /* Make sure each of the original sockets was seen exactly once. */
+       check_n_were_seen_once(listen_socks, listen_socks_len, listen_socks_len,
+                              counts, counts_len);
+       check_n_were_seen_once(established_socks, established_socks_len,
+                              established_socks_len, counts, counts_len);
+done:
+       free_fds(new_socks, established_socks_len);
+}
+
 static void force_realloc(int family, int sock_type, const char *addr,
                          __u16 port, int *socks, int socks_len,
                          int *established_socks, int established_socks_len,
@@ -368,6 +580,24 @@ done:
        free_fds(new_socks, socks_len);
 }
 
+static void force_realloc_established(int family, int sock_type,
+                                     const char *addr, __u16 port,
+                                     int *listen_socks, int listen_socks_len,
+                                     int *established_socks,
+                                     int established_socks_len,
+                                     struct sock_count *counts, int counts_len,
+                                     struct bpf_link *link, int iter_fd)
+{
+       /* Iterate through all sockets to trigger a realloc. */
+       read_n(iter_fd, -1, counts, counts_len);
+
+       /* Make sure each socket was seen exactly once. */
+       check_n_were_seen_once(listen_socks, listen_socks_len, listen_socks_len,
+                              counts, counts_len);
+       check_n_were_seen_once(established_socks, established_socks_len,
+                              established_socks_len, counts, counts_len);
+}
+
 struct test_case {
        void (*test)(int family, int sock_type, const char *addr, __u16 port,
                     int *socks, int socks_len, int *established_socks,
@@ -477,6 +707,69 @@ static struct test_case resume_tests[] = {
                .family = AF_INET6,
                .test = force_realloc,
        },
+       {
+               .description = "tcp: resume after removing a seen socket (established)",
+               /* Force all established sockets into one bucket */
+               .ehash_buckets = 1,
+               .connections = nr_soreuse,
+               .init_socks = nr_soreuse,
+               /* Room for connect()ed and accept()ed sockets */
+               .max_socks = nr_soreuse * 3,
+               .sock_type = SOCK_STREAM,
+               .family = AF_INET6,
+               .test = remove_seen_established,
+       },
+       {
+               .description = "tcp: resume after removing one unseen socket (established)",
+               /* Force all established sockets into one bucket */
+               .ehash_buckets = 1,
+               .connections = nr_soreuse,
+               .init_socks = nr_soreuse,
+               /* Room for connect()ed and accept()ed sockets */
+               .max_socks = nr_soreuse * 3,
+               .sock_type = SOCK_STREAM,
+               .family = AF_INET6,
+               .test = remove_unseen_established,
+       },
+       {
+               .description = "tcp: resume after removing all unseen sockets (established)",
+               /* Force all established sockets into one bucket */
+               .ehash_buckets = 1,
+               .connections = nr_soreuse,
+               .init_socks = nr_soreuse,
+               /* Room for connect()ed and accept()ed sockets */
+               .max_socks = nr_soreuse * 3,
+               .sock_type = SOCK_STREAM,
+               .family = AF_INET6,
+               .test = remove_all_established,
+       },
+       {
+               .description = "tcp: resume after adding a few sockets (established)",
+               /* Force all established sockets into one bucket */
+               .ehash_buckets = 1,
+               .connections = nr_soreuse,
+               .init_socks = nr_soreuse,
+               /* Room for connect()ed and accept()ed sockets */
+               .max_socks = nr_soreuse * 3,
+               .sock_type = SOCK_STREAM,
+               .family = AF_INET6,
+               .test = add_some_established,
+       },
+       {
+               .description = "tcp: force a realloc to occur (established)",
+               /* Force all established sockets into one bucket */
+               .ehash_buckets = 1,
+               /* Bucket size will need to double when going from listening to
+                * established sockets.
+                */
+               .connections = init_batch_size,
+               .init_socks = nr_soreuse,
+               /* Room for connect()ed and accept()ed sockets */
+               .max_socks = nr_soreuse + (init_batch_size * 2),
+               .sock_type = SOCK_STREAM,
+               .family = AF_INET6,
+               .test = force_realloc_established,
+       },
 };
 
 static void do_resume_test(struct test_case *tc)