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
 19static const char *cg_path[] = {
 20	"/", "/parent", "/parent/child1", "/parent/child2"
 21};
 22
 23static int cg_fd[] = {-1, -1, -1, -1};
 24static unsigned long long cg_id[] = {0, 0, 0, 0};
 25static char expected_output[64];
 26
 27static 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
 42static void cleanup_cgroups(void)
 43{
 44	int i;
 45
 46	for (i = 0; i < NUM_CGROUPS; i++)
 47		close(cg_fd[i]);
 48}
 49
 50static 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);
 90free_link:
 91	bpf_link__destroy(link);
 92}
 93
 94/* Invalid cgroup. */
 95static 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. */
112static 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. */
130static 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. */
141static 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. */
152static 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. */
168static 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. */
183static 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
192static void test_walk_dead_self_only(struct cgroup_iter *skel)
193{
194	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
195	char expected_output[128], buf[128];
196	const char *cgrp_name = "/dead";
197	union bpf_iter_link_info linfo;
198	int len, cgrp_fd, iter_fd;
199	struct bpf_link *link;
200	size_t left;
201	char *p;
202
203	cgrp_fd = create_and_get_cgroup(cgrp_name);
204	if (!ASSERT_GE(cgrp_fd, 0, "create cgrp"))
205		return;
206
207	/* The cgroup will be dead during read() iteration, so it only has
208	 * epilogue in the output
209	 */
210	snprintf(expected_output, sizeof(expected_output), EPILOGUE);
211
212	memset(&linfo, 0, sizeof(linfo));
213	linfo.cgroup.cgroup_fd = cgrp_fd;
214	linfo.cgroup.order = BPF_CGROUP_ITER_SELF_ONLY;
215	opts.link_info = &linfo;
216	opts.link_info_len = sizeof(linfo);
217
218	link = bpf_program__attach_iter(skel->progs.cgroup_id_printer, &opts);
219	if (!ASSERT_OK_PTR(link, "attach_iter"))
220		goto close_cgrp;
221
222	iter_fd = bpf_iter_create(bpf_link__fd(link));
223	if (!ASSERT_GE(iter_fd, 0, "iter_create"))
224		goto free_link;
225
226	/* Close link fd and cgroup fd */
227	bpf_link__destroy(link);
228	close(cgrp_fd);
229
230	/* Remove cgroup to mark it as dead */
231	remove_cgroup(cgrp_name);
232
233	/* Two kern_sync_rcu() and usleep() pairs are used to wait for the
234	 * releases of cgroup css, and the last kern_sync_rcu() and usleep()
235	 * pair is used to wait for the free of cgroup itself.
236	 */
237	kern_sync_rcu();
238	usleep(8000);
239	kern_sync_rcu();
240	usleep(8000);
241	kern_sync_rcu();
242	usleep(1000);
243
244	memset(buf, 0, sizeof(buf));
245	left = ARRAY_SIZE(buf);
246	p = buf;
247	while ((len = read(iter_fd, p, left)) > 0) {
248		p += len;
249		left -= len;
250	}
251
252	ASSERT_STREQ(buf, expected_output, "dead cgroup output");
253
254	/* read() after iter finishes should be ok. */
255	if (len == 0)
256		ASSERT_OK(read(iter_fd, buf, sizeof(buf)), "second_read");
257
258	close(iter_fd);
259	return;
260free_link:
261	bpf_link__destroy(link);
262close_cgrp:
263	close(cgrp_fd);
264}
265
266void test_cgroup_iter(void)
267{
268	struct cgroup_iter *skel = NULL;
269
270	if (setup_cgroup_environment())
271		return;
272
273	if (setup_cgroups())
274		goto out;
275
276	skel = cgroup_iter__open_and_load();
277	if (!ASSERT_OK_PTR(skel, "cgroup_iter__open_and_load"))
278		goto out;
279
280	if (test__start_subtest("cgroup_iter__invalid_cgroup"))
281		test_invalid_cgroup(skel);
282	if (test__start_subtest("cgroup_iter__invalid_cgroup_spec"))
283		test_invalid_cgroup_spec(skel);
284	if (test__start_subtest("cgroup_iter__preorder"))
285		test_walk_preorder(skel);
286	if (test__start_subtest("cgroup_iter__postorder"))
287		test_walk_postorder(skel);
288	if (test__start_subtest("cgroup_iter__ancestors_up_walk"))
289		test_walk_ancestors_up(skel);
290	if (test__start_subtest("cgroup_iter__early_termination"))
291		test_early_termination(skel);
292	if (test__start_subtest("cgroup_iter__self_only"))
293		test_walk_self_only(skel);
294	if (test__start_subtest("cgroup_iter__dead_self_only"))
295		test_walk_dead_self_only(skel);
296out:
297	cgroup_iter__destroy(skel);
298	cleanup_cgroups();
299	cleanup_cgroup_environment();
300}