Add explicit checks to all allocators in the NRT.

This patch:

* Updates nrt.c to ensure that allocation routines return NULL if
  the allocation fails. cf. malloc.
* Updates the nrt.context to add a check that wraps any allocation
  routine to check if the allocation failed and if so raises a
  MemoryError.
* Fixes the runtime system allocators to catch invalid allocation
  requests and handle failed allocation requests.
* Adds tests for the above.
* Fixes a storage classifier issue in the C code in the nrt tests.
* Adds the option `--debuginfo` to the `build_ext` command in
  setup.py, on unix systems this will ensure the DSOs Numba
  produces are built with debug symbols.
* Adds an update to the notes in nrt_external.h about the expected
  behaviour of an externally supplied allocation routine.

Fixes numba#8015
(the segfault part).

numba/core/runtime/context.py

@@ -16,6 +16,18 @@ def _require_nrt(self):
         if not self._enabled:
             raise errors.NumbaRuntimeError("NRT required but not enabled")
 
+    def _check_null_result(func):
+        def wrap(self, builder, *args):
+            memptr = func(self, builder, *args)
+            with builder.if_then(cgutils.is_null(builder, memptr),
+                                 likely=False):
+                msg = "Allocation failed (probably too large)."
+                self._context.call_conv.return_user_exc(builder, MemoryError,
+                                                        (msg,))
+            return memptr
+        return wrap
+
+    @_check_null_result
     def allocate(self, builder, size):
         """
         Low-level allocate a new memory area of `size` bytes.
@@ -39,6 +51,7 @@ def free(self, builder, ptr):
         fn = cgutils.get_or_insert_function(mod, fnty, "NRT_Free")
         return builder.call(fn, [ptr])
 
+    @_check_null_result
     def meminfo_alloc(self, builder, size):
         """
         Allocate a new MemInfo with a data payload of `size` bytes.
@@ -53,6 +66,7 @@ def meminfo_alloc(self, builder, size):
         fn.return_value.add_attribute("noalias")
         return builder.call(fn, [size])
 
+    @_check_null_result
     def meminfo_alloc_dtor(self, builder, size, dtor):
         self._require_nrt()
 
@@ -65,6 +79,7 @@ def meminfo_alloc_dtor(self, builder, size, dtor):
         return builder.call(fn, [size,
                                  builder.bitcast(dtor, cgutils.voidptr_t)])
 
+    @_check_null_result
     def meminfo_alloc_aligned(self, builder, size, align):
         """
         Allocate a new MemInfo with an aligned data payload of `size` bytes.
@@ -87,6 +102,7 @@ def meminfo_alloc_aligned(self, builder, size, align):
             assert align.type == u32, "align must be a uint32"
         return builder.call(fn, [size, align])
 
+    @_check_null_result
     def meminfo_new_varsize(self, builder, size):
         """
         Allocate a MemInfo pointing to a variable-sized data area.  The area
@@ -104,6 +120,7 @@ def meminfo_new_varsize(self, builder, size):
         fn.return_value.add_attribute("noalias")
         return builder.call(fn, [size])
 
+    @_check_null_result
     def meminfo_new_varsize_dtor(self, builder, size, dtor):
         """
         Like meminfo_new_varsize() but also set the destructor for
@@ -118,6 +135,7 @@ def meminfo_new_varsize_dtor(self, builder, size, dtor):
             mod, fnty, "NRT_MemInfo_new_varsize_dtor")
         return builder.call(fn, [size, dtor])
 
+    @_check_null_result
     def meminfo_varsize_alloc(self, builder, meminfo, size):
         """
         Allocate a new data area for a MemInfo created by meminfo_new_varsize().
@@ -132,6 +150,7 @@ def meminfo_varsize_alloc(self, builder, meminfo, size):
         return self._call_varsize_alloc(builder, meminfo, size,
                                         "NRT_MemInfo_varsize_alloc")
 
