112bb6ca4SYonghong Song // SPDX-License-Identifier: GPL-2.0
212bb6ca4SYonghong Song /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates.*/
312bb6ca4SYonghong Song 
412bb6ca4SYonghong Song #define _GNU_SOURCE
512bb6ca4SYonghong Song #include <unistd.h>
612bb6ca4SYonghong Song #include <sys/syscall.h>
712bb6ca4SYonghong Song #include <sys/types.h>
812bb6ca4SYonghong Song #include <test_progs.h>
912bb6ca4SYonghong Song #include "cgrp_ls_tp_btf.skel.h"
1012bb6ca4SYonghong Song #include "cgrp_ls_recursion.skel.h"
1112bb6ca4SYonghong Song #include "cgrp_ls_attach_cgroup.skel.h"
1212bb6ca4SYonghong Song #include "cgrp_ls_negative.skel.h"
13*41d76c72SYonghong Song #include "cgrp_ls_sleepable.skel.h"
1412bb6ca4SYonghong Song #include "network_helpers.h"
15*41d76c72SYonghong Song #include "cgroup_helpers.h"
1612bb6ca4SYonghong Song 
1712bb6ca4SYonghong Song struct socket_cookie {
1812bb6ca4SYonghong Song 	__u64 cookie_key;
1912bb6ca4SYonghong Song 	__u32 cookie_value;
2012bb6ca4SYonghong Song };
2112bb6ca4SYonghong Song 
2212bb6ca4SYonghong Song static void test_tp_btf(int cgroup_fd)
2312bb6ca4SYonghong Song {
2412bb6ca4SYonghong Song 	struct cgrp_ls_tp_btf *skel;
2512bb6ca4SYonghong Song 	long val1 = 1, val2 = 0;
2612bb6ca4SYonghong Song 	int err;
2712bb6ca4SYonghong Song 
2812bb6ca4SYonghong Song 	skel = cgrp_ls_tp_btf__open_and_load();
2912bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
3012bb6ca4SYonghong Song 		return;
3112bb6ca4SYonghong Song 
3212bb6ca4SYonghong Song 	/* populate a value in map_b */
3312bb6ca4SYonghong Song 	err = bpf_map_update_elem(bpf_map__fd(skel->maps.map_b), &cgroup_fd, &val1, BPF_ANY);
3412bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_update_elem"))
3512bb6ca4SYonghong Song 		goto out;
3612bb6ca4SYonghong Song 
3712bb6ca4SYonghong Song 	/* check value */
3812bb6ca4SYonghong Song 	err = bpf_map_lookup_elem(bpf_map__fd(skel->maps.map_b), &cgroup_fd, &val2);
3912bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_lookup_elem"))
4012bb6ca4SYonghong Song 		goto out;
4112bb6ca4SYonghong Song 	if (!ASSERT_EQ(val2, 1, "map_lookup_elem, invalid val"))
4212bb6ca4SYonghong Song 		goto out;
4312bb6ca4SYonghong Song 
4412bb6ca4SYonghong Song 	/* delete value */
4512bb6ca4SYonghong Song 	err = bpf_map_delete_elem(bpf_map__fd(skel->maps.map_b), &cgroup_fd);
4612bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_delete_elem"))
4712bb6ca4SYonghong Song 		goto out;
4812bb6ca4SYonghong Song 
4912bb6ca4SYonghong Song 	skel->bss->target_pid = syscall(SYS_gettid);
5012bb6ca4SYonghong Song 
5112bb6ca4SYonghong Song 	err = cgrp_ls_tp_btf__attach(skel);
5212bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "skel_attach"))
5312bb6ca4SYonghong Song 		goto out;
5412bb6ca4SYonghong Song 
5512bb6ca4SYonghong Song 	syscall(SYS_gettid);
5612bb6ca4SYonghong Song 	syscall(SYS_gettid);
5712bb6ca4SYonghong Song 
5812bb6ca4SYonghong Song 	skel->bss->target_pid = 0;
5912bb6ca4SYonghong Song 
6012bb6ca4SYonghong Song 	/* 3x syscalls: 1x attach and 2x gettid */
6112bb6ca4SYonghong Song 	ASSERT_EQ(skel->bss->enter_cnt, 3, "enter_cnt");
6212bb6ca4SYonghong Song 	ASSERT_EQ(skel->bss->exit_cnt, 3, "exit_cnt");
6312bb6ca4SYonghong Song 	ASSERT_EQ(skel->bss->mismatch_cnt, 0, "mismatch_cnt");
6412bb6ca4SYonghong Song out:
6512bb6ca4SYonghong Song 	cgrp_ls_tp_btf__destroy(skel);
6612bb6ca4SYonghong Song }
6712bb6ca4SYonghong Song 
6812bb6ca4SYonghong Song static void test_attach_cgroup(int cgroup_fd)
6912bb6ca4SYonghong Song {
7012bb6ca4SYonghong Song 	int server_fd = 0, client_fd = 0, err = 0;
7112bb6ca4SYonghong Song 	socklen_t addr_len = sizeof(struct sockaddr_in6);
7212bb6ca4SYonghong Song 	struct cgrp_ls_attach_cgroup *skel;
7312bb6ca4SYonghong Song 	__u32 cookie_expected_value;
7412bb6ca4SYonghong Song 	struct sockaddr_in6 addr;
7512bb6ca4SYonghong Song 	struct socket_cookie val;
7612bb6ca4SYonghong Song 
7712bb6ca4SYonghong Song 	skel = cgrp_ls_attach_cgroup__open_and_load();
7812bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open"))
7912bb6ca4SYonghong Song 		return;
8012bb6ca4SYonghong Song 
8112bb6ca4SYonghong Song 	skel->links.set_cookie = bpf_program__attach_cgroup(
8212bb6ca4SYonghong Song 		skel->progs.set_cookie, cgroup_fd);
8312bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel->links.set_cookie, "prog_attach"))
8412bb6ca4SYonghong Song 		goto out;
8512bb6ca4SYonghong Song 
8612bb6ca4SYonghong Song 	skel->links.update_cookie_sockops = bpf_program__attach_cgroup(
8712bb6ca4SYonghong Song 		skel->progs.update_cookie_sockops, cgroup_fd);
8812bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel->links.update_cookie_sockops, "prog_attach"))
8912bb6ca4SYonghong Song 		goto out;
9012bb6ca4SYonghong Song 
9112bb6ca4SYonghong Song 	skel->links.update_cookie_tracing = bpf_program__attach(
9212bb6ca4SYonghong Song 		skel->progs.update_cookie_tracing);
9312bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel->links.update_cookie_tracing, "prog_attach"))
9412bb6ca4SYonghong Song 		goto out;
9512bb6ca4SYonghong Song 
9612bb6ca4SYonghong Song 	server_fd = start_server(AF_INET6, SOCK_STREAM, "::1", 0, 0);
9712bb6ca4SYonghong Song 	if (!ASSERT_GE(server_fd, 0, "start_server"))
9812bb6ca4SYonghong Song 		goto out;
9912bb6ca4SYonghong Song 
10012bb6ca4SYonghong Song 	client_fd = connect_to_fd(server_fd, 0);
10112bb6ca4SYonghong Song 	if (!ASSERT_GE(client_fd, 0, "connect_to_fd"))
10212bb6ca4SYonghong Song 		goto close_server_fd;
10312bb6ca4SYonghong Song 
10412bb6ca4SYonghong Song 	err = bpf_map_lookup_elem(bpf_map__fd(skel->maps.socket_cookies),
10512bb6ca4SYonghong Song 				  &cgroup_fd, &val);
10612bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "map_lookup(socket_cookies)"))
10712bb6ca4SYonghong Song 		goto close_client_fd;
10812bb6ca4SYonghong Song 
10912bb6ca4SYonghong Song 	err = getsockname(client_fd, (struct sockaddr *)&addr, &addr_len);
11012bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "getsockname"))
11112bb6ca4SYonghong Song 		goto close_client_fd;
11212bb6ca4SYonghong Song 
11312bb6ca4SYonghong Song 	cookie_expected_value = (ntohs(addr.sin6_port) << 8) | 0xFF;
11412bb6ca4SYonghong Song 	ASSERT_EQ(val.cookie_value, cookie_expected_value, "cookie_value");
11512bb6ca4SYonghong Song 
11612bb6ca4SYonghong Song close_client_fd:
11712bb6ca4SYonghong Song 	close(client_fd);
11812bb6ca4SYonghong Song close_server_fd:
11912bb6ca4SYonghong Song 	close(server_fd);
12012bb6ca4SYonghong Song out:
12112bb6ca4SYonghong Song 	cgrp_ls_attach_cgroup__destroy(skel);
12212bb6ca4SYonghong Song }
12312bb6ca4SYonghong Song 
12412bb6ca4SYonghong Song static void test_recursion(int cgroup_fd)
12512bb6ca4SYonghong Song {
12612bb6ca4SYonghong Song 	struct cgrp_ls_recursion *skel;
12712bb6ca4SYonghong Song 	int err;
12812bb6ca4SYonghong Song 
12912bb6ca4SYonghong Song 	skel = cgrp_ls_recursion__open_and_load();
13012bb6ca4SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open_and_load"))
13112bb6ca4SYonghong Song 		return;
13212bb6ca4SYonghong Song 
13312bb6ca4SYonghong Song 	err = cgrp_ls_recursion__attach(skel);
13412bb6ca4SYonghong Song 	if (!ASSERT_OK(err, "skel_attach"))
13512bb6ca4SYonghong Song 		goto out;
13612bb6ca4SYonghong Song 
13712bb6ca4SYonghong Song 	/* trigger sys_enter, make sure it does not cause deadlock */
13812bb6ca4SYonghong Song 	syscall(SYS_gettid);
13912bb6ca4SYonghong Song 
14012bb6ca4SYonghong Song out:
14112bb6ca4SYonghong Song 	cgrp_ls_recursion__destroy(skel);
14212bb6ca4SYonghong Song }
14312bb6ca4SYonghong Song 
14412bb6ca4SYonghong Song static void test_negative(void)
14512bb6ca4SYonghong Song {
14612bb6ca4SYonghong Song 	struct cgrp_ls_negative *skel;
14712bb6ca4SYonghong Song 
14812bb6ca4SYonghong Song 	skel = cgrp_ls_negative__open_and_load();
14912bb6ca4SYonghong Song 	if (!ASSERT_ERR_PTR(skel, "skel_open_and_load")) {
15012bb6ca4SYonghong Song 		cgrp_ls_negative__destroy(skel);
15112bb6ca4SYonghong Song 		return;
15212bb6ca4SYonghong Song 	}
15312bb6ca4SYonghong Song }
15412bb6ca4SYonghong Song 
155*41d76c72SYonghong Song static void test_cgroup_iter_sleepable(int cgroup_fd, __u64 cgroup_id)
156*41d76c72SYonghong Song {
157*41d76c72SYonghong Song 	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
158*41d76c72SYonghong Song 	union bpf_iter_link_info linfo;
159*41d76c72SYonghong Song 	struct cgrp_ls_sleepable *skel;
160*41d76c72SYonghong Song 	struct bpf_link *link;
161*41d76c72SYonghong Song 	int err, iter_fd;
162*41d76c72SYonghong Song 	char buf[16];
163*41d76c72SYonghong Song 
164*41d76c72SYonghong Song 	skel = cgrp_ls_sleepable__open();
165*41d76c72SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open"))
166*41d76c72SYonghong Song 		return;
167*41d76c72SYonghong Song 
168*41d76c72SYonghong Song 	bpf_program__set_autoload(skel->progs.cgroup_iter, true);
169*41d76c72SYonghong Song 	err = cgrp_ls_sleepable__load(skel);
170*41d76c72SYonghong Song 	if (!ASSERT_OK(err, "skel_load"))
171*41d76c72SYonghong Song 		goto out;
172*41d76c72SYonghong Song 
173*41d76c72SYonghong Song 	memset(&linfo, 0, sizeof(linfo));
174*41d76c72SYonghong Song 	linfo.cgroup.cgroup_fd = cgroup_fd;
175*41d76c72SYonghong Song 	linfo.cgroup.order = BPF_CGROUP_ITER_SELF_ONLY;
176*41d76c72SYonghong Song 	opts.link_info = &linfo;
177*41d76c72SYonghong Song 	opts.link_info_len = sizeof(linfo);
178*41d76c72SYonghong Song 	link = bpf_program__attach_iter(skel->progs.cgroup_iter, &opts);
179*41d76c72SYonghong Song 	if (!ASSERT_OK_PTR(link, "attach_iter"))
180*41d76c72SYonghong Song 		goto out;
181*41d76c72SYonghong Song 
182*41d76c72SYonghong Song 	iter_fd = bpf_iter_create(bpf_link__fd(link));
183*41d76c72SYonghong Song 	if (!ASSERT_GE(iter_fd, 0, "iter_create"))
184*41d76c72SYonghong Song 		goto out;
185*41d76c72SYonghong Song 
186*41d76c72SYonghong Song 	/* trigger the program run */
187*41d76c72SYonghong Song 	(void)read(iter_fd, buf, sizeof(buf));
188*41d76c72SYonghong Song 
189*41d76c72SYonghong Song 	ASSERT_EQ(skel->bss->cgroup_id, cgroup_id, "cgroup_id");
190*41d76c72SYonghong Song 
191*41d76c72SYonghong Song 	close(iter_fd);
192*41d76c72SYonghong Song out:
193*41d76c72SYonghong Song 	cgrp_ls_sleepable__destroy(skel);
194*41d76c72SYonghong Song }
195*41d76c72SYonghong Song 
196*41d76c72SYonghong Song static void test_no_rcu_lock(__u64 cgroup_id)
197*41d76c72SYonghong Song {
198*41d76c72SYonghong Song 	struct cgrp_ls_sleepable *skel;
199*41d76c72SYonghong Song 	int err;
200*41d76c72SYonghong Song 
201*41d76c72SYonghong Song 	skel = cgrp_ls_sleepable__open();
202*41d76c72SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open"))
203*41d76c72SYonghong Song 		return;
204*41d76c72SYonghong Song 
205*41d76c72SYonghong Song 	skel->bss->target_pid = syscall(SYS_gettid);
206*41d76c72SYonghong Song 
207*41d76c72SYonghong Song 	bpf_program__set_autoload(skel->progs.no_rcu_lock, true);
208*41d76c72SYonghong Song 	err = cgrp_ls_sleepable__load(skel);
209*41d76c72SYonghong Song 	if (!ASSERT_OK(err, "skel_load"))
210*41d76c72SYonghong Song 		goto out;
211*41d76c72SYonghong Song 
212*41d76c72SYonghong Song 	err = cgrp_ls_sleepable__attach(skel);
213*41d76c72SYonghong Song 	if (!ASSERT_OK(err, "skel_attach"))
214*41d76c72SYonghong Song 		goto out;
215*41d76c72SYonghong Song 
216*41d76c72SYonghong Song 	syscall(SYS_getpgid);
217*41d76c72SYonghong Song 
218*41d76c72SYonghong Song 	ASSERT_EQ(skel->bss->cgroup_id, cgroup_id, "cgroup_id");
219*41d76c72SYonghong Song out:
220*41d76c72SYonghong Song 	cgrp_ls_sleepable__destroy(skel);
221*41d76c72SYonghong Song }
222*41d76c72SYonghong Song 
223*41d76c72SYonghong Song static void test_rcu_lock(void)
224*41d76c72SYonghong Song {
225*41d76c72SYonghong Song 	struct cgrp_ls_sleepable *skel;
226*41d76c72SYonghong Song 	int err;
227*41d76c72SYonghong Song 
228*41d76c72SYonghong Song 	skel = cgrp_ls_sleepable__open();
229*41d76c72SYonghong Song 	if (!ASSERT_OK_PTR(skel, "skel_open"))
230*41d76c72SYonghong Song 		return;
231*41d76c72SYonghong Song 
232*41d76c72SYonghong Song 	bpf_program__set_autoload(skel->progs.yes_rcu_lock, true);
233*41d76c72SYonghong Song 	err = cgrp_ls_sleepable__load(skel);
234*41d76c72SYonghong Song 	ASSERT_ERR(err, "skel_load");
235*41d76c72SYonghong Song 
236*41d76c72SYonghong Song 	cgrp_ls_sleepable__destroy(skel);
237*41d76c72SYonghong Song }
238*41d76c72SYonghong Song 
23912bb6ca4SYonghong Song void test_cgrp_local_storage(void)
24012bb6ca4SYonghong Song {
241*41d76c72SYonghong Song 	__u64 cgroup_id;
24212bb6ca4SYonghong Song 	int cgroup_fd;
24312bb6ca4SYonghong Song 
24412bb6ca4SYonghong Song 	cgroup_fd = test__join_cgroup("/cgrp_local_storage");
24512bb6ca4SYonghong Song 	if (!ASSERT_GE(cgroup_fd, 0, "join_cgroup /cgrp_local_storage"))
24612bb6ca4SYonghong Song 		return;
24712bb6ca4SYonghong Song 
248*41d76c72SYonghong Song 	cgroup_id = get_cgroup_id("/cgrp_local_storage");
24912bb6ca4SYonghong Song 	if (test__start_subtest("tp_btf"))
25012bb6ca4SYonghong Song 		test_tp_btf(cgroup_fd);
25112bb6ca4SYonghong Song 	if (test__start_subtest("attach_cgroup"))
25212bb6ca4SYonghong Song 		test_attach_cgroup(cgroup_fd);
25312bb6ca4SYonghong Song 	if (test__start_subtest("recursion"))
25412bb6ca4SYonghong Song 		test_recursion(cgroup_fd);
25512bb6ca4SYonghong Song 	if (test__start_subtest("negative"))
25612bb6ca4SYonghong Song 		test_negative();
257*41d76c72SYonghong Song 	if (test__start_subtest("cgroup_iter_sleepable"))
258*41d76c72SYonghong Song 		test_cgroup_iter_sleepable(cgroup_fd, cgroup_id);
259*41d76c72SYonghong Song 	if (test__start_subtest("no_rcu_lock"))
260*41d76c72SYonghong Song 		test_no_rcu_lock(cgroup_id);
261*41d76c72SYonghong Song 	if (test__start_subtest("rcu_lock"))
262*41d76c72SYonghong Song 		test_rcu_lock();
26312bb6ca4SYonghong Song 
26412bb6ca4SYonghong Song 	close(cgroup_fd);
26512bb6ca4SYonghong Song }
266