Merge pull request #88 from ngoldbaum/api-cleanup

ngoldbaum · web-flow · commit 3240d939dd5d · 2023-09-25T14:09:11.000-06:00
Clean up error handling and int overflow
diff --git a/stringdtype/stringdtype/src/casts.c b/stringdtype/stringdtype/src/casts.c
@@ -207,7 +207,8 @@ unicode_to_string(PyArrayMethod_Context *context, char *const data[],
         npy_packed_static_string *out_pss = (npy_packed_static_string *)out;
         npy_string_free(out_pss);
         if (npy_string_newemptysize(out_num_bytes, out_pss) < 0) {
-            gil_error(PyExc_MemoryError, "npy_string_newemptysize failed");
+            gil_error(PyExc_MemoryError,
+                      "Failed to allocate string in unicode to string cast");
             return -1;
         }
         npy_static_string out_ss = {0, NULL};
@@ -336,7 +337,7 @@ string_to_unicode(PyArrayMethod_Context *context, char *const data[],
     while (N--) {
         const npy_packed_static_string *ps = (npy_packed_static_string *)in;
         npy_static_string s = {0, NULL};
-        npy_static_string name = {0, NULL};
+        npy_static_string name;
         unsigned char *this_string = NULL;
         size_t n_bytes;
         int is_null = npy_load_string(ps, &s);
@@ -481,23 +482,28 @@ bool_to_string(PyArrayMethod_Context *NPY_UNUSED(context), char *const data[],
     while (N--) {
         npy_packed_static_string *out_pss = (npy_packed_static_string *)out;
         npy_string_free(out_pss);
+        char *ret_val = NULL;
+        size_t size = 0;
         if ((npy_bool)(*in) == 1) {
-            if (npy_string_newsize("True", 4, out_pss) < 0) {
-                gil_error(PyExc_MemoryError, "npy_string_newsize failed");
-                return -1;
-            }
+            ret_val = "True";
+            size = 4;
         }
         else if ((npy_bool)(*in) == 0) {
-            if (npy_string_newsize("False", 5, out_pss) < 0) {
-                gil_error(PyExc_MemoryError, "npy_string_newsize failed");
-                return -1;
-            }
+            ret_val = "False";
+            size = 5;
         }
         else {
             gil_error(PyExc_RuntimeError,
                       "invalid value encountered in bool to string cast");
             return -1;
         }
+        if (npy_string_newsize(ret_val, size, out_pss) < 0) {
+            // execution should never get here because this will be a small
+            // string on all platforms
+            gil_error(PyExc_MemoryError,
+                      "Failed to allocate string in bool to string cast");
+            return -1;
+        }
         in += in_stride;
         out += out_stride;
     }
@@ -593,7 +599,9 @@ pyobj_to_string(PyObject *obj, char *out)
     npy_packed_static_string *out_ss = (npy_packed_static_string *)out;
     npy_string_free(out_ss);
     if (npy_string_newsize(cstr_val, length, out_ss) < 0) {
-        PyErr_SetString(PyExc_MemoryError, "npy_string_newsize failed");
+        PyErr_SetString(PyExc_MemoryError,
+                        "Failed to allocate numpy string when converting from "
+                        "python string.");
         Py_DECREF(pystr_val);
         return -1;
     }
@@ -1087,7 +1095,8 @@ datetime_to_string(PyArrayMethod_Context *context, char *const data[],
             if (npy_string_newsize(datetime_buf, strlen(datetime_buf),
                                    out_pss) < 0) {
                 PyErr_SetString(PyExc_MemoryError,
-                                "npy_string_newsize failed");
+                                "Failed to allocate string when converting "
+                                "from a datetime.");
                 return -1;
             }
         }
diff --git a/stringdtype/stringdtype/src/dtype.c b/stringdtype/stringdtype/src/dtype.c
@@ -31,19 +31,12 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
             has_string_na = 1;
             Py_ssize_t size = 0;
             const char *buf = PyUnicode_AsUTF8AndSize(na_object, &size);
-            int res = npy_string_newsize(buf, (size_t)size,
-                                         &packed_default_string);
-            if (res == -1) {
+            if (npy_string_newsize(buf, (size_t)size, &packed_default_string) <
+                0) {
                 PyErr_NoMemory();
                 Py_DECREF(new);
                 return NULL;
             }
-            else if (res == -2) {
-                // this should never happen
-                assert(0);
-                Py_DECREF(new);
-                return NULL;
-            }
         }
         else {
             // treat as nan-like if != comparison returns a object whose truth
@@ -72,17 +65,10 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
 
         Py_ssize_t size = 0;
         const char *utf8_ptr = PyUnicode_AsUTF8AndSize(na_pystr, &size);
-        // discard const to initialize buffer
-        int res = npy_string_newsize(utf8_ptr, (size_t)size, &packed_na_name);
-        if (res == -1) {
+        if (npy_string_newsize(utf8_ptr, (size_t)size, &packed_na_name)) {
             PyErr_NoMemory();
             Py_DECREF(new);
-            return NULL;
-        }
-        else if (res == -2) {
-            // this should never happen
-            assert(0);
-            Py_DECREF(new);
+            Py_DECREF(na_pystr);
             return NULL;
         }
         Py_DECREF(na_pystr);
@@ -235,19 +221,11 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
             return -1;
         }
 
-        int res = npy_string_newsize(val, length, sdata);
-
-        if (res == -1) {
+        if (npy_string_newsize(val, length, sdata) < 0) {
             PyErr_NoMemory();
             Py_DECREF(val_obj);
             return -1;
         }
-        else if (res == -2) {
-            // this should never happen
-            assert(0);
-            Py_DECREF(val_obj);
-            return -1;
-        }
     }
 
     return 0;
diff --git a/stringdtype/stringdtype/src/static_string.c b/stringdtype/stringdtype/src/static_string.c
@@ -48,9 +48,6 @@ typedef union _npy_static_string_u {
 #define NPY_SHORT_STRING_MAX_SIZE \
     (sizeof(npy_static_string) - 1)  // 15 or 7 depending on arch
 
-// one byte in size is reserved for flags and small string optimization
-#define NPY_MAX_STRING_SIZE (1 << (sizeof(size_t) - 1)) - 1
-
 // Since this has no flags set, technically this is a heap-allocated string
 // with size zero. Practically, that doesn't matter because we always do size
 // checks before accessing heap data, but that may be confusing. The nice part
@@ -118,15 +115,15 @@ int
 npy_string_newsize(const char *init, size_t size,
                    npy_packed_static_string *to_init)
 {
-    if (to_init == NULL || size > NPY_MAX_STRING_SIZE) {
-        return -2;
-    }
-
     if (size == 0) {
         *to_init = *NPY_EMPTY_STRING;
         return 0;
     }
 
+    if (size > NPY_MAX_STRING_SIZE) {
+        return -1;
+    }
+
     _npy_static_string_u *to_init_u = ((_npy_static_string_u *)to_init);
 
     if (size > NPY_SHORT_STRING_MAX_SIZE) {
@@ -157,15 +154,15 @@ npy_string_newsize(const char *init, size_t size,
 int
 npy_string_newemptysize(size_t size, npy_packed_static_string *out)
 {
-    if (out == NULL || size > NPY_MAX_STRING_SIZE) {
-        return -2;
-    }
-
     if (size == 0) {
         *out = *NPY_EMPTY_STRING;
         return 0;
     }
 
+    if (size > NPY_MAX_STRING_SIZE) {
+        return -1;
+    }
+
     _npy_static_string_u *out_u = (_npy_static_string_u *)out;
 
     if (size > NPY_SHORT_STRING_MAX_SIZE) {
diff --git a/stringdtype/stringdtype/src/static_string.h b/stringdtype/stringdtype/src/static_string.h
@@ -21,13 +21,20 @@ extern const npy_packed_static_string *NPY_EMPTY_STRING;
 // working with it.
 extern const npy_packed_static_string *NPY_NULL_STRING;
 
+// one byte in size is reserved for flags and small string optimization
+#define NPY_MAX_STRING_SIZE (1 << (sizeof(size_t) - 1)) - 1
+
 // Allocates a new buffer for *to_init*, which must be set to NULL before
 // calling this function, filling the newly allocated buffer with the copied
 // contents of the first *size* entries in *init*, which must be valid and
 // initialized beforehand. Calling npy_string_free on *to_init* before calling
-// this function on an existing string is sufficient to initialize it. Returns
-// -1 if malloc fails and -2 if the internal buffer in *to_init* is not NULL
-// to indicate a programming error. Returns 0 on success.
+// this function on an existing string or copying the contents of
+// NPY_EMPTY_STRING into *to_init* is sufficient to initialize it. Does not
+// check if *to_init* is NULL or if the internal buffer is non-NULL, undefined
+// behavior or memory leaks are possible if this function is passed a pointer
+// to a an unintialized struct, a NULL pointer, or an existing heap-allocated
+// string.  Returns -1 if allocating the string would exceed the maximum
+// allowed string size or exhaust available memory. Returns 0 on success.
 int
 npy_string_newsize(const char *init, size_t size,
                    npy_packed_static_string *to_init);
@@ -47,8 +54,15 @@ npy_string_dup(const npy_packed_static_string *in,
 
 // Allocates a new string buffer for *out* with enough capacity to store
 // *size* bytes of text. Does not do any initialization, the caller must
-// initialize the string buffer. Returns -1 if malloc fails and -2 if *out* is
-// not NULL. Returns 0 on success.
+// initialize the string buffer after this function returns. Calling
+// npy_string_free on *to_init* before calling this function on an existing
+// string or copying the contents of NPY_EMPTY_STRING into *to_init* is
+// sufficient to initialize it.Does not check if *to_init* is NULL or if the
+// internal buffer is non-NULL, undefined behavior or memory leaks are
+// possible if this function is passed a pointer to a an unintialized struct,
+// a NULL pointer, or an existing heap-allocated string.  Returns 0 on
+// success. Returns -1 if allocating the string would exceed the maximum
+// allowed string size or exhaust available memory. Returns 0 on success.
 int
 npy_string_newemptysize(size_t size, npy_packed_static_string *out);
 
diff --git a/stringdtype/stringdtype/src/umath.c b/stringdtype/stringdtype/src/umath.c
@@ -68,19 +68,21 @@ multiply_resolve_descriptors(
             }                                                                \
             npy_##shortname factor = *(npy_##shortname *)iin;                \
             size_t cursize = is.size;                                        \
-            /* FIXME: check for overflow? */                                 \
             size_t newsize = cursize * factor;                               \
-                                                                             \
-            if (npy_string_newemptysize(newsize, ops) < 0) {                 \
+            /* newsize can only be less than cursize if there is overflow */ \
+            if (((newsize < cursize) ||                                      \
+                 npy_string_newemptysize(newsize, ops) < 0)) {               \
                 gil_error(PyExc_MemoryError,                                 \
-                          "npy_string_newemptysize failed");                 \
+                          "Failed to allocate string in string mutiply");    \
                 return -1;                                                   \
             }                                                                \
                                                                              \
             npy_static_string os = {0, NULL};                                \
             npy_load_string(ops, &os);                                       \
             for (size_t i = 0; i < (size_t)factor; i++) {                    \
                 /* excplicitly discard const; initializing new buffer */     \
+                /* multiply can't overflow because cursize * factor */       \
+                /* has already been checked and doesn't overflow */          \
                 memcpy((char *)os.buf + i * cursize, is.buf, cursize);       \
             }                                                                \
                                                                              \
@@ -245,6 +247,12 @@ add_strided_loop(PyArrayMethod_Context *context, char *const data[],
             }
         }
 
+        if ((s1.size + s2.size) < s1.size) {
+            // overflow
+            gil_error(PyExc_MemoryError,
+                      "Failed to allocate string in string add");
+        }
+
         if (npy_string_newemptysize(s1.size + s2.size, ops) < 0) {
             return -1;
         }
diff --git a/stringdtype/tests/test_stringdtype.py b/stringdtype/tests/test_stringdtype.py
@@ -700,3 +700,9 @@ def test_null_roundtripping(dtype):
     arr = np.array(data, dtype=dtype)
     assert data[0] == arr[0]
     assert data[1] == arr[1]
+
+
+def test_string_too_large_error():
+    arr = np.array(["a", "b", "c"], dtype=StringDType())
+    with pytest.raises(MemoryError):
+        arr * (2**63 - 2)