+    @_check_null_result
     def meminfo_varsize_realloc(self, builder, meminfo, size):
         """
         Reallocate a data area allocated by meminfo_new_varsize().

numba/core/runtime/nrt.c

@@ -189,8 +189,10 @@ NRT_MemInfo *NRT_MemInfo_new(void *data, size_t size,
                              NRT_dtor_function dtor, void *dtor_info)
 {
     NRT_MemInfo *mi = NRT_Allocate(sizeof(NRT_MemInfo));
-    NRT_Debug(nrt_debug_print("NRT_MemInfo_new mi=%p\n", mi));
-    NRT_MemInfo_init(mi, data, size, dtor, dtor_info, NULL);
+    if (mi != NULL) {
+        NRT_Debug(nrt_debug_print("NRT_MemInfo_new mi=%p\n", mi));
+        NRT_MemInfo_init(mi, data, size, dtor, dtor_info, NULL);
+    }
     return mi;
 }
 
@@ -212,9 +214,10 @@ void nrt_internal_dtor_safe(void *ptr, size_t size, void *info) {
 
 static
 void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out, NRT_ExternalAllocator *allocator) {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     NRT_Debug(nrt_debug_print("nrt_allocate_meminfo_and_data %p\n", allocator));
     char *base = NRT_Allocate_External(sizeof(NRT_MemInfo) + size, allocator);
+    if (base == NULL) return NULL; /* return early if alloc fails */
     mi = (NRT_MemInfo *) base;
     *mi_out = mi;
     return base + sizeof(NRT_MemInfo);
@@ -235,16 +238,18 @@ void nrt_internal_custom_dtor_safe(void *ptr, size_t size, void *info) {
 
 
 NRT_MemInfo *NRT_MemInfo_alloc(size_t size) {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     void *data = nrt_allocate_meminfo_and_data(size, &mi, NULL);
+    if (data == NULL) return NULL; /* return early if allocation failed */
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc %p\n", data));
     NRT_MemInfo_init(mi, data, size, NULL, NULL, NULL);
     return mi;
 }
 
 NRT_MemInfo *NRT_MemInfo_alloc_external(size_t size, NRT_ExternalAllocator *allocator) {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     void *data = nrt_allocate_meminfo_and_data(size, &mi, allocator);
+    if (data == NULL) return NULL; /* return early if allocation failed */
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc %p\n", data));
     NRT_MemInfo_init(mi, data, size, NULL, NULL, allocator);
     return mi;
@@ -255,8 +260,9 @@ NRT_MemInfo *NRT_MemInfo_alloc_safe(size_t size) {
 }
 
 NRT_MemInfo* NRT_MemInfo_alloc_dtor_safe(size_t size, NRT_dtor_function dtor) {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     void *data = nrt_allocate_meminfo_and_data(size, &mi, NULL);
+    if (data == NULL) return NULL; /* return early if allocation failed */
     /* Only fill up a couple cachelines with debug markers, to minimize
        overhead. */
     memset(data, 0xCB, MIN(size, 256));
@@ -270,9 +276,10 @@ static
 void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align,
                                           NRT_MemInfo **mi, NRT_ExternalAllocator *allocator)
 {
-    size_t offset, intptr, remainder;
+    size_t offset = 0, intptr = 0, remainder = 0;
     NRT_Debug(nrt_debug_print("nrt_allocate_meminfo_and_data_align %p\n", allocator));
     char *base = nrt_allocate_meminfo_and_data(size + 2 * align, mi, allocator);
+    if (base == NULL) return NULL; /* return early if allocation failed */
     intptr = (size_t) base;
     /* See if we are aligned */
     remainder = intptr % align;
@@ -285,16 +292,18 @@ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align,
 }
 
 NRT_MemInfo *NRT_MemInfo_alloc_aligned(size_t size, unsigned align) {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi, NULL);
+    if (data == NULL) return NULL; /* return early if allocation failed */
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc_aligned %p\n", data));
     NRT_MemInfo_init(mi, data, size, NULL, NULL, NULL);
     return mi;
 }
 
 NRT_MemInfo *NRT_MemInfo_alloc_safe_aligned(size_t size, unsigned align) {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi, NULL);
+    if (data == NULL) return NULL; /* return early if allocation failed */
     /* Only fill up a couple cachelines with debug markers, to minimize
        overhead. */
     memset(data, 0xCB, MIN(size, 256));
@@ -305,9 +314,10 @@ NRT_MemInfo *NRT_MemInfo_alloc_safe_aligned(size_t size, unsigned align) {
 }
 
 NRT_MemInfo *NRT_MemInfo_alloc_safe_aligned_external(size_t size, unsigned align, NRT_ExternalAllocator *allocator) {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc_safe_aligned_external %p\n", allocator));
     void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi, allocator);
+    if (data == NULL) return NULL; /* return early if allocation failed */
     /* Only fill up a couple cachelines with debug markers, to minimize
        overhead. */
     memset(data, 0xCB, MIN(size, 256));
@@ -397,10 +407,9 @@ nrt_varsize_dtor(void *ptr, size_t size, void *info) {
 
 NRT_MemInfo *NRT_MemInfo_new_varsize(size_t size)
 {
-    NRT_MemInfo *mi;
+    NRT_MemInfo *mi = NULL;
     void *data = NRT_Allocate(size);
-    if (data == NULL)
-        return NULL;
+    if (data == NULL) return NULL; /* return early if allocation failed */
 
     mi = NRT_MemInfo_new(data, size, nrt_varsize_dtor, NULL);
     NRT_Debug(nrt_debug_print("NRT_MemInfo_new_varsize size=%zu "

numba/core/runtime/nrt.py

@@ -92,10 +92,16 @@ def meminfo_alloc(self, size, safe=False):
         See `NRT_MemInfo_alloc_safe()` in "nrt.h" for details.
         """
         self._init_guard()
