Add kernel heap allocation tracing debug facility.

kernel/include/sortix/kernel/decl.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2013, 2014 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2013, 2014, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -36,4 +36,35 @@ typedef uintptr_t addr_t;
 	#define CPU X64
 #endif
 
+#ifdef __TRACE_ALLOCATION_SITES
+
+struct kernel_allocation_site
+{
+	struct __allocation_site allocation_site;
+	struct kernel_allocation_site* next;
+	bool registered;
+};
+
+extern struct kernel_allocation_site* first_kernel_allocation_site;
+
+#undef ALLOCATION_SITE
+#define ALLOCATION_SITE \
+({ \
+	static struct kernel_allocation_site site = \
+		{ { __FILE__, __LINE__, __func__, 0, 0 }, NULL, false }; \
+	if ( !site.registered ) \
+	{ \
+		site.registered = true; \
+		site.next = first_kernel_allocation_site; \
+		first_kernel_allocation_site = &site; \
+	} \
+	&site.allocation_site; \
+})
+
+#include <stddef.h>
+void* operator new(size_t size, struct __allocation_site* allocation_site);
+void* operator new[](size_t size, struct __allocation_site* allocation_site);
+#define new new (ALLOCATION_SITE)
+#endif
+
 #endif

kernel/kernel.cpp

@@ -102,6 +102,8 @@
 #include "x86-family/vbox.h"
 #endif
 
+struct kernel_allocation_site* first_kernel_allocation_site = NULL;
+
 // Keep the stack size aligned with $CPU/boot.s
 const size_t STACK_SIZE = 64*1024;
 extern "C" { __attribute__((aligned(16))) size_t stack[STACK_SIZE / sizeof(size_t)]; }

kernel/kernelinfo.cpp

@@ -50,6 +50,51 @@ ssize_t sys_kernelinfo(const char* user_req, char* user_resp, size_t resplen)
 	char* req = GetStringFromUser(user_req);
 	if ( !req )
 		return -1;
+#ifdef __TRACE_ALLOCATION_SITES
+	if ( !strcmp(req, "allocations") )
+	{
+		delete[] req;
+		bool exhausted = false;
+		size_t total_needed = 0;
+		for ( struct kernel_allocation_site* site = first_kernel_allocation_site;
+			  site;
+			  site = site->next )
+		{
+			char str[256];
+			snprintf(str, sizeof(str), "%20zu B %zu allocations %s:%zu %s %c",
+			         site->allocation_site.current_size,
+			         site->allocation_site.allocations,
+			         site->allocation_site.file,
+			         site->allocation_site.line,
+			         site->allocation_site.function,
+			         site->next ? '\n' : '\0');
+			size_t stringlen = strlen(str);
+			total_needed += stringlen;
+			if ( exhausted )
+				continue;
+			if ( resplen < stringlen )
+			{
+				exhausted = true;
+				continue;
+			}
+			if ( !CopyToUser(user_resp, str, sizeof(char) * stringlen) )
+				return -1;
+			user_resp += stringlen;
+			resplen -= stringlen;
+		}
+		if ( !exhausted && !resplen )
+			exhausted = true;
+		if ( !exhausted )
+		{
+			char zero = '\0';
+			if ( !CopyToUser(user_resp, &zero, 1) )
+				return -1;
+		}
+		if ( exhausted )
+			return errno = ERANGE, (ssize_t) total_needed;
+		return 0;
+	}
+#endif
 	const char* str = KernelInfo(req);
 	delete[] req;
 	if ( !str )

kernel/op-new.cpp

@@ -24,6 +24,19 @@
 #warning "security: -fcheck-new might not work on clang"
 #endif
 
