1*12bb6ca4SYonghong Song // SPDX-License-Identifier: GPL-2.0
2*12bb6ca4SYonghong Song /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates.*/
3*12bb6ca4SYonghong Song 
4*12bb6ca4SYonghong Song #define _GNU_SOURCE
5*12bb6ca4SYonghong Song #include <unistd.h>
6*12bb6ca4SYonghong Song #include <sys/syscall.h>
7*12bb6ca4SYonghong Song #include <sys/types.h>
8*12bb6ca4SYonghong Song #include <test_progs.h>
9*12bb6ca4SYonghong Song #include "cgrp_ls_tp_btf.skel.h"
10*12bb6ca4SYonghong Song #include "cgrp_ls_recursion.skel.h"
11*12bb6ca4SYonghong Song #include "cgrp_ls_attach_cgroup.skel.h"
12*12bb6ca4SYonghong Song #include "cgrp_ls_negative.skel.h"
13*12bb6ca4SYonghong Song #include "network_helpers.h"
14*12bb6ca4SYonghong Song 
15*12bb6ca4SYonghong Song struct socket_cookie {
16*12bb6ca4SYonghong Song 	__u64 cookie_key;
17*12bb6ca4SYonghong Song 	__u32 cookie_value;
18*12bb6ca4SYonghong Song };
19*12bb6ca4SYonghong Song 
20*12bb6ca4SYonghong Song static void test_tp_btf(int cgroup_fd)
21*12bb6ca4SYonghong Song {
22*12bb6ca4SYonghong Song 	struct cgrp_ls_tp_btf *skel;
23*12bb6ca4SYonghong Song 	long val1 = 1, val2 = 0;
24*12bb6ca4SYonghong Song 	int err;
25*12bb6ca4SYonghong Song 
26*12bb6ca4SYonghong Song 	skel = cgrp_ls_tp_btf__open_and_load();
27*12bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
28*12bb6ca4SYonghong Song 		return;
29*12bb6ca4SYonghong Song 
30*12bb6ca4SYonghong Song 	/* populate a value in map_b */
31*12bb6ca4SYonghong Song 	err = bpf_map_update_elem(bpf_map__fd(skel->maps.map_b), &cgroup_fd, &val1, BPF_ANY);
32*12bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_update_elem"))
33*12bb6ca4SYonghong Song 		goto out;
34*12bb6ca4SYonghong Song 
35*12bb6ca4SYonghong Song 	/* check value */
36*12bb6ca4SYonghong Song 	err = bpf_map_lookup_elem(bpf_map__fd(skel->maps.map_b), &cgroup_fd, &val2);
37*12bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_lookup_elem"))
38*12bb6ca4SYonghong Song 		goto out;
39*12bb6ca4SYonghong Song 	if (!ASSERT_EQ(val2, 1, "map_lookup_elem, invalid val"))
40*12bb6ca4SYonghong Song 		goto out;
41*12bb6ca4SYonghong Song 
42*12bb6ca4SYonghong Song 	/* delete value */
43*12bb6ca4SYonghong Song 	err = bpf_map_delete_elem(bpf_map__fd(skel->maps.map_b), &cgroup_fd);
44*12bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_delete_elem"))
45*12bb6ca4SYonghong Song 		goto out;
46*12bb6ca4SYonghong Song 
47*12bb6ca4SYonghong Song 	skel->bss->target_pid = syscall(SYS_gettid);
48*12bb6ca4SYonghong Song 
49*12bb6ca4SYonghong Song 	err = cgrp_ls_tp_btf__attach(skel);
50*12bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "skel_attach"))
51*12bb6ca4SYonghong Song 		goto out;
52*12bb6ca4SYonghong Song 
53*12bb6ca4SYonghong Song 	syscall(SYS_gettid);
54*12bb6ca4SYonghong Song 	syscall(SYS_gettid);
55*12bb6ca4SYonghong Song 
56*12bb6ca4SYonghong Song 	skel->bss->target_pid = 0;
57*12bb6ca4SYonghong Song 
58*12bb6ca4SYonghong Song 	/* 3x syscalls: 1x attach and 2x gettid */
59*12bb6ca4SYonghong Song 	ASSERT_EQ(skel->bss->enter_cnt, 3, "enter_cnt");
60*12bb6ca4SYonghong Song 	ASSERT_EQ(skel->bss->exit_cnt, 3, "exit_cnt");
61*12bb6ca4SYonghong Song 	ASSERT_EQ(skel->bss->mismatch_cnt, 0, "mismatch_cnt");
62*12bb6ca4SYonghong Song out:
63*12bb6ca4SYonghong Song 	cgrp_ls_tp_btf__destroy(skel);
64*12bb6ca4SYonghong Song }
65*12bb6ca4SYonghong Song 
66*12bb6ca4SYonghong Song static void test_attach_cgroup(int cgroup_fd)
67*12bb6ca4SYonghong Song {
68*12bb6ca4SYonghong Song 	int server_fd = 0, client_fd = 0, err = 0;
69*12bb6ca4SYonghong Song 	socklen_t addr_len = sizeof(struct sockaddr_in6);
70*12bb6ca4SYonghong Song 	struct cgrp_ls_attach_cgroup *skel;
71*12bb6ca4SYonghong Song 	__u32 cookie_expected_value;
72*12bb6ca4SYonghong Song 	struct sockaddr_in6 addr;
73*12bb6ca4SYonghong Song 	struct socket_cookie val;
74*12bb6ca4SYonghong Song 
75*12bb6ca4SYonghong Song 	skel = cgrp_ls_attach_cgroup__open_and_load();
76*12bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open"))
77*12bb6ca4SYonghong Song 		return;
78*12bb6ca4SYonghong Song 
79*12bb6ca4SYonghong Song 	skel->links.set_cookie = bpf_program__attach_cgroup(
80*12bb6ca4SYonghong Song 		skel->progs.set_cookie, cgroup_fd);
81*12bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel->links.set_cookie, "prog_attach"))
82*12bb6ca4SYonghong Song 		goto out;
83*12bb6ca4SYonghong Song 
84*12bb6ca4SYonghong Song 	skel->links.update_cookie_sockops = bpf_program__attach_cgroup(
85*12bb6ca4SYonghong Song 		skel->progs.update_cookie_sockops, cgroup_fd);
86*12bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel->links.update_cookie_sockops, "prog_attach"))
87*12bb6ca4SYonghong Song 		goto out;
88*12bb6ca4SYonghong Song 
89*12bb6ca4SYonghong Song 	skel->links.update_cookie_tracing = bpf_program__attach(
90*12bb6ca4SYonghong Song 		skel->progs.update_cookie_tracing);
91*12bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel->links.update_cookie_tracing, "prog_attach"))
92*12bb6ca4SYonghong Song 		goto out;
93*12bb6ca4SYonghong Song 
94*12bb6ca4SYonghong Song 	server_fd = start_server(AF_INET6, SOCK_STREAM, "::1", 0, 0);
95*12bb6ca4SYonghong Song 	if (!ASSERT_GE(server_fd, 0, "start_server"))
96*12bb6ca4SYonghong Song 		goto out;
97*12bb6ca4SYonghong Song 
98*12bb6ca4SYonghong Song 	client_fd = connect_to_fd(server_fd, 0);
99*12bb6ca4SYonghong Song 	if (!ASSERT_GE(client_fd, 0, "connect_to_fd"))
100*12bb6ca4SYonghong Song 		goto close_server_fd;
101*12bb6ca4SYonghong Song 
102*12bb6ca4SYonghong Song 	err = bpf_map_lookup_elem(bpf_map__fd(skel->maps.socket_cookies),
103*12bb6ca4SYonghong Song 				  &cgroup_fd, &val);
104*12bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_lookup(socket_cookies)"))
105*12bb6ca4SYonghong Song 		goto close_client_fd;
106*12bb6ca4SYonghong Song 
107*12bb6ca4SYonghong Song 	err = getsockname(client_fd, (struct sockaddr *)&addr, &addr_len);
108*12bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "getsockname"))
109*12bb6ca4SYonghong Song 		goto close_client_fd;
110*12bb6ca4SYonghong Song 
111*12bb6ca4SYonghong Song 	cookie_expected_value = (ntohs(addr.sin6_port) << 8) | 0xFF;
112*12bb6ca4SYonghong Song 	ASSERT_EQ(val.cookie_value, cookie_expected_value, "cookie_value");
113*12bb6ca4SYonghong Song 
114*12bb6ca4SYonghong Song close_client_fd:
115*12bb6ca4SYonghong Song 	close(client_fd);
116*12bb6ca4SYonghong Song close_server_fd:
117*12bb6ca4SYonghong Song 	close(server_fd);
118*12bb6ca4SYonghong Song out:
119*12bb6ca4SYonghong Song 	cgrp_ls_attach_cgroup__destroy(skel);
120*12bb6ca4SYonghong Song }
121*12bb6ca4SYonghong Song 
122*12bb6ca4SYonghong Song static void test_recursion(int cgroup_fd)
123*12bb6ca4SYonghong Song {
124*12bb6ca4SYonghong Song 	struct cgrp_ls_recursion *skel;
125*12bb6ca4SYonghong Song 	int err;
126*12bb6ca4SYonghong Song 
127*12bb6ca4SYonghong Song 	skel = cgrp_ls_recursion__open_and_load();
128*12bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
129*12bb6ca4SYonghong Song 		return;
130*12bb6ca4SYonghong Song 
131*12bb6ca4SYonghong Song 	err = cgrp_ls_recursion__attach(skel);
132*12bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "skel_attach"))
133*12bb6ca4SYonghong Song 		goto out;
134*12bb6ca4SYonghong Song 
135*12bb6ca4SYonghong Song 	/* trigger sys_enter, make sure it does not cause deadlock */
136*12bb6ca4SYonghong Song 	syscall(SYS_gettid);
137*12bb6ca4SYonghong Song 
138*12bb6ca4SYonghong Song out:
139*12bb6ca4SYonghong Song 	cgrp_ls_recursion__destroy(skel);
140*12bb6ca4SYonghong Song }
141*12bb6ca4SYonghong Song 
142*12bb6ca4SYonghong Song static void test_negative(void)
143*12bb6ca4SYonghong Song {
144*12bb6ca4SYonghong Song 	struct cgrp_ls_negative *skel;
145*12bb6ca4SYonghong Song 
146*12bb6ca4SYonghong Song 	skel = cgrp_ls_negative__open_and_load();
147*12bb6ca4SYonghong Song 	if (!ASSERT_ERR_PTR(skel, "skel_open_and_load")) {
148*12bb6ca4SYonghong Song 		cgrp_ls_negative__destroy(skel);
149*12bb6ca4SYonghong Song 		return;
150*12bb6ca4SYonghong Song 	}
151*12bb6ca4SYonghong Song }
152*12bb6ca4SYonghong Song 
153*12bb6ca4SYonghong Song void test_cgrp_local_storage(void)
154*12bb6ca4SYonghong Song {
155*12bb6ca4SYonghong Song 	int cgroup_fd;
156*12bb6ca4SYonghong Song 
157*12bb6ca4SYonghong Song 	cgroup_fd = test__join_cgroup("/cgrp_local_storage");
158*12bb6ca4SYonghong Song 	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup /cgrp_local_storage"))
159*12bb6ca4SYonghong Song 		return;
160*12bb6ca4SYonghong Song 
161*12bb6ca4SYonghong Song 	if (test__start_subtest("tp_btf"))
162*12bb6ca4SYonghong Song 		test_tp_btf(cgroup_fd);
163*12bb6ca4SYonghong Song 	if (test__start_subtest("attach_cgroup"))
164*12bb6ca4SYonghong Song 		test_attach_cgroup(cgroup_fd);
165*12bb6ca4SYonghong Song 	if (test__start_subtest("recursion"))
166*12bb6ca4SYonghong Song 		test_recursion(cgroup_fd);
167*12bb6ca4SYonghong Song 	if (test__start_subtest("negative"))
168*12bb6ca4SYonghong Song 		test_negative();
169*12bb6ca4SYonghong Song 
170*12bb6ca4SYonghong Song 	close(cgroup_fd);
171*12bb6ca4SYonghong Song }
172