+        if size < 0:
+            msg = f"Cannot allocate a negative number of bytes: {size}."
+            raise ValueError(msg)
         if safe:
             mi = _nrt.meminfo_alloc_safe(size)
         else:
             mi = _nrt.meminfo_alloc(size)
+        if mi == 0: # alloc failed or size was 0 and alloc returned NULL.
+            msg = f"Requested allocation of {size} bytes failed."
+            raise MemoryError(msg)
         return MemInfo(mi)
 
     def get_allocation_stats(self):

numba/core/runtime/nrt_external.h

@@ -34,8 +34,11 @@ typedef struct {
     */
     NRT_MemInfo* (*allocate)(size_t nbytes);
     /* Allocates memory using an external allocator but still using Numba's MemInfo.
-
-    */
+     *
+     * NOTE: An externally provided allocator must behave the same way as C99
+     *       stdlib.h's "malloc" function with respect to return value. i.e. on
+     *       error the allocator must return NULL.
+     */
     NRT_MemInfo* (*allocate_external)(size_t nbytes, NRT_ExternalAllocator *allocator);
 
     /* Convert externally allocated memory into a MemInfo.

numba/tests/test_dyn_array.py

@@ -416,6 +416,20 @@ def wrapped(n, input, out):
         # The following checks can discover a reference count error
         self.assertEqual(expected_refct, sys.getrefcount(input))
 
+    def test_invalid_size_array(self):
+
+        @njit
+        def foo(x):
+            np.empty(x)
+
+        # Exceptions leak references
+        self.disable_leak_check()
+
+        with self.assertRaises(MemoryError) as raises:
+            foo(types.size_t.maxval // 8 // 2)
+
+        self.assertIn("Allocation failed", str(raises.exception))
+
     def test_swap(self):
 
         def pyfunc(x, y, t):

numba/tests/test_nrt.py

@@ -219,6 +219,25 @@ def test_buffer(self):
         # We can't check this deterministically because the memory could be
         # consumed by another thread.
 
+    def test_allocate_invalid_size(self):
+        # Checks that attempting to allocate too big a region fails gracefully.
+        size = types.size_t.maxval // 8 // 2
+        for pred in (True, False):
+            with self.assertRaises(MemoryError) as raises:
+                rtsys.meminfo_alloc(size, safe=pred)
+            self.assertIn(f"Requested allocation of {size} bytes failed.",
+                          str(raises.exception))
+
+    def test_allocate_negative_size(self):
+        # Checks that attempting to allocate negative number of bytes fails
+        # gracefully.
+        size = -10
+        for pred in (True, False):
+            with self.assertRaises(ValueError) as raises:
+                rtsys.meminfo_alloc(size, safe=pred)
+            msg = f"Cannot allocate a negative number of bytes: {size}."
+            self.assertIn(msg, str(raises.exception))
+
 
 class TestTracemalloc(unittest.TestCase):
     """
@@ -610,7 +629,7 @@ def test_manage_memory(self):
         """
         cdef = """
 void* test_nrt_api(void *nrt);
-int status;
+extern int status;
         """
 
         ffi, mod = self.compile_cffi_module(name, source, cdef)

setup.py

@@ -76,19 +76,22 @@ def run(self):
     ('werror', None, 'Build extensions with -Werror'),
     ('wall', None, 'Build extensions with -Wall'),
     ('noopt', None, 'Build extensions without optimization'),
+    ('debuginfo', None, 'Build extensions with debug symbols'),
 ]
 
 
 class NumbaBuildExt(build_ext):
 
     user_options = build_ext.user_options + numba_be_user_options
-    boolean_options = build_ext.boolean_options + ['werror', 'wall', 'noopt']
+    debug_opts = ['werror', 'wall', 'noopt', 'debuginfo']
+    boolean_options = build_ext.boolean_options + debug_opts
 
     def initialize_options(self):
         super().initialize_options()
         self.werror = 0
         self.wall = 0
         self.noopt = 0
+        self.debuginfo = 0
 
     def run(self):
         extra_compile_args = []
@@ -97,6 +100,9 @@ def run(self):
                 extra_compile_args.append('/Od')
             else:
                 extra_compile_args.append('-O0')
+        if self.debuginfo:
+            if sys.platform != 'win32':
+                extra_compile_args.append('-g')
         if self.werror:
             extra_compile_args.append('-Werror')
         if self.wall: