From c9113fd3f5fcd78e9e560dbac75ed5aae359eb2d Mon Sep 17 00:00:00 2001
From: Jochen Hoenicke <hoenicke@gmail.com>
Date: Mon, 18 Jun 2018 21:01:54 +0100
Subject: [PATCH] firmware: fix message processing, typos in recovery

---
 firmware/messages.c | 13 ++++++++-----
 firmware/messages.h |  3 +--
 firmware/recovery.c |  6 ++++--
 firmware/storage.c  |  2 ++
 4 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/firmware/messages.c b/firmware/messages.c
index 4638057a88c..b388067fd6b 100644
--- a/firmware/messages.c
+++ b/firmware/messages.c
@@ -233,7 +233,7 @@ void msg_process(char type, uint16_t msg_id, const pb_field_t *fields, uint8_t *
 	}
 }
 
-void msg_read_common(char type, const uint8_t *buf, int len)
+void msg_read_common(char type, const uint8_t *buf, uint32_t len)
 {
 	static char read_state = READSTATE_IDLE;
 	static CONFIDENTIAL uint8_t msg_in[MSG_IN_SIZE];
@@ -271,8 +271,12 @@ void msg_read_common(char type, const uint8_t *buf, int len)
 			read_state = READSTATE_IDLE;
 			return;
 		}
-		memcpy(msg_in + msg_pos, buf + 1, len - 1);
-		msg_pos += len - 1;
+		/* raw data starts at buf + 1 with len - 1 bytes */
+		buf++;
+		len = MIN(len - 1, MSG_IN_SIZE - msg_pos);
+
+		memcpy(msg_in + msg_pos, buf, len);
+		msg_pos += len;
 	}
 
 	if (msg_pos >= msg_size) {
@@ -329,8 +333,7 @@ void msg_read_tiny(const uint8_t *buf, int len)
 	}
 
 	const pb_field_t *fields = 0;
-	// upstream nanopb is missing const qualifier, so we have to cast :-/
-	pb_istream_t stream = pb_istream_from_buffer((uint8_t *)buf + 9, msg_size);
+	pb_istream_t stream = pb_istream_from_buffer(buf + 9, msg_size);
 
 	switch (msg_id) {
 		case MessageType_MessageType_PinMatrixAck:
diff --git a/firmware/messages.h b/firmware/messages.h
index ace4a44480c..030af0e698c 100644
--- a/firmware/messages.h
+++ b/firmware/messages.h
@@ -42,11 +42,10 @@ const uint8_t *msg_debug_out_data(void);
 
 #endif
 
-void msg_read_common(char type, const uint8_t *buf, int len);
+void msg_read_common(char type, const uint8_t *buf, uint32_t len);
 bool msg_write_common(char type, uint16_t msg_id, const void *msg_ptr);
 
 void msg_read_tiny(const uint8_t *buf, int len);
-void msg_debug_read_tiny(const uint8_t *buf, int len);
 extern uint8_t msg_tiny[128];
 extern uint16_t msg_tiny_id;
 
diff --git a/firmware/recovery.c b/firmware/recovery.c
index 8990f29092d..be5661af241 100644
--- a/firmware/recovery.c
+++ b/firmware/recovery.c
@@ -295,7 +295,7 @@ static void display_choices(bool twoColumn, char choices[9][12], int num)
 
 	/* avoid picking out of range numbers */
 	for (int i = 0; i < displayedChoices; i++) {
-		if (word_matrix[i] > num)
+		if (word_matrix[i] >= num)
 			word_matrix[i] = 0;
 	}
 	/* two column layout: middle column = right column */
@@ -405,11 +405,13 @@ static void recovery_digit(const char digit) {
 		/* received final word */
 
 		/* Mark the chosen word for 250 ms */
-		int y = 54 - ((digit - '1')/3)*11;
+		int y = 54 - ((digit - '1') / 3) * 11;
 		int x = 64 * (((digit - '1') % 3) > 0);
 		oledInvert(x + 1, y, x + 62, y + 9);
 		oledRefresh();
+		usbTiny(1);
 		usbSleep(250);
+		usbTiny(0);
 
 		/* index of the chosen word */
 		int idx = TABLE2(TABLE1(word_pincode / 9) + (word_pincode % 9)) + choice;
diff --git a/firmware/storage.c b/firmware/storage.c
index ca77f3c5397..429a82e93f4 100644
--- a/firmware/storage.c
+++ b/firmware/storage.c
@@ -565,6 +565,7 @@ bool storage_getRootNode(HDNode *node, const char *curve, bool usePassphrase)
 			// decrypt hd node
 			uint8_t secret[64];
 			PBKDF2_HMAC_SHA512_CTX pctx;
+			char oldTiny = usbTiny(1);
 			pbkdf2_hmac_sha512_Init(&pctx, (const uint8_t *)sessionPassphrase, strlen(sessionPassphrase), (const uint8_t *)"TREZORHD", 8);
 			get_root_node_callback(0, BIP39_PBKDF2_ROUNDS);
 			for (int i = 0; i < 8; i++) {
@@ -572,6 +573,7 @@ bool storage_getRootNode(HDNode *node, const char *curve, bool usePassphrase)
 				get_root_node_callback((i + 1) * BIP39_PBKDF2_ROUNDS / 8, BIP39_PBKDF2_ROUNDS);
 			}
 			pbkdf2_hmac_sha512_Final(&pctx, secret);
+			usbTiny(oldTiny);
 			aes_decrypt_ctx ctx;
 			aes_decrypt_key256(secret, &ctx);
 			aes_cbc_decrypt(node->chain_code, node->chain_code, 32, secret + 32, &ctx);