From 8af061ade39c07cbb26f921c42217a7bfdd1b6ba Mon Sep 17 00:00:00 2001 From: Louis Verhaard Date: Fri, 22 Jan 2021 14:03:11 +0100 Subject: MLBEDSW-3832: Search allocator: improve C API - Removed unnecessary casts - Added more error handling Change-Id: Ic822877544f67452339a20dca4addddc050d195c Signed-off-by: Louis Verhaard --- ethosu/tensor_allocator/tensor_allocatormodule.cpp | 25 ++++++++++++++++------ .../tensor_allocator/test/test_tensor_allocator.py | 14 ++++++++++++ 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/ethosu/tensor_allocator/tensor_allocatormodule.cpp b/ethosu/tensor_allocator/tensor_allocatormodule.cpp index 02488add..52f1c690 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 +#include #include #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(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 input; std::vector 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 5350fde0..1011279c 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) -- cgit v1.2.1