Resguardo inicial

This commit is contained in:
lastdragon
2026-09-06 12:53:21 -06:00
parent 662efe3c2a
commit ad2088d178
94853 changed files with 43200404 additions and 0 deletions
@@ -0,0 +1,553 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <vmlinux.h>
#include <libarena/common.h>
#include <libarena/asan.h>
enum {
/*
* Is the access checked by check_region_inline
* a read or a write?
*/
ASAN_READ = 0x0U,
ASAN_WRITE = 0x1U,
};
/*
* Address sanitizer (ASAN) for arena-based BPF programs, inspired
* by KASAN.
*
* The API
* -------
*
* The implementation includes two kinds of components: Implementation
* of ASAN hooks injected by LLVM into the program, and API calls that
* allocators use to mark memory as valid or invalid. The full list is:
*
* LLVM stubs:
*
* void __asan_{load, store}<size>(intptr_t addr)
* Checks whether an access is valid. All variations covered
* by check_region_inline().
*
* void __asan_{store, load}((intptr_t addr, ssize_t size)
*
* void __asan_report_{load, store}<size>(intptr_t addr)
* Report an access violation for the program. Used when LLVM
* uses direct code generation for shadow map checks.
*
* void *__asan_memcpy(void *d, const void *s, size_t n)
* void *__asan_memmove(void *d, const void *s, size_t n)
* void *__asan_memset(void *p, int c, size_t n)
* Hooks for ASAN instrumentation of the LLVM mem* builtins.
* Currently unimplemented just like the builtins themselves.
*
* API methods:
*
* asan_init()
* Initialize the ASAN map for the arena.
*
* asan_poison()
* Mark a region of memory as poisoned. Accessing poisoned memory
* causes asan_report() to fire. Invoked during free().
*
* asan_unpoison()
* Mark a region as unpoisoned after alloc().
*
* asan_shadow_set()
* Check a byte's validity directly.
*
* The Algorithm In Brief
* ----------------------
* Each group of 8 bytes is mapped to a "granule" in the shadow map. This
* granule is the size of the byte and describes which bytes are valid.
* Possible values are:
*
* 0: All bytes are valid. Makes checks in the middle of an allocated region
* (most of them) fast.
* (0, 7]: How many consecutive bytes are valid, starting from the lowest one.
* The tradeoff is that we can't poison individual bytes in the middle of a
* valid region.
* [0x80, 0xff]: Special poison values, can be used to denote specific error
* modes (e.g., recently freed vs uninitialized memory).
*
* The mapping between a memory location and its shadow is:
* shadow_addr = shadow_base + (addr >> 3). We retain the 8:1 data:shadow
* ratio of existing ASAN implementations as a compromise between tracking
* granularity and space usage/scan overhead.
*/
#ifdef BPF_ARENA_ASAN
#pragma clang attribute push(__attribute__((no_sanitize("address"))), \
apply_to = function)
#define SHADOW_ALL_ZEROES ((u64)-1)
/*
* Canary variable for ASAN violations. Set to the offending address.
*/
volatile u64 asan_violated = 0;
/*
* Shadow map occupancy map.
*/
volatile u64 __asan_shadow_memory_dynamic_address;
volatile u32 asan_reported = false;
volatile bool asan_inited = false;
/*
* Set during program load.
*/
volatile bool asan_report_once = false;
/*
* BPF does not currently support the memset/memcpy/memcmp intrinsics.
* For large sequential copies, or assignments of large data structures,
* the frontend will generate an intrinsic that causes the BPF backend
* to exit due to a missing implementation. Provide a simple implementation
* just for memset to use it for poisoning/unpoisoning the map.
*/
__weak int asan_memset(s8 __arena *dst, s8 val, size_t size)
{
size_t i;
for (i = zero; i < size && can_loop; i++)
dst[i] = val;
return 0;
}
/* Validate a 1-byte access, always within a single byte. */
static __always_inline bool memory_is_poisoned_1(s8 __arena *addr)
{
s8 shadow_value = *(s8 __arena *)mem_to_shadow(addr);
/* Byte is 0, access is valid. */
if (likely(!shadow_value))
return false;
/*
* Byte is non-zero. Access is valid if granule offset in [0, shadow_value),
* so the memory is poisoned if shadow_value is negative or smaller than
* the granule's value.
*/
return ASAN_GRANULE(addr) >= shadow_value;
}
/* Validate a 2- 4-, 8-byte access, shadow spans up to 2 bytes. */
static __always_inline bool memory_is_poisoned_2_4_8(s8 __arena *addr, u64 size)
{
u64 end = (u64)addr + size - 1;
/*
* Region fully within a single byte (addition didn't
* overflow above ASAN_GRANULE).
*/
if (likely(ASAN_GRANULE(end) >= size - 1))
return memory_is_poisoned_1((s8 __arena *)end);
/*
* Otherwise first byte must be fully unpoisoned, and second byte
* must be unpoisoned up to the end of the accessed region.
*/
return *(s8 __arena *)mem_to_shadow(addr) || memory_is_poisoned_1((s8 __arena *)end);
}
__weak bool asan_shadow_set(void __arena *addr)
{
return memory_is_poisoned_1(addr);
}
static __always_inline u64 first_nonzero_byte(u64 addr, size_t size)
{
while (size && can_loop) {
if (unlikely(*(s8 __arena *)addr))
return addr;
addr += 1;
size -= 1;
}
return SHADOW_ALL_ZEROES;
}
static __always_inline bool memory_is_poisoned_n(s8 __arena *addr, u64 size)
{
u64 ret;
u64 start;
u64 end;
/* Size of [start, end] is end - start + 1. */
start = (u64)mem_to_shadow(addr);
end = (u64)mem_to_shadow(addr + size - 1);
ret = first_nonzero_byte(start, (end - start) + 1);
if (likely(ret == SHADOW_ALL_ZEROES))
return false;
return unlikely(ret != end || ASAN_GRANULE(addr + size - 1) >= *(s8 __arena *)end);
}
__weak int asan_report(s8 __arena *addr, size_t sz, u32 flags)
{
u32 reported = __sync_val_compare_and_swap(&asan_reported, false, true);
/* Only report the first ASAN violation. */
if (reported && asan_report_once)
return 0;
asan_violated = (u64)addr;
arena_stderr("Memory violation for address %p (0x%lx) for %s of size %ld\n",
addr, (u64)addr,
(flags & ASAN_WRITE) ? "write" : "read",
sz);
bpf_stream_print_stack(BPF_STDERR);
return 0;
}
static __always_inline bool check_asan_args(s8 __arena *addr, size_t size,
bool *result)
{
bool valid = true;
/* Size 0 accesses are valid even if the address is invalid. */
if (unlikely(size == 0))
goto confirmed_valid;
/*
* Wraparound is possible for values close to the the edge of the
* 4GiB boundary of the arena (last valid address is 1UL << 32 - 1).
*
*
* The wraparound detection below works for small sizes. check_asan_args is
* always called from the builtin ASAN checks, so 1 <= size <= 64. Even
* for storeN/loadN that we do not expect to encounter the intrinsics will
* not have a large enough size that:
*
* - addr + size > MAX_U32
* - (u32)(addr + size) > (u32) addr
*
* which would defeat wraparound detection.
*/
if (unlikely((u32)(u64)(addr + size) < (u32)(u64)addr))
goto confirmed_invalid;
return false;
confirmed_invalid:
valid = false;
/* FALLTHROUGH */
confirmed_valid:
*result = valid;
return true;
}
static __always_inline bool check_region_inline(intptr_t ptr, size_t size,
u32 flags)
{
s8 __arena *addr = (s8 __arena *)(u64)ptr;
bool is_poisoned, is_valid;
if (check_asan_args(addr, size, &is_valid)) {
if (!is_valid)
asan_report(addr, size, flags);
return is_valid;
}
switch (size) {
case 1:
is_poisoned = memory_is_poisoned_1(addr);
break;
case 2:
case 4:
case 8:
is_poisoned = memory_is_poisoned_2_4_8(addr, size);
break;
default:
is_poisoned = memory_is_poisoned_n(addr, size);
}
if (is_poisoned) {
asan_report(addr, size, flags);
return false;
}
return true;
}
/*
* __alias is not supported for BPF so define *__noabort() variants as wrappers.
*/
#define DEFINE_ASAN_LOAD_STORE(size) \
__hidden void __asan_store##size(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_WRITE); \
} \
__hidden void __asan_store##size##_noabort(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_WRITE); \
} \
__hidden void __asan_load##size(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_READ); \
} \
__hidden void __asan_load##size##_noabort(intptr_t addr) \
{ \
check_region_inline(addr, size, ASAN_READ); \
} \
__hidden void __asan_report_store##size(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_WRITE); \
} \
__hidden void __asan_report_store##size##_noabort(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_WRITE); \
} \
__hidden void __asan_report_load##size(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_READ); \
} \
__hidden void __asan_report_load##size##_noabort(intptr_t addr) \
{ \
asan_report((s8 __arena *)addr, size, ASAN_READ); \
}
DEFINE_ASAN_LOAD_STORE(1);
DEFINE_ASAN_LOAD_STORE(2);
DEFINE_ASAN_LOAD_STORE(4);
DEFINE_ASAN_LOAD_STORE(8);
void __asan_storeN(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_WRITE);
}
void __asan_storeN_noabort(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_WRITE);
}
void __asan_loadN(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_READ);
}
void __asan_loadN_noabort(intptr_t addr, ssize_t size)
{
check_region_inline(addr, size, ASAN_READ);
}
/*
* We currently do not sanitize globals.
*/
void __asan_register_globals(intptr_t globals, size_t n)
{
}
void __asan_unregister_globals(intptr_t globals, size_t n)
{
}
/*
* We do not currently have memcpy/memmove/memset intrinsics
* in LLVM. Do not implement sanitization.
*/
void *__asan_memcpy(void *d, const void *s, size_t n)
{
arena_stderr("ASAN: Unexpected %s call", __func__);
return NULL;
}
void *__asan_memmove(void *d, const void *s, size_t n)
{
arena_stderr("ASAN: Unexpected %s call", __func__);
return NULL;
}
void *__asan_memset(void *p, int c, size_t n)
{
arena_stderr("ASAN: Unexpected %s call", __func__);
return NULL;
}
/*
* Poisoning code, used when we add more freed memory to the allocator by:
* a) pulling memory from the arena segment using bpf_arena_alloc_pages()
* b) freeing memory from application code
*/
__hidden __noasan int asan_poison(void __arena *addr, s8 val, size_t size)
{
s8 __arena *shadow;
size_t len;
/*
* Poisoning from a non-granule address makes no sense: We can only allocate
* memory to the application that has a granule-aligned starting address,
* and bpf_arena_alloc_pages returns page-aligned memory. A non-aligned
* addr then implies we're freeing a different address than the one we
* allocated.
*/
if (unlikely((u64)addr & ASAN_GRANULE_MASK))
return -EINVAL;
/*
* We cannot free an unaligned region because it'd be possible that we
* cannot describe the resulting poisoning state of the granule in
* the ASAN encoding.
*
* Every granule represents a region of memory that looks like the
* following (P for poisoned bytes, C for clear):
*
* <Clear> <Poisoned>
* [ C C C ... P P ]
*
* The value of the granule's shadow map is the number of clear bytes in
* it. We cannot represent granules with the following state:
*
* [ P P ... C C ... P P ]
*
* That would be possible if we could free unaligned regions, so prevent that.
*/
if (unlikely(size & ASAN_GRANULE_MASK))
return -EINVAL;
shadow = mem_to_shadow(addr);
len = size >> ASAN_SHADOW_SHIFT;
asan_memset(shadow, val, len);
return 0;
}
/*
* Unpoisoning code for marking memory as valid during allocation calls.
*
* Very similar to asan_poison, except we need to round up instead of
* down, then partially poison the last granule if necessary.
*
* Partial poisoning is useful for keeping the padding poisoned. Allocations
* are granule-aligned, so we we're reserving granule-aligned sizes for the
* allocation. However, we want to still treat accesses to the padding as
* invalid. Partial poisoning takes care of that. Freeing and poisoning the
* memory is still done in granule-aligned sizes and repoisons the already
* poisoned padding.
*/
__hidden __noasan int asan_unpoison(void __arena *addr, size_t size)
{
size_t partial = size & ASAN_GRANULE_MASK;
s8 __arena *shadow;
size_t len;
/*
* We cannot allocate in the middle of the granule. The ASAN shadow
* map encoding only describes regions of memory where every granule
* follows this format (P for poisoned, C for clear):
*
* <Clear> <Poisoned>
* [ C C C ... P P ]
*
* This is so we can use a single number in [0, ASAN_SHADOW_SCALE)
* to represent the poison state of the granule.
*/
if (unlikely((u64)addr & ASAN_GRANULE_MASK))
return -EINVAL;
shadow = mem_to_shadow(addr);
len = size >> ASAN_SHADOW_SHIFT;
asan_memset(shadow, 0, len);
/*
* If we are allocating a non-granule aligned region, we need to adjust
* the last byte of the shadow map to list how many bytes in the granule
* are unpoisoned. If the region is aligned, then the memset call above
* was enough.
*/
if (partial)
shadow[len] = partial;
return 0;
}
/*
* Initialize ASAN state when necessary. Triggered from userspace before
* allocator startup.
*/
SEC("syscall")
__weak __noasan int asan_init(struct asan_init_args *args)
{
u64 globals_pages = args->arena_globals_pages;
u64 all_pages = args->arena_all_pages;
u64 shadow_map, shadow_pgoff;
u64 shadow_pages;
if (asan_inited)
return 0;
/*
* Round up the shadow map size to the nearest page.
*/
shadow_pages = all_pages >> ASAN_SHADOW_SHIFT;
if ((all_pages & ((1 << ASAN_SHADOW_SHIFT) - 1)))
shadow_pages += 1;
if (all_pages > (1ULL << 32) / __PAGE_SIZE) {
arena_stderr("error: arena size %lx too large", all_pages);
return -EINVAL;
}
if (globals_pages > all_pages) {
arena_stderr("error: globals %lx do not fit in arena %lx",
globals_pages, all_pages);
return -EINVAL;
}
if (globals_pages + shadow_pages >= all_pages) {
arena_stderr("error: globals %lx do not leave room for shadow map %lx "
"(arena pages %lx)",
globals_pages, shadow_pages, all_pages);
return -EINVAL;
}
shadow_pgoff = all_pages - shadow_pages - globals_pages;
__asan_shadow_memory_dynamic_address = shadow_pgoff * __PAGE_SIZE;
/*
* Allocate the last (1/ASAN_SHADOW_SCALE)th of an arena's pages for the map
* We find the offset and size from the arena map.
*
* The allocated map pages are zeroed out, meaning all memory is marked as valid
* even if it's not allocated already. This is expected: Since the actual memory
* pages are not allocated, accesses to it will trigger page faults and will be
* reported through BPF streams. Any pages allocated through bpf_arena_alloc_pages
* should be poisoned by the allocator right after the call succeeds.
*/
shadow_map = (u64)bpf_arena_alloc_pages(
&arena, (void __arena *)__asan_shadow_memory_dynamic_address,
shadow_pages, NUMA_NO_NODE, 0);
if (!shadow_map) {
arena_stderr("Could not allocate shadow map\n");
__asan_shadow_memory_dynamic_address = 0;
return -ENOMEM;
}
asan_inited = true;
return 0;
}
#pragma clang attribute pop
#endif /* BPF_ARENA_ASAN */
__weak char _license[] SEC("license") = "GPL";
@@ -0,0 +1,903 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/buddy.h>
/*
* Buddy allocator arena-based implementation.
*
* Memory is organized into chunks. These chunks
* cannot be coalesced or split. Allocating
* chunks allocates their memory eagerly.
*
* Internally, each chunk is organized into blocks.
* Blocks _can_ be coalesced/split, but only inside
* the chunk. Each block can be allocated or
* unallocated. If allocated, the entire block holds
* user data. If unallocated, the block is mostly
* invalid memory, with the exception of a header
* used for freelist tracking.
*
* The header is placed at an offset inside the block
* to prevent off-by-one errors from the previous block
* from trivially overwriting the header. Such an error
* is also not catchable by ASAN, since the header remains
* valid memory even after the block is freed. It is still
* theoretically possible for the header to be corrupted
* without being caught by ASAN, but harder.
*
* Since the allocator needs to track order information for
* both allocated and free blocks, and allocated blocks cannot
* store a header, the allocator also stores per-chunk order
* information in a reserved region at the beginning of the
* chunk. The header includes a bitmap with the order of blocks
* and their allocation state. It also includes the freelist
* heads for the allocation itself.
*/
enum {
BUDDY_POISONED = (s8)0xef,
/* Number of pages to be allocated per chunk. */
BUDDY_CHUNK_PAGES = BUDDY_CHUNK_BYTES / __PAGE_SIZE
};
static inline int buddy_lock(struct buddy __arena *buddy)
{
return arena_spin_lock(&buddy->lock);
}
static inline void buddy_unlock(struct buddy __arena *buddy)
{
arena_spin_unlock(&buddy->lock);
}
/*
* Reserve part of the arena address space for the allocator. We use
* this to get aligned addresses for the chunks, since the arena
* page alloc kfuncs do not support aligning to a boundary (in this
* case 1 MiB, see buddy.h on how this is derived).
*/
static int buddy_reserve_arena_vaddr(struct buddy __arena *buddy)
{
buddy->vaddr = 0;
return bpf_arena_reserve_pages(&arena,
(void __arena *)BUDDY_VADDR_OFFSET,
BUDDY_VADDR_SIZE / __PAGE_SIZE);
}
/*
* Free up any unused address space. Used only during teardown.
*/
static void buddy_unreserve_arena_vaddr(struct buddy __arena *buddy)
{
bpf_arena_free_pages(
&arena, (void __arena *)(BUDDY_VADDR_OFFSET + buddy->vaddr),
(BUDDY_VADDR_SIZE - buddy->vaddr) / __PAGE_SIZE);
buddy->vaddr = 0;
}
/*
* Carve out part of the reserved address space and hand it over
* to the buddy allocator.
*
* We are assuming the buddy allocator is the only allocator in the
* system, so there is no race between this function reserving a
* page range and some other allocator actually making the BPF call
* to really create and reserve it.
*
* However, bump allocation must still be atomic because this function
* is called without the buddy lock from multiple threads concurrently.
*/
__weak int buddy_alloc_arena_vaddr(struct buddy __arena *buddy, u64 *vaddrp)
{
u64 vaddr, old, new;
if (!buddy || !vaddrp)
return -EINVAL;
do {
vaddr = buddy->vaddr;
new = vaddr + BUDDY_CHUNK_BYTES;
if (new > BUDDY_VADDR_SIZE)
return -EINVAL;
old = __sync_val_compare_and_swap(&buddy->vaddr, vaddr, new);
} while (old != vaddr && can_loop);
if (old != vaddr)
return -EINVAL;
*vaddrp = BUDDY_VADDR_OFFSET + vaddr;
return 0;
}
static u64 arena_next_pow2(__u64 n)
{
n--;
n |= n >> 1;
n |= n >> 2;
n |= n >> 4;
n |= n >> 8;
n |= n >> 16;
n |= n >> 32;
n++;
return n;
}
__weak
int idx_set_allocated(struct buddy_chunk __arena *chunk, u64 idx, bool allocated)
{
bool already_allocated;
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("setting state of invalid idx (%ld, max %d)\n", idx,
BUDDY_CHUNK_ITEMS);
return -EINVAL;
}
already_allocated = chunk->allocated[idx / 8] & (1 << (idx % 8));
if (unlikely(already_allocated == allocated)) {
arena_stderr("Double %s of idx %ld for chunk %p",
allocated ? "alloc" : "free",
idx, chunk);
return -EINVAL;
}
if (allocated)
chunk->allocated[idx / 8] |= 1 << (idx % 8);
else
chunk->allocated[idx / 8] &= ~(1 << (idx % 8));
return 0;
}
static int idx_is_allocated(struct buddy_chunk __arena *chunk, u64 idx, bool *allocated)
{
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("getting state of invalid idx (%llu, max %d)\n", idx,
BUDDY_CHUNK_ITEMS);
return -EINVAL;
}
*allocated = chunk->allocated[idx / 8] & (1 << (idx % 8));
return 0;
}
__weak
int idx_set_order(struct buddy_chunk __arena *chunk, u64 idx, u8 order)
{
u8 prev_order;
if (unlikely(order >= BUDDY_CHUNK_NUM_ORDERS)) {
arena_stderr("setting invalid order %u\n", order);
return -EINVAL;
}
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("setting order of invalid idx (%d, max %d)\n", idx,
BUDDY_CHUNK_ITEMS);
return -EINVAL;
}
/*
* We store two order instances per byte, one per nibble.
* Retain the existing nibble.
*/
prev_order = chunk->orders[idx / 2];
if (idx & 0x1) {
order &= 0xf;
order |= (prev_order & 0xf0);
} else {
order <<= 4;
order |= (prev_order & 0xf);
}
chunk->orders[idx / 2] = order;
return 0;
}
static u8 idx_get_order(struct buddy_chunk __arena *chunk, u64 idx)
{
u8 result;
_Static_assert(BUDDY_CHUNK_NUM_ORDERS <= 16,
"order must fit in 4 bits");
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("getting order of invalid idx %u\n", idx);
return BUDDY_CHUNK_NUM_ORDERS;
}
result = chunk->orders[idx / 2];
return (idx & 0x1) ? (result & 0xf) : (result >> 4);
}
static void __arena *idx_to_addr(struct buddy_chunk __arena *chunk, size_t idx)
{
u64 address;
if (unlikely(idx >= BUDDY_CHUNK_ITEMS)) {
arena_stderr("translating invalid idx %u\n", idx);
return NULL;
}
/*
* The data blocks start in the chunk after the metadata block.
* We find the actual address by indexing into the region at an
* BUDDY_MIN_ALLOC_BYTES granularity, the minimum allowed.
* The index number already accounts for the fact that the first
* blocks in the chunk are occupied by the metadata, so we do
* not need to offset it.
*/
address = (u64)chunk + (idx * BUDDY_MIN_ALLOC_BYTES);
return (void __arena *)address;
}
static struct buddy_header __arena *idx_to_header(struct buddy_chunk __arena *chunk, size_t idx)
{
bool allocated;
u64 address;
if (unlikely(idx_is_allocated(chunk, idx, &allocated))) {
arena_stderr("accessing invalid idx 0x%lx\n", idx);
return NULL;
}
if (unlikely(allocated)) {
arena_stderr("accessing allocated idx 0x%lx as header\n", idx);
return NULL;
}
address = (u64)idx_to_addr(chunk, idx);
if (!address)
return NULL;
/*
* Offset the header within the block. This avoids accidental overwrites
* to the header because of off-by-one errors when using adjacent blocks.
*
* The offset has been chosen as a compromise between ASAN effectiveness
* and allocator granularity:
* 1) ASAN dictates valid data runs are 8-byte aligned.
* 2) We want to keep a low minimum allocation size (currently 16).
*
* As a result, we have only two possible positions for the header: Bytes
* 0 and 8. Keeping the header in byte 0 means off-by-ones from the previous
* block touch the header, and, since the header must be accessible, ASAN
* will not trigger. Keeping the header on byte 8 means off-by-one errors from
* the previous block are caught by ASAN. Negative offsets are rarer, so
* while accesses into the block from the next block are possible, they are
* less probable.
*/
return (struct buddy_header __arena *)(address + BUDDY_HEADER_OFF);
}
static void header_add_freelist(struct buddy_chunk __arena *chunk, struct buddy_header __arena *header,
u64 idx, u8 order)
{
struct buddy_header __arena *tmp_header;
idx_set_order(chunk, idx, order);
header->next_index = chunk->freelists[order];
header->prev_index = BUDDY_CHUNK_ITEMS;
if (header->next_index != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, header->next_index);
tmp_header->prev_index = idx;
}
chunk->freelists[order] = idx;
}
static void header_remove_freelist(struct buddy_chunk __arena *chunk,
struct buddy_header __arena *header, u8 order)
{
struct buddy_header __arena *tmp_header;
if (header->prev_index != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, header->prev_index);
tmp_header->next_index = header->next_index;
}
if (header->next_index != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, header->next_index);
tmp_header->prev_index = header->prev_index;
}
/* Pop off the list head if necessary. */
if (idx_to_header(chunk, chunk->freelists[order]) == header)
chunk->freelists[order] = header->next_index;
header->prev_index = BUDDY_CHUNK_ITEMS;
header->next_index = BUDDY_CHUNK_ITEMS;
}
static u64 size_to_order(size_t size)
{
u64 order;
/*
* Legal sizes are [1, 4GiB] (the biggest possible arena).
* Of course, sizes close to GiB are practically impossible
* to fulfill and allocation will fail, but that's taken care
* of by the caller.
*/
if (unlikely(size == 0 || size > (1UL << 32))) {
arena_stderr("illegal size request %lu\n", size);
return 64;
}
/*
* To find the order of the allocation we find the first power of two
* >= the requested size, take the log2, then adjust it for the minimum
* allocation size by removing the minimum shift from it. Requests
* smaller than the minimum allocation size are rounded up.
*/
order = arena_fls(arena_next_pow2(size)) - 1;
if (order < BUDDY_MIN_ALLOC_SHIFT)
return 0;
return order - BUDDY_MIN_ALLOC_SHIFT;
}
__weak
int add_leftovers_to_freelist(struct buddy_chunk __arena *chunk, u32 cur_idx,
u64 min_order, u64 max_order)
{
struct buddy_header __arena *header;
u64 ord;
u32 idx;
for (ord = min_order; ord < max_order && can_loop; ord++) {
/* Mark the buddy as free and add it to the freelists. */
idx = cur_idx + (1 << ord);
header = idx_to_header(chunk, idx);
if (unlikely(!header)) {
arena_stderr("idx %u has no header", idx);
return -EINVAL;
}
asan_unpoison(header, sizeof(*header));
header_add_freelist(chunk, header, idx, ord);
}
return 0;
}
static struct buddy_chunk __arena *buddy_chunk_get(struct buddy __arena *buddy)
{
u64 order, ord, min_order, max_order;
struct buddy_chunk __arena *chunk;
size_t left;
int power2;
u64 vaddr;
u32 idx;
int ret;
/*
* Step 1: Allocate a properly aligned chunk, and
* prep it for insertion into the buddy allocator.
* We don't need the allocator lock until step 2.
*/
ret = buddy_alloc_arena_vaddr(buddy, &vaddr);
if (ret)
return NULL;
/* Addresses must be aligned to the chunk boundary. */
if (vaddr % BUDDY_CHUNK_BYTES)
return NULL;
/* Unreserve the address space. */
bpf_arena_free_pages(&arena, (void __arena *)vaddr,
BUDDY_CHUNK_PAGES);
chunk = bpf_arena_alloc_pages(&arena, (void __arena *)vaddr,
BUDDY_CHUNK_PAGES, NUMA_NO_NODE, 0);
if (!chunk) {
arena_stderr("[ALLOC FAILED]");
return NULL;
}
if (buddy_lock(buddy)) {
/*
* We cannot reclaim the vaddr space, but that is ok - this
* operation should always succeed. The error path is to catch
* accidental deadlocks that will cause -ENOMEMs to the program as
* the allocator fails to refill itself, in which case vaddr usage
* is the least of our worries.
*/
bpf_arena_free_pages(&arena, (void __arena *)vaddr, BUDDY_CHUNK_PAGES);
return NULL;
}
asan_poison(chunk, BUDDY_POISONED, BUDDY_CHUNK_PAGES * __PAGE_SIZE);
/* Unpoison the chunk itself. */
asan_unpoison(chunk, sizeof(*chunk));
/* Mark all freelists as empty. */
for (ord = zero; ord < BUDDY_CHUNK_NUM_ORDERS && can_loop; ord++)
chunk->freelists[ord] = BUDDY_CHUNK_ITEMS;
/*
* Initialize the chunk by carving out a page range to hold the metadata
* struct above, then dumping the rest of the pages into the allocator.
*/
_Static_assert(BUDDY_CHUNK_PAGES * __PAGE_SIZE >=
BUDDY_MIN_ALLOC_BYTES *
BUDDY_CHUNK_ITEMS,
"chunk must fit within the allocation");
/*
* Step 2: Reserve a chunk for the chunk metadata, then breaks
* the rest of the full allocation into the different buckets.
* We allocating the memory by grabbing blocks of progressively
* smaller sizes from the allocator, which are guaranteed to be
* continuous.
*
* This operation also populates the allocator.
*
* Algorithm:
*
* - max_order: The last order allocation we made
* - left: How many bytes are left to allocate
* - cur_index: Current index into the top-level block we are
* allocating from.
*
* Step 3:
* - Find the largest power-of-2 allocation still smaller than left (infimum)
* - Reserve a chunk of that size, along with its buddy
* - For every order from [infimum + 1, last order), carve out a block
* and put it into the allocator.
*
* Example: Chunk size 0b1010000 (80 bytes)
*
* Step 1:
*
* idx infimum 1 << max_order
* 0 64 128 1 << 20
* |________|_________|______________________|
*
* Blocks set aside:
* [0, 64) - Completely allocated
* [64, 128) - Will be further split in the next iteration
*
* Blocks added to the allocator:
* [128, 256)
* [256, 512)
* ...
* [1 << 18, 1 << 19)
* [1 << 19, 1 << 20)
*
* Step 2:
*
* idx infimum idx + 1 << max_order
* 64 80 96 64 + 1 << 6 = 128
* |________|_________|______________________|
*
* Blocks set aside:
* [64, 80) - Completely allocated
*
* Blocks added to the allocator:
* [80, 96) - left == 0 so the buddy is unused and marked as freed
* [96, 128)
*/
max_order = BUDDY_CHUNK_NUM_ORDERS;
left = sizeof(*chunk);
idx = 0;
while (left && can_loop) {
power2 = arena_fls(left) - 1;
/*
* Note: The condition below only triggers to catch serious bugs
* early. There is no sane way to undo any block insertions from
* the allocated chunk, so just leak any leftover allocations,
* emit a diagnostic, unlock and exit.
*
*/
if (unlikely(power2 >= BUDDY_CHUNK_NUM_ORDERS)) {
arena_stderr(
"buddy chunk metadata require allocation of order %d\n",
power2);
arena_stderr(
"chunk has size of 0x%lx bytes (left %lx bytes)\n",
sizeof(*chunk), left);
buddy_unlock(buddy);
return NULL;
}
/* Round up allocations that are too small. */
left -= (power2 >= BUDDY_MIN_ALLOC_SHIFT) ? 1 << power2 : left;
order = (power2 >= BUDDY_MIN_ALLOC_SHIFT) ? power2 - BUDDY_MIN_ALLOC_SHIFT : 0;
if (idx_set_allocated(chunk, idx, true)) {
buddy_unlock(buddy);
return NULL;
}
/*
* Starting an order above the one we allocated, populate
* the allocator with free blocks. If this is the last
* allocation (left == 0), also mark the buddy as free.
*
* See comment above about error handling: The error path
* is only there as a way to mitigate deeply buggy allocator
* states by emitting a diagnostic in add_leftovers_to_freelist()
* and leaking any memory not added in the freelists.
*/
min_order = left ? order + 1 : order;
if (add_leftovers_to_freelist(chunk, idx, min_order, max_order)) {
buddy_unlock(buddy);
return NULL;
}
/* Adjust the index. */
idx += 1 << order;
max_order = order;
}
buddy_unlock(buddy);
return chunk;
}
__weak int buddy_init(struct buddy __arena *buddy)
{
struct buddy_chunk __arena *chunk;
int ret;
if (!asan_ready())
return -EINVAL;
/* Reserve enough address space to ensure allocations are aligned. */
ret = buddy_reserve_arena_vaddr(buddy);
if (ret)
return ret;
_Static_assert(BUDDY_CHUNK_PAGES > 0,
"chunk must use one or more pages");
chunk = buddy_chunk_get(buddy);
if (buddy_lock(buddy)) {
bpf_arena_free_pages(&arena, chunk, BUDDY_CHUNK_PAGES);
return -EINVAL;
}
/* Chunk is already properly unpoisoned if allocated. */
if (chunk)
chunk->next = buddy->first_chunk;
/* Put the chunk at the beginning of the list. */
buddy->first_chunk = chunk;
buddy_unlock(buddy);
return chunk ? 0 : -ENOMEM;
}
/*
* Destroy the allocator. This does not check whether there are any allocations
* currently in use, so any pages being accessed will start taking arena faults.
* We do not take a lock because we are freeing arena pages, and nobody should
* be using the allocator at that point in the execution.
*/
__weak int buddy_destroy(struct buddy __arena *buddy)
{
struct buddy_chunk __arena *chunk, *next;
if (!buddy)
return -EINVAL;
/*
* Traverse all buddy chunks and free them back to the arena
* with the same granularity they were allocated with.
*/
for (chunk = buddy->first_chunk; chunk && can_loop; chunk = next) {
next = chunk->next;
/* Wholesale poison the entire block. */
asan_poison(chunk, BUDDY_POISONED,
BUDDY_CHUNK_PAGES * __PAGE_SIZE);
bpf_arena_free_pages(&arena, chunk, BUDDY_CHUNK_PAGES);
}
/* Free up any part of the address space that did not get used. */
buddy_unreserve_arena_vaddr(buddy);
/* Clear all fields. */
buddy->first_chunk = NULL;
return 0;
}
__weak u64 buddy_chunk_alloc(struct buddy_chunk __arena *chunk, int order_req)
{
struct buddy_header __arena *header, *tmp_header, *next_header;
u32 idx, tmpidx, retidx;
u64 address;
u64 order = 0;
u64 i;
for (order = order_req; order < BUDDY_CHUNK_NUM_ORDERS && can_loop; order++) {
if (chunk->freelists[order] != BUDDY_CHUNK_ITEMS)
break;
}
if (order >= BUDDY_CHUNK_NUM_ORDERS)
return (u64)NULL;
retidx = chunk->freelists[order];
header = idx_to_header(chunk, retidx);
if (unlikely(!header))
return (u64) NULL;
chunk->freelists[order] = header->next_index;
if (header->next_index != BUDDY_CHUNK_ITEMS) {
next_header = idx_to_header(chunk, header->next_index);
next_header->prev_index = BUDDY_CHUNK_ITEMS;
}
header->prev_index = BUDDY_CHUNK_ITEMS;
header->next_index = BUDDY_CHUNK_ITEMS;
if (idx_set_order(chunk, retidx, order_req))
return (u64)NULL;
if (idx_set_allocated(chunk, retidx, true))
return (u64)NULL;
/*
* Do not unpoison the address yet, will be done by the caller
* because the caller has the exact allocation size requested.
*/
address = (u64)idx_to_addr(chunk, retidx);
if (!address)
return (u64)NULL;
/* If we allocated from a larger-order chunk, split the buddies. */
for (i = order_req; i < order && can_loop; i++) {
/*
* Flip the bit for the current order (the bit is guaranteed
* to be 0, so just add 1 << i).
*/
idx = retidx + (1 << i);
/* Add the buddy of the allocation to the free list. */
header = idx_to_header(chunk, idx);
/* Unpoison the buddy header */
asan_unpoison(header, sizeof(*header));
if (idx_set_order(chunk, idx, i))
return (u64)NULL;
/* Push the header to the beginning of the freelists list. */
tmpidx = chunk->freelists[i];
header->prev_index = BUDDY_CHUNK_ITEMS;
header->next_index = tmpidx;
if (tmpidx != BUDDY_CHUNK_ITEMS) {
tmp_header = idx_to_header(chunk, tmpidx);
tmp_header->prev_index = idx;
}
chunk->freelists[i] = idx;
}
return address;
}
/* Scan the existing chunks for available memory. */
static u64 buddy_alloc_from_existing_chunks(struct buddy __arena *buddy, int order)
{
struct buddy_chunk __arena *chunk;
u64 address;
for (chunk = buddy->first_chunk; chunk != NULL && can_loop;
chunk = chunk->next) {
address = buddy_chunk_alloc(chunk, order);
if (address)
return address;
}
return (u64)NULL;
}
/*
* Try an allocation from a newly allocated chunk. Also
* incorporate the chunk into the linked list.
*/
static u64 buddy_alloc_from_new_chunk(struct buddy __arena *buddy, struct buddy_chunk __arena *chunk, int order)
{
u64 address;
if (buddy_lock(buddy))
return (u64)NULL;
/*
* Add the chunk into the allocator and try
* to allocate specifically from that chunk.
*/
chunk->next = buddy->first_chunk;
buddy->first_chunk = chunk;
address = buddy_chunk_alloc(buddy->first_chunk, order);
buddy_unlock(buddy);
return (u64)address;
}
__weak
void __arena *buddy_alloc(struct buddy __arena *buddy, size_t size)
{
void __arena *address = NULL;
struct buddy_chunk __arena *chunk;
int order;
if (!buddy)
return NULL;
order = size_to_order(size);
if (order >= BUDDY_CHUNK_NUM_ORDERS || order < 0) {
arena_stderr("invalid order %d (sz %lu)\n", order, size);
return NULL;
}
if (buddy_lock(buddy))
return NULL;
address = (u8 __arena *)buddy_alloc_from_existing_chunks(buddy, order);
buddy_unlock(buddy);
if (address)
goto done;
/* Get a new chunk. */
chunk = buddy_chunk_get(buddy);
if (chunk)
address = (u8 __arena *)buddy_alloc_from_new_chunk(buddy, chunk, order);
done:
/* If we failed to allocate memory, return NULL. */
if (!address)
return NULL;
/*
* Unpoison exactly the amount of bytes requested. If the
* data is smaller than the header, we must poison any
* unused bytes that were part of the header.
*/
if (size < BUDDY_HEADER_OFF + sizeof(struct buddy_header __arena))
asan_poison(address + BUDDY_HEADER_OFF, BUDDY_POISONED,
sizeof(struct buddy_header __arena));
asan_unpoison(address, size);
return address;
}
static __always_inline int buddy_free_unlocked(struct buddy __arena *buddy, u64 addr)
{
struct buddy_header __arena *header, *buddy_header;
u64 idx, buddy_idx, tmp_idx;
struct buddy_chunk __arena *chunk;
bool allocated;
u8 order;
int ret;
if (!buddy)
return -EINVAL;
if (addr & (BUDDY_MIN_ALLOC_BYTES - 1)) {
arena_stderr("Freeing unaligned address %llx\n", addr);
return -EINVAL;
}
/* Get (chunk, idx) out of the address. */
chunk = (void __arena *)(addr & ~BUDDY_CHUNK_OFFSET_MASK);
idx = (addr & BUDDY_CHUNK_OFFSET_MASK) / BUDDY_MIN_ALLOC_BYTES;
/* Mark the block as unallocated so we can access the header. */
ret = idx_set_allocated(chunk, idx, false);
if (ret)
return ret;
order = idx_get_order(chunk, idx);
header = idx_to_header(chunk, idx);
/* The header is in the block itself, keep it unpoisoned. */
asan_poison((u8 __arena *)addr, BUDDY_POISONED,
BUDDY_MIN_ALLOC_BYTES << order);
asan_unpoison(header, sizeof(*header));
/*
* Coalescing loop. Merge with free buddies of equal order.
* For every coalescing step, keep the left buddy and
* drop the right buddy's header.
*/
for (; order < BUDDY_CHUNK_NUM_ORDERS && can_loop; order++) {
buddy_idx = idx ^ (1 << order);
/* Check if the buddy is actually free. */
idx_is_allocated(chunk, buddy_idx, &allocated);
if (allocated)
break;
/*
* If buddy is not the same order as the chunk
* being freed, then we're done coalescing.
*/
if (idx_get_order(chunk, buddy_idx) != order)
break;
buddy_header = idx_to_header(chunk, buddy_idx);
header_remove_freelist(chunk, buddy_header, order);
/* Keep the left header out of the two buddies, drop the other one. */
if (buddy_idx < idx) {
tmp_idx = idx;
idx = buddy_idx;
buddy_idx = tmp_idx;
}
/* Remove the buddy from the freelists so that we can merge it. */
idx_set_order(chunk, buddy_idx, order);
buddy_header = idx_to_header(chunk, buddy_idx);
asan_poison(buddy_header, BUDDY_POISONED,
sizeof(*buddy_header));
}
/* Header properly freed but not in any freelists yet .*/
idx_set_order(chunk, idx, order);
header = idx_to_header(chunk, idx);
header_add_freelist(chunk, header, idx, order);
return 0;
}
__weak int buddy_free(struct buddy __arena *buddy, void __arena *addr)
{
int ret;
if (!buddy)
return -EINVAL;
/* Freeing NULL is a valid no-op. */
if (!addr)
return 0;
ret = buddy_lock(buddy);
if (ret)
return ret;
ret = buddy_free_unlocked(buddy, (u64)addr);
buddy_unlock(buddy);
return ret;
}
__weak char _license[] SEC("license") = "GPL";
@@ -0,0 +1,52 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/buddy.h>
const volatile u32 zero = 0;
struct buddy __arena buddy;
int arena_fls(__u64 word)
{
if (!word)
return 0;
return 64 - __builtin_clzll(word);
}
SEC("syscall")
__weak int arena_get_info(struct arena_get_info_args *args)
{
args->arena_base = arena_base(&arena);
return 0;
}
SEC("syscall")
__weak int arena_alloc_reserve(struct arena_alloc_reserve_args *args)
{
return bpf_arena_reserve_pages(&arena, NULL, args->nr_pages);
}
SEC("syscall")
__weak int arena_buddy_reset(void)
{
buddy_destroy(&buddy);
return buddy_init(&buddy);
}
__weak void __arena *arena_malloc(size_t size)
{
return buddy_alloc(&buddy, size);
}
__weak void arena_free(void __arena *ptr)
{
buddy_free(&buddy, ptr);
}
char _license[] SEC("license") = "GPL";
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,234 @@
// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/*
* Copyright (c) 2025-2026 Meta Platforms, Inc. and affiliates.
* Copyright (c) 2025-2026 Emil Tsalapatis <etsal@meta.com>
*/
#include <bpf_atomic.h>
#include <libarena/common.h>
#include <libarena/asan.h>
#include <libarena/spmc.h>
static inline
u64 spmc_arr_size(volatile struct spmc_arr __arena *spmc_arr)
{
return SPMC_ARR_BASESZ << spmc_arr->order;
}
static inline
u64 spmc_arr_get(volatile struct spmc_arr __arena *spmc_arr, u64 ind)
{
u64 ret = READ_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)]);
return ret;
}
static inline
void spmc_arr_put(volatile struct spmc_arr __arena *spmc_arr, u64 ind, u64 value)
{
WRITE_ONCE(spmc_arr->data[ind % spmc_arr_size(spmc_arr)], value);
}
static inline
void spmc_arr_copy(volatile struct spmc_arr __arena *dst,
volatile struct spmc_arr __arena *src, u64 b, u64 t)
{
u64 i;
for (i = t; i < b && can_loop; i++)
spmc_arr_put(dst, i, spmc_arr_get(src, i));
}
static inline
int spmc_order_init(struct spmc __arena *spmc, int order)
{
volatile struct spmc_arr __arena *arr = &spmc->arr[order];
if (unlikely(!spmc))
return -EINVAL;
if (order >= SPMC_ARR_ORDERS)
return -E2BIG;
/* Already allocated? */
if (arr->data)
return 0;
arr->data = arena_malloc((SPMC_ARR_BASESZ << order) * sizeof(*arr->data));
if (!arr->data)
return -ENOMEM;
return 0;
}
__weak
int spmc_owned_add(struct spmc __arena *spmc, u64 val)
{
volatile struct spmc_arr __arena *newarr;
volatile struct spmc_arr __arena *arr;
ssize_t sz;
u64 b, t;
int ret;
if (unlikely(!spmc))
return -EINVAL;
/*
* Bottom must always be read first, also
* see spmc_steal().
*/
b = smp_load_acquire(&spmc->bottom);
t = READ_ONCE(spmc->top);
arr = READ_ONCE(spmc->cur);
sz = b - t;
if (sz >= spmc_arr_size(arr) - 1) {
ret = spmc_order_init(spmc, arr->order + 1);
if (ret)
return ret;
newarr = &spmc->arr[arr->order + 1];
spmc_arr_copy(newarr, arr, b, t);
smp_store_release(&spmc->cur, newarr);
arr = newarr;
}
spmc_arr_put(arr, b, val);
smp_store_release(&spmc->bottom, b + 1);
return 0;
}
__weak
int spmc_owned_remove(struct spmc __arena *spmc, u64 *val)
{
volatile struct spmc_arr __arena *arr;
int ret = 0;
ssize_t sz;
u64 value;
u64 b, t;
if (unlikely(!spmc || !val))
return -EINVAL;
b = READ_ONCE(spmc->bottom) - 1;
WRITE_ONCE(spmc->bottom, b);
smp_mb();
t = READ_ONCE(spmc->top);
arr = READ_ONCE(spmc->cur);
sz = b - t;
if (sz < 0) {
WRITE_ONCE(spmc->bottom, t);
return -ENOENT;
}
value = spmc_arr_get(arr, b);
if (sz > 0) {
*val = value;
return 0;
}
if (cmpxchg(&spmc->top, t, t + 1) != t)
ret = -EAGAIN;
WRITE_ONCE(spmc->bottom, t + 1);
if (ret)
return ret;
*val = value;
return 0;
}
__weak
int spmc_steal(struct spmc __arena *spmc, u64 *val)
{
volatile struct spmc_arr __arena *arr;
ssize_t sz;
u64 value;
u64 b, t;
if (unlikely(!spmc || !val))
return -EINVAL;
/*
* It is important that t is read before b for
* stealers to avoid racing with the owner.
* Races between stealers are dealt with using
* CAS to increment the top value below.
*/
t = smp_load_acquire(&spmc->top);
b = smp_load_acquire(&spmc->bottom);
sz = b - t;
if (sz <= 0)
return -ENOENT;
arr = smp_load_acquire(&spmc->cur);
value = spmc_arr_get(arr, t);
if (cmpxchg(&spmc->top, t, t + 1) != t)
return -EAGAIN;
*val = value;
return 0;
}
__weak
struct spmc __arena *spmc_create(void)
{
/*
* Marked as volatile because otherwise the array
* reference in the internal loop gets demoted to
* scalar and the program fails verification.
*/
struct spmc __arena *volatile spmc;
int ret, i;
spmc = arena_malloc(sizeof(*spmc));
if (!spmc)
return NULL;
spmc->bottom = 0;
spmc->top = 0;
for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++) {
spmc->arr[i].data = NULL;
spmc->arr[i].order = i;
}
ret = spmc_order_init((struct spmc __arena *)spmc, 0);
if (ret) {
arena_free(spmc);
return NULL;
}
spmc->cur = &spmc->arr[0];
return (struct spmc __arena *)spmc;
}
__weak
int spmc_destroy(struct spmc __arena *spmc)
{
int i;
if (unlikely(!spmc))
return -EINVAL;
for (i = 0; i < SPMC_ARR_ORDERS && can_loop; i++)
arena_free(spmc->arr[i].data);
arena_free(spmc);
return 0;
}