From 8979066420892a5c559bc5ba8eb135e10e846871 Mon Sep 17 00:00:00 2001
From: BlackMark <me@blackmark.me>
Date: Sun, 12 Apr 2020 11:47:11 +0200
Subject: [PATCH] Fix receiveByte timeout and add global timeout

---
 stk500v2/main.cpp | 50 +++++++++++++++++++++++++++++++++++------------
 1 file changed, 37 insertions(+), 13 deletions(-)

diff --git a/stk500v2/main.cpp b/stk500v2/main.cpp
index 4513bf1..405ce4a 100644
--- a/stk500v2/main.cpp
+++ b/stk500v2/main.cpp
@@ -2,6 +2,8 @@
 
 #include "uart/uart.hpp"
 
+#include <math.h>
+
 #include <avr/boot.h>
 #include <avr/interrupt.h>
 #include <avr/io.h>
@@ -9,7 +11,10 @@
 
 #include "command.hpp"
 
-using uart_interface = uart::Hardware0<uart::Config<115200>, uart::Driven::BLOCKING>;
+static constexpr auto TIMEOUT = 5000;
+static constexpr auto BAUD_RATE = 115200;
+
+using uart_interface = uart::Hardware0<uart::Config<BAUD_RATE>, uart::Driven::BLOCKING>;
 uart::Uart<uart_interface> serial;
 
 struct Message {
@@ -21,11 +26,26 @@ struct Message {
 	uint8_t checksum;
 };
 
-static inline bool receiveByte(uint8_t &data, uint32_t timeout = 100000)
+static inline bool receiveByte(uint8_t &data, uint16_t &timeout)
 {
-	for (uint32_t i = 0; i < timeout; ++i) {
-		if (serial.rxByte(data))
+	constexpr auto MICROSECOND = 1000.0 * 1000;
+	constexpr auto SYMBOL_SIZE = 9;
+	constexpr auto BYTE_DELAY_US = (SYMBOL_SIZE * MICROSECOND) / BAUD_RATE;
+	constexpr auto NUM_MS_DELAY_STEPS = static_cast<uint16_t>(round(1000 / BYTE_DELAY_US));
+	uint16_t msDelay = NUM_MS_DELAY_STEPS;
+
+	while (timeout) {
+		if (serial.rxByte(data)) {
+			timeout = TIMEOUT;
 			return true;
+		}
+
+		_delay_us(BYTE_DELAY_US);
+
+		if (--msDelay == 0) {
+			msDelay = NUM_MS_DELAY_STEPS;
+			--timeout;
+		}
 	}
 
 	return false;
@@ -42,23 +62,23 @@ static inline uint8_t calcChecksum(const Message &msg)
 	return checksum;
 }
 
-static inline bool receiveMessage(Message &msg)
+static inline bool receiveMessage(Message &msg, uint16_t &timeout)
 {
-	if (!receiveByte(msg.start) || msg.start != MESSAGE_START)
+	if (!receiveByte(msg.start, timeout) || msg.start != MESSAGE_START)
 		return false;
-	if (!receiveByte(msg.number))
+	if (!receiveByte(msg.number, timeout))
 		return false;
-	if (!receiveByte(*(reinterpret_cast<uint8_t *>(&msg.size) + 1)))
+	if (!receiveByte(*(reinterpret_cast<uint8_t *>(&msg.size) + 1), timeout))
 		return false;
-	if (!receiveByte(*reinterpret_cast<uint8_t *>(&msg.size)) || msg.size > sizeof(msg.body))
+	if (!receiveByte(*reinterpret_cast<uint8_t *>(&msg.size), timeout) || msg.size > sizeof(msg.body))
 		return false;
-	if (!receiveByte(msg.token) || msg.token != TOKEN)
+	if (!receiveByte(msg.token, timeout) || msg.token != TOKEN)
 		return false;
 	for (uint16_t i = 0; i < msg.size; ++i) {
-		if (!receiveByte(msg.body[i]))
+		if (!receiveByte(msg.body[i], timeout))
 			return false;
 	}
-	if (!receiveByte(msg.checksum) || msg.checksum != calcChecksum(msg))
+	if (!receiveByte(msg.checksum, timeout) || msg.checksum != calcChecksum(msg))
 		return false;
 
 	return true;
@@ -551,11 +571,15 @@ int main()
 	Message msg;
 	uint32_t addr = 0x0000;
 	ChipEraseState chipEraseFlag = ChipEraseState::NONE;
+	uint16_t timeout = TIMEOUT;
 
 	while (true) {
-		if (receiveMessage(msg)) {
+		if (receiveMessage(msg, timeout)) {
 			handleMessage(msg, addr, chipEraseFlag);
 		}
+
+		if (timeout == 0)
+			timeout = TIMEOUT;
 	}
 
 	return 0;