diff options
Diffstat (limited to 'py/modstruct.c')
| -rw-r--r-- | py/modstruct.c | 85 |
1 files changed, 43 insertions, 42 deletions
diff --git a/py/modstruct.c b/py/modstruct.c index 1daa33338..8617a8e0d 100644 --- a/py/modstruct.c +++ b/py/modstruct.c @@ -82,26 +82,10 @@ STATIC mp_uint_t get_fmt_num(const char **p) { return val; } -STATIC uint calcsize_items(const char *fmt) { - uint cnt = 0; - while (*fmt) { - int num = 1; - if (unichar_isdigit(*fmt)) { - num = get_fmt_num(&fmt); - if (*fmt == 's') { - num = 1; - } - } - cnt += num; - fmt++; - } - return cnt; -} - -STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) { - const char *fmt = mp_obj_str_get_str(fmt_in); +STATIC size_t calc_size_items(const char *fmt, size_t *total_sz) { char fmt_type = get_fmt_type(&fmt); - mp_uint_t size; + size_t total_cnt = 0; + size_t size; for (size = 0; *fmt; fmt++) { mp_uint_t cnt = 1; if (unichar_isdigit(*fmt)) { @@ -109,8 +93,10 @@ STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) { } if (*fmt == 's') { + total_cnt += 1; size += cnt; } else { + total_cnt += cnt; mp_uint_t align; size_t sz = mp_binary_get_size(fmt_type, *fmt, &align); while (cnt--) { @@ -120,6 +106,14 @@ STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) { } } } + *total_sz = size; + return total_cnt; +} + +STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) { + const char *fmt = mp_obj_str_get_str(fmt_in); + size_t size; + calc_size_items(fmt, &size); return MP_OBJ_NEW_SMALL_INT(size); } MP_DEFINE_CONST_FUN_OBJ_1(struct_calcsize_obj, struct_calcsize); @@ -130,8 +124,9 @@ STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t *args) { // Since we implement unpack and unpack_from using the same function // we relax the "exact" requirement, and only implement "big enough". const char *fmt = mp_obj_str_get_str(args[0]); + size_t total_sz; + size_t num_items = calc_size_items(fmt, &total_sz); char fmt_type = get_fmt_type(&fmt); - uint num_items = calcsize_items(fmt); mp_obj_tuple_t *res = MP_OBJ_TO_PTR(mp_obj_new_tuple(num_items, NULL)); mp_buffer_info_t bufinfo; mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_READ); @@ -152,21 +147,23 @@ STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t *args) { p += offset; } - for (uint i = 0; i < num_items;) { - mp_uint_t sz = 1; + // Check that the input buffer is big enough to unpack all the values + if (p + total_sz > end_p) { + mp_raise_ValueError("buffer too small"); + } + + for (size_t i = 0; i < num_items;) { + mp_uint_t cnt = 1; if (unichar_isdigit(*fmt)) { - sz = get_fmt_num(&fmt); - } - if (p + sz > end_p) { - mp_raise_ValueError("buffer too small"); + cnt = get_fmt_num(&fmt); } mp_obj_t item; if (*fmt == 's') { - item = mp_obj_new_bytes(p, sz); - p += sz; + item = mp_obj_new_bytes(p, cnt); + p += cnt; res->items[i++] = item; } else { - while (sz--) { + while (cnt--) { item = mp_binary_get_val(fmt_type, *fmt, &p); res->items[i++] = item; } @@ -177,36 +174,35 @@ STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t *args) { } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_unpack_from_obj, 2, 3, struct_unpack_from); -STATIC void struct_pack_into_internal(mp_obj_t fmt_in, byte *p, byte* end_p, size_t n_args, const mp_obj_t *args) { +// This function assumes there is enough room in p to store all the values +STATIC void struct_pack_into_internal(mp_obj_t fmt_in, byte *p, size_t n_args, const mp_obj_t *args) { const char *fmt = mp_obj_str_get_str(fmt_in); char fmt_type = get_fmt_type(&fmt); size_t i; for (i = 0; i < n_args;) { - mp_uint_t sz = 1; + mp_uint_t cnt = 1; if (*fmt == '\0') { // more arguments given than used by format string; CPython raises struct.error here break; } if (unichar_isdigit(*fmt)) { - sz = get_fmt_num(&fmt); - } - if (p + sz > end_p) { - mp_raise_ValueError("buffer too small"); + cnt = get_fmt_num(&fmt); } if (*fmt == 's') { mp_buffer_info_t bufinfo; mp_get_buffer_raise(args[i++], &bufinfo, MP_BUFFER_READ); - mp_uint_t to_copy = sz; + mp_uint_t to_copy = cnt; if (bufinfo.len < to_copy) { to_copy = bufinfo.len; } memcpy(p, bufinfo.buf, to_copy); - memset(p + to_copy, 0, sz - to_copy); - p += sz; + memset(p + to_copy, 0, cnt - to_copy); + p += cnt; } else { - while (sz--) { + // If we run out of args then we just finish; CPython would raise struct.error + while (cnt-- && i < n_args) { mp_binary_set_val(fmt_type, *fmt, args[i++], &p); } } @@ -221,8 +217,7 @@ STATIC mp_obj_t struct_pack(size_t n_args, const mp_obj_t *args) { vstr_init_len(&vstr, size); byte *p = (byte*)vstr.buf; memset(p, 0, size); - byte *end_p = &p[size]; - struct_pack_into_internal(args[0], p, end_p, n_args - 1, &args[1]); + struct_pack_into_internal(args[0], p, n_args - 1, &args[1]); return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr); } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_pack_obj, 1, MP_OBJ_FUN_ARGS_MAX, struct_pack); @@ -242,7 +237,13 @@ STATIC mp_obj_t struct_pack_into(size_t n_args, const mp_obj_t *args) { byte *end_p = &p[bufinfo.len]; p += offset; - struct_pack_into_internal(args[0], p, end_p, n_args - 3, &args[3]); + // Check that the output buffer is big enough to hold all the values + mp_int_t sz = MP_OBJ_SMALL_INT_VALUE(struct_calcsize(args[0])); + if (p + sz > end_p) { + mp_raise_ValueError("buffer too small"); + } + + struct_pack_into_internal(args[0], p, n_args - 3, &args[3]); return mp_const_none; } MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_pack_into_obj, 3, MP_OBJ_FUN_ARGS_MAX, struct_pack_into); |