| 1 | // SPDX-License-Identifier: GPL-2.0 |
|---|
| 2 | #include <test_progs.h> |
|---|
| 3 | #include "udp_limit.skel.h" |
|---|
| 4 | |
|---|
| 5 | #include <sys/types.h> |
|---|
| 6 | #include <sys/socket.h> |
|---|
| 7 | |
|---|
| 8 | void test_udp_limit(void) |
|---|
| 9 | { |
|---|
| 10 | struct udp_limit *skel; |
|---|
| 11 | int fd1 = -1, fd2 = -1; |
|---|
| 12 | int cgroup_fd; |
|---|
| 13 | |
|---|
| 14 | cgroup_fd = test__join_cgroup("/udp_limit" ); |
|---|
| 15 | if (!ASSERT_GE(cgroup_fd, 0, "cg-join" )) |
|---|
| 16 | return; |
|---|
| 17 | |
|---|
| 18 | skel = udp_limit__open_and_load(); |
|---|
| 19 | if (!ASSERT_OK_PTR(skel, "skel-load" )) |
|---|
| 20 | goto close_cgroup_fd; |
|---|
| 21 | |
|---|
| 22 | skel->links.sock = bpf_program__attach_cgroup(skel->progs.sock, cgroup_fd); |
|---|
| 23 | if (!ASSERT_OK_PTR(skel->links.sock, "cg_attach_sock" )) |
|---|
| 24 | goto close_skeleton; |
|---|
| 25 | skel->links.sock_release = bpf_program__attach_cgroup(skel->progs.sock_release, cgroup_fd); |
|---|
| 26 | if (!ASSERT_OK_PTR(skel->links.sock_release, "cg_attach_sock_release" )) |
|---|
| 27 | goto close_skeleton; |
|---|
| 28 | |
|---|
| 29 | /* BPF program enforces a single UDP socket per cgroup, |
|---|
| 30 | * verify that. |
|---|
| 31 | */ |
|---|
| 32 | fd1 = socket(AF_INET, SOCK_DGRAM, 0); |
|---|
| 33 | if (!ASSERT_GE(fd1, 0, "socket(fd1)" )) |
|---|
| 34 | goto close_skeleton; |
|---|
| 35 | |
|---|
| 36 | fd2 = socket(AF_INET, SOCK_DGRAM, 0); |
|---|
| 37 | if (!ASSERT_LT(fd2, 0, "socket(fd2)" )) |
|---|
| 38 | goto close_skeleton; |
|---|
| 39 | |
|---|
| 40 | /* We can reopen again after close. */ |
|---|
| 41 | close(fd1); |
|---|
| 42 | fd1 = -1; |
|---|
| 43 | |
|---|
| 44 | fd1 = socket(AF_INET, SOCK_DGRAM, 0); |
|---|
| 45 | if (!ASSERT_GE(fd1, 0, "socket(fd1-again)" )) |
|---|
| 46 | goto close_skeleton; |
|---|
| 47 | |
|---|
| 48 | /* Make sure the program was invoked the expected |
|---|
| 49 | * number of times: |
|---|
| 50 | * - open fd1 - BPF_CGROUP_INET_SOCK_CREATE |
|---|
| 51 | * - attempt to openfd2 - BPF_CGROUP_INET_SOCK_CREATE |
|---|
| 52 | * - close fd1 - BPF_CGROUP_INET_SOCK_RELEASE |
|---|
| 53 | * - open fd1 again - BPF_CGROUP_INET_SOCK_CREATE |
|---|
| 54 | */ |
|---|
| 55 | if (!ASSERT_EQ(skel->bss->invocations, 4, "bss-invocations" )) |
|---|
| 56 | goto close_skeleton; |
|---|
| 57 | |
|---|
| 58 | /* We should still have a single socket in use */ |
|---|
| 59 | if (!ASSERT_EQ(skel->bss->in_use, 1, "bss-in_use" )) |
|---|
| 60 | goto close_skeleton; |
|---|
| 61 | |
|---|
| 62 | close_skeleton: |
|---|
| 63 | if (fd1 >= 0) |
|---|
| 64 | close(fd1); |
|---|
| 65 | if (fd2 >= 0) |
|---|
| 66 | close(fd2); |
|---|
| 67 | udp_limit__destroy(skel); |
|---|
| 68 | close_cgroup_fd: |
|---|
| 69 | close(cgroup_fd); |
|---|
| 70 | } |
|---|
| 71 | |
|---|
