Crash nicely on float input to bytearray_malloc

sabnzbd · Dec 17, 2023 · 876b03c · 876b03c
1 parent 8f92f90
commit 876b03c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/src/utils.cc b/src/utils.cc
@@ -19,5 +19,9 @@
 #include "utils.h"
 
 PyObject* bytearray_malloc(PyObject* self, PyObject* Py_input_size) {
+    if(!PyLong_Check(Py_input_size)) {
+        PyErr_SetString(PyExc_TypeError, "Expected type 'int'.");
+        return NULL;
+    }
     return PyByteArray_FromStringAndSize(NULL, PyLong_AsSsize_t(Py_input_size));
 }
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -8,7 +8,9 @@ def test_bytearray_malloc():
 
 
 def test_bytearray_malloc_bad_inputs():
+    with pytest.raises(TypeError):
+        sabctools.bytearray_malloc(10.0)
     with pytest.raises(SystemError):
         sabctools.bytearray_malloc(-1)
-    with pytest.raises(SystemError):
+    with pytest.raises(TypeError):
         sabctools.bytearray_malloc("foo")