From 98d6c9720f4936769a8053d6035f0303fa72b8fe Mon Sep 17 00:00:00 2001
From: BlackMark <me@blackmark.me>
Date: Sat, 11 Apr 2020 10:35:52 +0200
Subject: [PATCH] Add flash reading functionality

---
 stk500v2/main.cpp | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/stk500v2/main.cpp b/stk500v2/main.cpp
index 4832936..375dc62 100644
--- a/stk500v2/main.cpp
+++ b/stk500v2/main.cpp
@@ -3,6 +3,7 @@
 #include "uart/uart.hpp"
 
 #include <avr/boot.h>
+#include <avr/pgmspace.h>
 
 #include "command.hpp"
 
@@ -110,6 +111,20 @@ static inline bool isReadLockIsp(const Message &msg)
 	return false;
 }
 
+static inline bool isLoadAddress(const Message &msg)
+{
+	if (msg.size == 5 && msg.body[0] == CMD_LOAD_ADDRESS)
+		return true;
+	return false;
+}
+
+static inline bool isReadFlashIsp(const Message &msg)
+{
+	if (msg.size == 4 && msg.body[0] == CMD_READ_FLASH_ISP)
+		return true;
+	return false;
+}
+
 static inline bool isLeaveProgmodeIsp(const Message &msg)
 {
 	if (msg.size == 3 && msg.body[0] == CMD_LEAVE_PROGMODE_ISP)
@@ -226,6 +241,28 @@ static inline void formatReadLockIspAnswer(Message &msg)
 	msg.checksum = calcChecksum(msg);
 }
 
+static inline void formatLoadAddressAnswer(Message &msg)
+{
+	msg.size = 2;
+	msg.body[1] = STATUS_CMD_OK;
+
+	msg.checksum = calcChecksum(msg);
+}
+
+static inline void formatReadFlashIspAnswer(Message &msg, uint32_t &addr)
+{
+	const uint16_t numBytes = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
+	msg.size = 3 + numBytes;
+	msg.body[1] = STATUS_CMD_OK;
+	for (uint16_t i = 0; i < numBytes; ++i) {
+		msg.body[i + 2] = pgm_read_byte(static_cast<uint16_t>(addr + i));
+	}
+	addr += numBytes;
+	msg.body[numBytes + 2] = STATUS_CMD_OK;
+
+	msg.checksum = calcChecksum(msg);
+}
+
 static inline void formatLeaveProgmodeIspAnswer(Message &msg)
 {
 	msg.size = 2;
@@ -245,6 +282,8 @@ static inline void formatErrorAnswer(Message &msg)
 
 static inline void handleMessage(Message &msg)
 {
+	static uint32_t s_address = 0;
+
 	if (isSignOn(msg))
 		formatSignOnAnswer(msg);
 	else if (isGetParameter(msg))
@@ -259,6 +298,14 @@ static inline void handleMessage(Message &msg)
 		formatReadFuseIspAnswer(msg);
 	else if (isReadLockIsp(msg))
 		formatReadLockIspAnswer(msg);
+	else if (isLoadAddress(msg)) {
+		s_address = msg.body[1];
+		s_address = (s_address << 8) | msg.body[2];
+		s_address = (s_address << 8) | msg.body[3];
+		s_address = (s_address << 8) | msg.body[4];
+		formatLoadAddressAnswer(msg);
+	} else if (isReadFlashIsp(msg))
+		formatReadFlashIspAnswer(msg, s_address);
 	else if (isLeaveProgmodeIsp(msg))
 		formatLeaveProgmodeIspAnswer(msg);
 	else