check for creating strings that are too large

Author: ngoldbaum

stringdtype/stringdtype/src/static_string.c

@@ -120,6 +120,10 @@ npy_string_newsize(const char *init, size_t size,
         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) {
@@ -155,6 +159,10 @@ npy_string_newemptysize(size_t size, npy_packed_static_string *out)
         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) {

stringdtype/stringdtype/src/static_string.h

@@ -33,7 +33,8 @@ extern const npy_packed_static_string *NPY_NULL_STRING;
 // 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 malloc fails. Returns 0 on success.
+// 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);
@@ -60,7 +61,8 @@ npy_string_dup(const npy_packed_static_string *in,
 // 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 malloc fails. Returns 0 on success.
+// 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);
 

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)