app: aboot: check battery voltage before erase or flash the whole image

Halfway stop to erase or flash image is a risky when the battery's
voltage is too low. So it's reasonable to check the battery voltage
before to erase or flash image.

Change-Id: Ia31b9133f02b29c75cdc3c9c47b4500b009e5c01
diff --git a/app/aboot/aboot.c b/app/aboot/aboot.c
index 98006bd..c4eec13 100644
--- a/app/aboot/aboot.c
+++ b/app/aboot/aboot.c
@@ -2385,13 +2385,6 @@
 
 void cmd_erase(const char *arg, void *data, unsigned sz)
 {
-#if CHECK_BAT_VOLTAGE
-	if (!target_battery_soc_ok()) {
-		fastboot_fail("Warning: battery's capacity is very low\n");
-		return;
-	}
-#endif
-
 #if VERIFIED_BOOT
 	if (target_build_variant_user())
 	{
@@ -2991,12 +2984,6 @@
 
 void cmd_flash(const char *arg, void *data, unsigned sz)
 {
-#if CHECK_BAT_VOLTAGE
-	if (!target_battery_soc_ok()) {
-		fastboot_fail("Warning: battery's capacity is very low\n");
-		return;
-	}
-#endif
 	if(target_is_emmc_boot())
 		cmd_flash_mmc(arg, data, sz);
 	else
diff --git a/app/aboot/fastboot.c b/app/aboot/fastboot.c
index f99ed03..205b74a 100644
--- a/app/aboot/fastboot.c
+++ b/app/aboot/fastboot.c
@@ -500,6 +500,10 @@
 {
 	struct fastboot_cmd *cmd;
 	int r;
+#if CHECK_BAT_VOLTAGE
+	boolean is_first_erase_flash = false;
+#endif
+
 	dprintf(INFO,"fastboot: processing commands\n");
 
 	uint8_t *buffer = (uint8_t *)memalign(CACHE_LINE, ROUNDUP(4096, CACHE_LINE));
@@ -523,6 +527,24 @@
 		buffer[r] = 0;
 		dprintf(INFO,"fastboot: %s\n", buffer);
 
+#if CHECK_BAT_VOLTAGE
+		/* check battery voltage before erase or flash image */
+		if (!strncmp((const char*) buffer, "getvar:partition-type", 21))
+			is_first_erase_flash = true;
+
+		if (is_first_erase_flash) {
+			if (!strncmp((const char*) buffer, "erase", 5) ||
+				!strncmp((const char*) buffer, "flash", 5)) {
+				if (!target_battery_soc_ok()) {
+					dprintf(INFO,"fastboot: battery voltage: %d\n",
+						target_get_battery_voltage());
+					fastboot_fail("Warning: battery's capacity is very low\n");
+					return;
+				}
+			}
+		}
+#endif
+
 		fastboot_state = STATE_COMMAND;
 
 		for (cmd = cmdlist; cmd; cmd = cmd->next) {
@@ -532,6 +554,15 @@
 				    (void*) download_base, download_size);
 			if (fastboot_state == STATE_COMMAND)
 				fastboot_fail("unknown reason");
+
+#if CHECK_BAT_VOLTAGE
+			if (!strncmp((const char*) buffer, "erase", 5) ||
+				!strncmp((const char*) buffer, "flash", 5)) {
+				if (is_first_erase_flash) {
+					is_first_erase_flash = false;
+				}
+			}
+#endif
 			goto again;
 		}