1// SPDX-License-Identifier: GPL-2.0
  2/* Copyright (c) 2024 Meta Platforms, Inc. and affiliates. */
  3
  4#define BPF_NO_KFUNC_PROTOTYPES
  5#include <vmlinux.h>
  6#include <bpf/bpf_helpers.h>
  7#include <bpf/bpf_tracing.h>
  8#include "bpf_misc.h"
  9#include "bpf_experimental.h"
 10#include "bpf_arena_common.h"
 11
 12#define ARENA_SIZE (1ull << 32)
 13
 14struct {
 15	__uint(type, BPF_MAP_TYPE_ARENA);
 16	__uint(map_flags, BPF_F_MMAPABLE);
 17	__uint(max_entries, ARENA_SIZE / PAGE_SIZE);
 18} arena SEC(".maps");
 19
 20SEC("syscall")
 21__success __retval(0)
 22int big_alloc1(void *ctx)
 23{
 24#if defined(__BPF_FEATURE_ADDR_SPACE_CAST)
 25	volatile char __arena *page1, *page2, *no_page, *page3;
 26	void __arena *base;
 27
 28	page1 = base = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
 29	if (!page1)
 30		return 1;
 31	*page1 = 1;
 32	page2 = bpf_arena_alloc_pages(&arena, base + ARENA_SIZE - PAGE_SIZE * 2,
 33				      1, NUMA_NO_NODE, 0);
 34	if (!page2)
 35		return 2;
 36	*page2 = 2;
 37	no_page = bpf_arena_alloc_pages(&arena, base + ARENA_SIZE - PAGE_SIZE,
 38					1, NUMA_NO_NODE, 0);
 39	if (no_page)
 40		return 3;
 41	if (*page1 != 1)
 42		return 4;
 43	if (*page2 != 2)
 44		return 5;
 45	bpf_arena_free_pages(&arena, (void __arena *)page1, 1);
 46	if (*page2 != 2)
 47		return 6;
 48	if (*page1 != 0) /* use-after-free should return 0 */
 49		return 7;
 50	page3 = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
 51	if (!page3)
 52		return 8;
 53	*page3 = 3;
 54	if (page1 != page3)
 55		return 9;
 56	if (*page2 != 2)
 57		return 10;
 58	if (*(page1 + PAGE_SIZE) != 0)
 59		return 11;
 60	if (*(page1 - PAGE_SIZE) != 0)
 61		return 12;
 62	if (*(page2 + PAGE_SIZE) != 0)
 63		return 13;
 64	if (*(page2 - PAGE_SIZE) != 0)
 65		return 14;
 66#endif
 67	return 0;
 68}
 69
 70#if defined(__BPF_FEATURE_ADDR_SPACE_CAST)
 71#define PAGE_CNT 100
 72__u8 __arena * __arena page[PAGE_CNT]; /* occupies the first page */
 73__u8 __arena *base;
 74
 75/*
 76 * Check that arena's range_tree algorithm allocates pages sequentially
 77 * on the first pass and then fills in all gaps on the second pass.
 78 */
 79__noinline int alloc_pages(int page_cnt, int pages_atonce, bool first_pass,
 80		int max_idx, int step)
 81{
 82	__u8 __arena *pg;
 83	int i, pg_idx;
 84
 85	for (i = 0; i < page_cnt; i++) {
 86		pg = bpf_arena_alloc_pages(&arena, NULL, pages_atonce,
 87					   NUMA_NO_NODE, 0);
 88		if (!pg)
 89			return step;
 90		pg_idx = (unsigned long) (pg - base) / PAGE_SIZE;
 91		if (first_pass) {
 92			/* Pages must be allocated sequentially */
 93			if (pg_idx != i)
 94				return step + 100;
 95		} else {
 96			/* Allocator must fill into gaps */
 97			if (pg_idx >= max_idx || (pg_idx & 1))
 98				return step + 200;
 99		}
100		*pg = pg_idx;
101		page[pg_idx] = pg;
102		cond_break;
103	}
104	return 0;
105}
106
107SEC("syscall")
108__success __retval(0)
109int big_alloc2(void *ctx)
110{
111	__u8 __arena *pg;
112	int i, err;
113
114	base = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
115	if (!base)
116		return 1;
117	bpf_arena_free_pages(&arena, (void __arena *)base, 1);
118
119	err = alloc_pages(PAGE_CNT, 1, true, PAGE_CNT, 2);
120	if (err)
121		return err;
122
123	/* Clear all even pages */
124	for (i = 0; i < PAGE_CNT; i += 2) {
125		pg = page[i];
126		if (*pg != i)
127			return 3;
128		bpf_arena_free_pages(&arena, (void __arena *)pg, 1);
129		page[i] = NULL;
130		cond_break;
131	}
132
133	/* Allocate into freed gaps */
134	err = alloc_pages(PAGE_CNT / 2, 1, false, PAGE_CNT, 4);
135	if (err)
136		return err;
137
138	/* Free pairs of pages */
139	for (i = 0; i < PAGE_CNT; i += 4) {
140		pg = page[i];
141		if (*pg != i)
142			return 5;
143		bpf_arena_free_pages(&arena, (void __arena *)pg, 2);
144		page[i] = NULL;
145		page[i + 1] = NULL;
146		cond_break;
147	}
148
149	/* Allocate 2 pages at a time into freed gaps */
150	err = alloc_pages(PAGE_CNT / 4, 2, false, PAGE_CNT, 6);
151	if (err)
152		return err;
153
154	/* Check pages without freeing */
155	for (i = 0; i < PAGE_CNT; i += 2) {
156		pg = page[i];
157		if (*pg != i)
158			return 7;
159		cond_break;
160	}
161
162	pg = bpf_arena_alloc_pages(&arena, NULL, 1, NUMA_NO_NODE, 0);
163
164	if (!pg)
165		return 8;
166	/*
167	 * The first PAGE_CNT pages are occupied. The new page
168	 * must be above.
169	 */
170	if ((pg - base) / PAGE_SIZE < PAGE_CNT)
171		return 9;
172	return 0;
173}
174#endif
175char _license[] SEC("license") = "GPL";