From 16955979a2a1d253e795cb52f9c1175dddc88325 Mon Sep 17 00:00:00 2001
From: BlackMark <me@blackmark.me>
Date: Sat, 11 Apr 2020 22:19:22 +0200
Subject: [PATCH] Implement chip erase

---
 stk500v2/main.cpp | 107 ++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 99 insertions(+), 8 deletions(-)

diff --git a/stk500v2/main.cpp b/stk500v2/main.cpp
index 95e819e..499b4ab 100644
--- a/stk500v2/main.cpp
+++ b/stk500v2/main.cpp
@@ -146,6 +146,13 @@ static inline bool isReadEepromIsp(const Message &msg)
 	return false;
 }
 
+static inline bool isChipEraseIsp(const Message &msg)
+{
+	if (msg.size == 7 && msg.body[0] == CMD_CHIP_ERASE_ISP)
+		return true;
+	return false;
+}
+
 static inline bool isProgramEepromIsp(const Message &msg)
 {
 	if (msg.body[0] == CMD_PROGRAM_EEPROM_ISP) {
@@ -325,6 +332,61 @@ void writeEepromByte(uint8_t *addr, uint8_t value)
 
 } // namespace
 
+static inline void performChipErase()
+{
+	constexpr auto getEepromEraseFuseBit = []() -> bool {
+		constexpr auto EESAVE = 3;
+		return boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS) & (1 << EESAVE);
+	};
+
+	constexpr auto getBootloaderSize = []() -> uint16_t {
+		const auto extendedFuse = boot_lock_fuse_bits_get(GET_HIGH_FUSE_BITS);
+		constexpr auto BOOTSZ0 = 1;
+		constexpr auto BOOTSZ1 = 2;
+
+		if (extendedFuse & (1 << BOOTSZ1 | 1 << BOOTSZ0))
+			return 256 * 2;
+		else if (extendedFuse & (1 << BOOTSZ1))
+			return 512 * 2;
+		else if (extendedFuse & (1 << BOOTSZ0))
+			return 1024 * 2;
+		return 2048 * 2;
+	};
+
+	constexpr auto eraseFlash = [getBootloaderSize]() {
+		const auto bootloaderSize = getBootloaderSize();
+		const auto flashSize = (FLASHEND - bootloaderSize + 1);
+
+		for (uint16_t i = 0; i < flashSize; i += SPM_PAGESIZE) {
+			boot_page_erase(i);
+			boot_spm_busy_wait();
+		}
+		boot_rww_enable();
+	};
+
+	constexpr auto eraseEeprom = [getEepromEraseFuseBit]() {
+		const auto eraseEeprom = getEepromEraseFuseBit();
+		if (eraseEeprom) {
+			constexpr auto EEPROM_SIZE = E2END + 1;
+			for (uint16_t i = 0; i < EEPROM_SIZE; ++i) {
+				writeEepromByte(reinterpret_cast<uint8_t *>(i), 0xFF);
+				waitEepromReady();
+			}
+		}
+	};
+
+	eraseFlash();
+	eraseEeprom();
+}
+
+static inline void formatChipEraseIspAnswer(Message &msg)
+{
+	msg.size = 2;
+	msg.body[1] = STATUS_CMD_OK;
+
+	msg.checksum = calcChecksum(msg);
+}
+
 static inline void formatReadEepromIspAnswer(Message &msg, uint32_t &addr)
 {
 	const uint16_t numBytes = static_cast<uint16_t>(msg.body[1]) << 8 | msg.body[2];
@@ -370,7 +432,25 @@ static inline void formatErrorAnswer(Message &msg)
 	msg.checksum = calcChecksum(msg);
 }
 
-static inline void handleMessage(Message &msg)
+enum class ChipEraseState {
+	NONE = 0b00,
+	REQUEST = 0b01,
+	RESPONSE = 0b10,
+	PERFORM = 0b11,
+};
+
+constexpr ChipEraseState operator|(const ChipEraseState &self, const ChipEraseState &other)
+{
+	return static_cast<ChipEraseState>(static_cast<uint8_t>(self) | static_cast<uint8_t>(other));
+}
+
+constexpr ChipEraseState &operator|=(ChipEraseState &self, const ChipEraseState &other)
+{
+	self = self | other;
+	return self;
+}
+
+static inline void handleMessage(Message &msg, ChipEraseState &chipEraseFlag)
 {
 	static uint32_t s_address;
 
@@ -398,24 +478,35 @@ static inline void handleMessage(Message &msg)
 		formatReadFlashIspAnswer(msg, s_address);
 	else if (isReadEepromIsp(msg))
 		formatReadEepromIspAnswer(msg, s_address);
-	else if (isProgramEepromIsp(msg))
+	else if (isChipEraseIsp(msg)) {
+		chipEraseFlag = ChipEraseState::REQUEST;
+		formatChipEraseIspAnswer(msg);
+	} else if (isProgramEepromIsp(msg))
 		formatProgramEepromIspAnswer(msg, s_address);
-	else if (isLeaveProgmodeIsp(msg))
+	else if (isLeaveProgmodeIsp(msg)) {
+		chipEraseFlag |= ChipEraseState::RESPONSE;
 		formatLeaveProgmodeIspAnswer(msg);
-	else
+	} else
 		formatErrorAnswer(msg);
+
+	transmitMessage(msg);
+
+	if (chipEraseFlag == ChipEraseState::PERFORM) {
+		performChipErase();
+		chipEraseFlag = ChipEraseState::NONE;
+	}
 }
 
 int main()
 {
 	serial.init();
 
-	static Message s_msg;
+	Message msg;
+	ChipEraseState chipEraseFlag = ChipEraseState::NONE;
 
 	while (true) {
-		if (receiveMessage(s_msg)) {
-			handleMessage(s_msg);
-			transmitMessage(s_msg);
+		if (receiveMessage(msg)) {
+			handleMessage(msg, chipEraseFlag);
 		}
 	}