From 3352904f1d08dcc76424c64b3b9bcc67a53356ed Mon Sep 17 00:00:00 2001
From: iceman1001 <herrmann1001@gmail.com>
Date: Sat, 31 Aug 2019 21:47:31 +0200
Subject: [PATCH] chg 'mem spiffs load' - use loadfile_safe

---
 client/cmdflashmemspiffs.c | 40 +++++++++++---------------------------
 1 file changed, 11 insertions(+), 29 deletions(-)

diff --git a/client/cmdflashmemspiffs.c b/client/cmdflashmemspiffs.c
index 5783e2705..0eef78259 100644
--- a/client/cmdflashmemspiffs.c
+++ b/client/cmdflashmemspiffs.c
@@ -330,12 +330,15 @@ static int CmdFlashMemSpiFFSLoad(const char *Cmd) {
                 if (param_getstr(Cmd, cmdp + 1, filename, FILE_PATH_SIZE) >= FILE_PATH_SIZE) {
                     PrintAndLogEx(FAILED, "Filename too long");
                     errors = true;
-                    break;
                 }
                 cmdp += 2;
                 break;
             case 'o':
                 param_getstr(Cmd, cmdp + 1, destfilename, 32);
+                if (strlen(destfilename) == 0) {
+                    PrintAndLogEx(FAILED, "Destination Filename missing or invalid");
+                    errors = true;
+                }
                 cmdp += 2;
                 break;
             default:
@@ -345,42 +348,20 @@ static int CmdFlashMemSpiFFSLoad(const char *Cmd) {
         }
     }
 
-    if (destfilename[0] == '\0') {
-        PrintAndLogEx(FAILED, "Destination Filename missing or invalid");
-        errors = true;
-    }
-
     // Validations
-    if (errors || cmdp == 0) {
-        usage_flashmemspiffs_load();
-        return PM3_EINVARG;
-    }
+    if (errors || cmdp == 0)
+        return usage_flashmemspiffs_load();
 
     size_t datalen = 0;
-    int res = 0;
-    uint8_t *data = calloc(FLASH_MEM_MAX_SIZE, sizeof(uint8_t));
+    uint8_t *data = NULL;
 
-    res = loadFile(filename, "", data, FLASH_MEM_MAX_SIZE, &datalen);
+    int res = loadFile_safe(filename, "", (void**)&data, &datalen);
     // int res = loadFileEML( filename, data, &datalen);
-    if (res) {
+    if (res != PM3_SUCCESS) {
         free(data);
         return PM3_EFILE;
     }
-
-    if (datalen > FLASH_MEM_MAX_SIZE) {
-        PrintAndLogEx(ERR, "error, filesize is larger than available memory");
-        free(data);
-        return PM3_EOVFLOW;
-    }
-
-    uint8_t *newdata = realloc(data, datalen);
-    if (newdata == NULL) {
-        free(data);
-        return PM3_EMALLOC;
-    } else {
-        data = newdata;
-    }
-
+    
     // We want to mount before multiple operation so the lazy writes/append will not
     // trigger a mount + umount each loop iteration (lazy ops device side)
     SendCommandNG(CMD_SPIFFS_MOUNT, NULL, 0);
@@ -425,6 +406,7 @@ static int CmdFlashMemSpiFFSLoad(const char *Cmd) {
         if (!isok) {
             conn.block_after_ACK = false;
             PrintAndLogEx(FAILED, "Flash write fail [offset %u]", bytes_sent);
+            free(data);
             return PM3_EFLASH;
         }
     }