Merge pull request adafruit#7772 from jepler/fix-struct-pack

struct: Check that argument counts match, similar to cpython3
jepler · Mar 22, 2023 · 4525714 · 4525714
2 parents e1f1647 + 98c546b
commit 4525714
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 23 deletions.
diff --git a/ports/unix/variants/coverage/mpconfigvariant.h b/ports/unix/variants/coverage/mpconfigvariant.h
@@ -62,6 +62,7 @@
 #define MICROPY_PY_FRAMEBUF            (1)
 #define MICROPY_PY_COLLECTIONS_NAMEDTUPLE__ASDICT (1)
 #define MICROPY_PY_COLLECTIONS_ORDEREDDICT (1)
+#define MICROPY_PY_STRUCT              (0) // uses shared-bindings struct
 #define MICROPY_PY_UCRYPTOLIB          (1)
 #define MICROPY_PY_UCRYPTOLIB_CTR      (1)
 #define MICROPY_PY_MICROPYTHON_HEAP_LOCKED (1)

diff --git a/ports/unix/variants/coverage/mpconfigvariant.mk b/ports/unix/variants/coverage/mpconfigvariant.mk
@@ -33,6 +33,7 @@ SRC_BITMAP := \
 	shared-bindings/bitmaptools/__init__.c \
 	shared-bindings/displayio/Bitmap.c \
 	shared-bindings/rainbowio/__init__.c \
+	shared-bindings/struct/__init__.c \
 	shared-bindings/traceback/__init__.c \
 	shared-bindings/util.c \
 	shared-bindings/zlib/__init__.c \
@@ -45,6 +46,7 @@ SRC_BITMAP := \
 	shared-module/displayio/ColorConverter.c \
 	shared-module/os/getenv.c \
 	shared-module/rainbowio/__init__.c \
+	shared-module/struct/__init__.c \
 	shared-module/traceback/__init__.c \
 	shared-module/zlib/__init__.c \
 
@@ -57,6 +59,7 @@ CFLAGS += \
 	-DCIRCUITPY_OS_GETENV=1 \
 	-DCIRCUITPY_GIFIO=1 \
 	-DCIRCUITPY_RAINBOWIO=1 \
+	-DCIRCUITPY_STRUCT=1 \
 	-DCIRCUITPY_TRACEBACK=1 \
 	-DCIRCUITPY_ZLIB=1
 

diff --git a/shared-module/struct/__init__.c b/shared-module/struct/__init__.c
@@ -129,43 +129,45 @@ void shared_modules_struct_pack_into(mp_obj_t fmt_in, byte *p, byte *end_p, size
         mp_raise_RuntimeError(translate("buffer too small"));
     }
 
-    size_t i;
+    size_t i = 0;
     byte *p_base = p;
-    for (i = 0; i < n_args;) {
+    while (*fmt) {
         mp_uint_t sz = 1;
-        if (*fmt == '\0') {
-            // more arguments given than used by format string; CPython raises struct.error here
-            mp_raise_RuntimeError(translate("too many arguments provided with the given format"));
-        }
         struct_validate_format(*fmt);
 
         if (unichar_isdigit(*fmt)) {
             sz = 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;
-            if (bufinfo.len < to_copy) {
-                to_copy = bufinfo.len;
+            if (i < n_args) {
+                mp_buffer_info_t bufinfo;
+                mp_get_buffer_raise(args[i], &bufinfo, MP_BUFFER_READ);
+                mp_uint_t to_copy = sz;
+                if (bufinfo.len < to_copy) {
+                    to_copy = bufinfo.len;
+                }
+                memcpy(p, bufinfo.buf, to_copy);
+                memset(p + to_copy, 0, sz - to_copy);
             }
-            memcpy(p, bufinfo.buf, to_copy);
-            memset(p + to_copy, 0, sz - to_copy);
+            i++;
             p += sz;
         } else {
             while (sz--) {
                 // Pad bytes don't have a corresponding argument.
                 if (*fmt == 'x') {
                     mp_binary_set_val(fmt_type, *fmt, MP_OBJ_NEW_SMALL_INT(0), p_base, &p);
                 } else {
-                    mp_binary_set_val(fmt_type, *fmt, args[i], p_base, &p);
+                    if (i < n_args) {
+                        mp_binary_set_val(fmt_type, *fmt, args[i], p_base, &p);
+                    }
                     i++;
                 }
             }
         }
         fmt++;
     }
+    (void)mp_arg_validate_length(n_args, i, MP_QSTR_args);
 }
 
 mp_obj_tuple_t *shared_modules_struct_unpack_from(mp_obj_t fmt_in, byte *p, byte *end_p, bool exact_size) {

diff --git a/tests/circuitpython/struct_nargs.py b/tests/circuitpython/struct_nargs.py
@@ -0,0 +1,38 @@
+import struct
+
+
+def do_pack(*args):
+    try:
+        print(struct.pack(*args))
+    except Exception as e:
+        print("ERROR")
+
+
+do_pack("H")
+do_pack("H", 1)
+do_pack("H", 1, 2)
+
+do_pack("2H")
+do_pack("2H", 1)
+do_pack("2H", 1, 2)
+
+do_pack("3H")
+do_pack("3H", 1)
+do_pack("3H", 1, 2)
+do_pack("3H", 1, 2, 3)
+do_pack("3H", 1, 2, 3, 4)
+
+do_pack("4sH", b"x")
+do_pack("4sH", b"x", 1)
+do_pack("4sH", b"x", 1, 2)
+
+do_pack("4s2H", b"x")
+do_pack("4s2H", b"x", 1)
+do_pack("4s2H", b"x", 1, 2)
+do_pack("4s2H", b"x", 1, 2, 3)
+
+do_pack("4s3H", b"x")
+do_pack("4s3H", b"x", 1)
+do_pack("4s3H", b"x", 1, 2)
+do_pack("4s3H", b"x", 1, 2, 3)
+do_pack("4s3H", b"x", 1, 2, 3, 4)
diff --git a/tests/unix/extra_coverage.py.exp b/tests/unix/extra_coverage.py.exp
@@ -35,15 +35,15 @@ bitmaptools     btree           cexample        cmath
 collections     cppexample      displayio       errno
 ffi             framebuf        gc              hashlib
 json            math            qrio            rainbowio
-re              sys             termios         traceback
-ubinascii       uctypes         uerrno          uheapq
-uio             ujson           ulab            ulab.numpy
-ulab.numpy.fft  ulab.numpy.linalg               ulab.scipy
-ulab.scipy.linalg               ulab.scipy.optimize
-ulab.scipy.signal               ulab.scipy.special
-ulab.utils      uos             urandom         ure
-uselect         ustruct         utime           utimeq
-uzlib           zlib
+re              struct          sys             termios
+traceback       ubinascii       uctypes         uerrno
+uheapq          uio             ujson           ulab
+ulab.numpy      ulab.numpy.fft  ulab.numpy.linalg
+ulab.scipy      ulab.scipy.linalg
+ulab.scipy.optimize             ulab.scipy.signal
+ulab.scipy.special              ulab.utils      uos
+urandom         ure             uselect         utime
+utimeq          uzlib           zlib
 ime
 
 utime           utimeq