py/binary,objint: Add overflow checks for int to bytes conversions.

Matt Wozniski · projectgus · commit 50bd33b2a4e8 · 2024-11-21T15:38:16.000+11:00
For both small and long integers, raise an exception if calling struct.pack, adding an element to an array.array, or formatting an int with int.to_bytes would overflow the requested size. (Cherry-picked from CircuitPython commit 095c844.) Signed-off-by: Angus Gratton <angus@redyak.com.au>
diff --git a/py/binary.c b/py/binary.c
@@ -441,15 +441,18 @@ void mp_binary_set_val(char struct_type, char val_type, mp_obj_t val_in, byte *p
             break;
         }
         #endif
-        default:
+        default: {
+            bool signed_type = is_signed(val_type);
             #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
             if (mp_obj_is_exact_type(val_in, &mp_type_int)) {
+                mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
                 mp_obj_int_to_bytes_impl(val_in, struct_type == '>', size, p);
                 return;
             }
             #endif
 
             val = mp_obj_get_int(val_in);
+            mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
             // zero/sign extend if needed
             if (MP_BYTES_PER_OBJ_WORD < 8 && size > sizeof(val)) {
                 int c = (mp_int_t)val < 0 ? 0xff : 0x00;
@@ -459,6 +462,7 @@ void mp_binary_set_val(char struct_type, char val_type, mp_obj_t val_in, byte *p
                 }
             }
             break;
+        }
     }
 
     mp_binary_set_int(MIN((size_t)size, sizeof(val)), struct_type == '>', p, val);
@@ -478,16 +482,24 @@ void mp_binary_set_val_array(char typecode, void *p, size_t index, mp_obj_t val_
         case 'O':
             ((mp_obj_t *)p)[index] = val_in;
             break;
-        default:
+        default: {
+            size_t size = mp_binary_get_size('@', typecode, NULL);
+            bool signed_type = is_signed(typecode);
             #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
             if (mp_obj_is_exact_type(val_in, &mp_type_int)) {
-                size_t size = mp_binary_get_size('@', typecode, NULL);
+                mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
                 mp_obj_int_to_bytes_impl(val_in, MP_ENDIANNESS_BIG,
                     size, (uint8_t *)p + index * size);
                 return;
             }
             #endif
-            mp_binary_set_val_array_from_int(typecode, p, index, mp_obj_get_int(val_in));
+            mp_int_t val = mp_obj_get_int(val_in);
+            if (val < 0 && typecode == BYTEARRAY_TYPECODE) {
+                val = val & 0xFF;
+            }
+            mp_obj_int_buffer_overflow_check(mp_obj_new_int(val), size, signed_type);
+            mp_binary_set_val_array_from_int(typecode, p, index, val);
+        }
     }
 }
 
diff --git a/py/objint.c b/py/objint.c
@@ -300,6 +300,48 @@ char *mp_obj_int_formatted(char **buf, size_t *buf_size, size_t *fmt_size, mp_co
     return b;
 }
 
+void mp_obj_int_buffer_overflow_check(mp_obj_t self_in, size_t nbytes, bool is_signed) {
+    if (is_signed) {
+        // edge = 1 << (nbytes * 8 - 1)
+        mp_obj_t edge = mp_binary_op(MP_BINARY_OP_INPLACE_LSHIFT,
+            mp_obj_new_int(1),
+            mp_obj_new_int(nbytes * 8 - 1));
+
+        // if self >= edge, we don't fit
+        if (mp_binary_op(MP_BINARY_OP_MORE_EQUAL, self_in, edge) == mp_const_true) {
+            goto raise;
+        }
+
+        // edge = -edge
+        edge = mp_unary_op(MP_UNARY_OP_NEGATIVE, edge);
+
+        // if self < edge, we don't fit
+        if (mp_binary_op(MP_BINARY_OP_LESS, self_in, edge) == mp_const_true) {
+            goto raise;
+        }
+    } else {
+        if (mp_obj_int_sign(self_in) < 0) {
+            // Negative numbers never fit in an unsigned value
+            goto raise;
+        }
+
+        // edge = 1 << (nbytes * 8)
+        mp_obj_t edge = mp_binary_op(MP_BINARY_OP_INPLACE_LSHIFT,
+            mp_obj_new_int(1),
+            mp_obj_new_int(nbytes * 8));
+
+        // if self >= edge, we don't fit
+        if (mp_binary_op(MP_BINARY_OP_MORE_EQUAL, self_in, edge) == mp_const_true) {
+            goto raise;
+        }
+    }
+
+    return;
+
+raise:
+    mp_raise_msg_varg(&mp_type_OverflowError, MP_ERROR_TEXT("value would overflow a %d byte buffer"), nbytes);
+}
+
 #if MICROPY_LONGINT_IMPL == MICROPY_LONGINT_IMPL_NONE
 
 int mp_obj_int_sign(mp_obj_t self_in) {
@@ -434,6 +476,8 @@ static mp_obj_t int_to_bytes(size_t n_args, const mp_obj_t *args) {
     vstr_init_len(&vstr, dlen);
     byte *data = (byte *)vstr.buf;
 
+    mp_obj_int_buffer_overflow_check(args[0], dlen, false);
+
     #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
     if (!mp_obj_is_small_int(args[0])) {
         overflow = !mp_obj_int_to_bytes_impl(args[0], big_endian, dlen, data);
diff --git a/py/objint.h b/py/objint.h
@@ -53,6 +53,7 @@ char *mp_obj_int_formatted(char **buf, size_t *buf_size, size_t *fmt_size, mp_co
     int base, const char *prefix, char base_char, char comma);
 char *mp_obj_int_formatted_impl(char **buf, size_t *buf_size, size_t *fmt_size, mp_const_obj_t self_in,
     int base, const char *prefix, char base_char, char comma);
+void mp_obj_int_buffer_overflow_check(mp_obj_t self_in, size_t nbytes, bool is_signed);
 mp_int_t mp_obj_int_hash(mp_obj_t self_in);
 mp_obj_t mp_obj_int_from_bytes_impl(bool big_endian, size_t len, const byte *buf);
 // Returns true if 'self_in' fit into 'len' bytes of 'buf' without overflowing, 'buf' is truncated otherwise.
