1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2021 Facebook */
3
4#include "vmlinux.h"
5#include <bpf/bpf_helpers.h>
6#include "bpf_misc.h"
7
8char _license[] SEC("license") = "GPL";
9
10struct callback_ctx {
11	int output;
12};
13
14struct {
15	__uint(type, BPF_MAP_TYPE_HASH);
16	__uint(max_entries, 32);
17	__type(key, int);
18	__type(value, int);
19} map1 SEC(".maps");
20
21/* These should be set by the user program */
22u32 nested_callback_nr_loops;
23u32 stop_index = -1;
24u32 nr_loops;
25int pid;
26int callback_selector;
27
28/* Making these global variables so that the userspace program
29 * can verify the output through the skeleton
30 */
31int nr_loops_returned;
32int g_output;
33int err;
34
35static int callback(__u32 index, void *data)
36{
37	struct callback_ctx *ctx = data;
38
39	if (index >= stop_index)
40		return 1;
41
42	ctx->output += index;
43
44	return 0;
45}
46
47static int empty_callback(__u32 index, void *data)
48{
49	return 0;
50}
51
52static int nested_callback2(__u32 index, void *data)
53{
54	nr_loops_returned += bpf_loop(nested_callback_nr_loops, callback, data, 0);
55
56	return 0;
57}
58
59static int nested_callback1(__u32 index, void *data)
60{
61	bpf_loop(nested_callback_nr_loops, nested_callback2, data, 0);
62	return 0;
63}
64
65SEC("fentry/" SYS_PREFIX "sys_nanosleep")
66int test_prog(void *ctx)
67{
68	struct callback_ctx data = {};
69
70	if (bpf_get_current_pid_tgid() >> 32 != pid)
71		return 0;
72
73	nr_loops_returned = bpf_loop(nr_loops, callback, &data, 0);
74
75	if (nr_loops_returned < 0)
76		err = nr_loops_returned;
77	else
78		g_output = data.output;
79
80	return 0;
81}
82
83SEC("fentry/" SYS_PREFIX "sys_nanosleep")
84int prog_null_ctx(void *ctx)
85{
86	if (bpf_get_current_pid_tgid() >> 32 != pid)
87		return 0;
88
89	nr_loops_returned = bpf_loop(nr_loops, empty_callback, NULL, 0);
90
91	return 0;
92}
93
94SEC("fentry/" SYS_PREFIX "sys_nanosleep")
95int prog_invalid_flags(void *ctx)
96{
97	struct callback_ctx data = {};
98
99	if (bpf_get_current_pid_tgid() >> 32 != pid)
100		return 0;
101
102	err = bpf_loop(nr_loops, callback, &data, 1);
103
104	return 0;
105}
106
107SEC("fentry/" SYS_PREFIX "sys_nanosleep")
108int prog_nested_calls(void *ctx)
109{
110	struct callback_ctx data = {};
111
112	if (bpf_get_current_pid_tgid() >> 32 != pid)
113		return 0;
114
115	nr_loops_returned = 0;
116	bpf_loop(nr_loops, nested_callback1, &data, 0);
117
118	g_output = data.output;
119
120	return 0;
121}
122
123static int callback_set_f0(int i, void *ctx)
124{
125	g_output = 0xF0;
126	return 0;
127}
128
129static int callback_set_0f(int i, void *ctx)
130{
131	g_output = 0x0F;
132	return 0;
133}
134
135/*
136 * non-constant callback is a corner case for bpf_loop inline logic
137 */
138SEC("fentry/" SYS_PREFIX "sys_nanosleep")
139int prog_non_constant_callback(void *ctx)
140{
141	if (bpf_get_current_pid_tgid() >> 32 != pid)
142		return 0;
143
144	int (*callback)(int i, void *ctx);
145
146	g_output = 0;
147
148	if (callback_selector == 0x0F)
149		callback = callback_set_0f;
150	else
151		callback = callback_set_f0;
152
153	bpf_loop(1, callback, NULL, 0);
154
155	return 0;
156}
157
158static int stack_check_inner_callback(void *ctx)
159{
160	return 0;
161}
162
163static int map1_lookup_elem(int key)
164{
165	int *val = bpf_map_lookup_elem(&map1, &key);
166
167	return val ? *val : -1;
168}
169
170static void map1_update_elem(int key, int val)
171{
172	bpf_map_update_elem(&map1, &key, &val, BPF_ANY);
173}
174
175static int stack_check_outer_callback(void *ctx)
176{
177	int a = map1_lookup_elem(1);
178	int b = map1_lookup_elem(2);
179	int c = map1_lookup_elem(3);
180	int d = map1_lookup_elem(4);
181	int e = map1_lookup_elem(5);
182	int f = map1_lookup_elem(6);
183
184	bpf_loop(1, stack_check_inner_callback, NULL, 0);
185
186	map1_update_elem(1, a + 1);
187	map1_update_elem(2, b + 1);
188	map1_update_elem(3, c + 1);
189	map1_update_elem(4, d + 1);
190	map1_update_elem(5, e + 1);
191	map1_update_elem(6, f + 1);
192
193	return 0;
194}
195
196/* Some of the local variables in stack_check and
197 * stack_check_outer_callback would be allocated on stack by
198 * compiler. This test should verify that stack content for these
199 * variables is preserved between calls to bpf_loop (might be an issue
200 * if loop inlining allocates stack slots incorrectly).
201 */
202SEC("fentry/" SYS_PREFIX "sys_nanosleep")
203int stack_check(void *ctx)
204{
205	if (bpf_get_current_pid_tgid() >> 32 != pid)
206		return 0;
207
208	int a = map1_lookup_elem(7);
209	int b = map1_lookup_elem(8);
210	int c = map1_lookup_elem(9);
211	int d = map1_lookup_elem(10);
212	int e = map1_lookup_elem(11);
213	int f = map1_lookup_elem(12);
214
215	bpf_loop(1, stack_check_outer_callback, NULL, 0);
216
217	map1_update_elem(7,  a + 1);
218	map1_update_elem(8, b + 1);
219	map1_update_elem(9, c + 1);
220	map1_update_elem(10, d + 1);
221	map1_update_elem(11, e + 1);
222	map1_update_elem(12, f + 1);
223
224	return 0;
225}
226