From e273f6d2e5ad9191b2c066b73c2af803637951b2 Mon Sep 17 00:00:00 2001
From: Angus Gratton <angus@redyak.com.au>
Date: Tue, 13 Feb 2024 09:22:45 +1100
Subject: [PATCH] WIP: Add Valgrind "max address space" option.

This option prevents the Python heap from reusing freed memory addresses
for new allocations (obviously only viable for some workloads!)

Advantage is that valgrind will be able to catch any use-after-free, that
might have been missed if the freed memory address was reallocated for
another use.

This work was funded through GitHub Sponsors.

Signed-off-by: Angus Gratton <angus@redyak.com.au>
---
 py/gc.c       | 9 +++++++++
 py/mpconfig.h | 6 ++++++
 2 files changed, 15 insertions(+)

diff --git a/py/gc.c b/py/gc.c
index ff34188885..3b54fc96c7 100644
--- a/py/gc.c
+++ b/py/gc.c
@@ -649,9 +649,11 @@ void gc_collect_end(void) {
     #if MICROPY_GC_SPLIT_HEAP
     MP_STATE_MEM(gc_last_free_area) = &MP_STATE_MEM(area);
     #endif
+    #if !MICROPY_DEBUG_VALGRIND_MAX_ADDRSPACE
     for (mp_state_mem_area_t *area = &MP_STATE_MEM(area); area != NULL; area = NEXT_AREA(area)) {
         area->gc_last_free_atb_index = 0;
     }
+    #endif
     MP_STATE_THREAD(gc_lock_depth)--;
     GC_EXIT();
 }
@@ -838,6 +840,9 @@ found:
         #endif
         area->gc_last_free_atb_index = (i + 1) / BLOCKS_PER_ATB;
     }
+    #if MICROPY_DEBUG_VALGRIND_MAX_ADDRSPACE
+    area->gc_last_free_atb_index = (i + 1) / BLOCKS_PER_ATB;
+    #endif
 
     area->gc_last_used_block = MAX(area->gc_last_used_block, end_block);
 
@@ -966,10 +971,12 @@ void gc_free(void *ptr) {
     }
     #endif
 
+    #if !MICROPY_DEBUG_VALGRIND_MAX_ADDRSPACE
     // set the last_free pointer to this block if it's earlier in the heap
     if (block / BLOCKS_PER_ATB < area->gc_last_free_atb_index) {
         area->gc_last_free_atb_index = block / BLOCKS_PER_ATB;
     }
+    #endif
 
     // free head and all of its tail blocks
     do {
@@ -1133,10 +1140,12 @@ void *gc_realloc(void *ptr_in, size_t n_bytes, bool allow_move) {
         }
         #endif
 
+        #if !MICROPY_DEBUG_VALGRIND_MAX_ADDRSPACE
         // set the last_free pointer to end of this block if it's earlier in the heap
         if ((block + new_blocks) / BLOCKS_PER_ATB < area->gc_last_free_atb_index) {
             area->gc_last_free_atb_index = (block + new_blocks) / BLOCKS_PER_ATB;
         }
+        #endif
 
         VALGRIND_MP_RESIZE_BLOCK(ptr, n_blocks, n_bytes);
 
diff --git a/py/mpconfig.h b/py/mpconfig.h
index af2480266b..cf966a1e53 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -540,6 +540,12 @@
 #define MICROPY_DEBUG_VALGRIND (0)
 #endif
 
+// Whether valgrind should always use new memory addresses for allocations,
+// making it easier to find use-after-free bugs.
+#ifndef MICROPY_DEBUG_VALGRIND_MAX_ADDRSPACE
+#define MICROPY_DEBUG_VALGRIND_MAX_ADDRSPACE (MICROPY_DEBUG_VALGRIND)
+#endif
+
 /*****************************************************************************/
 /* Optimisations                                                             */