tools/testing/selftests/bpf/prog_tests/sock_iter_batch.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 // Copyright (c) 2024 Meta
   3
   4 #include <test_progs.h>
   5 #include "network_helpers.h"
   6 #include "sock_iter_batch.skel.h"
   7
   8 #define TEST_NS "sock_iter_batch_netns"
   9
  10 static const int nr_soreuse = 4;
  11
  12 static void do_test(int sock_type, bool onebyone)
  13 {
  14         int err, i, nread, to_read, total_read, iter_fd = -1;
  15         int first_idx, second_idx, indices[nr_soreuse];
  16         struct bpf_link *link = NULL;
  17         struct sock_iter_batch *skel;
  18         int *fds[2] = {};
  19
  20         skel = sock_iter_batch__open();
  21         if (!ASSERT_OK_PTR(skel, "sock_iter_batch__open"))
  22                 return;
  23
  24         /* Prepare 2 buckets of sockets in the kernel hashtable */
  25         for (i = 0; i < ARRAY_SIZE(fds); i++) {
  26                 int local_port;
  27
  28                 fds[i] = start_reuseport_server(AF_INET6, sock_type, "::1", 0, 0,
  29                                                 nr_soreuse);
  30                 if (!ASSERT_OK_PTR(fds[i], "start_reuseport_server"))
  31                         goto done;
  32                 local_port = get_socket_local_port(*fds[i]);
  33                 if (!ASSERT_GE(local_port, 0, "get_socket_local_port"))
  34                         goto done;
  35                 skel->rodata->ports[i] = ntohs(local_port);
  36         }
  37
  38         err = sock_iter_batch__load(skel);
  39         if (!ASSERT_OK(err, "sock_iter_batch__load"))
  40                 goto done;
  41
  42         link = bpf_program__attach_iter(sock_type == SOCK_STREAM ?
  43                                         skel->progs.iter_tcp_soreuse :
  44                                         skel->progs.iter_udp_soreuse,
  45                                         NULL);
  46         if (!ASSERT_OK_PTR(link, "bpf_program__attach_iter"))
  47                 goto done;
  48
  49         iter_fd = bpf_iter_create(bpf_link__fd(link));
  50         if (!ASSERT_GE(iter_fd, 0, "bpf_iter_create"))
  51                 goto done;
  52
  53         /* Test reading a bucket (either from fds[0] or fds[1]).
  54          * Only read "nr_soreuse - 1" number of sockets
  55          * from a bucket and leave one socket out from
  56          * that bucket on purpose.
  57          */
  58         to_read = (nr_soreuse - 1) * sizeof(*indices);
  59         total_read = 0;
  60         first_idx = -1;
  61         do {
  62                 nread = read(iter_fd, indices, onebyone ? sizeof(*indices) : to_read);
  63                 if (nread <= 0 || nread % sizeof(*indices))
  64                         break;
  65                 total_read += nread;
  66
  67                 if (first_idx == -1)
  68                         first_idx = indices[0];
  69                 for (i = 0; i < nread / sizeof(*indices); i++)
  70                         ASSERT_EQ(indices[i], first_idx, "first_idx");
  71         } while (total_read < to_read);
  72         ASSERT_EQ(nread, onebyone ? sizeof(*indices) : to_read, "nread");
  73         ASSERT_EQ(total_read, to_read, "total_read");
  74
  75         free_fds(fds[first_idx], nr_soreuse);
  76         fds[first_idx] = NULL;
  77
  78         /* Read the "whole" second bucket */
  79         to_read = nr_soreuse * sizeof(*indices);
  80         total_read = 0;
  81         second_idx = !first_idx;
  82         do {
  83                 nread = read(iter_fd, indices, onebyone ? sizeof(*indices) : to_read);
  84                 if (nread <= 0 || nread % sizeof(*indices))
  85                         break;
  86                 total_read += nread;
  87
  88                 for (i = 0; i < nread / sizeof(*indices); i++)
  89                         ASSERT_EQ(indices[i], second_idx, "second_idx");
  90         } while (total_read <= to_read);
  91         ASSERT_EQ(nread, 0, "nread");
  92         /* Both so_reuseport ports should be in different buckets, so
  93          * total_read must equal to the expected to_read.
  94          *
  95          * For a very unlikely case, both ports collide at the same bucket,
  96          * the bucket offset (i.e. 3) will be skipped and it cannot
  97          * expect the to_read number of bytes.
  98          */
  99         if (skel->bss->bucket[0] != skel->bss->bucket[1])
 100                 ASSERT_EQ(total_read, to_read, "total_read");
 101
 102 done:
 103         for (i = 0; i < ARRAY_SIZE(fds); i++)
 104                 free_fds(fds[i], nr_soreuse);
 105         if (iter_fd < 0)
 106                 close(iter_fd);
 107         bpf_link__destroy(link);
 108         sock_iter_batch__destroy(skel);
 109 }
 110
 111 void test_sock_iter_batch(void)
 112 {
 113         struct nstoken *nstoken = NULL;
 114
 115         SYS_NOFAIL("ip netns del " TEST_NS);
 116         SYS(done, "ip netns add %s", TEST_NS);
 117         SYS(done, "ip -net %s link set dev lo up", TEST_NS);
 118
 119         nstoken = open_netns(TEST_NS);
 120         if (!ASSERT_OK_PTR(nstoken, "open_netns"))
 121                 goto done;
 122
 123         if (test__start_subtest("tcp")) {
 124                 do_test(SOCK_STREAM, true);
 125                 do_test(SOCK_STREAM, false);
 126         }
 127         if (test__start_subtest("udp")) {
 128                 do_test(SOCK_DGRAM, true);
 129                 do_test(SOCK_DGRAM, false);
 130         }
 131         close_netns(nstoken);
 132
 133 done:
 134         SYS_NOFAIL("ip netns del " TEST_NS);
 135 }