+#ifdef __TRACE_ALLOCATION_SITES
+#undef new
+
+void* operator new(size_t size, struct __allocation_site* allocation_site)
+{
+	return malloc_trace(allocation_site, size);
+}
+
+void* operator new[](size_t size, struct __allocation_site* allocation_site)
+{
+	return malloc_trace(allocation_site, size);
+}
+#else
 void* operator new(size_t size)
 {
 	return malloc(size);
@@ -33,6 +46,7 @@ void* operator new[](size_t size)
 {
 	return malloc(size);
 }
+#endif
 
 void operator delete(void* addr)
 {

libc/Makefile

@@ -60,20 +60,6 @@ malloc/__heap_verify.o \
 netinet/if_ether/etheraddr_broadcast.o \
 netinet/in/in6addr_any.o \
 netinet/in/in6addr_loopback.o \
-regex/regcomp.o \
-regex/regerror.o \
-regex/regexec.o \
-regex/regfree.o \
-sha2/sha224hl.o \
-sha2/sha224.o \
-sha2/sha256hl.o \
-sha2/sha256.o \
-sha2/sha384hl.o \
-sha2/sha384.o \
-sha2/sha512_256hl.o \
-sha2/sha512_256.o \
-sha2/sha512hl.o \
-sha2/sha512.o \
 signal/sigaddset.o \
 signal/sigandset.o \
 signal/sigdelset.o \
@@ -86,7 +72,6 @@ signal/sigorset.o \
 ssp/__stack_chk_fail.o \
 stdio/asprintf.o \
 stdio/cbprintf.o \
-stdio/cbscanf.o \
 stdio/clearerr.o \
 stdio/clearerr_unlocked.o \
 stdio/dprintf.o \
@@ -116,8 +101,6 @@ stdio/fgets.o \
 stdio/fgets_unlocked.o \
 stdio/fileno_unlocked.o \
 stdio/flockfile.o \
-stdio/fmemopen.o \
-stdio/fnewfile.o \
 stdio/fparsemode.o \
 stdio/fprintf_unlocked.o \
 stdio/fputc.o \
@@ -128,8 +111,6 @@ stdio/fread.o \
 stdio/fread_unlocked.o \
 stdio/fregister.o \
 stdio/fresetfile.o \
-stdio/fscanf.o \
-stdio/fscanf_unlocked.o \
 stdio/fseek.o \
 stdio/fseeko.o \
 stdio/fseeko_unlocked.o \
@@ -142,28 +123,20 @@ stdio/funlockfile.o \
 stdio/funregister.o \
 stdio/fwrite.o \
 stdio/fwrite_unlocked.o \
-stdio/getdelim.o \
-stdio/getline.o \
-stdio/open_memstream.o \
 stdio/rewind.o \
 stdio/setbuf.o \
 stdio/setvbuf.o \
 stdio/setvbuf_unlocked.o \
 stdio/snprintf.o \
 stdio/sprintf.o \
-stdio/sscanf.o \
 stdio/ungetc.o \
 stdio/ungetc_unlocked.o \
 stdio/vasprintf.o \
 stdio/vcbprintf.o \
 stdio/vdprintf.o \
 stdio/vfprintf_unlocked.o \
-stdio/vfscanf.o \
-stdio/vfscanf_unlocked.o \
-stdio/vcbscanf.o \
 stdio/vsnprintf.o \
 stdio/vsprintf.o \
-stdio/vsscanf.o \
 stdlib/abort.o \
 stdlib/abs.o \
 stdlib/arc4random_buf.o \
@@ -278,7 +251,6 @@ wchar/wcscmp.o \
 wchar/wcscoll.o \
 wchar/wcscpy.o \
 wchar/wcscspn.o \
-wchar/wcsdup.o \
 wchar/wcsftime.o \
 wchar/wcslcat.o \
 wchar/wcslcpy.o \
@@ -329,6 +301,34 @@ wctype/towupper.o \
 wctype/wctype.o \
 
 HOSTEDOBJS=\
+regex/regcomp.o \
+regex/regerror.o \
+regex/regexec.o \
+regex/regfree.o \
+sha2/sha224hl.o \
+sha2/sha224.o \
+sha2/sha256hl.o \
+sha2/sha256.o \
+sha2/sha384hl.o \
+sha2/sha384.o \
+sha2/sha512_256hl.o \
+sha2/sha512_256.o \
+sha2/sha512hl.o \
+sha2/sha512.o \
+stdio/cbscanf.o \
+stdio/fscanf.o \
+stdio/fscanf_unlocked.o \
+stdio/sscanf.o \
+stdio/vfscanf.o \
+stdio/vfscanf_unlocked.o \
+stdio/vcbscanf.o \
+stdio/vsscanf.o \
+stdio/fmemopen.o \
+stdio/open_memstream.o \
+stdio/getdelim.o \
+stdio/getline.o \
+stdio/fnewfile.o \
+wchar/wcsdup.o \
 blf/blowfish.o \
 $(CPUDIR)/fork.o \
 $(CPUDIR)/setjmp.o \

libc/include/malloc.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2012, 2013, 2014, 2015, 2016 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2012, 2013, 2014, 2015, 2016, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -117,6 +117,16 @@ struct heap_alloc
 #warning "You need to implement HEAP_CHUNK_MAGIC for your native word width"
 #endif
 
+#ifdef __TRACE_ALLOCATION_SITES
+#define MAGIC_IS_ALLOCATION_SITE(magic)( (((size_t) (magic)) & 0x3) == 0x2)
+#define ALLOCATION_SITE_OF_MAGIC(magic) ((struct __allocation_site*) ((magic) & ~0x3UL))
+#define MAGIC_OF_ALLOCATION_SITE(magic) (((size_t) (magic)) | 0x2)
+#else
+#define MAGIC_IS_ALLOCATION_SITE(magic) 0
+#define ALLOCATION_SITE_OF_MAGIC(magic) NULL
+#define MAGIC_OF_ALLOCATION_SITE(magic) NULL
+#endif
+
 /* The heap is split into a number of parts that each consists of a number of
    of chunks (used and unused). The heap normally is just a single part, but if
    the address space gets too fragmented, it may not be possible to extend the
@@ -346,7 +356,8 @@ bool heap_chunk_is_used(struct heap_chunk* chunk)
 {
 	assert(HEAP_IS_POINTER_ALIGNED(chunk, chunk->chunk_size));
 
-	return chunk->chunk_magic == HEAP_CHUNK_MAGIC;
+	return chunk->chunk_magic == HEAP_CHUNK_MAGIC ||
+	       MAGIC_IS_ALLOCATION_SITE(chunk->chunk_magic);
 }
 
 /* Returns the trailing structure following the given chunk. */

libc/include/stdio.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2013, 2014, 2015 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2013, 2014, 2015, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -242,8 +242,14 @@ int vdprintf(int fildes, const char* __restrict format, __gnuc_va_list ap)
 
 /* Functions copied from elsewhere. */
 #if __USE_SORTIX
+#ifdef __TRACE_ALLOCATION_SITES
+int asprintf_trace(struct __allocation_site*, char** __restrict, const char* __restrict, ...)
+	__attribute__((__format__ (printf, 3, 4)));
+#define asprintf(a, ...) asprintf_trace(ALLOCATION_SITE, (a), __VA_ARGS__)
+#else
 int asprintf(char** __restrict, const char* __restrict, ...)
 	__attribute__((__format__ (printf, 2, 3)));
+#endif
 void clearerr_unlocked(FILE* stream);
 int feof_unlocked(FILE* stream);
 int ferror_unlocked(FILE* stream);
@@ -255,9 +261,15 @@ int fputc_unlocked(int c, FILE* stream);
 int fputs_unlocked(const char* __restrict, FILE* __restrict stream);
 size_t fread_unlocked(void* __restrict ptr, size_t size, size_t nitems, FILE* __restrict stream);
 size_t fwrite_unlocked(const void* __restrict ptr, size_t size, size_t nitems, FILE* __restrict stream);
+#ifdef __TRACE_ALLOCATION_SITES
+int vasprintf_trace(struct __allocation_site*, char** __restrict, const char* __restrict, __gnuc_va_list)
+	__attribute__((__format__ (printf, 3, 0)));
+#define vasprintf(a, ...) vasprintf_trace(ALLOCATION_SITE, (a), __VA_ARGS__)
+#else
 int vasprintf(char** __restrict, const char* __restrict, __gnuc_va_list)
 	__attribute__((__format__ (printf, 2, 0)));
 #endif
+#endif
 
 /* Functions that are Sortix extensions. */
 #if __USE_SORTIX

libc/include/stdlib.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011-2017 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011-2017, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -88,7 +88,12 @@ double atof(const char* value);
 int atoi(const char*);
 long atol(const char*);
 void* bsearch(const void*, const void*, size_t, size_t, int (*)(const void*, const void*));
+#ifdef __TRACE_ALLOCATION_SITES
+void* calloc_trace(struct __allocation_site*, size_t, size_t);
+#define calloc(a, b) calloc_trace(ALLOCATION_SITE, (a), (b))
+#else
 void* calloc(size_t, size_t);
+#endif
 char* canonicalize_file_name(const char* path);
 char* canonicalize_file_name_at(int dirfd, const char* path);
 int clearenv(void);
@@ -99,7 +104,12 @@ void free(void*);
 char* getenv(const char*);
 long labs(long);
 ldiv_t ldiv(long, long);
+#ifdef __TRACE_ALLOCATION_SITES
+void* malloc_trace(struct __allocation_site*, size_t);
+#define malloc(a) malloc_trace(ALLOCATION_SITE, (a))
+#else
 void* malloc(size_t);
+#endif
 int mblen(const char*, size_t);
 size_t mbstowcs(wchar_t* __restrict, const char* __restrict, size_t);
 int mbtowc(wchar_t *__restrict, const char* __restrict, size_t);
@@ -116,7 +126,12 @@ void qsort_r(void*, size_t, size_t, int (*)(const void*, const void*, void*), vo
 __attribute__((__warning__("rand() isn't random, use arc4random()")))
 #endif
 int rand(void);
+#ifdef __TRACE_ALLOCATION_SITES
+void* realloc_trace(struct __allocation_site*, void*, size_t);
+#define realloc(a, b) realloc_trace(ALLOCATION_SITE, (a), (b))
+#else
 void* realloc(void*, size_t);
+#endif
 char* realpath(const char* __restrict, char* __restrict);
 int setenv(const char*, const char*, int);
 #if !defined(__is_sortix_libc) /* not a warning inside libc */
@@ -189,7 +204,12 @@ int posix_openpt(int);
 uint32_t arc4random(void);
 void arc4random_buf(void*, size_t);
 uint32_t arc4random_uniform(uint32_t);
+#ifdef __TRACE_ALLOCATION_SITES
+void* reallocarray_trace(struct __allocation_site*, void*, size_t, size_t);
+#define reallocarray(a, b, c) reallocarray_trace(ALLOCATION_SITE, (a), (b), (c))
+#else
 void* reallocarray(void*, size_t, size_t);
+#endif
 int ptsname_r(int, char*, size_t);
 #endif
 

libc/include/string.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2013, 2014, 2017 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2013, 2014, 2017, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -91,8 +91,13 @@ void* memccpy(void* __restrict, const void* __restrict, int, size_t);
 
 /* Functions from XOPEN 420 gone into POSIX 2008. */
 #if __USE_SORTIX || 420 <= __USE_XOPEN || 200809L <= __USE_POSIX
+#ifdef __TRACE_ALLOCATION_SITES
+char* strdup_trace(struct __allocation_site*, const char*);
+#define strdup(a) strdup_trace(ALLOCATION_SITE, (a))
+#else
 char* strdup(const char*);
 #endif
+#endif
 
 /* Functions from POSIX 2001. */
 #if __USE_SORTIX || 200112L <= __USE_POSIX
@@ -111,7 +116,12 @@ char* strtok_r(char* __restrict, const char* __restrict, char** __restrict);
 char* stpcpy(char* __restrict, const char* __restrict);
 char* stpncpy(char* __restrict, const char* __restrict, size_t);
 int strcoll_l(const char*, const char*, locale_t);
+#ifdef __TRACE_ALLOCATION_SITES
+char* strndup_trace(struct __allocation_site*, const char*, size_t);
+#define strndup(a, b) strndup_trace(ALLOCATION_SITE, (a), (b))
+#else
 char* strndup(const char*, size_t);
+#endif
 size_t strnlen(const char*, size_t);
 #if __USE_SORTIX && __SORTIX_STDLIB_REDIRECTS
 const char* strsignal(int signum) __asm__ ("sortix_strsignal");

libc/include/sys/cdefs.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2013, 2015 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2013, 2015, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -128,4 +128,24 @@
          is updated to not rely on this macro. */
 #undef __SORTIX_HAS_GETSERVBYNAME__
 
+#if (defined(__is_sortix_libk) || defined(__is_sortix_kernel)) && \
+    defined(__TRACE_KMALLOC)
+#undef __TRACE_ALLOCATION_SITES
+#define __TRACE_ALLOCATION_SITES
+#endif
+
+#ifdef __TRACE_ALLOCATION_SITES
+#include <stddef.h>
+struct __allocation_site
+{
+	const char* file;
+	size_t line;
+	const char* function;
+	size_t current_size;
+	size_t allocations;
+};
+#define __TRACE_ALLOCATION_SITES
+#define ALLOCATION_SITE 0
+#endif
+
 #endif

libc/stdio/asprintf.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2014, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -20,13 +20,24 @@
 #include <stdarg.h>
 #include <stdio.h>
 
+#ifdef __TRACE_ALLOCATION_SITES
+int asprintf_trace(struct __allocation_site* allocation_site,
+                   char** restrict result_ptr,
+                   const char* restrict format,
+                   ...)
+#else
 int asprintf(char** restrict result_ptr,
              const char* restrict format,
              ...)
+#endif
 {
 	va_list list;
 	va_start(list, format);
+#ifdef __TRACE_ALLOCATION_SITES
+	int result = vasprintf_trace(allocation_site, result_ptr, format, list);
+#else
 	int result = vasprintf(result_ptr, format, list);
+#endif
 	va_end(list);
 	return result;
 }

libc/stdio/vasprintf.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014, 2015 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2014, 2015, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -27,6 +27,9 @@ struct vasprintf
 	char* buffer;
 	size_t used;
 	size_t size;
+#ifdef __TRACE_ALLOCATION_SITES
+	struct __allocation_site* allocation_site;
+#endif
 };
 
 static size_t vasprintf_callback(void* ctx, const char* string, size_t length)
@@ -39,7 +42,12 @@ static size_t vasprintf_callback(void* ctx, const char* string, size_t length)
 		size_t new_size = 2 * state->size;
 		if ( new_size < needed_size )
 			new_size = needed_size;
+#ifdef __TRACE_ALLOCATION_SITES
+		char* new_buffer = (char*) realloc_trace(state->allocation_site,
+		                                         state->buffer, new_size);
+#else
 		char* new_buffer = (char*) realloc(state->buffer, new_size);
+#endif
 		if ( !new_buffer )
 		{
 			free(state->buffer);
@@ -54,14 +62,26 @@ static size_t vasprintf_callback(void* ctx, const char* string, size_t length)
 	return length;
 }
 
+#ifdef __TRACE_ALLOCATION_SITES
+int vasprintf_trace(struct __allocation_site* allocation_site,
+                    char** restrict result_ptr,
+                    const char* restrict format,
+                    va_list list)
+#else
 int vasprintf(char** restrict result_ptr,
               const char* restrict format,
               va_list list)
+#endif
 {
 	struct vasprintf state;
 	state.used = 0;
 	state.size = 32;
+#ifdef __TRACE_ALLOCATION_SITES
+	state.allocation_site = allocation_site;
+	if ( !(state.buffer = (char*) malloc_trace(allocation_site, state.size)) )
+#else
 	if ( !(state.buffer = (char*) malloc(state.size)) )
+#endif
 		return -1;
 	int result = vcbprintf(&state, vasprintf_callback, format, list);
 	if ( !state.buffer )

libc/stdlib/calloc.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2014 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2014, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -22,12 +22,21 @@
 #include <stdlib.h>
 #include <string.h>
 
+#ifdef __TRACE_ALLOCATION_SITES
+void* calloc_trace(struct __allocation_site* allocation_site,
+                   size_t nmemb, size_t size)
+#else
 void* calloc(size_t nmemb, size_t size)
+#endif
 {
 	if ( size && nmemb && SIZE_MAX / size < nmemb )
 		return errno = ENOMEM, (void*) NULL;
 	size_t total = nmemb * size;
+#ifdef __TRACE_ALLOCATION_SITES
+	void* result = malloc_trace(allocation_site, total);
+#else
 	void* result = malloc(total);
+#endif
 	if ( !result )
 		return NULL;
 	memset(result, 0, total);

libc/stdlib/free.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2013, 2014, 2015 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2013, 2014, 2015, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -39,6 +39,16 @@ void free(void* addr)
 	// Retrieve the chunk that contains this allocation.
 	struct heap_chunk* chunk = heap_data_to_chunk((uint8_t*) addr);
 
+#ifdef __TRACE_ALLOCATION_SITES
+	if ( MAGIC_IS_ALLOCATION_SITE(chunk->chunk_magic) )
+	{
+		struct __allocation_site* allocation_site =
+			ALLOCATION_SITE_OF_MAGIC(chunk->chunk_magic);
+		allocation_site->current_size -= chunk->chunk_size;
+		allocation_site->allocations--;
+	}
+#endif
+
 	// Return the chunk to the heap.
 	heap_insert_chunk(chunk);
 

libc/stdlib/malloc.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2013, 2014, 2015 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2013, 2014, 2015, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -31,7 +31,12 @@
 
 #if !defined(HEAP_GUARD_DEBUG)
 
+#ifdef __TRACE_ALLOCATION_SITES
+void* malloc_trace(struct __allocation_site* allocation_site,
+                   size_t original_size)
+#else
 void* malloc(size_t original_size)
+#endif
 {
 	if ( !heap_size_has_bin(original_size) )
 		return errno = ENOMEM, (void*) NULL;
@@ -89,6 +94,12 @@ void* malloc(size_t original_size)
 	if ( heap_can_split_chunk(result_chunk, chunk_size) )
 		heap_split_chunk(result_chunk, chunk_size);
 
+#ifdef __TRACE_ALLOCATION_SITES
+	allocation_site->current_size += result_chunk->chunk_size;
+	allocation_site->allocations++;
+	result_chunk->chunk_magic = MAGIC_OF_ALLOCATION_SITE(allocation_site);
+#endif
+
 	__heap_verify();
 	__heap_unlock();
 

libc/stdlib/realloc.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2013, 2014, 2015 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2013, 2014, 2015, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -31,10 +31,20 @@
 
 #if !defined(HEAP_GUARD_DEBUG)
 
+#ifdef __TRACE_ALLOCATION_SITES
+void* realloc_trace(struct __allocation_site* new_allocation_site,
+                    void* ptr,
+                    size_t requested_size)
+#else
 void* realloc(void* ptr, size_t requested_size)
+#endif
 {
 	if ( !ptr )
+#ifdef __TRACE_ALLOCATION_SITES
+		return malloc_trace(new_allocation_site, requested_size);
+#else
 		return malloc(requested_size);
+#endif
 
 	if ( !heap_size_has_bin(requested_size) )
 		return errno = ENOMEM, (void*) NULL;
@@ -55,7 +65,14 @@ void* realloc(void* ptr, size_t requested_size)
 	// Retrieve the chunk that contains this allocation.
 	struct heap_chunk* chunk = heap_data_to_chunk((uint8_t*) ptr);
 
-	assert(chunk->chunk_magic == HEAP_CHUNK_MAGIC);
+#ifdef __TRACE_ALLOCATION_SITES
+	assert(MAGIC_IS_ALLOCATION_SITE(chunk->chunk_magic));
+	struct __allocation_site* allocation_site =
+		ALLOCATION_SITE_OF_MAGIC(chunk->chunk_magic);
+#endif
+
+	assert(chunk->chunk_magic == HEAP_CHUNK_MAGIC ||
+	       MAGIC_IS_ALLOCATION_SITE(chunk->chunk_magic));
 	assert(heap_chunk_to_post(chunk)->chunk_magic == HEAP_CHUNK_MAGIC);
 	assert(heap_chunk_to_post(chunk)->chunk_size == chunk->chunk_size);
 
@@ -72,8 +89,17 @@ void* realloc(void* ptr, size_t requested_size)
 	if ( requested_chunk_size < chunk->chunk_size )
 	{
 		assert(requested_chunk_size <= chunk->chunk_size);
+#ifdef __TRACE_ALLOCATION_SITES
+		allocation_site->current_size -= chunk->chunk_size;
+		allocation_site->allocations--;
+#endif
 		if ( heap_can_split_chunk(chunk, requested_chunk_size) )
 			heap_split_chunk(chunk, requested_chunk_size);
+#ifdef __TRACE_ALLOCATION_SITES
+		allocation_site->current_size += chunk->chunk_size;
+		allocation_site->allocations++;
+		chunk->chunk_magic = MAGIC_OF_ALLOCATION_SITE(allocation_site);
+#endif
 		__heap_verify();
 		__heap_unlock();
 		return heap_chunk_to_data(chunk);
@@ -88,11 +114,20 @@ void* realloc(void* ptr, size_t requested_size)
 	     !heap_chunk_is_used(right) &&
 	     requested_chunk_size <= chunk->chunk_size + right->chunk_size )
 	{
+#ifdef __TRACE_ALLOCATION_SITES
+		allocation_site->current_size -= chunk->chunk_size;
+		allocation_site->allocations--;
+#endif
 		heap_remove_chunk(right);
 		heap_chunk_format((uint8_t*) chunk, chunk->chunk_size + right->chunk_size);
 		assert(requested_chunk_size <= chunk->chunk_size);
 		if ( heap_can_split_chunk(chunk, requested_chunk_size) )
 			heap_split_chunk(chunk, requested_chunk_size);
+#ifdef __TRACE_ALLOCATION_SITES
+		allocation_site->current_size += chunk->chunk_size;
+		allocation_site->allocations++;
+		chunk->chunk_magic = MAGIC_OF_ALLOCATION_SITE(allocation_site);
+#endif
 		__heap_verify();
 		__heap_unlock();
 		return heap_chunk_to_data(chunk);
@@ -108,7 +143,11 @@ void* realloc(void* ptr, size_t requested_size)
 
 	assert(orignal_ptr_size < requested_size);
 
+#ifdef __TRACE_ALLOCATION_SITES
+	void* result = malloc_trace(allocation_site, requested_size);
+#else
 	void* result = malloc(requested_size);
+#endif
 	if ( !result )
 		return (void*) NULL;
 

libc/stdlib/reallocarray.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2014, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -21,9 +21,18 @@
 #include <stdint.h>
 #include <stdlib.h>
 
+#ifdef __TRACE_ALLOCATION_SITES
+void* reallocarray_trace(struct __allocation_site* allocation_site,
+                         void* ptr, size_t nmemb, size_t size)
+#else
 void* reallocarray(void* ptr, size_t nmemb, size_t size)
+#endif
 {
 	if ( size && nmemb && SIZE_MAX / size < nmemb )
 		return errno = ENOMEM, (void*) NULL;
+#ifdef __TRACE_ALLOCATION_SITES
+	return realloc_trace(allocation_site, ptr, nmemb * size);
+#else
 	return realloc(ptr, nmemb * size);
+#endif
 }

libc/string/strdup.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2014 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2014, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -20,10 +20,18 @@
 #include <stdlib.h>
 #include <string.h>
 
+#ifdef __TRACE_ALLOCATION_SITES
+char* strdup_trace(struct __allocation_site* allocation_site, const char* input)
+#else
 char* strdup(const char* input)
+#endif
 {
 	size_t input_length = strlen(input);
+#ifdef __TRACE_ALLOCATION_SITES
+	char* result = (char*) malloc_trace(allocation_site, input_length + 1);
+#else
 	char* result = (char*) malloc(input_length + 1);
+#endif
 	if ( !result )
 		return NULL;
 	memcpy(result, input, input_length + 1);

libc/string/strndup.c

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2012, 2014 Jonas 'Sortie' Termansen.
+ * Copyright (c) 2011, 2012, 2014, 2022 Jonas 'Sortie' Termansen.
  *
  * Permission to use, copy, modify, and distribute this software for any
  * purpose with or without fee is hereby granted, provided that the above
@@ -20,10 +20,19 @@
 #include <stdlib.h>
 #include <string.h>
 
+#ifdef __TRACE_ALLOCATION_SITES
+char* strndup_trace(struct __allocation_site* allocation_site,
+                    const char* input, size_t n)
+#else
 char* strndup(const char* input, size_t n)
+#endif
 {
 	size_t input_size = strnlen(input, n);
+#ifdef __TRACE_ALLOCATION_SITES
+	char* result = (char*) malloc_trace(allocation_site, input_size + 1);
+#else
 	char* result = (char*) malloc(input_size + 1);
+#endif
 	if ( !result )
 		return NULL;
 	memcpy(result, input, input_size);