4 files changed, 42 insertions, 11 deletions

diff --git a/src/flash/mflash.c b/src/flash/mflash.c
index 272127be..ba344222 100644
--- a/src/flash/mflash.c
+++ b/src/flash/mflash.c
@@ -720,14 +720,20 @@ COMMAND_HANDLER(mg_write_cmd)
 	if (ret != ERROR_OK)
 		return ret;
 
+	int filesize;
 	buffer = malloc(MG_FILEIO_CHUNK);
 	if (!buffer) {
 		fileio_close(&fileio);
 		return ERROR_FAIL;
 	}
+	int retval = fileio_size(&fileio, &filesize);
+	if (retval != ERROR_OK) {
+		fileio_close(&fileio);
+		return retval;
+	}
 
-	cnt = fileio_size(&fileio) / MG_FILEIO_CHUNK;
-	res = fileio_size(&fileio) % MG_FILEIO_CHUNK;
+	cnt = filesize / MG_FILEIO_CHUNK;
+	res = filesize % MG_FILEIO_CHUNK;
 
 	struct duration bench;
 	duration_start(&bench);
@@ -752,8 +758,8 @@ COMMAND_HANDLER(mg_write_cmd)
 	if (duration_measure(&bench) == ERROR_OK)
 	{
 		command_print(CMD_CTX, "wrote %ld bytes from file %s "
-				"in %fs (%0.3f kB/s)", (long)fileio_size(&fileio), CMD_ARGV[1],
-				duration_elapsed(&bench), duration_kbps(&bench, fileio_size(&fileio)));
+				"in %fs (%0.3f kB/s)", (long)filesize, CMD_ARGV[1],
+				duration_elapsed(&bench), duration_kbps(&bench, filesize));
 	}
 
 	free(buffer);
diff --git a/src/flash/nand/fileio.c b/src/flash/nand/fileio.c
index 0a006fc0..c7515e2f 100644
--- a/src/flash/nand/fileio.c
+++ b/src/flash/nand/fileio.c
@@ -180,7 +180,13 @@ COMMAND_HELPER(nand_fileio_parse_args, struct nand_fileio_state *state,
 		return retval;
 
 	if (!need_size)
-		state->size = fileio_size(&state->fileio);
+	{
+		int filesize;
+		retval = fileio_size(&state->fileio, &filesize);
+		if (retval != ERROR_OK)
+			return retval;
+		state->size = filesize;
+	}
 
 	*dev = nand;
 
diff --git a/src/flash/nand/tcl.c b/src/flash/nand/tcl.c
index a54f8ea8..15cf1797 100644
--- a/src/flash/nand/tcl.c
+++ b/src/flash/nand/tcl.c
@@ -388,9 +388,14 @@ COMMAND_HANDLER(handle_nand_dump_command)
 
 	if (nand_fileio_finish(&s) == ERROR_OK)
 	{
+		int filesize;
+		retval = fileio_size(&s.fileio, &filesize);
+		if (retval != ERROR_OK)
+			return retval;
+
 		command_print(CMD_CTX, "dumped %ld bytes in %fs (%0.3f KiB/s)",
-				(long)fileio_size(&s.fileio), duration_elapsed(&s.bench),
-				duration_kbps(&s.bench, fileio_size(&s.fileio)));
+				(long)filesize, duration_elapsed(&s.bench),
+				duration_kbps(&s.bench, filesize));
 	}
 	return ERROR_OK;
 }
diff --git a/src/flash/nor/tcl.c b/src/flash/nor/tcl.c
index f36ab7da..142f31f7 100644
--- a/src/flash/nor/tcl.c
+++ b/src/flash/nor/tcl.c
@@ -604,9 +604,23 @@ COMMAND_HANDLER(handle_flash_write_bank_command)
 		return ERROR_OK;
 	}
 
-	buffer = malloc(fileio_size(&fileio));
+	int filesize;
+	retval = fileio_size(&fileio, &filesize);
+	if (retval != ERROR_OK)
+	{
+		fileio_close(&fileio);
+		return retval;
+	}
+
+	buffer = malloc(filesize);
+	if (buffer == NULL)
+	{
+		fileio_close(&fileio);
+		LOG_ERROR("Out of memory");
+		return ERROR_FAIL;
+	}
 	size_t buf_cnt;
-	if (fileio_read(&fileio, fileio_size(&fileio), buffer, &buf_cnt) != ERROR_OK)
+	if (fileio_read(&fileio, filesize, buffer, &buf_cnt) != ERROR_OK)
 	{
 		free(buffer);
 		fileio_close(&fileio);
@@ -622,8 +636,8 @@ COMMAND_HANDLER(handle_flash_write_bank_command)
 	{
 		command_print(CMD_CTX, "wrote %ld bytes from file %s to flash bank %u"
 				" at offset 0x%8.8" PRIx32 " in %fs (%0.3f KiB/s)",
-				(long)fileio_size(&fileio), CMD_ARGV[1], p->bank_number, offset,
-				duration_elapsed(&bench), duration_kbps(&bench, fileio_size(&fileio)));
+				(long)filesize, CMD_ARGV[1], p->bank_number, offset,
+				duration_elapsed(&bench), duration_kbps(&bench, filesize));
 	}
 
 	fileio_close(&fileio);