Fix: Harden Python binding error paths and shape-overflow guards

Add nk_size_mul_checked_ to types.h and route the shape-product and cdist size
computations in python/{each,distance}.c through it, so a buffer reporting an
overflowing shape raises OverflowError instead of wrapping into an undersized
allocation walked at the true extent.

The elementwise scalar-array paths returned NULL without setting an exception
when the buffer dtype is unsupported (a CPython protocol violation); they now
raise TypeError. The packed and symmetric matrix verbs leaked the owned output
tensor on an invalid row range; they now Py_DECREF it. The DLPack importer
propagates a PyObject_IsTrue error and nulls owner->managed if PyCapsule_SetName
fails, avoiding a double-free of the producer tensor.

Author: ashvardanian

include/numkong/types.h

@@ -1748,6 +1748,14 @@ NK_PUBLIC nk_size_t nk_size_round_up_to_multiple_(nk_size_t number, nk_size_t di
     return nk_size_divide_round_up_(number, divisor) * divisor;
 }
 
+/** @brief Multiplies two sizes with overflow detection. Writes the product and returns 1 on success;
+ *         returns 0 (leaving @p product unchanged) when @p a * @p b would overflow `nk_size_t`. */
+NK_PUBLIC int nk_size_mul_checked_(nk_size_t a, nk_size_t b, nk_size_t *product) NK_STREAMING_COMPATIBLE_ {
+    if (b != 0 && a > NK_SIZE_MAX / b) return 0;
+    *product = a * b;
+    return 1;
+}
+
 NK_INTERNAL nk_f32_t nk_f32_abs_(nk_f32_t x) { return x < 0 ? -x : x; }
 NK_INTERNAL nk_f64_t nk_f64_abs_(nk_f64_t x) { return x < 0 ? -x : x; }
 NK_INTERNAL nk_i64_t nk_i64_abs_(nk_i64_t x) { return x < 0 ? -x : x; }

python/distance.c

@@ -815,7 +815,11 @@ static PyObject *implement_cdist(                        //
         goto cleanup;
     }
 
-    size_t const count_pairs = a_parsed.rows * b_parsed.rows;
+    size_t count_pairs;
+    if (!nk_size_mul_checked_(a_parsed.rows, b_parsed.rows, &count_pairs)) {
+        PyErr_SetString(PyExc_OverflowError, "cdist result size (a_rows * b_rows) overflows size_t");
+        goto cleanup;
+    }
     char *distances_start = NULL;
     size_t distances_rows_stride_bytes = 0;
     size_t distances_cols_stride_bytes = 0;

python/dlpack_interop.c

@@ -271,10 +271,16 @@ static int nk_fill_dl_tensor(Tensor *tensor, DLTensor *out, nk_dlpack_export_ctx
         if (!versioned && (tensor->dtype == nk_e2m3_k || tensor->dtype == nk_e3m2_k || tensor->dtype == nk_e2m1_k)) {
             PyErr_SetString( //
                 PyExc_TypeError,
-                "Sub-byte FP4/FP6 (e2m1/e2m3/e3m2) DLPack export requires max_version >= (1, 0) so " "the "
-                                                                                                     "IS_SUBBYTE_TYPE_"
-                                                                                                     "PADDED flag can "
-                                                                                                     "be set");
+                "Sub-byte FP4/FP6 (e2m1/e2m3/e3m2) DLPack export requires max_version >= (1, 0) so " "the " "IS_"
+                                                                                                            "SUBBYTE_"
+                                                                                                            "TYPE_" "PA"
+                                                                                                                    "DD"
+                                                                                                                    "ED"
+                                                                                                                    " f"
+                                                                                                                    "la"
+                                                                                                                    "g "
+                                                                                                                    "ca"
+                                                                                                                    "n " "be set");
         }
         else { PyErr_Format(PyExc_TypeError, "dtype %d has no DLPack mapping", (int)tensor->dtype); }
         return -1;
