Quelle  sock_iter_batch.c   Sprache: C

// SPDX-License-Identifier: GPL-2.0
// Copyright (c) 2024 Meta

#include <poll.h>
#include <test_progs.h>
#include "network_helpers.h"
#include "sock_iter_batch.skel.h"

#define TEST_NS "sock_iter_batch_netns"
#define TEST_CHILD_NS "sock_iter_batch_child_netns"

static const int init_batch_size = 16;
static const int nr_soreuse = 4;

struct iter_out {
 int idx;
 __u64 cookie;
} __packed;

struct sock_count {
 __u64 cookie;
 int count;
};

static int insert(__u64 cookie, struct sock_count counts[], int counts_len)
{
 int insert = -1;
 int i = 0;

 for (; i < counts_len; i++) {
  if (!counts[i].cookie) {
   insert = i;
  } else if (counts[i].cookie == cookie) {
   insert = i;
   break;
  }
 }
 if (insert < 0)
  return insert;

 counts[insert].cookie = cookie;
 counts[insert].count++;

 return counts[insert].count;
}

static int read_n(int iter_fd, int n, struct sock_count counts[],
    int counts_len)
{
 struct iter_out out;
 int nread = 1;
 int i = 0;

 for (; nread > 0 && (n < 0 || i < n); i++) {
  nread = read(iter_fd, &out, sizeof(out));
  if (!nread || !ASSERT_EQ(nread, sizeof(out), "nread"))
   break;
  ASSERT_GE(insert(out.cookie, counts, counts_len), 0, "insert");
 }

 ASSERT_TRUE(n < 0 || i == n, "n < 0 || i == n");

 return i;
}

static __u64 socket_cookie(int fd)
{
 __u64 cookie;
 socklen_t cookie_len = sizeof(cookie);

 if (!ASSERT_OK(getsockopt(fd, SOL_SOCKET, SO_COOKIE, &cookie,
      &cookie_len), "getsockopt(SO_COOKIE)"))
  return 0;
 return cookie;
}

static bool was_seen(int fd, struct sock_count counts[], int counts_len)
{
 __u64 cookie = socket_cookie(fd);
 int i = 0;

 for (; cookie && i < counts_len; i++)
  if (cookie == counts[i].cookie)
   return true;

 return false;
}

static int get_seen_socket(int *fds, struct sock_count counts[], int n)
{
 int i = 0;

 for (; i < n; i++)
  if (was_seen(fds[i], counts, n))
   return i;
 return -1;
}

static int get_nth_socket(int *fds, int fds_len, struct bpf_link *link, int n)
{
 int i, nread, iter_fd;
 int nth_sock_idx = -1;
 struct iter_out out;

 iter_fd = bpf_iter_create(bpf_link__fd(link));
 if (!ASSERT_OK_FD(iter_fd, "bpf_iter_create"))
  return -1;

 for (; n >= 0; n--) {
  nread = read(iter_fd, &out, sizeof(out));
  if (!nread || !ASSERT_GE(nread, 1, "nread"))
   goto done;
 }

 for (i = 0; i < fds_len && nth_sock_idx < 0; i++)
  if (fds[i] >= 0 && socket_cookie(fds[i]) == out.cookie)
   nth_sock_idx = i;
done:
 close(iter_fd);
 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);
 int count = 0;
 int i = 0;

 for (; cookie && !count && i < n; i++)
  if (cookie == counts[i].cookie)
   count = counts[i].count;

 return count;
}

static void check_n_were_seen_once(int *fds, int fds_len, int n,
       struct sock_count counts[], int counts_len)
{
 int seen_once = 0;
 int seen_cnt;
 int i = 0;

 for (; i < fds_len; i++) {
  /* Skip any sockets that were closed or that weren't seen
 * exactly once.
 */
  if (fds[i] < 0)
   continue;
  seen_cnt = get_seen_count(fds[i], counts, counts_len);
  if (seen_cnt && ASSERT_EQ(seen_cnt, 1, "seen_cnt"))
   seen_once++;
 }

