py/modstruct: Check and prevent buffer-write overflow in struct packing.

Prior to this patch, the size of the buffer given to pack_into() was checked
for being too small by using the count of the arguments, not their actual
size.  For example, a format spec of '4I' would only check that there was 4
bytes available, not 16; and 'I' would check for 1 byte, not 4.

The pack() function is ok because its buffer is created to be exactly the
correct size.

The fix in this patch calculates the total size of the format spec at the
start of pack_into() and verifies that the buffer is large enough.  This
adds some computational overhead, to iterate through the whole format spec.
The alternative is to check during the packing, but that requires extra
code to handle alignment, and the check is anyway not needed for pack().
So to maintain minimal code size the check is done using struct_calcsize.

py/modstruct.c

@@ -174,37 +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);
+            cnt = get_fmt_num(&fmt);
-        }
-        if (p + sz > end_p) {
-            mp_raise_ValueError("buffer too small");
         }
 
         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);
+            memset(p + to_copy, 0, cnt - to_copy);
-            p += sz;
+            p += cnt;
         } else {
             // If we run out of args then we just finish; CPython would raise struct.error
-            while (sz-- && i < n_args) {
+            while (cnt-- && i < n_args) {
                 mp_binary_set_val(fmt_type, *fmt, args[i++], &p);
             }
         }
@@ -219,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, n_args - 1, &args[1]);
-    struct_pack_into_internal(args[0], p, end_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);
@@ -240,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);

tests/basics/struct1.py

@@ -69,6 +69,12 @@ print(buf)
 struct.pack_into('<bbb', buf, -6, 0x44, 0x45, 0x46)
 print(buf)
 
+# check that we get an error if the buffer is too small
+try:
+    struct.pack_into('I', bytearray(1), 0, 0)
+except:
+    print('struct.error')
+
 try:
     struct.pack_into('<bbb', buf, 7, 0x41, 0x42, 0x43)
 except:

tests/basics/struct2.py

@@ -30,6 +30,10 @@ try:
     struct.unpack('2H', b'\x00\x00')
 except:
     print('Exception')
+try:
+    struct.pack_into('2I', bytearray(4), 0, 0)
+except:
+    print('Exception')
 
 # check that unknown types raise an exception
 try: