ports/nrf: Upmerging port with upstream master

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);