Improve RGBMatrix allocation tracking

This prevents leaks but not all use-after-free issues.

Author: tannewt

ports/espressif/boards/adafruit_matrixportal_s3/pins.c

@@ -5,10 +5,10 @@ STATIC const mp_rom_obj_tuple_t matrix_addr_tuple = {
     {&mp_type_tuple},
     5,
     {
-        MP_ROM_PTR(&pin_GPIO35),
+        MP_ROM_PTR(&pin_GPIO45),
         MP_ROM_PTR(&pin_GPIO36),
         MP_ROM_PTR(&pin_GPIO48),
-        MP_ROM_PTR(&pin_GPIO45),
+        MP_ROM_PTR(&pin_GPIO35),
         MP_ROM_PTR(&pin_GPIO21),
     }
 };

ports/espressif/common-hal/rgbmatrix/RGBMatrix.c

@@ -61,6 +61,9 @@ void common_hal_rgbmatrix_timer_enable(void *ptr) {
 }
 
 void common_hal_rgbmatrix_timer_disable(void *ptr) {
+    if (ptr == NULL) {
+        return;
+    }
     timer_index_t *timer = (timer_index_t *)ptr;
     if (timer->idx == TIMER_MAX) {
         return;

shared-bindings/rgbmatrix/RGBMatrix.c

@@ -376,9 +376,7 @@ STATIC const mp_rom_map_elem_t rgbmatrix_rgbmatrix_locals_dict_table[] = {
 STATIC MP_DEFINE_CONST_DICT(rgbmatrix_rgbmatrix_locals_dict, rgbmatrix_rgbmatrix_locals_dict_table);
 
 STATIC void rgbmatrix_rgbmatrix_get_bufinfo(mp_obj_t self_in, mp_buffer_info_t *bufinfo) {
-    rgbmatrix_rgbmatrix_obj_t *self = (rgbmatrix_rgbmatrix_obj_t *)self_in;
-
-    *bufinfo = self->bufinfo;
+    common_hal_rgbmatrix_rgbmatrix_get_bufinfo(self_in, bufinfo);
 }
 
 // These version exists so that the prototype matches the protocol,
@@ -442,7 +440,7 @@ STATIC mp_int_t rgbmatrix_rgbmatrix_get_buffer(mp_obj_t self_in, mp_buffer_info_
     if ((flags & MP_BUFFER_WRITE) && !(self->bufinfo.typecode & MP_OBJ_ARRAY_TYPECODE_FLAG_RW)) {
         return 1;
     }
-    *bufinfo = self->bufinfo;
+    common_hal_rgbmatrix_rgbmatrix_get_bufinfo(self_in, bufinfo);
     bufinfo->typecode = 'H';
     return 0;
 }

shared-bindings/rgbmatrix/RGBMatrix.h

@@ -34,6 +34,7 @@ extern const mp_obj_type_t rgbmatrix_RGBMatrix_type;
 void common_hal_rgbmatrix_rgbmatrix_construct(rgbmatrix_rgbmatrix_obj_t *self, int width, int bit_depth, uint8_t rgb_count, uint8_t *rgb_pins, uint8_t addr_count, uint8_t *addr_pins, uint8_t clock_pin, uint8_t latch_pin, uint8_t oe_pin, bool doublebuffer, mp_obj_t framebuffer, int8_t tile, bool serpentine, void *timer);
 void common_hal_rgbmatrix_rgbmatrix_deinit(rgbmatrix_rgbmatrix_obj_t *);
 void rgbmatrix_rgbmatrix_collect_ptrs(rgbmatrix_rgbmatrix_obj_t *);
+void common_hal_rgbmatrix_rgbmatrix_get_bufinfo(rgbmatrix_rgbmatrix_obj_t *self, mp_buffer_info_t *bufinfo);
 void common_hal_rgbmatrix_rgbmatrix_reconstruct(rgbmatrix_rgbmatrix_obj_t *self);
 void common_hal_rgbmatrix_rgbmatrix_set_paused(rgbmatrix_rgbmatrix_obj_t *self, bool paused);
 bool common_hal_rgbmatrix_rgbmatrix_get_paused(rgbmatrix_rgbmatrix_obj_t *self);

shared-module/rgbmatrix/RGBMatrix.c

@@ -79,8 +79,12 @@ STATIC void common_hal_rgbmatrix_rgbmatrix_construct1(rgbmatrix_rgbmatrix_obj_t
         }
         // verify that the matrix is big enough
         mp_get_index(mp_obj_get_type(self->framebuffer), self->bufinfo.len, MP_OBJ_NEW_SMALL_INT(self->bufsize - 1), false);
+        self->allocation = NULL;
     } else {
-        self->bufinfo.buf = common_hal_rgbmatrix_allocator_impl(self->bufsize);
+        // The supervisor allocation can move memory by changing self->allocation->ptr.
+        // So we hold onto it and update bufinfo every time we use it.
+        self->allocation = allocate_memory(align32_size(self->bufsize), false, true);
+        self->bufinfo.buf = self->allocation->ptr;
         self->bufinfo.len = self->bufsize;
         self->bufinfo.typecode = 'H' | MP_OBJ_ARRAY_TYPECODE_FLAG_RW;
     }
@@ -146,7 +150,9 @@ STATIC void free_pin_seq(uint8_t *seq, int count) {
 
 extern int pm_row_count;
 STATIC void common_hal_rgbmatrix_rgbmatrix_deinit1(rgbmatrix_rgbmatrix_obj_t *self) {
-    common_hal_rgbmatrix_timer_disable(self->timer);
+    if (self->timer != NULL) {
+        common_hal_rgbmatrix_timer_disable(self->timer);
+    }
 
     if (_PM_protoPtr == &self->protomatter) {
         _PM_protoPtr = NULL;
@@ -160,15 +166,15 @@ STATIC void common_hal_rgbmatrix_rgbmatrix_deinit1(rgbmatrix_rgbmatrix_obj_t *se
 
     // If it was supervisor-allocated, it is supervisor-freed and the pointer
     // is zeroed, otherwise the pointer is just zeroed
-    if (self->bufinfo.buf) {
-        common_hal_rgbmatrix_free_impl(self->bufinfo.buf);
-        self->bufinfo.buf = NULL;
+    if (self->allocation != NULL) {
+        free_memory(self->allocation);
     }
 
-
     // If a framebuffer was passed in to the constructor, clear the reference
     // here so that it will become GC'able
     self->framebuffer = mp_const_none;
+
+    self->bufinfo.buf = NULL;
 }
 
 void common_hal_rgbmatrix_rgbmatrix_deinit(rgbmatrix_rgbmatrix_obj_t *self) {
@@ -187,6 +193,13 @@ void common_hal_rgbmatrix_rgbmatrix_deinit(rgbmatrix_rgbmatrix_obj_t *self) {
     self->base.type = &mp_type_NoneType;
 }
 
+void common_hal_rgbmatrix_rgbmatrix_get_bufinfo(rgbmatrix_rgbmatrix_obj_t *self, mp_buffer_info_t *bufinfo) {
+    if (self->allocation != NULL) {
+        self->bufinfo.buf = self->allocation->ptr;
+    }
+    *bufinfo = self->bufinfo;
+}
+
 void common_hal_rgbmatrix_rgbmatrix_reconstruct(rgbmatrix_rgbmatrix_obj_t *self) {
     if (self->framebuffer != mp_const_none) {
         memset(&self->bufinfo, 0, sizeof(self->bufinfo));
@@ -196,11 +209,6 @@ void common_hal_rgbmatrix_rgbmatrix_reconstruct(rgbmatrix_rgbmatrix_obj_t *self)
     common_hal_rgbmatrix_rgbmatrix_deinit1(self);
     common_hal_rgbmatrix_rgbmatrix_construct1(self, mp_const_none);
     #endif
-    if (self->bufinfo.buf == NULL) {
-        self->bufinfo.buf = common_hal_rgbmatrix_allocator_impl(self->bufsize);
-        self->bufinfo.len = self->bufsize;
-        self->bufinfo.typecode = 'H' | MP_OBJ_ARRAY_TYPECODE_FLAG_RW;
-    }
     memset(self->bufinfo.buf, 0, self->bufinfo.len);
     common_hal_rgbmatrix_rgbmatrix_set_paused(self, false);
 }
@@ -214,6 +222,9 @@ void common_hal_rgbmatrix_rgbmatrix_set_paused(rgbmatrix_rgbmatrix_obj_t *self,
         _PM_stop(&self->protomatter);
     } else if (!paused && self->paused) {
         _PM_resume(&self->protomatter);
+        if (self->allocation) {
+            self->bufinfo.buf = self->allocation->ptr;
+        }
         _PM_convert_565(&self->protomatter, self->bufinfo.buf, self->width);
         _PM_swapbuffer_maybe(&self->protomatter);
     }
@@ -226,6 +237,9 @@ bool common_hal_rgbmatrix_rgbmatrix_get_paused(rgbmatrix_rgbmatrix_obj_t *self)
 
 void common_hal_rgbmatrix_rgbmatrix_refresh(rgbmatrix_rgbmatrix_obj_t *self) {
     if (!self->paused) {
+        if (self->allocation != NULL) {
+            self->bufinfo.buf = self->allocation->ptr;
+        }
         _PM_convert_565(&self->protomatter, self->bufinfo.buf, self->width);
         _PM_swapbuffer_maybe(&self->protomatter);
     }
@@ -240,11 +254,43 @@ int common_hal_rgbmatrix_rgbmatrix_get_height(rgbmatrix_rgbmatrix_obj_t *self) {
     return computed_height;
 }
 
+// Track the returned pointers and their matching allocation so that we can free
+// them even when the memory was moved by the supervisor. This prevents leaks
+// but doesn't protect against the memory being used after its been freed! The
+// long term fix is to utilize a permanent heap that can be shared with MP's
+// split heap.
+typedef struct matrix_allocation {
+    void *original_pointer;
+    supervisor_allocation *allocation;
+} matrix_allocation_t;
+
+// Four should be more than we ever need. ProtoMatter does 3 allocations currently.
+static matrix_allocation_t allocations[4];
+
 void *common_hal_rgbmatrix_allocator_impl(size_t sz) {
     supervisor_allocation *allocation = allocate_memory(align32_size(sz), false, true);
-    return allocation ? allocation->ptr : NULL;
+    if (allocation == NULL) {
+        return NULL;
+    }
+    for (size_t i = 0; i < sizeof(allocations); i++) {
+        matrix_allocation_t *matrix_allocation = &allocations[i];
+        if (matrix_allocation->original_pointer == NULL) {
+            matrix_allocation->original_pointer = allocation->ptr;
+            matrix_allocation->allocation = allocation;
+            return allocation->ptr;
+        }
+    }
+    return NULL;
 }
 
 void common_hal_rgbmatrix_free_impl(void *ptr_in) {
-    free_memory(allocation_from_ptr(ptr_in));
+    for (size_t i = 0; i < sizeof(allocations); i++) {
+        matrix_allocation_t *matrix_allocation = &allocations[i];
+        if (matrix_allocation->original_pointer == ptr_in) {
+            matrix_allocation->original_pointer = NULL;
+            free_memory(matrix_allocation->allocation);
+            matrix_allocation->allocation = NULL;
+            return;
+        }
+    }
 }

shared-module/rgbmatrix/RGBMatrix.h

@@ -28,12 +28,14 @@
 
 #include "py/obj.h"
 #include "lib/protomatter/src/core.h"
+#include "supervisor/memory.h"
 
 extern const mp_obj_type_t rgbmatrix_RGBMatrix_type;
 typedef struct {
     mp_obj_base_t base;
     mp_obj_t framebuffer;
     mp_buffer_info_t bufinfo;
+    supervisor_allocation *allocation;
     Protomatter_core protomatter;
     void *timer;
     uint16_t bufsize, width;