1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
3
4#include <vmlinux.h>
5#include <bpf/bpf_helpers.h>
6#include <bpf/bpf_tracing.h>
7#include "bpf_misc.h"
8#include "bpf_experimental.h"
9#include "bpf_arena_common.h"
10
11struct {
12	__uint(type, BPF_MAP_TYPE_ARENA);
13	__uint(map_flags, BPF_F_MMAPABLE);
14	__uint(max_entries, 2); /* arena of two pages close to 32-bit boundary*/
15#ifdef __TARGET_ARCH_arm64
16        __ulong(map_extra, (1ull << 32) | (~0u - __PAGE_SIZE * 2 + 1)); /* start of mmap() region */
17#else
18        __ulong(map_extra, (1ull << 44) | (~0u - __PAGE_SIZE * 2 + 1)); /* start of mmap() region */
19#endif
20} arena SEC(".maps");
21
22SEC("syscall")
23__success __retval(0)
24int basic_alloc1(void *ctx)
25{
26#if defined(__BPF_FEATURE_ADDR_SPACE_CAST)
27	volatile int __arena *page1, *page2, *no_page, *page3;
28
29	page1 = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
30	if (!page1)
31		return 1;
32	*page1 = 1;
33	page2 = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
34	if (!page2)
35		return 2;
36	*page2 = 2;
37	no_page = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
38	if (no_page)
39		return 3;
40	if (*page1 != 1)
41		return 4;
42	if (*page2 != 2)
43		return 5;
44	bpf_arena_free_pages(&arena, (void __arena *)page2, 1);
45	if (*page1 != 1)
46		return 6;
47	if (*page2 != 0) /* use-after-free should return 0 */
48		return 7;
49	page3 = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
50	if (!page3)
51		return 8;
52	*page3 = 3;
53	if (page2 != page3)
54		return 9;
55	if (*page1 != 1)
56		return 10;
57#endif
58	return 0;
59}
60
61SEC("syscall")
62__success __retval(0)
63int basic_alloc2(void *ctx)
64{
65#if defined(__BPF_FEATURE_ADDR_SPACE_CAST)
66	volatile char __arena *page1, *page2, *page3, *page4;
67
68	page1 = bpf_arena_alloc_pages(&arena, NULL, 2, NUMA_NO_NODE, 0);
69	if (!page1)
70		return 1;
71	page2 = page1 + __PAGE_SIZE;
72	page3 = page1 + __PAGE_SIZE * 2;
73	page4 = page1 - __PAGE_SIZE;
74	*page1 = 1;
75	*page2 = 2;
76	*page3 = 3;
77	*page4 = 4;
78	if (*page1 != 1)
79		return 1;
80	if (*page2 != 2)
81		return 2;
82	if (*page3 != 0)
83		return 3;
84	if (*page4 != 0)
85		return 4;
86	bpf_arena_free_pages(&arena, (void __arena *)page1, 2);
87	if (*page1 != 0)
88		return 5;
89	if (*page2 != 0)
90		return 6;
91	if (*page3 != 0)
92		return 7;
93	if (*page4 != 0)
94		return 8;
95#endif
96	return 0;
97}
98
99struct bpf_arena___l {
100        struct bpf_map map;
101} __attribute__((preserve_access_index));
102
103SEC("syscall")
104__success __retval(0) __log_level(2)
105int basic_alloc3(void *ctx)
106{
107	struct bpf_arena___l *ar = (struct bpf_arena___l *)&arena;
108	volatile char __arena *pages;
109
110	pages = bpf_arena_alloc_pages(&ar->map, NULL, ar->map.max_entries, NUMA_NO_NODE, 0);
111	if (!pages)
112		return 1;
113	return 0;
114}
115
116SEC("iter.s/bpf_map")
117__success __log_level(2)
118int iter_maps1(struct bpf_iter__bpf_map *ctx)
119{
120	struct bpf_map *map = ctx->map;
121
122	if (!map)
123		return 0;
124	bpf_arena_alloc_pages(map, NULL, map->max_entries, 0, 0);
125	return 0;
126}
127
128SEC("iter.s/bpf_map")
129__failure __msg("expected pointer to STRUCT bpf_map")
130int iter_maps2(struct bpf_iter__bpf_map *ctx)
131{
132	struct seq_file *seq = ctx->meta->seq;
133
134	bpf_arena_alloc_pages((void *)seq, NULL, 1, 0, 0);
135	return 0;
136}
137
138SEC("iter.s/bpf_map")
139__failure __msg("untrusted_ptr_bpf_map")
140int iter_maps3(struct bpf_iter__bpf_map *ctx)
141{
142	struct bpf_map *map = ctx->map;
143
144	if (!map)
145		return 0;
146	bpf_arena_alloc_pages(map->inner_map_meta, NULL, map->max_entries, 0, 0);
147	return 0;
148}
149
150char _license[] SEC("license") = "GPL";
151