diff --git a/micropython/bootloader/bootloader.h b/micropython/bootloader/bootloader.h deleted file mode 100644 index 7dfacf8783..0000000000 --- a/micropython/bootloader/bootloader.h +++ /dev/null @@ -1,4 +0,0 @@ -#define UPLOAD_CHUNK_SIZE (128*1024) - -#define USB_PACKET_SIZE 64 -#define USB_IFACE_NUM 0 diff --git a/micropython/bootloader/main.c b/micropython/bootloader/main.c index 0393cec31c..a16ac5432c 100644 --- a/micropython/bootloader/main.c +++ b/micropython/bootloader/main.c @@ -13,7 +13,6 @@ #include "version.h" #include "mini_printf.h" -#include "bootloader.h" #include "messages.h" #define IMAGE_MAGIC 0x465A5254 // TRZF @@ -160,7 +159,7 @@ void mainloop(void) for (;;) { int r = usb_hid_read_blocking(USB_IFACE_NUM, buf, USB_PACKET_SIZE, 100); - if (r <= 0) { + if (r != USB_PACKET_SIZE) { continue; } assert(r == USB_PACKET_SIZE); @@ -173,23 +172,23 @@ void mainloop(void) switch (msg_id) { case 0: // Initialize display_printf("received Initialize\n"); - process_msg_Initialize(USB_IFACE_NUM); + process_msg_Initialize(USB_IFACE_NUM, msg_size, buf); break; case 1: // Ping display_printf("received Ping\n"); - process_msg_Ping(USB_IFACE_NUM); + process_msg_Ping(USB_IFACE_NUM, msg_size, buf); break; case 6: // FirmwareErase display_printf("received FirmwareErase\n"); - process_msg_FirmwareErase(USB_IFACE_NUM); + process_msg_FirmwareErase(USB_IFACE_NUM, msg_size, buf); break; case 7: // FirmwareUpload display_printf("received FirmwareUpload\n"); - process_msg_FirmwareUpload(USB_IFACE_NUM); + process_msg_FirmwareUpload(USB_IFACE_NUM, msg_size, buf); break; default: display_printf("received unknown message\n"); - process_msg_unknown(USB_IFACE_NUM); + process_msg_unknown(USB_IFACE_NUM, msg_size, buf); break; } } diff --git a/micropython/bootloader/messages.c b/micropython/bootloader/messages.c index aa837cadcc..0b6d876807 100644 --- a/micropython/bootloader/messages.c +++ b/micropython/bootloader/messages.c @@ -3,7 +3,6 @@ #include "usb.h" #include "version.h" -#include "bootloader.h" #include "messages.h" @@ -54,7 +53,7 @@ static bool _usb_write(pb_ostream_t *stream, const pb_byte_t *buf, size_t count) state->packet_index++; memset(state->buf, 0, USB_PACKET_SIZE); state->buf[0] = '?'; - state->packet_pos = 1; + state->packet_pos = MSG_HEADER2_LEN; } } @@ -89,7 +88,7 @@ static bool _send_msg(uint8_t iface_num, uint16_t msg_id, const pb_field_t field usb_write_state state = { .iface_num = iface_num, .packet_index = 0, - .packet_pos = MSG_HEADER_LEN, + .packet_pos = MSG_HEADER1_LEN, .buf = { '?', '#', '#', (msg_id >> 8) & 0xFF, msg_id & 0xFF, @@ -127,8 +126,27 @@ static bool _encode_string(pb_ostream_t *stream, const pb_field_t *field, void * #define MSG_ASSIGN_STRING(FIELD, VALUE) do { msg.FIELD.funcs.encode = &_encode_string; msg.FIELD.arg = VALUE; } while (0) #define MSG_SEND(TYPE) do { _send_msg(iface_num, MessageType_MessageType_##TYPE, TYPE##_fields, &msg); } while (0) -void process_msg_Initialize(uint8_t iface_num) +static void consume_message(uint8_t iface_num, uint32_t msg_size, uint8_t *buf, void (* consumer)(const uint8_t *buf, uint8_t len)) { + if (consumer) { + consumer(buf + MSG_HEADER1_LEN, USB_PACKET_SIZE - MSG_HEADER1_LEN); + } + int remaining_chunks = (msg_size - (USB_PACKET_SIZE - MSG_HEADER1_LEN)) / (USB_PACKET_SIZE - MSG_HEADER2_LEN); + for (int i = 0; i < remaining_chunks; i++) { + int r = usb_hid_read_blocking(USB_IFACE_NUM, buf, USB_PACKET_SIZE, 100); + if (r != USB_PACKET_SIZE) { + break; + } + if (consumer) { + consumer(buf + MSG_HEADER2_LEN, USB_PACKET_SIZE - MSG_HEADER2_LEN); + } + } +} + +void process_msg_Initialize(uint8_t iface_num, uint32_t msg_size, uint8_t *buf) +{ + consume_message(iface_num, msg_size, buf, NULL); + MSG_INIT(Features); MSG_ASSIGN_STRING(vendor, "trezor.io"); MSG_ASSIGN_VALUE(major_version, VERSION_MAJOR); @@ -140,16 +158,20 @@ void process_msg_Initialize(uint8_t iface_num) MSG_SEND(Features); } -void process_msg_Ping(uint8_t iface_num) +void process_msg_Ping(uint8_t iface_num, uint32_t msg_size, uint8_t *buf) { + consume_message(iface_num, msg_size, buf, NULL); + MSG_INIT(Success); // TODO: read message from Ping MSG_ASSIGN_STRING(message, "PONG!"); MSG_SEND(Success); } -void process_msg_FirmwareErase(uint8_t iface_num) +void process_msg_FirmwareErase(uint8_t iface_num, uint32_t msg_size, uint8_t *buf) { + consume_message(iface_num, msg_size, buf, NULL); + // TODO: implement MSG_INIT(Failure); MSG_ASSIGN_VALUE(code, FailureType_Failure_FirmwareError); @@ -157,8 +179,10 @@ void process_msg_FirmwareErase(uint8_t iface_num) MSG_SEND(Failure); } -void process_msg_FirmwareUpload(uint8_t iface_num) +void process_msg_FirmwareUpload(uint8_t iface_num, uint32_t msg_size, uint8_t *buf) { + consume_message(iface_num, msg_size, buf, NULL); + // TODO: implement MSG_INIT(Failure); MSG_ASSIGN_VALUE(code, FailureType_Failure_FirmwareError); @@ -166,8 +190,10 @@ void process_msg_FirmwareUpload(uint8_t iface_num) MSG_SEND(Failure); } -void process_msg_unknown(uint8_t iface_num) +void process_msg_unknown(uint8_t iface_num, uint32_t msg_size, uint8_t *buf) { + consume_message(iface_num, msg_size, buf, NULL); + MSG_INIT(Failure); MSG_ASSIGN_VALUE(code, FailureType_Failure_UnexpectedMessage); MSG_ASSIGN_STRING(message, "Unexpected message"); diff --git a/micropython/bootloader/messages.h b/micropython/bootloader/messages.h index 1ba6d7b317..a25fd39845 100644 --- a/micropython/bootloader/messages.h +++ b/micropython/bootloader/messages.h @@ -6,14 +6,20 @@ #include -#define MSG_HEADER_LEN 9 +#define MSG_HEADER1_LEN 9 +#define MSG_HEADER2_LEN 1 + +#define USB_PACKET_SIZE 64 +#define USB_IFACE_NUM 0 + +#define UPLOAD_CHUNK_SIZE (128*1024) bool msg_parse_header(const uint8_t *buf, uint16_t *msg_id, uint32_t *msg_size); -void process_msg_Initialize(uint8_t iface_num); -void process_msg_Ping(uint8_t iface_num); -void process_msg_FirmwareErase(uint8_t iface_num); -void process_msg_FirmwareUpload(uint8_t iface_num); -void process_msg_unknown(uint8_t iface_num); +void process_msg_Initialize(uint8_t iface_num, uint32_t msg_size, uint8_t *buf); +void process_msg_Ping(uint8_t iface_num, uint32_t msg_size, uint8_t *buf); +void process_msg_FirmwareErase(uint8_t iface_num, uint32_t msg_size, uint8_t *buf); +void process_msg_FirmwareUpload(uint8_t iface_num, uint32_t msg_size, uint8_t *buf); +void process_msg_unknown(uint8_t iface_num, uint32_t msg_size, uint8_t *buf); #endif