bootloader: more work on firmware update

micropython/bootloader/messages.c

@@ -186,22 +186,8 @@ static bool _recv_msg(uint8_t iface_num, uint32_t msg_size, uint8_t *buf, const
     return true;
 }
 
-/*
-static bool _decode(pb_istream_t *stream, const pb_field_t *field, void **arg)
-{
-    pb_byte_t *buf = *arg;
-    memset(buf, 0, 1024);
-    if (stream->bytes_left > 1024 - 1) {
-        return false;
-    }
-    if (!pb_read(stream, buf, stream->bytes_left)) {
-        return false;
-    }
-    return true;
-}
-*/
-
 #define MSG_RECV_INIT(TYPE) TYPE msg_recv = TYPE##_init_default
+#define MSG_RECV_CALLBACK(FIELD, CALLBACK) do { msg_recv.FIELD.funcs.decode = &CALLBACK; } while (0)
 #define MSG_RECV(TYPE) do { _recv_msg(iface_num, msg_size, buf, TYPE##_fields, &msg_recv); } while(0)
 
 void process_msg_Initialize(uint8_t iface_num, uint32_t msg_size, uint8_t *buf)
@@ -230,29 +216,73 @@ void process_msg_Ping(uint8_t iface_num, uint32_t msg_size, uint8_t *buf)
     MSG_SEND(Success);
 }
 
+static uint32_t firmware_size, firmware_flashed, chunk_requested;
+
 void process_msg_FirmwareErase(uint8_t iface_num, uint32_t msg_size, uint8_t *buf)
 {
+    firmware_size = 0;
+    firmware_flashed = 0;
+    chunk_requested = 0;
+
     MSG_RECV_INIT(FirmwareErase);
     MSG_RECV(FirmwareErase);
 
-    MSG_SEND_INIT(FirmwareRequest);
+    firmware_size = msg_recv.has_length ? msg_recv.length : 0;
-    MSG_SEND_ASSIGN_VALUE(offset, 0);
+    if (firmware_size > 0 && firmware_size % 4 == 0) {
-    MSG_SEND_ASSIGN_VALUE(length, FIRMWARE_CHUNK_SIZE);
+        chunk_requested = (firmware_size > FIRMWARE_CHUNK_SIZE) ? FIRMWARE_CHUNK_SIZE : firmware_size;
-    MSG_SEND(FirmwareRequest);
+        MSG_SEND_INIT(FirmwareRequest);
+        MSG_SEND_ASSIGN_VALUE(offset, 0);
+        MSG_SEND_ASSIGN_VALUE(length, chunk_requested);
+        MSG_SEND(FirmwareRequest);
+    } else {
+        MSG_SEND_INIT(Failure);
+        MSG_SEND_ASSIGN_VALUE(code, FailureType_Failure_DataError);
+        MSG_SEND_ASSIGN_STRING(message, "Wrong firmware size");
+        MSG_SEND(Failure);
+    }
+}
+
+static uint32_t chunk_size = 0;
+
+static bool _read_payload(pb_istream_t *stream, const pb_field_t *field, void **arg)
+{
+#define BUFSIZE 1024
+    pb_byte_t buf[BUFSIZE];
+    chunk_size = stream->bytes_left;
+    while (stream->bytes_left) {
+        if (!pb_read(stream, buf, (stream->bytes_left > BUFSIZE) ? BUFSIZE : stream->bytes_left)) {
+            return false;
+        }
+    }
+    return true;
 }
 
 void process_msg_FirmwareUpload(uint8_t iface_num, uint32_t msg_size, uint8_t *buf)
 {
     MSG_RECV_INIT(FirmwareUpload);
+    MSG_RECV_CALLBACK(payload, _read_payload);
     MSG_RECV(FirmwareUpload);
-/*
+
-    MSG_SEND_INIT(FirmwareRequest);
+    if (chunk_size != chunk_requested) {
-    MSG_SEND_ASSIGN_VALUE(offset, FIRMWARE_CHUNK_SIZE);
+        MSG_SEND_INIT(Failure);
-    MSG_SEND_ASSIGN_VALUE(length, FIRMWARE_CHUNK_SIZE);
+        MSG_SEND_ASSIGN_VALUE(code, FailureType_Failure_DataError);
-    MSG_SEND(FirmwareRequest);
+        MSG_SEND_ASSIGN_STRING(message, "Invalid chunk size");
-*/
+        MSG_SEND(Failure);
-    MSG_SEND_INIT(Success);
+    }
-    MSG_SEND(Success);
+
+    firmware_size -= chunk_requested;
+    firmware_flashed += chunk_requested;
+
+    if (firmware_size > 0) {
+        chunk_requested = (firmware_size > FIRMWARE_CHUNK_SIZE) ? FIRMWARE_CHUNK_SIZE : firmware_size;
+        MSG_SEND_INIT(FirmwareRequest);
+        MSG_SEND_ASSIGN_VALUE(offset, firmware_flashed);
+        MSG_SEND_ASSIGN_VALUE(length, chunk_requested);
+        MSG_SEND(FirmwareRequest);
+    } else {
+        MSG_SEND_INIT(Success);
+        MSG_SEND(Success);
+    }
 }
 
 void process_msg_unknown(uint8_t iface_num, uint32_t msg_size, uint8_t *buf)

micropython/bootloader/protob/messages.pb.c

@@ -68,7 +68,8 @@ const pb_field_t ButtonAck_fields[1] = {
     PB_LAST_FIELD
 };
 
