From 68aa9d631dc1b448eb9ad1311133871dd00f01a8 Mon Sep 17 00:00:00 2001
From: slurdge <slurdge@slurdge.org>
Date: Fri, 19 Jul 2019 18:21:20 +0200
Subject: [PATCH] Flasher support for versionning

Only bootrom with version > 1.0.0 will allow 512K writes
---
 bootrom/bootrom.c |  2 +-
 client/flash.c    | 24 +++++++++++++++++++++++-
 include/pm3_cmd.h |  5 ++++-
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/bootrom/bootrom.c b/bootrom/bootrom.c
index 491cafec7..030074177 100644
--- a/bootrom/bootrom.c
+++ b/bootrom/bootrom.c
@@ -143,7 +143,7 @@ void UsbPacketReceived(uint8_t *packet, int len) {
 
         case CMD_BL_VERSION: {
             dont_ack = 1;
-            arg0 = VERSION_1_0_0;
+            arg0 = BL_VERSION_1_0_0;
             reply_old(CMD_BL_VERSION, arg0, 0, 0, 0, 0);
         }
 
diff --git a/client/flash.c b/client/flash.c
index abd81a9de..2c990a80a 100644
--- a/client/flash.c
+++ b/client/flash.c
@@ -356,8 +356,30 @@ int flash_start_flashing(int enable_bl_writes, char *serial_port_name, uint32_t
         *chipinfo = resp.oldarg[0];
     }
 
+    int version = BL_VERSION_INVALID;
+    if (state & DEVICE_INFO_FLAG_UNDERSTANDS_VERSION) {
+        SendCommandBL(CMD_BL_VERSION, 0, 0, 0, NULL, 0);
+        PacketResponseNG resp;
+        WaitForResponse(CMD_BL_VERSION, &resp);
+        version = resp.oldarg[0];
+        if ((version < BL_VERSION_FIRST || version > BL_VERSION_LAST)) {
+            version = BL_VERSION_INVALID;
+        }
+    } else {
+        PrintAndLogEx(ERR, _RED_("Note: Your bootloader does not understand the new CMD_BL_VERSION command"));
+        PrintAndLogEx(ERR, _RED_("It is recommended that you update your bootloader") "\n");
+    }
+
+    bool allow_512k_writes = false;
+    if (version == BL_VERSION_INVALID) {
+        PrintAndLogEx(ERR, _RED_("Note: Your bootloader reported an invalid version number"));
+        PrintAndLogEx(ERR, _RED_("It is recommended that you update your bootloader") "\n");
+    } else if (version >= BL_VERSION_1_0_0) {
+        allow_512k_writes = true;
+    }
+
     uint32_t flash_end = FLASH_START + AT91C_IFLASH_PAGE_SIZE * AT91C_IFLASH_NB_OF_PAGES / 2;
-    if (((*chipinfo & 0xF00) >> 8) > 9) {
+    if ((((*chipinfo & 0xF00) >> 8) > 9) && allow_512k_writes) {
         flash_end = FLASH_START + AT91C_IFLASH_PAGE_SIZE * AT91C_IFLASH_NB_OF_PAGES;
     }
 
diff --git a/include/pm3_cmd.h b/include/pm3_cmd.h
index 55204e1a4..640a3e2bf 100644
--- a/include/pm3_cmd.h
+++ b/include/pm3_cmd.h
@@ -541,7 +541,10 @@ typedef struct {
 #define DEVICE_INFO_FLAG_UNDERSTANDS_VERSION         (1<<6)
 
 // Different versions here. Each version should increse the number
-#define VERSION_1_0_0 1
+#define BL_VERSION_INVALID  0
+#define BL_VERSION_1_0_0    1
+#define BL_VERSION_FIRST    BL_VERSION_1_0_0
+#define BL_VERSION_LAST     BL_VERSION_1_0_0
 
 /* CMD_START_FLASH may have three arguments: start of area to flash,
    end of area to flash, optional magic.