From 6f9d8bae423ac4c37d7c4fab7badb71626c55f6d Mon Sep 17 00:00:00 2001
From: BlackMark <me@blackmark.me>
Date: Sun, 12 Apr 2020 00:15:26 +0200
Subject: [PATCH] Implement writing flash memory

---
 stk500v2/main.cpp | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

diff --git a/stk500v2/main.cpp b/stk500v2/main.cpp
index 1a2e870..4513bf1 100644
--- a/stk500v2/main.cpp
+++ b/stk500v2/main.cpp
@@ -153,6 +153,16 @@ static inline bool isChipEraseIsp(const Message &msg)
 	return false;
 }
 
+static inline bool isProgramFlashIsp(const Message &msg)
+{
+	if (msg.body[0] == CMD_PROGRAM_FLASH_ISP) {
+		const auto dataSize = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
+		if (msg.size == (dataSize + 10) && dataSize == SPM_PAGESIZE)
+			return true;
+	}
+	return false;
+}
+
 static inline bool isProgramEepromIsp(const Message &msg)
 {
 	if (msg.body[0] == CMD_PROGRAM_EEPROM_ISP) {
@@ -330,6 +340,24 @@ void writeEepromByte(uint8_t *addr, uint8_t value)
 	EECR |= (1 << EEPE);
 }
 
+//////////////////////////////////////////////////////////////////////////
+
+void writeFlashPage(uint32_t pageAddress, const uint8_t *data)
+{
+	boot_page_erase(pageAddress);
+	boot_spm_busy_wait();
+
+	for (uint16_t i = 0; i < SPM_PAGESIZE; i += 2) {
+		uint16_t dataWord = *data++;
+		dataWord |= (*data++) << 8;
+		boot_page_fill(pageAddress + i, dataWord);
+	}
+
+	boot_page_write(pageAddress);
+	boot_spm_busy_wait();
+	boot_rww_enable();
+}
+
 } // namespace
 
 static inline uint16_t getBootloaderSize()
@@ -407,6 +435,19 @@ static inline void formatReadEepromIspAnswer(Message &msg, uint32_t &addr)
 	msg.checksum = calcChecksum(msg);
 }
 
+static inline void formatProgramFlashIspAnswer(Message &msg, uint32_t &addr)
+{
+	if (addr < getFlashSize())
+		writeFlashPage(addr, msg.body + 10);
+	const uint16_t numBytes = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
+	addr += numBytes;
+
+	msg.size = 2;
+	msg.body[1] = STATUS_CMD_OK;
+
+	msg.checksum = calcChecksum(msg);
+}
+
 static inline void formatProgramEepromIspAnswer(Message &msg, uint32_t &addr)
 {
 	const uint16_t numBytes = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
@@ -485,7 +526,9 @@ static inline void handleMessage(Message &msg, uint32_t &addr, ChipEraseState &c
 	else if (isChipEraseIsp(msg)) {
 		chipEraseFlag = ChipEraseState::REQUEST;
 		formatChipEraseIspAnswer(msg);
-	} else if (isProgramEepromIsp(msg))
+	} else if (isProgramFlashIsp(msg))
+		formatProgramFlashIspAnswer(msg, addr);
+	else if (isProgramEepromIsp(msg))
 		formatProgramEepromIspAnswer(msg, addr);
 	else if (isLeaveProgmodeIsp(msg)) {
 		chipEraseFlag |= ChipEraseState::RESPONSE;