MLBEDSW-3832: Search allocator: improve C API

- Removed unnecessary casts
- Added more error handling

Change-Id: Ic822877544f67452339a20dca4addddc050d195c
Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>

2 files changed, 33 insertions, 6 deletions

diff --git a/ethosu/tensor_allocator/tensor_allocatormodule.cpp b/ethosu/tensor_allocator/tensor_allocatormodule.cpp
index 02488ad..52f1c69 100644
--- a/ethosu/tensor_allocator/tensor_allocatormodule.cpp
+++ b/ethosu/tensor_allocator/tensor_allocatormodule.cpp
@@ -18,6 +18,7 @@
 
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
+#include <cstdint>
 
 #include <vector>
 #include "search_allocator.h"
@@ -45,23 +46,35 @@ static PyObject *method_allocate (PyObject *self, PyObject *args)
     int available_size = 0;
 
     /* Arguments to the method are delivered as a tuple, unpack the
-      * tuple to get the individual arguments, note the second is
-      * optional.
-      */
+     * tuple to get the individual arguments, note the second is
+     * optional.
+     */
     if (!PyArg_ParseTuple(args, "O|i", &input_list_object, &available_size)) {
         return NULL;
     }
 
     /* Unpack the length of the input integer list. */
-    int input_length = static_cast<int>(PyObject_Length (input_list_object));
+    auto input_length = PyObject_Length(input_list_object);
     if (input_length < 0) {
-        input_length = 0;
+        return NULL;
+    }
+    if (input_length % 3 != 0) {
+        PyErr_SetString(PyExc_ValueError, "Input length must be multiple of 3");
+        return NULL;
     }
     std::vector<uint32_t> input;
     std::vector<uint32_t> output;
     for (int i = 0; i < input_length; ++i) {
         PyObject *obj = PyList_GetItem(input_list_object, i);
-        uint32_t value = (uint32_t)PyLong_AsLong(obj);
+        if (!PyLong_Check(obj)) {
+            PyErr_SetString(PyExc_ValueError, "Illegal value in input");
+            return NULL;
+        }
+        auto value = PyLong_AsLong(obj);
+        if (value < 0 || value > UINT32_MAX) {
+            PyErr_SetString(PyExc_ValueError, "Input value out of bounds");
+            return NULL;
+        }
         input.push_back(value);
     }
     allocate(input, available_size, output);
diff --git a/ethosu/tensor_allocator/test/test_tensor_allocator.py b/ethosu/tensor_allocator/test/test_tensor_allocator.py
index 5350fde..1011279 100644
--- a/ethosu/tensor_allocator/test/test_tensor_allocator.py
+++ b/ethosu/tensor_allocator/test/test_tensor_allocator.py
@@ -47,3 +47,17 @@ def test_allocate(lrs, expected_size):
     res = tensor_allocator.allocate(input, 0)
     assert len(res) == len(lrs)
     assert max(addr + lr[2] for addr, lr in zip(res, lrs)) == expected_size
+
+
+def test_allocate_empty_input():
+    assert [] == tensor_allocator.allocate([], 0)
+
+
+invalid_input_test_data = [None, 3, [1, 2, 16, 9, 15], [1, 5, None], [-1, 0, 16], [1, 2, 10000000000]]
+
+
+@pytest.mark.parametrize("input", invalid_input_test_data)
+def test_allocate_invalid_input(input):
+    """Tests the search allocator with invalid input data"""
+    with pytest.raises(Exception):
+        tensor_allocator.allocate(input, 0)