Experiment to replace malloc(3)/calloc(3)/realloc(3)/free(3) with custom allocator 32 byte aligned.

mem.c

// SPDX-License-Identifier: GPL-2.0
/*
 * Experiment to replace malloc(3)/calloc(3)/realloc(3)/free(3) with custom
 * allocator 32 byte aligned.
 *
 * @author Ammar Faizi <ammarfaizi2@gmail.com> https://www.facebook.com/ammarfaizi2
 * @license GPL-2.0
 *
 * gcc -Wall -Wextra -O3 -shared -fPIC mem.c -o mem.so
 * LD_PRELOAD=$(pwd)/mem.so your_program
 *
 */

#include <errno.h>
#include <stdlib.h>
#include <stdint.h>
#include <sys/mman.h>


static inline void *my_memcpy(void *dest, void const *src, size_t n)
{
	void *orig = dest;
	size_t rcx = n >> 3;

	asm volatile(
		"rep movsq"
		: "+D"(dest), "+S"(src), "+c"(rcx)
		:
		: "memory"
	);

	rcx = n & 7;
	asm volatile(
		"rep movsb"
		: "+D"(dest), "+S"(src), "+c"(rcx)
		:
		: "memory"
	);

	return orig;
}


static inline void *my_mmap(void *addr, size_t length, int prot, int flags,
			    int fd, off_t offset)
{
	void *ret;
	register int _flags asm("r10") = flags;
	register int _fd asm("r8") = fd;
	register off_t _offset asm("r9") = offset;

	asm volatile(
		"syscall"
		: "=a"(ret)
		: "a"(9), "D"(addr), "S"(length), "d"(prot),
		  "r"(_flags), "r"(_fd), "r"(_offset)
		: "memory", "r11", "rcx"
	);
	return ret;
}


static inline int my_munmap(void *addr, size_t length)
{
	int ret;

	asm volatile(
		"syscall"
		: "=a"(ret)
		: "a"(11), "D"(addr), "S"(length)
		: "memory", "r11", "rcx"
	);
	return ret;
}

#define unlikely(EXPR) __builtin_expect(EXPR, 0)

void * __attribute__((noinline)) malloc(size_t len)
{
	void *start_map;
	uintptr_t user_ptr, cmperr;
	size_t add_req = 0;

	add_req += sizeof(size_t);
	add_req += sizeof(uint8_t);
	add_req += 0x1full;

	start_map = my_mmap(NULL, add_req + len, PROT_READ | PROT_WRITE,
			    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

	cmperr = 0xffffffffffffff00ull;
	if (unlikely(((uintptr_t)start_map & cmperr) == cmperr)) {
		errno = ENOMEM;
		return NULL;
	}

	/* Align 32-byte and take space to save the length and diff */
	user_ptr = ((uintptr_t)start_map + add_req) & ~0x1full;
	*(size_t  *)(user_ptr - 8) = len;
	*(uint8_t *)(user_ptr - 9) = (uint8_t)(user_ptr - (uintptr_t)start_map);

	return (void *)user_ptr;
}


void free(void *__user_ptr)
{
	size_t len;
	uint8_t diff;
	uintptr_t user_ptr = (uintptr_t)__user_ptr;

	if (__user_ptr == NULL)
		return;

	len  = *(size_t  *)(user_ptr - 8);
	diff = *(uint8_t *)(user_ptr - 9);
	my_munmap((void *)(user_ptr - diff), len);
}


void *realloc(void *__user_ptr, size_t new_len)
{
	void *new_mem;
	size_t len;
	uint8_t diff;
	uintptr_t user_ptr = (uintptr_t)__user_ptr;

	new_mem = malloc(new_len);
	if (unlikely(__user_ptr == NULL || new_mem == NULL))
		return new_mem;

	len  = *(size_t  *)(user_ptr - 8);
	diff = *(uint8_t *)(user_ptr - 9);

	my_memcpy(new_mem, __user_ptr, (new_len < len) ? new_len : len);
	my_munmap((void *)(user_ptr - diff), len);
	return new_mem;
}


void *calloc(size_t nmemb, size_t len)
{
	size_t x = nmemb * len;
	if (unlikely(nmemb != 0 && x / nmemb != len)) {
		errno = EOVERFLOW;
		return NULL;
	}
	return malloc(x);
}

#if 0
/* For testing only */
#include <stdio.h>
#include <string.h
int main(void)
{
	char *test = malloc(1);

	for (size_t i = 2; i <= (1024 * 1024); i++) {
		test = realloc(test, i);
		memset(test, 'q', i);
	}

	free(test);
}
#endif