-const pb_field_t FirmwareErase_fields[1] = {
+const pb_field_t FirmwareErase_fields[2] = {
+    PB_FIELD(  1, UINT32  , OPTIONAL, STATIC  , FIRST, FirmwareErase, length, length, 0),
     PB_LAST_FIELD
 };
 

micropython/bootloader/protob/messages.pb.h

@@ -40,11 +40,6 @@ typedef struct _ButtonAck {
 /* @@protoc_insertion_point(struct:ButtonAck) */
 } ButtonAck;
 
-typedef struct _FirmwareErase {
-    char dummy_field;
-/* @@protoc_insertion_point(struct:FirmwareErase) */
-} FirmwareErase;
-
 typedef struct _GetFeatures {
     char dummy_field;
 /* @@protoc_insertion_point(struct:GetFeatures) */
@@ -111,6 +106,12 @@ typedef struct _Features {
 /* @@protoc_insertion_point(struct:Features) */
 } Features;
 
+typedef struct _FirmwareErase {
+    bool has_length;
+    uint32_t length;
+/* @@protoc_insertion_point(struct:FirmwareErase) */
+} FirmwareErase;
+
 typedef struct _FirmwareRequest {
     bool has_offset;
     uint32_t offset;
@@ -156,7 +157,7 @@ typedef struct _Success {
 #define Failure_init_default                     {false, (FailureType)0, false, ""}
 #define ButtonRequest_init_default               {false, (ButtonRequestType)0, false, ""}
 #define ButtonAck_init_default                   {0}
-#define FirmwareErase_init_default               {0}
+#define FirmwareErase_init_default               {false, 0}
 #define FirmwareRequest_init_default             {false, 0, false, 0}
 #define FirmwareUpload_init_default              {{{NULL}, NULL}, false, {0, {0}}}
 #define Initialize_init_zero                     {0}
@@ -167,7 +168,7 @@ typedef struct _Success {
 #define Failure_init_zero                        {false, (FailureType)0, false, ""}
 #define ButtonRequest_init_zero                  {false, (ButtonRequestType)0, false, ""}
 #define ButtonAck_init_zero                      {0}
-#define FirmwareErase_init_zero                  {0}
+#define FirmwareErase_init_zero                  {false, 0}
 #define FirmwareRequest_init_zero                {false, 0, false, 0}
 #define FirmwareUpload_init_zero                 {{{NULL}, NULL}, false, {0, {0}}}
 
@@ -193,6 +194,7 @@ typedef struct _Success {
 #define Features_pin_cached_tag                  16
 #define Features_passphrase_cached_tag           17
 #define Features_firmware_present_tag            18
+#define FirmwareErase_length_tag                 1
 #define FirmwareRequest_offset_tag               1
 #define FirmwareRequest_length_tag               2
 #define FirmwareUpload_payload_tag               1
@@ -212,7 +214,7 @@ extern const pb_field_t Success_fields[2];
 extern const pb_field_t Failure_fields[3];
 extern const pb_field_t ButtonRequest_fields[3];
 extern const pb_field_t ButtonAck_fields[1];
-extern const pb_field_t FirmwareErase_fields[1];
+extern const pb_field_t FirmwareErase_fields[2];
 extern const pb_field_t FirmwareRequest_fields[3];
 extern const pb_field_t FirmwareUpload_fields[3];
 
@@ -225,7 +227,7 @@ extern const pb_field_t FirmwareUpload_fields[3];
 #define Failure_size                             270
 #define ButtonRequest_size                       270
 #define ButtonAck_size                           0
-#define FirmwareErase_size                       0
+#define FirmwareErase_size                       6
 #define FirmwareRequest_size                     12
 /* FirmwareUpload_size depends on runtime parameters */
 

micropython/bootloader/protob/messages.proto

@@ -107,6 +107,7 @@ message ButtonAck {
  * @next Failure
  */
 message FirmwareErase {
+	optional uint32 length = 1;			// length of new firmware
 }
 
 /**

micropython/bootloader/protob/types.pb.h

@@ -18,13 +18,13 @@ extern "C" {
 typedef enum _FailureType {
     FailureType_Failure_UnexpectedMessage = 1,
     FailureType_Failure_ButtonExpected = 2,
-    FailureType_Failure_SyntaxError = 3,
+    FailureType_Failure_DataError = 3,
     FailureType_Failure_ActionCancelled = 4,
     FailureType_Failure_PinExpected = 5,
     FailureType_Failure_PinCancelled = 6,
     FailureType_Failure_PinInvalid = 7,
     FailureType_Failure_InvalidSignature = 8,
-    FailureType_Failure_Other = 9,
+    FailureType_Failure_ProcessError = 9,
     FailureType_Failure_NotEnoughFunds = 10,
     FailureType_Failure_NotInitialized = 11,
     FailureType_Failure_FirmwareError = 99

micropython/bootloader/protob/types.proto

@@ -5,13 +5,13 @@
 enum FailureType {
 	Failure_UnexpectedMessage = 1;
 	Failure_ButtonExpected = 2;
-	Failure_SyntaxError = 3;
+	Failure_DataError = 3;
 	Failure_ActionCancelled = 4;
 	Failure_PinExpected = 5;
 	Failure_PinCancelled = 6;
 	Failure_PinInvalid = 7;
 	Failure_InvalidSignature = 8;
-	Failure_Other = 9;
+	Failure_ProcessError = 9;
 	Failure_NotEnoughFunds = 10;
 	Failure_NotInitialized = 11;
 	Failure_FirmwareError = 99;

vendor/trezor-common

@@ -1 +1 @@
-Subproject commit 73b9ffc36a9b54c56bbfa643b4735c63cf309ca2
+Subproject commit 0001cb18c065d7f3aacd6f8d7c5be42cd7477a94