From 6ead44933ac50a263c0e36c9b312e350d6147202 Mon Sep 17 00:00:00 2001
From: Philippe Teuwen <phil@teuwen.org>
Date: Wed, 16 Feb 2022 03:32:59 +0100
Subject: [PATCH] Flasher: factorize and prompt also wehn we can't check the
 version

---
 client/src/flash.c | 82 +++++++++++++++++++++++++---------------------
 1 file changed, 45 insertions(+), 37 deletions(-)

diff --git a/client/src/flash.c b/client/src/flash.c
index 843597b50..4231b82a9 100644
--- a/client/src/flash.c
+++ b/client/src/flash.c
@@ -235,6 +235,34 @@ static int check_segs(flash_file_t *ctx, int can_write_bl, uint32_t flash_size)
     return PM3_SUCCESS;
 }
 
+static int ask_confirmation(void) {
+    PrintAndLogEx(INFO,  "Make sure to flash a correct and up-to-date version");
+    PrintAndLogEx(NORMAL,  "Do you want to flash the current image? (yes/no)");
+    char answer[10];
+    if ((fgets (answer, sizeof(answer), stdin) == NULL) || (strncmp(answer, "yes", 3) != 0)) {
+        return PM3_EOPABORTED;
+    } else {
+        return PM3_SUCCESS;
+    }
+}
+
+static int print_and_validate_version(struct version_information_t *vi) {
+    if (vi->magic != VERSION_INFORMATION_MAGIC)
+        return PM3_EFILE;
+    char temp[PM3_CMD_DATA_SIZE - 12]; // same limit as for ARM image
+    FormatVersionInformation(temp, sizeof(temp), "", vi);
+    PrintAndLogEx(SUCCESS, _CYAN_("ELF file version") _YELLOW_(" %s"), temp);
+    if (strlen(g_version_information.armsrc) == 9) {
+        if (strncmp(vi->armsrc, g_version_information.armsrc, 9) != 0) {
+            PrintAndLogEx(WARNING, _RED_("ARM firmware does not match the source at the time the client was compiled"));
+            return ask_confirmation();
+        } else {
+            return PM3_SUCCESS;
+        }
+    }
+    return PM3_EUNDEF;
+}
+
 // Load an ELF file for flashing
 int flash_load(flash_file_t *ctx) {
     FILE *fd;
@@ -316,53 +344,33 @@ int flash_load(flash_file_t *ctx) {
     shdrs = (Elf32_Shdr_t *)(ctx->elf + le32(ehdr->e_shoff));
     shstr = ctx->elf + le32(shdrs[ehdr->e_shstrndx].sh_offset);
 
-    char temp[PM3_CMD_DATA_SIZE - 12]; // same limit as for ARM image
     for (uint16_t i = 0; i < le16(ehdr->e_shnum); i++) {
         if (strcmp(((char *)shstr) + shdrs[i].sh_name, ".version_information") == 0) {
             vi = (struct version_information_t *)(ctx->elf + le32(shdrs[i].sh_offset));
-            FormatVersionInformation(temp, sizeof(temp), "", vi);
-            PrintAndLogEx(SUCCESS, _CYAN_("ELF file version") _YELLOW_(" %s"), temp);
-            if (strlen(g_version_information.armsrc) == 9) {
-                if (strncmp(vi->armsrc, g_version_information.armsrc, 9) != 0) {
-                    PrintAndLogEx(WARNING, _RED_("ARM firmware does not match the source at the time the client was compiled"));
-                    PrintAndLogEx(INFO,  "Make sure to flash a correct and up-to-date version");
-                    PrintAndLogEx(NORMAL,  "Do you want to flash the current image? (yes/no)");
-                    char answer[10];
-                    if ((fgets (answer, sizeof(answer), stdin) == NULL) || (strncmp(answer, "yes", 3) != 0)) {
-                        res = PM3_EOPABORTED;
-                        goto fail;
-                    }
-                }
-            }
+            res = print_and_validate_version(vi);
+            break;
         }
         if (strcmp(((char *)shstr) + shdrs[i].sh_name, ".bootphase1") == 0) {
             uint32_t offset = *(uint32_t*)(ctx->elf + le32(shdrs[i].sh_offset) + le32(shdrs[i].sh_size) - 4);
-            if (offset < le32(shdrs[i].sh_addr))
-                continue;
-            offset -= le32(shdrs[i].sh_addr);
-            if (offset >= le32(shdrs[i].sh_size))
-                continue;
-            vi = (struct version_information_t *)(ctx->elf + le32(shdrs[i].sh_offset) + offset);
-            if (vi->magic != VERSION_INFORMATION_MAGIC)
-                continue;
-            FormatVersionInformation(temp, sizeof(temp), "", vi);
-            PrintAndLogEx(SUCCESS, _CYAN_("ELF file version") _YELLOW_(" %s"), temp);
-            if (strlen(g_version_information.armsrc) == 9) {
-                if (strncmp(vi->armsrc, g_version_information.armsrc, 9) != 0) {
-                    PrintAndLogEx(WARNING, _RED_("ARM firmware does not match the source at the time the client was compiled"));
-                    PrintAndLogEx(INFO,  "Make sure to flash a correct and up-to-date version");
-                    PrintAndLogEx(NORMAL,  "Do you want to flash the current image? (yes/no)");
-                    char answer[10];
-                    if ((fgets (answer, sizeof(answer), stdin) == NULL) || (strncmp(answer, "yes", 3) != 0)) {
-                        res = PM3_EOPABORTED;
-                        goto fail;
-                    }
+            if (offset >= le32(shdrs[i].sh_addr)) {
+                offset -= le32(shdrs[i].sh_addr);
+                if (offset < le32(shdrs[i].sh_size)) {
+                    vi = (struct version_information_t *)(ctx->elf + le32(shdrs[i].sh_offset) + offset);
+                    res = print_and_validate_version(vi);
                 }
             }
+            break;
         }
     }
-    return PM3_SUCCESS;
-
+    if (res == PM3_SUCCESS)
+        return res;
+    if (res == PM3_EOPABORTED)
+        goto fail;
+    // We could not find proper version_information, so we ask for confirmation
+    PrintAndLogEx(WARNING, "Unable to check version_information");
+    res = ask_confirmation();
+    if (res == PM3_SUCCESS)
+        return res;
 fail:
     flash_free(ctx);
     return res;