 ASSERT_EQ(seen_once, n, "seen_once");
}

static int accept_from_one(struct pollfd *server_poll_fds,
      int server_poll_fds_len)
{
 static const int poll_timeout_ms = 5000; /* 5s */
 int ret;
 int i;

 ret = poll(server_poll_fds, server_poll_fds_len, poll_timeout_ms);
 if (!ASSERT_EQ(ret, 1, "poll"))
  return -1;

 for (i = 0; i < server_poll_fds_len; i++)
  if (server_poll_fds[i].revents & POLLIN)
   return accept(server_poll_fds[i].fd, NULL, NULL);

 return -1;
}

static int *connect_to_server(int family, int sock_type, const char *addr,
         __u16 port, int nr_connects, int *server_fds,
         int server_fds_len)
{
 struct pollfd *server_poll_fds = NULL;
 int *established_socks = NULL;
 int i;

 server_poll_fds = calloc(server_fds_len, sizeof(*server_poll_fds));
 if (!ASSERT_OK_PTR(server_poll_fds, "server_poll_fds"))
  return NULL;

 for (i = 0; i < server_fds_len; i++) {
  server_poll_fds[i].fd = server_fds[i];
  server_poll_fds[i].events = POLLIN;
 }

 i = 0;

 established_socks = malloc(sizeof(*established_socks) * nr_connects*2);
 if (!ASSERT_OK_PTR(established_socks, "established_socks"))
  goto error;

 while (nr_connects--) {
  established_socks[i] = connect_to_addr_str(family, sock_type,
          addr, port, NULL);
  if (!ASSERT_OK_FD(established_socks[i], "connect_to_addr_str"))
   goto error;
  i++;
  established_socks[i] = accept_from_one(server_poll_fds,
             server_fds_len);
  if (!ASSERT_OK_FD(established_socks[i], "accept_from_one"))
   goto error;
  i++;
 }

 free(server_poll_fds);
 return established_socks;
error:
 free_fds(established_socks, i);
 free(server_poll_fds);
 return NULL;
}

