From 17fa41620e9ac85dadb2348e3d00436d036ed7d2 Mon Sep 17 00:00:00 2001
From: Andrew Kozlik <andrew.kozlik@gmail.com>
Date: Mon, 23 Sep 2019 13:01:22 +0200
Subject: [PATCH] legacy: refactor readprotobufint()

---
 legacy/bootloader/usb.c | 42 ++++++++++++++++++++++++++++++++++++++---
 legacy/util.c           | 22 ---------------------
 legacy/util.h           |  3 ---
 3 files changed, 39 insertions(+), 28 deletions(-)

diff --git a/legacy/bootloader/usb.c b/legacy/bootloader/usb.c
index a73ffc099..1d6d77bb1 100644
--- a/legacy/bootloader/usb.c
+++ b/legacy/bootloader/usb.c
@@ -30,6 +30,7 @@
 #include "memzero.h"
 #include "oled.h"
 #include "rng.h"
+#include "secbool.h"
 #include "secp256k1.h"
 #include "sha2.h"
 #include "signatures.h"
@@ -117,6 +118,34 @@ static void check_and_write_chunk(void) {
   chunk_idx++;
 }
 
+// read protobuf integer and advance pointer
+static secbool readprotobufint(const uint8_t **ptr, uint32_t *result) {
+  *result = 0;
+
+  for (int i = 0; i <= 3; ++i) {
+    *result += (**ptr & 0x7F) << (7 * i);
+    if ((**ptr & 0x80) == 0) {
+      (*ptr)++;
+      return sectrue;
+    }
+    (*ptr)++;
+  }
+
+  if (**ptr & 0xF0) {
+    // result does not fit into uint32_t
+    *result = 0;
+
+    // skip over the rest of the integer
+    while (**ptr & 0x80) (*ptr)++;
+    (*ptr)++;
+    return secfalse;
+  }
+
+  *result += (uint32_t)(**ptr) << 28;
+  (*ptr)++;
+  return sectrue;
+}
+
 static void rx_callback(usbd_device *dev, uint8_t ep) {
   (void)ep;
   static uint16_t msg_id = 0xFFFF;
@@ -226,7 +255,12 @@ static void rx_callback(usbd_device *dev, uint8_t ep) {
       }
       // read payload length
       const uint8_t *p = buf + 10;
-      flash_len = readprotobufint(&p);
+      if (readprotobufint(&p, &flash_len) != sectrue) {  // integer too large
+        send_msg_failure(dev);
+        flash_state = STATE_END;
+        show_halt("Firmware is too big.", NULL);
+        return;
+      }
       if (flash_len <= FLASH_FWHEADER_LEN) {  // firmware is too small
         send_msg_failure(dev);
         flash_state = STATE_END;
@@ -254,7 +288,8 @@ static void rx_callback(usbd_device *dev, uint8_t ep) {
       chunk_idx = 0;
       w = 0;
       while (p < buf + 64) {
-        w = (w >> 8) | (*p << 24);  // assign byte to first byte of uint32_t w
+        // assign byte to first byte of uint32_t w
+        w = (w >> 8) | (((uint32_t)*p) << 24);
         wi++;
         if (wi == 4) {
           FW_HEADER[flash_pos / 4] = w;
@@ -285,7 +320,8 @@ static void rx_callback(usbd_device *dev, uint8_t ep) {
 
     const uint8_t *p = buf + 1;
     while (p < buf + 64 && flash_pos < flash_len) {
-      w = (w >> 8) | (*p << 24);  // assign byte to first byte of uint32_t w
+      // assign byte to first byte of uint32_t w
+      w = (w >> 8) | (((uint32_t)*p) << 24);
       wi++;
       if (wi == 4) {
         if (flash_pos < FLASH_FWHEADER_LEN) {
diff --git a/legacy/util.c b/legacy/util.c
index 1876e31cd..b0b94c804 100644
--- a/legacy/util.c
+++ b/legacy/util.c
@@ -40,25 +40,3 @@ void data2hex(const void *data, uint32_t len, char *str) {
   }
   str[len * 2] = 0;
 }
-
-uint32_t readprotobufint(const uint8_t **ptr) {
-  uint32_t result = (**ptr & 0x7F);
-  if (**ptr & 0x80) {
-    (*ptr)++;
-    result += (**ptr & 0x7F) * 128;
-    if (**ptr & 0x80) {
-      (*ptr)++;
-      result += (**ptr & 0x7F) * 128 * 128;
-      if (**ptr & 0x80) {
-        (*ptr)++;
-        result += (**ptr & 0x7F) * 128 * 128 * 128;
-        if (**ptr & 0x80) {
-          (*ptr)++;
-          result += (**ptr & 0x7F) * 128 * 128 * 128 * 128;
-        }
-      }
-    }
-  }
-  (*ptr)++;
-  return result;
-}
diff --git a/legacy/util.h b/legacy/util.h
index 9b382bb5e..593dc8ff1 100644
--- a/legacy/util.h
+++ b/legacy/util.h
@@ -58,9 +58,6 @@ void uint32hex(uint32_t num, char *str);
 // converts data to hexa
 void data2hex(const void *data, uint32_t len, char *str);
 
-// read protobuf integer and advance pointer
-uint32_t readprotobufint(const uint8_t **ptr);
-
 // defined in startup.s (or setup.c for emulator)
 extern void __attribute__((noreturn)) shutdown(void);