1e5a9df51SDavid Vernet // SPDX-License-Identifier: GPL-2.0
2e5a9df51SDavid Vernet /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
3e5a9df51SDavid Vernet 
4e5a9df51SDavid Vernet #include <linux/bpf.h>
5e5a9df51SDavid Vernet #include <bpf/bpf_helpers.h>
6e5a9df51SDavid Vernet #include "bpf_misc.h"
7e5a9df51SDavid Vernet #include "test_user_ringbuf.h"
8e5a9df51SDavid Vernet 
9e5a9df51SDavid Vernet char _license[] SEC("license") = "GPL";
10e5a9df51SDavid Vernet 
11e5a9df51SDavid Vernet struct {
12e5a9df51SDavid Vernet 	__uint(type, BPF_MAP_TYPE_USER_RINGBUF);
13e5a9df51SDavid Vernet } user_ringbuf SEC(".maps");
14e5a9df51SDavid Vernet 
15e5a9df51SDavid Vernet struct {
16e5a9df51SDavid Vernet 	__uint(type, BPF_MAP_TYPE_RINGBUF);
17e5a9df51SDavid Vernet } kernel_ringbuf SEC(".maps");
18e5a9df51SDavid Vernet 
19e5a9df51SDavid Vernet /* inputs */
20e5a9df51SDavid Vernet int pid, err, val;
21e5a9df51SDavid Vernet 
22e5a9df51SDavid Vernet int read = 0;
23e5a9df51SDavid Vernet 
24e5a9df51SDavid Vernet /* Counter used for end-to-end protocol test */
25e5a9df51SDavid Vernet __u64 kern_mutated = 0;
26e5a9df51SDavid Vernet __u64 user_mutated = 0;
27e5a9df51SDavid Vernet __u64 expected_user_mutated = 0;
28e5a9df51SDavid Vernet 
29e5a9df51SDavid Vernet static int
is_test_process(void)30e5a9df51SDavid Vernet is_test_process(void)
31e5a9df51SDavid Vernet {
32e5a9df51SDavid Vernet 	int cur_pid = bpf_get_current_pid_tgid() >> 32;
33e5a9df51SDavid Vernet 
34e5a9df51SDavid Vernet 	return cur_pid == pid;
35e5a9df51SDavid Vernet }
36e5a9df51SDavid Vernet 
37e5a9df51SDavid Vernet static long
record_sample(struct bpf_dynptr * dynptr,void * context)38e5a9df51SDavid Vernet record_sample(struct bpf_dynptr *dynptr, void *context)
39e5a9df51SDavid Vernet {
40e5a9df51SDavid Vernet 	const struct sample *sample = NULL;
41e5a9df51SDavid Vernet 	struct sample stack_sample;
42e5a9df51SDavid Vernet 	int status;
43e5a9df51SDavid Vernet 	static int num_calls;
44e5a9df51SDavid Vernet 
45e5a9df51SDavid Vernet 	if (num_calls++ % 2 == 0) {
46e5a9df51SDavid Vernet 		status = bpf_dynptr_read(&stack_sample, sizeof(stack_sample), dynptr, 0, 0);
47e5a9df51SDavid Vernet 		if (status) {
48e5a9df51SDavid Vernet 			bpf_printk("bpf_dynptr_read() failed: %d\n", status);
49e5a9df51SDavid Vernet 			err = 1;
506e44b9f3SDavid Vernet 			return 1;
51e5a9df51SDavid Vernet 		}
52e5a9df51SDavid Vernet 	} else {
53e5a9df51SDavid Vernet 		sample = bpf_dynptr_data(dynptr, 0, sizeof(*sample));
54e5a9df51SDavid Vernet 		if (!sample) {
55e5a9df51SDavid Vernet 			bpf_printk("Unexpectedly failed to get sample\n");
56e5a9df51SDavid Vernet 			err = 2;
576e44b9f3SDavid Vernet 			return 1;
58e5a9df51SDavid Vernet 		}
59e5a9df51SDavid Vernet 		stack_sample = *sample;
60e5a9df51SDavid Vernet 	}
61e5a9df51SDavid Vernet 
62e5a9df51SDavid Vernet 	__sync_fetch_and_add(&read, 1);
63e5a9df51SDavid Vernet 	return 0;
64e5a9df51SDavid Vernet }
65e5a9df51SDavid Vernet 
66e5a9df51SDavid Vernet static void
handle_sample_msg(const struct test_msg * msg)67e5a9df51SDavid Vernet handle_sample_msg(const struct test_msg *msg)
68e5a9df51SDavid Vernet {
69e5a9df51SDavid Vernet 	switch (msg->msg_op) {
70e5a9df51SDavid Vernet 	case TEST_MSG_OP_INC64:
71e5a9df51SDavid Vernet 		kern_mutated += msg->operand_64;
72e5a9df51SDavid Vernet 		break;
73e5a9df51SDavid Vernet 	case TEST_MSG_OP_INC32:
74e5a9df51SDavid Vernet 		kern_mutated += msg->operand_32;
75e5a9df51SDavid Vernet 		break;
76e5a9df51SDavid Vernet 	case TEST_MSG_OP_MUL64:
77e5a9df51SDavid Vernet 		kern_mutated *= msg->operand_64;
78e5a9df51SDavid Vernet 		break;
79e5a9df51SDavid Vernet 	case TEST_MSG_OP_MUL32:
80e5a9df51SDavid Vernet 		kern_mutated *= msg->operand_32;
81e5a9df51SDavid Vernet 		break;
82e5a9df51SDavid Vernet 	default:
83e5a9df51SDavid Vernet 		bpf_printk("Unrecognized op %d\n", msg->msg_op);
84e5a9df51SDavid Vernet 		err = 2;
85e5a9df51SDavid Vernet 	}
86e5a9df51SDavid Vernet }
87e5a9df51SDavid Vernet 
88e5a9df51SDavid Vernet static long
read_protocol_msg(struct bpf_dynptr * dynptr,void * context)89e5a9df51SDavid Vernet read_protocol_msg(struct bpf_dynptr *dynptr, void *context)
90e5a9df51SDavid Vernet {
91e5a9df51SDavid Vernet 	const struct test_msg *msg = NULL;
92e5a9df51SDavid Vernet 
93e5a9df51SDavid Vernet 	msg = bpf_dynptr_data(dynptr, 0, sizeof(*msg));
94e5a9df51SDavid Vernet 	if (!msg) {
95e5a9df51SDavid Vernet 		err = 1;
96e5a9df51SDavid Vernet 		bpf_printk("Unexpectedly failed to get msg\n");
97e5a9df51SDavid Vernet 		return 0;
98e5a9df51SDavid Vernet 	}
99e5a9df51SDavid Vernet 
100e5a9df51SDavid Vernet 	handle_sample_msg(msg);
101e5a9df51SDavid Vernet 
102e5a9df51SDavid Vernet 	return 0;
103e5a9df51SDavid Vernet }
104e5a9df51SDavid Vernet 
publish_next_kern_msg(__u32 index,void * context)105e5a9df51SDavid Vernet static int publish_next_kern_msg(__u32 index, void *context)
106e5a9df51SDavid Vernet {
107e5a9df51SDavid Vernet 	struct test_msg *msg = NULL;
108e5a9df51SDavid Vernet 	int operand_64 = TEST_OP_64;
109e5a9df51SDavid Vernet 	int operand_32 = TEST_OP_32;
110e5a9df51SDavid Vernet 
111e5a9df51SDavid Vernet 	msg = bpf_ringbuf_reserve(&kernel_ringbuf, sizeof(*msg), 0);
112e5a9df51SDavid Vernet 	if (!msg) {
113e5a9df51SDavid Vernet 		err = 4;
114e5a9df51SDavid Vernet 		return 1;
115e5a9df51SDavid Vernet 	}
116e5a9df51SDavid Vernet 
117e5a9df51SDavid Vernet 	switch (index % TEST_MSG_OP_NUM_OPS) {
118e5a9df51SDavid Vernet 	case TEST_MSG_OP_INC64:
119e5a9df51SDavid Vernet 		msg->operand_64 = operand_64;
120e5a9df51SDavid Vernet 		msg->msg_op = TEST_MSG_OP_INC64;
121e5a9df51SDavid Vernet 		expected_user_mutated += operand_64;
122e5a9df51SDavid Vernet 		break;
123e5a9df51SDavid Vernet 	case TEST_MSG_OP_INC32:
124e5a9df51SDavid Vernet 		msg->operand_32 = operand_32;
125e5a9df51SDavid Vernet 		msg->msg_op = TEST_MSG_OP_INC32;
126e5a9df51SDavid Vernet 		expected_user_mutated += operand_32;
127e5a9df51SDavid Vernet 		break;
128e5a9df51SDavid Vernet 	case TEST_MSG_OP_MUL64:
129e5a9df51SDavid Vernet 		msg->operand_64 = operand_64;
130e5a9df51SDavid Vernet 		msg->msg_op = TEST_MSG_OP_MUL64;
131e5a9df51SDavid Vernet 		expected_user_mutated *= operand_64;
132e5a9df51SDavid Vernet 		break;
133e5a9df51SDavid Vernet 	case TEST_MSG_OP_MUL32:
134e5a9df51SDavid Vernet 		msg->operand_32 = operand_32;
135e5a9df51SDavid Vernet 		msg->msg_op = TEST_MSG_OP_MUL32;
136e5a9df51SDavid Vernet 		expected_user_mutated *= operand_32;
137e5a9df51SDavid Vernet 		break;
138e5a9df51SDavid Vernet 	default:
139e5a9df51SDavid Vernet 		bpf_ringbuf_discard(msg, 0);
140e5a9df51SDavid Vernet 		err = 5;
141e5a9df51SDavid Vernet 		return 1;
142e5a9df51SDavid Vernet 	}
143e5a9df51SDavid Vernet 
144e5a9df51SDavid Vernet 	bpf_ringbuf_submit(msg, 0);
145e5a9df51SDavid Vernet 
146e5a9df51SDavid Vernet 	return 0;
147e5a9df51SDavid Vernet }
148e5a9df51SDavid Vernet 
149e5a9df51SDavid Vernet static void
publish_kern_messages(void)150e5a9df51SDavid Vernet publish_kern_messages(void)
151e5a9df51SDavid Vernet {
152e5a9df51SDavid Vernet 	if (expected_user_mutated != user_mutated) {
153e5a9df51SDavid Vernet 		bpf_printk("%lu != %lu\n", expected_user_mutated, user_mutated);
154e5a9df51SDavid Vernet 		err = 3;
155e5a9df51SDavid Vernet 		return;
156e5a9df51SDavid Vernet 	}
157e5a9df51SDavid Vernet 
158e5a9df51SDavid Vernet 	bpf_loop(8, publish_next_kern_msg, NULL, 0);
159e5a9df51SDavid Vernet }
160e5a9df51SDavid Vernet 
161e5a9df51SDavid Vernet SEC("fentry/" SYS_PREFIX "sys_prctl")
test_user_ringbuf_protocol(void * ctx)162e5a9df51SDavid Vernet int test_user_ringbuf_protocol(void *ctx)
163e5a9df51SDavid Vernet {
164e5a9df51SDavid Vernet 	long status = 0;
165e5a9df51SDavid Vernet 
166e5a9df51SDavid Vernet 	if (!is_test_process())
167e5a9df51SDavid Vernet 		return 0;
168e5a9df51SDavid Vernet 
169e5a9df51SDavid Vernet 	status = bpf_user_ringbuf_drain(&user_ringbuf, read_protocol_msg, NULL, 0);
170e5a9df51SDavid Vernet 	if (status < 0) {
171e5a9df51SDavid Vernet 		bpf_printk("Drain returned: %ld\n", status);
172e5a9df51SDavid Vernet 		err = 1;
173e5a9df51SDavid Vernet 		return 0;
174e5a9df51SDavid Vernet 	}
175e5a9df51SDavid Vernet 
176e5a9df51SDavid Vernet 	publish_kern_messages();
177e5a9df51SDavid Vernet 
178e5a9df51SDavid Vernet 	return 0;
179e5a9df51SDavid Vernet }
180e5a9df51SDavid Vernet 
181e5a9df51SDavid Vernet SEC("fentry/" SYS_PREFIX "sys_getpgid")
test_user_ringbuf(void * ctx)182e5a9df51SDavid Vernet int test_user_ringbuf(void *ctx)
183e5a9df51SDavid Vernet {
184e5a9df51SDavid Vernet 	if (!is_test_process())
185e5a9df51SDavid Vernet 		return 0;
186e5a9df51SDavid Vernet 
187e5a9df51SDavid Vernet 	err = bpf_user_ringbuf_drain(&user_ringbuf, record_sample, NULL, 0);
188e5a9df51SDavid Vernet 
189e5a9df51SDavid Vernet 	return 0;
190e5a9df51SDavid Vernet }
191e5a9df51SDavid Vernet 
192e5a9df51SDavid Vernet static long
do_nothing_cb(struct bpf_dynptr * dynptr,void * context)193e5a9df51SDavid Vernet do_nothing_cb(struct bpf_dynptr *dynptr, void *context)
194e5a9df51SDavid Vernet {
195e5a9df51SDavid Vernet 	__sync_fetch_and_add(&read, 1);
196e5a9df51SDavid Vernet 	return 0;
197e5a9df51SDavid Vernet }
198e5a9df51SDavid Vernet 
199*84c22fa8STiezhu Yang SEC("fentry/" SYS_PREFIX "sys_prlimit64")
test_user_ringbuf_epoll(void * ctx)200e5a9df51SDavid Vernet int test_user_ringbuf_epoll(void *ctx)
201e5a9df51SDavid Vernet {
202e5a9df51SDavid Vernet 	long num_samples;
203e5a9df51SDavid Vernet 
204e5a9df51SDavid Vernet 	if (!is_test_process())
205e5a9df51SDavid Vernet 		return 0;
206e5a9df51SDavid Vernet 
207e5a9df51SDavid Vernet 	num_samples = bpf_user_ringbuf_drain(&user_ringbuf, do_nothing_cb, NULL, 0);
208e5a9df51SDavid Vernet 	if (num_samples <= 0)
209e5a9df51SDavid Vernet 		err = 1;
210e5a9df51SDavid Vernet 
211e5a9df51SDavid Vernet 	return 0;
212e5a9df51SDavid Vernet }
213