Merge pull request #1860 from godlygeek/int_to_bytes_overflow_checks

Add overflow checks for int to bytes conversions

5 files changed, 58 insertions, 4 deletions

diff --git a/py/binary.c b/py/binary.c
index 9c3a49e8f..0e4a1d5f6 100644
--- a/py/binary.c
+++ b/py/binary.c
@@ -304,15 +304,18 @@ void mp_binary_set_val(char struct_type, char val_type, mp_obj_t val_in, byte **
             break;
         }
 #endif
-        default:
+        default: {
+            bool signed_type = is_signed(val_type);
             #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
             if (MP_OBJ_IS_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;
             } else
             #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 (BYTES_PER_WORD < 8 && size > sizeof(val)) {
                     int c = (is_signed(val_type) && (mp_int_t)val < 0) ? 0xff : 0x00;
@@ -322,6 +325,7 @@ void mp_binary_set_val(char struct_type, char val_type, mp_obj_t val_in, byte **
                     }
                 }
             }
+        }
     }
 
     mp_binary_set_int(MIN((size_t)size, sizeof(val)), struct_type == '>', p, val);
@@ -343,16 +347,22 @@ void mp_binary_set_val_array(char typecode, void *p, mp_uint_t index, mp_obj_t v
             ((mp_obj_t*)p)[index] = val_in;
             break;
 #endif
-        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_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);
+            mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
+            mp_binary_set_val_array_from_int(typecode, p, index, val);
+        }
     }
 }
 
diff --git a/py/objint.c b/py/objint.c
index fd746d331..80f6dfe55 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -300,6 +300,38 @@ 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) {
+        // self must be < 2**(bits - 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 (mp_binary_op(MP_BINARY_OP_LESS, self_in, edge) == mp_const_true) {
+            // and >= -2**(bits - 1)
+            edge = mp_unary_op(MP_UNARY_OP_NEGATIVE, edge);
+            if (mp_binary_op(MP_BINARY_OP_MORE_EQUAL, self_in, edge) == mp_const_true) {
+                return;
+            }
+        }
+    } else {
+        // self must be >= 0
+        if (mp_obj_int_sign(self_in) >= 0) {
+            // and < 2**(bits)
+            mp_obj_t edge = mp_binary_op(MP_BINARY_OP_INPLACE_LSHIFT,
+                                         mp_obj_new_int(1),
+                                         mp_obj_new_int(nbytes * 8));
+
+            if (mp_binary_op(MP_BINARY_OP_LESS, self_in, edge) == mp_const_true) {
+                return;
+            }
+        }
+    }
+
+    mp_raise_OverflowError_varg(translate("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) {
@@ -435,6 +467,8 @@ STATIC mp_obj_t int_to_bytes(size_t n_args, const mp_obj_t *args) {
     byte *data = (byte*)vstr.buf;
     memset(data, 0, len);
 
+    mp_obj_int_buffer_overflow_check(args[0], len, false);
+
     #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
     if (!MP_OBJ_IS_SMALL_INT(args[0])) {
         mp_obj_int_to_bytes_impl(args[0], big_endian, len, data);
diff --git a/py/objint.h b/py/objint.h
index 4b95acde9..654404ac9 100644
--- 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);
 void mp_obj_int_to_bytes_impl(mp_obj_t self_in, bool big_endian, size_t len, byte *buf);
diff --git a/py/runtime.c b/py/runtime.c
index 1a75a3870..9968f26e0 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -1598,6 +1598,14 @@ NORETURN void mp_raise_NotImplementedError_varg(const compressed_string_t *fmt,
     nlr_raise(exception);
 }
 
+NORETURN void mp_raise_OverflowError_varg(const compressed_string_t *fmt, ...) {
+    va_list argptr;
+    va_start(argptr,fmt);
+    mp_obj_t exception = mp_obj_new_exception_msg_vlist(&mp_type_OverflowError, fmt, argptr);
+    va_end(argptr);
+    nlr_raise(exception);
+}
+
 NORETURN void mp_raise_MpyError(const compressed_string_t *msg) {
     mp_raise_msg(&mp_type_MpyError, msg);
 }
diff --git a/py/runtime.h b/py/runtime.h
index c5a94b0b1..e8398cf0e 100644
--- a/py/runtime.h
+++ b/py/runtime.h
@@ -163,6 +163,7 @@ NORETURN void mp_raise_OSError_msg(const compressed_string_t *msg);
 NORETURN void mp_raise_OSError_msg_varg(const compressed_string_t *fmt, ...);
 NORETURN void mp_raise_NotImplementedError(const compressed_string_t *msg);
 NORETURN void mp_raise_NotImplementedError_varg(const compressed_string_t *fmt, ...);
+NORETURN void mp_raise_OverflowError_varg(const compressed_string_t *fmt, ...);
 NORETURN void mp_raise_MpyError(const compressed_string_t *msg);
 NORETURN void mp_raise_recursion_depth(void);