1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2022 Google */
3 
4 #include <test_progs.h>
5 #include <bpf/libbpf.h>
6 #include <bpf/btf.h>
7 #include "cgroup_iter.skel.h"
8 #include "cgroup_helpers.h"
9 
10 #define ROOT           0
11 #define PARENT         1
12 #define CHILD1         2
13 #define CHILD2         3
14 #define NUM_CGROUPS    4
15 
16 #define PROLOGUE       "prologue\n"
17 #define EPILOGUE       "epilogue\n"
18 
19 static const char *cg_path[] = {
20 	"/", "/parent", "/parent/child1", "/parent/child2"
21 };
22 
23 static int cg_fd[] = {-1, -1, -1, -1};
24 static unsigned long long cg_id[] = {0, 0, 0, 0};
25 static char expected_output[64];
26 
27 static int setup_cgroups(void)
28 {
29 	int fd, i = 0;
30 
31 	for (i = 0; i < NUM_CGROUPS; i++) {
32 		fd = create_and_get_cgroup(cg_path[i]);
33 		if (fd < 0)
34 			return fd;
35 
36 		cg_fd[i] = fd;
37 		cg_id[i] = get_cgroup_id(cg_path[i]);
38 	}
39 	return 0;
40 }
41 
42 static void cleanup_cgroups(void)
43 {
44 	int i;
45 
46 	for (i = 0; i < NUM_CGROUPS; i++)
47 		close(cg_fd[i]);
48 }
49 
50 static void read_from_cgroup_iter(struct bpf_program *prog, int cgroup_fd,
51 				  int order, const char *testname)
52 {
53 	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
54 	union bpf_iter_link_info linfo;
55 	struct bpf_link *link;
56 	int len, iter_fd;
57 	static char buf[128];
58 	size_t left;
59 	char *p;
60 
61 	memset(&linfo, 0, sizeof(linfo));
62 	linfo.cgroup.cgroup_fd = cgroup_fd;
63 	linfo.cgroup.order = order;
64 	opts.link_info = &linfo;
65 	opts.link_info_len = sizeof(linfo);
66 
67 	link = bpf_program__attach_iter(prog, &opts);
68 	if (!ASSERT_OK_PTR(link, "attach_iter"))
69 		return;
70 
71 	iter_fd = bpf_iter_create(bpf_link__fd(link));
72 	if (iter_fd < 0)
73 		goto free_link;
74 
75 	memset(buf, 0, sizeof(buf));
76 	left = ARRAY_SIZE(buf);
77 	p = buf;
78 	while ((len = read(iter_fd, p, left)) > 0) {
79 		p += len;
80 		left -= len;
81 	}
82 
83 	ASSERT_STREQ(buf, expected_output, testname);
84 
85 	/* read() after iter finishes should be ok. */
86 	if (len == 0)
87 		ASSERT_OK(read(iter_fd, buf, sizeof(buf)), "second_read");
88 
89 	close(iter_fd);
90 free_link:
91 	bpf_link__destroy(link);
92 }
93 
94 /* Invalid cgroup. */
95 static void test_invalid_cgroup(struct cgroup_iter *skel)
96 {
97 	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
98 	union bpf_iter_link_info linfo;
99 	struct bpf_link *link;
100 
101 	memset(&linfo, 0, sizeof(linfo));
102 	linfo.cgroup.cgroup_fd = (__u32)-1;
103 	opts.link_info = &linfo;
104 	opts.link_info_len = sizeof(linfo);
105 
106 	link = bpf_program__attach_iter(skel->progs.cgroup_id_printer, &opts);
107 	ASSERT_ERR_PTR(link, "attach_iter");
108 	bpf_link__destroy(link);
109 }
110 
111 /* Specifying both cgroup_fd and cgroup_id is invalid. */
112 static void test_invalid_cgroup_spec(struct cgroup_iter *skel)
113 {
114 	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
115 	union bpf_iter_link_info linfo;
116 	struct bpf_link *link;
117 
118 	memset(&linfo, 0, sizeof(linfo));
119 	linfo.cgroup.cgroup_fd = (__u32)cg_fd[PARENT];
120 	linfo.cgroup.cgroup_id = (__u64)cg_id[PARENT];
121 	opts.link_info = &linfo;
122 	opts.link_info_len = sizeof(linfo);
123 
124 	link = bpf_program__attach_iter(skel->progs.cgroup_id_printer, &opts);
125 	ASSERT_ERR_PTR(link, "attach_iter");
126 	bpf_link__destroy(link);
127 }
128 
129 /* Preorder walk prints parent and child in order. */
130 static void test_walk_preorder(struct cgroup_iter *skel)
131 {
132 	snprintf(expected_output, sizeof(expected_output),
133 		 PROLOGUE "%8llu\n%8llu\n%8llu\n" EPILOGUE,
134 		 cg_id[PARENT], cg_id[CHILD1], cg_id[CHILD2]);
135 
136 	read_from_cgroup_iter(skel->progs.cgroup_id_printer, cg_fd[PARENT],
137 			      BPF_CGROUP_ITER_DESCENDANTS_PRE, "preorder");
138 }
139 
140 /* Postorder walk prints child and parent in order. */
141 static void test_walk_postorder(struct cgroup_iter *skel)
142 {
143 	snprintf(expected_output, sizeof(expected_output),
144 		 PROLOGUE "%8llu\n%8llu\n%8llu\n" EPILOGUE,
145 		 cg_id[CHILD1], cg_id[CHILD2], cg_id[PARENT]);
146 
147 	read_from_cgroup_iter(skel->progs.cgroup_id_printer, cg_fd[PARENT],
148 			      BPF_CGROUP_ITER_DESCENDANTS_POST, "postorder");
149 }
150 
151 /* Walking parents prints parent and then root. */
152 static void test_walk_ancestors_up(struct cgroup_iter *skel)
153 {
154 	/* terminate the walk when ROOT is met. */
155 	skel->bss->terminal_cgroup = cg_id[ROOT];
156 
157 	snprintf(expected_output, sizeof(expected_output),
158 		 PROLOGUE "%8llu\n%8llu\n" EPILOGUE,
159 		 cg_id[PARENT], cg_id[ROOT]);
160 
161 	read_from_cgroup_iter(skel->progs.cgroup_id_printer, cg_fd[PARENT],
162 			      BPF_CGROUP_ITER_ANCESTORS_UP, "ancestors_up");
163 
164 	skel->bss->terminal_cgroup = 0;
165 }
166 
167 /* Early termination prints parent only. */
168 static void test_early_termination(struct cgroup_iter *skel)
169 {
170 	/* terminate the walk after the first element is processed. */
171 	skel->bss->terminate_early = 1;
172 
173 	snprintf(expected_output, sizeof(expected_output),
174 		 PROLOGUE "%8llu\n" EPILOGUE, cg_id[PARENT]);
175 
176 	read_from_cgroup_iter(skel->progs.cgroup_id_printer, cg_fd[PARENT],
177 			      BPF_CGROUP_ITER_DESCENDANTS_PRE, "early_termination");
178 
179 	skel->bss->terminate_early = 0;
180 }
181 
182 /* Waling self prints self only. */
183 static void test_walk_self_only(struct cgroup_iter *skel)
184 {
185 	snprintf(expected_output, sizeof(expected_output),
186 		 PROLOGUE "%8llu\n" EPILOGUE, cg_id[PARENT]);
187 
188 	read_from_cgroup_iter(skel->progs.cgroup_id_printer, cg_fd[PARENT],
189 			      BPF_CGROUP_ITER_SELF_ONLY, "self_only");
190 }
191 
192 void test_cgroup_iter(void)
193 {
194 	struct cgroup_iter *skel = NULL;
195 
196 	if (setup_cgroup_environment())
197 		return;
198 
199 	if (setup_cgroups())
200 		goto out;
201 
202 	skel = cgroup_iter__open_and_load();
203 	if (!ASSERT_OK_PTR(skel, "cgroup_iter__open_and_load"))
204 		goto out;
205 
206 	if (test__start_subtest("cgroup_iter__invalid_cgroup"))
207 		test_invalid_cgroup(skel);
208 	if (test__start_subtest("cgroup_iter__invalid_cgroup_spec"))
209 		test_invalid_cgroup_spec(skel);
210 	if (test__start_subtest("cgroup_iter__preorder"))
211 		test_walk_preorder(skel);
212 	if (test__start_subtest("cgroup_iter__postorder"))
213 		test_walk_postorder(skel);
214 	if (test__start_subtest("cgroup_iter__ancestors_up_walk"))
215 		test_walk_ancestors_up(skel);
216 	if (test__start_subtest("cgroup_iter__early_termination"))
217 		test_early_termination(skel);
218 	if (test__start_subtest("cgroup_iter__self_only"))
219 		test_walk_self_only(skel);
220 out:
221 	cgroup_iter__destroy(skel);
222 	cleanup_cgroups();
223 	cleanup_cgroup_environment();
224 }
225