// 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);
}