@@ -393,7 +399,9 @@ PyObject *Tensor_dlpack(PyObject *self, PyObject *args, PyObject *kwargs) {
 
     Tensor *tensor = (Tensor *)self;
 
-    if (copy != Py_None && PyObject_IsTrue(copy)) {
+    int const requested_copy = (copy != Py_None) ? PyObject_IsTrue(copy) : 0;
+    if (requested_copy < 0) return NULL; // bool(copy) raised
+    if (requested_copy) {
         PyErr_SetString(PyExc_BufferError, "NumKong DLPack exporter only supports zero-copy (copy=False)");
         return NULL;
     }
@@ -630,13 +638,24 @@ PyObject *api_from_dlpack(PyObject *self, PyObject *obj) {
     nk_dlpack_owner_t *owner = PyObject_New(nk_dlpack_owner_t, &DLPackOwnerType);
     if (!owner) goto fail;
     owner->is_versioned = is_versioned;
+    // Rename the capsule so the producer's destructor becomes a no-op and ownership transfers to `owner`.
+    // If the rename fails, leave the capsule's original name so the producer still frees it, and null
+    // `owner->managed` so our dealloc does not free it a second time.
     if (is_versioned) {
         owner->managed = PyCapsule_GetPointer(capsule, "dltensor_versioned");
-        PyCapsule_SetName(capsule, "used_dltensor_versioned");
+        if (PyCapsule_SetName(capsule, "used_dltensor_versioned") != 0) {
+            owner->managed = NULL;
+            Py_DECREF(owner);
+            goto fail;
+        }
     }
     else {
         owner->managed = PyCapsule_GetPointer(capsule, "dltensor");
-        PyCapsule_SetName(capsule, "used_dltensor");
+        if (PyCapsule_SetName(capsule, "used_dltensor") != 0) {
+            owner->managed = NULL;
+            Py_DECREF(owner);
+            goto fail;
+        }
     }
 
     // Build the Tensor view directly (parallels Tensor_view_object in tensor.c,

python/each.c

@@ -536,7 +536,11 @@ static PyObject *add_scalar_array(PyObject *array_obj, PyObject *scalar_obj, PyO
 
     nk_dtype_t dtype = resolve_nk_dtype_in_py_buffer(&a_buffer);
     if (out_dtype_obj) { dtype = py_object_to_nk_dtype(out_dtype_obj); }
-    if (dtype == nk_dtype_unknown_k) goto cleanup;
+    if (dtype == nk_dtype_unknown_k) {
+        if (!PyErr_Occurred())
+            PyErr_SetString(PyExc_TypeError, "unsupported buffer dtype for the requested elementwise operation");
+        goto cleanup;
+    }
 
     nk_each_scale_punned_t scale_kernel = NULL;
     nk_capability_t capability = nk_cap_serial_k;
@@ -555,7 +559,11 @@ static PyObject *add_scalar_array(PyObject *array_obj, PyObject *scalar_obj, PyO
 
     size_t const element_size = nk_dtype_bytes_per_value(dtype);
     size_t total_elements = 1;
-    for (int dim = 0; dim < a_buffer.ndim; dim++) total_elements *= (size_t)a_buffer.shape[dim];
+    for (int dim = 0; dim < a_buffer.ndim; dim++)
+        if (!nk_size_mul_checked_(total_elements, (size_t)a_buffer.shape[dim], &total_elements)) {
+            PyErr_SetString(PyExc_OverflowError, "tensor element count overflows size_t");
+            goto cleanup;
+        }
 
     char *result_data = NULL;
     Py_ssize_t result_strides[NK_TENSOR_MAX_RANK];
@@ -655,7 +663,11 @@ static PyObject *add_array_array(PyObject *a_obj, PyObject *b_obj, PyObject *out
     }
 
     if (out_dtype_obj) { dtype = py_object_to_nk_dtype(out_dtype_obj); }
-    if (dtype == nk_dtype_unknown_k) goto cleanup;
+    if (dtype == nk_dtype_unknown_k) {
+        if (!PyErr_Occurred())
+            PyErr_SetString(PyExc_TypeError, "unsupported buffer dtype for the requested elementwise operation");
+        goto cleanup;
+    }
 
     nk_each_sum_punned_t sum_kernel = NULL;
     nk_capability_t capability = nk_cap_serial_k;
@@ -668,7 +680,11 @@ static PyObject *add_array_array(PyObject *a_obj, PyObject *b_obj, PyObject *out
 
     int const num_dims = a_buffer.ndim;
     size_t total_elements = 1;
-    for (int dim = 0; dim < num_dims; dim++) total_elements *= (size_t)a_buffer.shape[dim];
+    for (int dim = 0; dim < num_dims; dim++)
+        if (!nk_size_mul_checked_(total_elements, (size_t)a_buffer.shape[dim], &total_elements)) {
+            PyErr_SetString(PyExc_OverflowError, "tensor element count overflows size_t");
+            goto cleanup;
+        }
 
     a_promoted = ensure_contiguous_buffer(a_buffer.buf, a_dtype, dtype, num_dims, a_buffer.shape, a_buffer.strides,
                                           total_elements, &a_needs_free);
@@ -827,7 +843,11 @@ static PyObject *multiply_scalar_array(PyObject *array_obj, PyObject *scalar_obj
 
     nk_dtype_t dtype = resolve_nk_dtype_in_py_buffer(&a_buffer);
     if (out_dtype_obj) { dtype = py_object_to_nk_dtype(out_dtype_obj); }
-    if (dtype == nk_dtype_unknown_k) goto cleanup;
+    if (dtype == nk_dtype_unknown_k) {
+        if (!PyErr_Occurred())
+            PyErr_SetString(PyExc_TypeError, "unsupported buffer dtype for the requested elementwise operation");
+        goto cleanup;
+    }
 
     nk_each_scale_punned_t scale_kernel = NULL;
     nk_capability_t capability = nk_cap_serial_k;
@@ -846,7 +866,11 @@ static PyObject *multiply_scalar_array(PyObject *array_obj, PyObject *scalar_obj
 
     size_t const element_size = nk_dtype_bytes_per_value(dtype);
     size_t total_elements = 1;
-    for (int dim = 0; dim < a_buffer.ndim; dim++) total_elements *= (size_t)a_buffer.shape[dim];
+    for (int dim = 0; dim < a_buffer.ndim; dim++)
+        if (!nk_size_mul_checked_(total_elements, (size_t)a_buffer.shape[dim], &total_elements)) {
+            PyErr_SetString(PyExc_OverflowError, "tensor element count overflows size_t");
+            goto cleanup;
+        }
 
     char *result_data = NULL;
     Py_ssize_t result_strides[NK_TENSOR_MAX_RANK];
@@ -946,7 +970,11 @@ static PyObject *multiply_array_array(PyObject *a_obj, PyObject *b_obj, PyObject
     }
 
     if (out_dtype_obj) { dtype = py_object_to_nk_dtype(out_dtype_obj); }
-    if (dtype == nk_dtype_unknown_k) goto cleanup;
+    if (dtype == nk_dtype_unknown_k) {
+        if (!PyErr_Occurred())
+            PyErr_SetString(PyExc_TypeError, "unsupported buffer dtype for the requested elementwise operation");
+        goto cleanup;
+    }
 
     nk_each_fma_punned_t fma_kernel = NULL;
     nk_capability_t capability = nk_cap_serial_k;
@@ -965,7 +993,11 @@ static PyObject *multiply_array_array(PyObject *a_obj, PyObject *b_obj, PyObject
 
     int const num_dims = a_buffer.ndim;
     size_t total_elements = 1;
-    for (int dim = 0; dim < num_dims; dim++) total_elements *= (size_t)a_buffer.shape[dim];
+    for (int dim = 0; dim < num_dims; dim++)
+        if (!nk_size_mul_checked_(total_elements, (size_t)a_buffer.shape[dim], &total_elements)) {
+            PyErr_SetString(PyExc_OverflowError, "tensor element count overflows size_t");
+            goto cleanup;
+        }
 
     a_promoted = ensure_contiguous_buffer(a_buffer.buf, a_dtype, dtype, num_dims, a_buffer.shape, a_buffer.strides,
                                           total_elements, &a_needs_free);

python/matrix.c

@@ -454,6 +454,7 @@ static PyObject *api_packed_common( //
     if (end_row < 0) end_row = (Py_ssize_t)height;
     if (start_row > (Py_ssize_t)height || end_row > (Py_ssize_t)height || start_row > end_row) {
         PyBuffer_Release(&a_buffer);
+        if (owns_result) Py_DECREF(result);
         PyErr_Format(PyExc_ValueError, "Invalid row range [%zd, %zd) for matrix with %zu rows", start_row, end_row,
                      (size_t)height);
         return NULL;
@@ -597,6 +598,7 @@ static PyObject *api_symmetric_common( //
         if (row_start > n_vectors || row_end > n_vectors || row_start > row_end) {
             PyErr_Format(PyExc_ValueError, "Invalid row range [%zu, %zu) for %zu vectors", (size_t)row_start,
                          (size_t)row_end, (size_t)n_vectors);
+            if (owns_result) Py_DECREF(result);
             goto cleanup;
         }
         nk_size_t row_count_val = (nk_size_t)(row_end - row_start);