diff --git a/UnityPyBoost/TypeTreeHelper.cpp b/UnityPyBoost/TypeTreeHelper.cpp index 65e4e304..12926ccf 100644 --- a/UnityPyBoost/TypeTreeHelper.cpp +++ b/UnityPyBoost/TypeTreeHelper.cpp @@ -15,6 +15,7 @@ typedef struct Reader uint8_t *ptr; uint8_t *end; uint8_t *start; + bool exhausted; } ReaderT; typedef struct TypeTreeReaderConfig @@ -75,6 +76,7 @@ inline PyObject *read_bool(ReaderT *reader) { if (reader->ptr + 1 > reader->end) { + reader->exhausted = true; PyErr_SetString(PyExc_ValueError, "read_bool out of bounds"); return NULL; } @@ -87,7 +89,8 @@ inline PyObject *read_bool_array(ReaderT *reader, int32_t count) { if (reader->ptr + count > reader->end) { - PyErr_SetString(PyExc_ValueError, "read_bool out of bounds"); + reader->exhausted = true; + PyErr_SetString(PyExc_ValueError, "read_bool_array out of bounds"); return NULL; } PyObject *list = PyList_New(count); @@ -104,6 +107,7 @@ inline PyObject *read_u8(ReaderT *reader) { if (reader->ptr + 1 > reader->end) { + reader->exhausted = true; PyErr_SetString(PyExc_ValueError, "read_u8 out of bounds"); return NULL; } @@ -114,7 +118,8 @@ inline PyObject *read_u8_array(ReaderT *reader, int32_t count) { if (reader->ptr + count > reader->end) { - PyErr_SetString(PyExc_ValueError, "read_u8 out of bounds"); + reader->exhausted = true; + PyErr_SetString(PyExc_ValueError, "read_u8_array out of bounds"); return NULL; } PyObject *list = PyList_New(count); @@ -129,6 +134,7 @@ inline PyObject *read_s8(ReaderT *reader) { if (reader->ptr + 1 > reader->end) { + reader->exhausted = true; PyErr_SetString(PyExc_ValueError, "read_s8 out of bounds"); return NULL; } @@ -139,7 +145,8 @@ inline PyObject *read_s8_array(ReaderT *reader, int32_t count) { if (reader->ptr + count > reader->end) { - PyErr_SetString(PyExc_ValueError, "read_s8 out of bounds"); + reader->exhausted = true; + PyErr_SetString(PyExc_ValueError, "read_s8_array out of bounds"); return NULL; } PyObject *list = PyList_New(count); @@ -159,6 +166,7 @@ inline PyObject *read_num(ReaderT *reader) if (reader->ptr + sizeof(T) > reader->end) { + reader->exhausted = true; return PyErr_Format(PyExc_ValueError, "read_%s out of bounds", typeid(T).name()); } T value = *(T *)reader->ptr; @@ -206,6 +214,7 @@ inline PyObject *read_num_array(ReaderT *reader, int32_t count) if (reader->ptr + sizeof(T) * count > reader->end) { + reader->exhausted = true; return PyErr_Format(PyExc_ValueError, "read_%s_array out of bounds", typeid(T).name()); } PyObject *list = PyList_New(count); @@ -260,6 +269,7 @@ inline bool _read_length(ReaderT *reader, int32_t *length) { if (reader->ptr + sizeof(int32_t) > reader->end) { + reader->exhausted = true; PyErr_SetString(PyExc_ValueError, "read_length out of bounds"); return false; } @@ -287,6 +297,7 @@ inline PyObject *read_str(ReaderT *reader) } if (reader->ptr + length > reader->end) { + reader->exhausted = true; PyErr_SetString(PyExc_ValueError, "read_str out of bounds"); return NULL; } @@ -306,6 +317,7 @@ inline PyObject *read_bytes(ReaderT *reader) } if (reader->ptr + length > reader->end) { + reader->exhausted = true; PyErr_SetString(PyExc_ValueError, "read_bytes out of bounds"); return NULL; } @@ -355,6 +367,10 @@ inline PyObject *read_pair_array(ReaderT *reader, TypeTreeNodeObject *node, Type TypeTreeNodeObject *second_child = (TypeTreeNodeObject *)PyList_GET_ITEM(node->m_Children, 1); PyObject *list = PyList_New(count); + if (list == NULL) + { + return NULL; + } for (auto i = 0; i < count; i++) { PyObject *first = read_typetree_value(reader, first_child, config); @@ -689,6 +705,13 @@ const NodeDataType SUPPORTED_VALUE_ARRAY_READ_TYPES[] = { template PyObject *read_typetree_value(ReaderT *reader, TypeTreeNodeObject *node, TypeTreeReaderConfigT *config) { + if (reader->exhausted) + { + if (!PyErr_Occurred()) + PyErr_SetString(PyExc_ValueError, "Read past end of typetree data"); + return NULL; + } + bool align = node->_align; PyObject *value = nullptr; @@ -827,6 +850,10 @@ PyObject *read_typetree_value(ReaderT *reader, TypeTreeNodeObject *node, TypeTre if (std::find(std::begin(SUPPORTED_VALUE_ARRAY_READ_TYPES), std::end(SUPPORTED_VALUE_ARRAY_READ_TYPES), child->_data_type) == std::end(SUPPORTED_VALUE_ARRAY_READ_TYPES)) { value = PyList_New(length); + if (value == NULL) + { + return NULL; + } for (int i = 0; i < length; i++) { PyObject *item = read_typetree_value(reader, child, config); @@ -869,6 +896,13 @@ PyObject *read_typetree_value(ReaderT *reader, TypeTreeNodeObject *node, TypeTre template PyObject *read_typetree_value_array(ReaderT *reader, TypeTreeNodeObject *node, TypeTreeReaderConfigT *config, int32_t count) { + if (reader->exhausted) + { + if (!PyErr_Occurred()) + PyErr_SetString(PyExc_ValueError, "Read past end of typetree data"); + return NULL; + } + bool align = node->_align; PyObject *value = nullptr; @@ -1012,7 +1046,7 @@ PyObject *read_typetree(PyObject *self, PyObject *args, PyObject *kwargs) } } - reader = {static_cast(view.buf), static_cast(view.buf) + view.len, static_cast(view.buf)}; + reader = {static_cast(view.buf), static_cast(view.buf) + view.len, static_cast(view.buf), false}; if (swap) {