static void remove_seen(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,
   int counts_len, struct bpf_link *link, int iter_fd)
{
 int close_idx;

 /* Iterate through the first socks_len - 1 sockets. */
 read_n(iter_fd, socks_len - 1, counts, counts_len);

 /* Make sure we saw socks_len - 1 sockets exactly once. */
 check_n_were_seen_once(socks, socks_len, socks_len - 1, counts,
          counts_len);

 /* Close a socket we've already seen to remove it from the bucket. */
 close_idx = get_seen_socket(socks, counts, counts_len);
 if (!ASSERT_GE(close_idx, 0, "close_idx"))
  return;
 close(socks[close_idx]);
 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(socks, socks_len, socks_len - 1, counts,
          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,
     struct sock_count *counts, int counts_len,
     struct bpf_link *link, int iter_fd)
{
 int close_idx;

 /* Iterate through the first socket. */
 read_n(iter_fd, 1, counts, counts_len);

 /* Make sure we saw a socket from fds. */
 check_n_were_seen_once(socks, 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(socks, socks_len, link, 1);
 if (!ASSERT_GE(close_idx, 0, "close_idx"))
  return;
 close(socks[close_idx]);
 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(socks, socks_len, socks_len - 1, counts,
          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,
         struct sock_count *counts, int counts_len,
         struct bpf_link *link, int iter_fd)
{
 int close_idx, i;

 /* Iterate through the first socket. */
 read_n(iter_fd, 1, counts, counts_len);

 /* Make sure we saw a socket from fds. */
 check_n_were_seen_once(socks, 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.
 */
 for (i = 0; i < socks_len - 1; i++) {
  close_idx = get_nth_socket(socks, socks_len, link, 1);
  if (!ASSERT_GE(close_idx, 0, "close_idx"))
   return;
  close(socks[close_idx]);
  socks[close_idx] = -1;
 }

 /* Make sure there are no more sockets returned */
 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,
       int counts_len, struct bpf_link *link, int iter_fd)
{
 int *new_socks = NULL;

 /* Iterate through the first socks_len - 1 sockets. */
 read_n(iter_fd, socks_len - 1, counts, counts_len);

 /* Make sure we saw socks_len - 1 sockets exactly once. */
 check_n_were_seen_once(socks, socks_len, socks_len - 1, counts,
          counts_len);

 /* Double the number of sockets in the bucket. */
 new_socks = start_reuseport_server(family, sock_type, addr, port, 0,
        socks_len);
 if (!ASSERT_OK_PTR(new_socks, "start_reuseport_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(socks, socks_len, socks_len, counts,
          counts_len);
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,
     struct sock_count *counts, int counts_len,
     struct bpf_link *link, int iter_fd)
{
 int *new_socks = NULL;

 /* Iterate through the first socket just to initialize the batch. */
 read_n(iter_fd, 1, counts, counts_len);

 /* Double the number of sockets in the bucket to force a realloc on the
 * next read.
 */
 new_socks = start_reuseport_server(family, sock_type, addr, port, 0,
        socks_len);
 if (!ASSERT_OK_PTR(new_socks, "start_reuseport_server"))
  goto done;

 /* Iterate through the rest of the sockets. */
 read_n(iter_fd, -1, counts, counts_len);

 /* Make sure each socket from the first set was seen exactly once. */
 check_n_were_seen_once(socks, socks_len, socks_len, counts,
          counts_len);
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,
       int established_socks_len, struct sock_count *counts,
       int counts_len, struct bpf_link *link, int iter_fd);
 const char *description;
 int ehash_buckets;
 int connections;
 int init_socks;
 int max_socks;
 int sock_type;
 int family;
};

static struct test_case resume_tests[] = {
 {
  .description = "udp: resume after removing a seen socket",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_DGRAM,
  .family = AF_INET6,
  .test = remove_seen,
 },
 {
  .description = "udp: resume after removing one unseen socket",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_DGRAM,
  .family = AF_INET6,
  .test = remove_unseen,
 },
 {
  .description = "udp: resume after removing all unseen sockets",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_DGRAM,
  .family = AF_INET6,
  .test = remove_all,
 },
 {
  .description = "udp: resume after adding a few sockets",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_DGRAM,
  /* Use AF_INET so that new sockets are added to the head of the
 * bucket's list.
 */
  .family = AF_INET,
  .test = add_some,
 },
 {
  .description = "udp: force a realloc to occur",
  .init_socks = init_batch_size,
  .max_socks = init_batch_size * 2,
  .sock_type = SOCK_DGRAM,
  /* Use AF_INET6 so that new sockets are added to the tail of the
 * bucket's list, needing to be added to the next batch to force
 * a realloc.
 */
  .family = AF_INET6,
  .test = force_realloc,
 },
 {
  .description = "tcp: resume after removing a seen socket (listening)",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_STREAM,
  .family = AF_INET6,
  .test = remove_seen,
 },
 {
  .description = "tcp: resume after removing one unseen socket (listening)",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_STREAM,
  .family = AF_INET6,
  .test = remove_unseen,
 },
 {
  .description = "tcp: resume after removing all unseen sockets (listening)",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_STREAM,
  .family = AF_INET6,
  .test = remove_all,
 },
 {
  .description = "tcp: resume after adding a few sockets (listening)",
  .init_socks = nr_soreuse,
  .max_socks = nr_soreuse,
  .sock_type = SOCK_STREAM,
  /* Use AF_INET so that new sockets are added to the head of the
 * bucket's list.
 */
  .family = AF_INET,
  .test = add_some,
 },
 {
  .description = "tcp: force a realloc to occur (listening)",
  .init_socks = init_batch_size,
  .max_socks = init_batch_size * 2,
  .sock_type = SOCK_STREAM,
  /* Use AF_INET6 so that new sockets are added to the tail of the
 * bucket's list, needing to be added to the next batch to force
 * a realloc.
 */
  .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)
{
 struct sock_iter_batch *skel = NULL;
 struct sock_count *counts = NULL;
 static const __u16 port = 10001;
 struct nstoken *nstoken = NULL;
 struct bpf_link *link = NULL;
 int *established_fds = NULL;
 int err, iter_fd = -1;
 const char *addr;
 int *fds = NULL;

 if (tc->ehash_buckets) {
  SYS_NOFAIL("ip netns del " TEST_CHILD_NS);
  SYS(done, "sysctl -wq net.ipv4.tcp_child_ehash_entries=%d",
      tc->ehash_buckets);
  SYS(done, "ip netns add %s", TEST_CHILD_NS);
  SYS(done, "ip -net %s link set dev lo up", TEST_CHILD_NS);
  nstoken = open_netns(TEST_CHILD_NS);
  if (!ASSERT_OK_PTR(nstoken, "open_child_netns"))
   goto done;
 }

 counts = calloc(tc->max_socks, sizeof(*counts));
 if (!ASSERT_OK_PTR(counts, "counts"))
  goto done;
 skel = sock_iter_batch__open();
 if (!ASSERT_OK_PTR(skel, "sock_iter_batch__open"))
  goto done;

 /* Prepare a bucket of sockets in the kernel hashtable */
 addr = tc->family == AF_INET6 ? "::1" : "127.0.0.1";
 fds = start_reuseport_server(tc->family, tc->sock_type, addr, port, 0,
         tc->init_socks);
 if (!ASSERT_OK_PTR(fds, "start_reuseport_server"))
  goto done;
 if (tc->connections) {
  established_fds = connect_to_server(tc->family, tc->sock_type,
          addr, port,
          tc->connections, fds,
          tc->init_socks);
  if (!ASSERT_OK_PTR(established_fds, "connect_to_server"))
   goto done;
 }
 skel->rodata->ports[0] = 0;
 skel->rodata->ports[1] = 0;
 skel->rodata->sf = tc->family;
 skel->rodata->ss = 0;

 err = sock_iter_batch__load(skel);
 if (!ASSERT_OK(err, "sock_iter_batch__load"))
  goto done;

 link = bpf_program__attach_iter(tc->sock_type == SOCK_STREAM ?
     skel->progs.iter_tcp_soreuse :
     skel->progs.iter_udp_soreuse,
     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;

 tc->test(tc->family, tc->sock_type, addr, port, fds, tc->init_socks,
   established_fds, tc->connections*2, counts, tc->max_socks,
   link, iter_fd);
done:
 close_netns(nstoken);
 SYS_NOFAIL("ip netns del " TEST_CHILD_NS);
 SYS_NOFAIL("sysctl -w net.ipv4.tcp_child_ehash_entries=0");
 free(counts);
 free_fds(fds, tc->init_socks);
 free_fds(established_fds, tc->connections*2);
 if (iter_fd >= 0)
  close(iter_fd);
 bpf_link__destroy(link);
 sock_iter_batch__destroy(skel);
}

static void do_resume_tests(void)
{
 int i;

 for (i = 0; i < ARRAY_SIZE(resume_tests); i++) {
  if (test__start_subtest(resume_tests[i].description)) {
   do_resume_test(&resume_tests[i]);
  }
 }
}

static void do_test(int sock_type, bool onebyone)
{
 int err, i, nread, to_read, total_read, iter_fd = -1;
 struct iter_out outputs[nr_soreuse];
 struct bpf_link *link = NULL;
 struct sock_iter_batch *skel;
 int first_idx, second_idx;
 int *fds[2] = {};

 skel = sock_iter_batch__open();
 if (!ASSERT_OK_PTR(skel, "sock_iter_batch__open"))
  return;

 /* Prepare 2 buckets of sockets in the kernel hashtable */
 for (i = 0; i < ARRAY_SIZE(fds); i++) {
  int local_port;

  fds[i] = start_reuseport_server(AF_INET6, sock_type, "::1", 0, 0,
      nr_soreuse);
  if (!ASSERT_OK_PTR(fds[i], "start_reuseport_server"))
   goto done;
  local_port = get_socket_local_port(*fds[i]);
  if (!ASSERT_GE(local_port, 0, "get_socket_local_port"))
   goto done;
  skel->rodata->ports[i] = ntohs(local_port);
 }
 skel->rodata->sf = AF_INET6;
 if (sock_type == SOCK_STREAM)
  skel->rodata->ss = TCP_LISTEN;

 err = sock_iter_batch__load(skel);
 if (!ASSERT_OK(err, "sock_iter_batch__load"))
  goto done;

 link = bpf_program__attach_iter(sock_type == SOCK_STREAM ?
     skel->progs.iter_tcp_soreuse :
     skel->progs.iter_udp_soreuse,
     NULL);
 if (!ASSERT_OK_PTR(link, "bpf_program__attach_iter"))
  goto done;

 iter_fd = bpf_iter_create(bpf_link__fd(link));
 if (!ASSERT_GE(iter_fd, 0, "bpf_iter_create"))
  goto done;

 /* Test reading a bucket (either from fds[0] or fds[1]).
 * Only read "nr_soreuse - 1" number of sockets
 * from a bucket and leave one socket out from
 * that bucket on purpose.
 */
 to_read = (nr_soreuse - 1) * sizeof(*outputs);
 total_read = 0;
 first_idx = -1;
 do {
  nread = read(iter_fd, outputs, onebyone ? sizeof(*outputs) : to_read);
  if (nread <= 0 || nread % sizeof(*outputs))
   break;
  total_read += nread;

  if (first_idx == -1)
   first_idx = outputs[0].idx;
  for (i = 0; i < nread / sizeof(*outputs); i++)
   ASSERT_EQ(outputs[i].idx, first_idx, "first_idx");
 } while (total_read < to_read);
 ASSERT_EQ(nread, onebyone ? sizeof(*outputs) : to_read, "nread");
 ASSERT_EQ(total_read, to_read, "total_read");

 free_fds(fds[first_idx], nr_soreuse);
 fds[first_idx] = NULL;

 /* Read the "whole" second bucket */
 to_read = nr_soreuse * sizeof(*outputs);
 total_read = 0;
 second_idx = !first_idx;
 do {
  nread = read(iter_fd, outputs, onebyone ? sizeof(*outputs) : to_read);
  if (nread <= 0 || nread % sizeof(*outputs))
   break;
  total_read += nread;

  for (i = 0; i < nread / sizeof(*outputs); i++)
   ASSERT_EQ(outputs[i].idx, second_idx, "second_idx");
 } while (total_read <= to_read);
 ASSERT_EQ(nread, 0, "nread");
 /* Both so_reuseport ports should be in different buckets, so
 * total_read must equal to the expected to_read.
 *
 * For a very unlikely case, both ports collide at the same bucket,
 * the bucket offset (i.e. 3) will be skipped and it cannot
 * expect the to_read number of bytes.
 */
 if (skel->bss->bucket[0] != skel->bss->bucket[1])
  ASSERT_EQ(total_read, to_read, "total_read");

done:
 for (i = 0; i < ARRAY_SIZE(fds); i++)
  free_fds(fds[i], nr_soreuse);
 if (iter_fd < 0)
  close(iter_fd);
 bpf_link__destroy(link);
 sock_iter_batch__destroy(skel);
}

void test_sock_iter_batch(void)
{
 struct nstoken *nstoken = NULL;

 SYS_NOFAIL("ip netns del " TEST_NS);
 SYS(done, "ip netns add %s", TEST_NS);
 SYS(done, "ip -net %s link set dev lo up", TEST_NS);

 nstoken = open_netns(TEST_NS);
 if (!ASSERT_OK_PTR(nstoken, "open_netns"))
  goto done;

 if (test__start_subtest("tcp")) {
  do_test(SOCK_STREAM, true);
  do_test(SOCK_STREAM, false);
 }
 if (test__start_subtest("udp")) {
  do_test(SOCK_DGRAM, true);
  do_test(SOCK_DGRAM, false);
 }
 do_resume_tests();
 close_netns(nstoken);

done:
 SYS_NOFAIL("ip netns del " TEST_NS);
}