From 095c844004bcd8680a4bf68901adbd9cac6a4302 Mon Sep 17 00:00:00 2001
From: Matt Wozniski <godlygeek+git@gmail.com>
Date: Wed, 8 May 2019 23:50:35 -0400
Subject: [PATCH] Add overflow checks for int to bytes conversions

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.
---
 py/binary.c | 18 ++++++++++++++----
 py/objint.c | 45 +++++++++++++++++++++++++++++++++++++++++++++
 py/objint.h |  1 +
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/py/binary.c b/py/binary.c
index 9c3a49e8f9..0e4a1d5f6d 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 fd746d3310..fc672b1123 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -300,6 +300,49 @@ 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_ValueError_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 +478,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 4b95acde9f..654404ac98 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);