3 files changed, 48 insertions, 30 deletions

diff --git a/device.h b/device.h
index dad8ccd56..f74b6b047 100644
--- a/device.h
+++ b/device.h
@@ -91,13 +91,16 @@ class Device {
     static const int kHighlightDown = -3;
     static const int kInvokeItem = -4;
 
-    // Called when we do a wipe data/factory reset operation (either via a
-    // reboot from the main system with the --wipe_data flag, or when the
-    // user boots into recovery manually and selects the option from the
-    // menu.)  Can perform whatever device-specific wiping actions are
-    // needed.  Return 0 on success.  The userdata and cache partitions
-    // are erased AFTER this returns (whether it returns success or not).
-    virtual int WipeData() { return 0; }
+    // Called before and after we do a wipe data/factory reset operation,
+    // either via a reboot from the main system with the --wipe_data flag,
+    // or when the user boots into recovery image manually and selects the
+    // option from the menu, to perform whatever device-specific wiping
+    // actions are needed.
+    // Return true on success; returning false from PreWipeData will prevent
+    // the regular wipe, and returning false from PostWipeData will cause
+    // the wipe to be considered a failure.
+    virtual bool PreWipeData() { return true; }
+    virtual bool PostWipeData() { return true; }
 
   private:
     RecoveryUI* ui_;
diff --git a/recovery.cpp b/recovery.cpp
index a7dc6ed9c..83ca5812d 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -421,8 +421,7 @@ typedef struct _saved_log_file {
     struct _saved_log_file* next;
 } saved_log_file;
 
-static int
-erase_volume(const char *volume) {
+static bool erase_volume(const char* volume) {
     bool is_cache = (strcmp(volume, CACHE_ROOT) == 0);
 
     ui->SetBackground(RecoveryUI::ERASING);
@@ -503,7 +502,7 @@ erase_volume(const char *volume) {
         copy_logs();
     }
 
-    return result;
+    return (result == 0);
 }
 
 static int
@@ -677,13 +676,13 @@ static bool wipe_data(int should_confirm, Device* device) {
     modified_flash = true;
 
     ui->Print("\n-- Wiping data...\n");
-    if (device->WipeData() == 0 && erase_volume("/data") == 0 && erase_volume("/cache") == 0) {
-        ui->Print("Data wipe complete.\n");
-        return true;
-    } else {
-        ui->Print("Data wipe failed.\n");
-        return false;
-    }
+    bool success =
+        device->PreWipeData() &&
+        erase_volume("/data") &&
+        erase_volume("/cache") &&
+        device->PostWipeData();
+    ui->Print("Data wipe %s.\n", success ? "complete" : "failed");
+    return success;
 }
 
 // Return true on success.
@@ -695,13 +694,9 @@ static bool wipe_cache(bool should_confirm, Device* device) {
     modified_flash = true;
 
     ui->Print("\n-- Wiping cache...\n");
-    if (erase_volume("/cache") == 0) {
-        ui->Print("Cache wipe complete.\n");
-        return true;
-    } else {
-        ui->Print("Cache wipe failed.\n");
-        return false;
-    }
+    bool success = erase_volume("/cache");
+    ui->Print("Cache wipe %s.\n", success ? "complete" : "failed");
+    return success;
 }
 
 static void choose_recovery_file(Device* device) {
diff --git a/updater/blockimg.c b/updater/blockimg.c
index dfba7e420..fc35abe4d 100644
--- a/updater/blockimg.c
+++ b/updater/blockimg.c
@@ -39,6 +39,11 @@
 
 #define BLOCKSIZE 4096
 
+// Set this to 0 to interpret 'erase' transfers to mean do a
+// BLKDISCARD ioctl (the normal behavior).  Set to 1 to interpret
+// erase to mean fill the region with zeroes.
+#define DEBUG_ERASE  0
+
 #ifndef BLKDISCARD
 #define BLKDISCARD _IO(0x12,119)
 #endif
@@ -1278,7 +1283,8 @@ static int PerformCommandZero(CommandParameters* params) {
     }
 
     if (params->cmdname[0] == 'z') {
-        // Update only for the zero command, as the erase command will call this
+        // Update only for the zero command, as the erase command will call
+        // this if DEBUG_ERASE is defined.
         params->written += tgt->size;
     }
 
@@ -1459,15 +1465,14 @@ pcdout:
 static int PerformCommandErase(CommandParameters* params) {
     char* range = NULL;
     int i;
+    int j;
     int rc = -1;
     RangeSet* tgt = NULL;
     struct stat st;
     uint64_t blocks[2];
 
-    // Always zero the blocks first to work around possibly flaky BLKDISCARD
-    // Bug: 20881595
-    if (PerformCommandZero(params) != 0) {
-        goto pceout;
+    if (DEBUG_ERASE) {
+        return PerformCommandZero(params);
     }
 
     if (!params) {
@@ -1487,7 +1492,7 @@ static int PerformCommandErase(CommandParameters* params) {
     range = strtok_r(NULL, " ", &params->cpos);
 
     if (range == NULL) {
-        fprintf(stderr, "missing target blocks for zero\n");
+        fprintf(stderr, "missing target blocks for erase\n");
         goto pceout;
     }
 
@@ -1496,7 +1501,22 @@ static int PerformCommandErase(CommandParameters* params) {
     if (params->canwrite) {
         fprintf(stderr, " erasing %d blocks\n", tgt->size);
 
+        allocate(BLOCKSIZE, &params->buffer, &params->bufsize);
+        memset(params->buffer, 0, BLOCKSIZE);
+
         for (i = 0; i < tgt->count; ++i) {
+            // Always zero the blocks first to work around possibly flaky BLKDISCARD
+            // Bug: 20881595
+            if (!check_lseek(params->fd, (off64_t) tgt->pos[i * 2] * BLOCKSIZE, SEEK_SET)) {
+                goto pceout;
+            }
+
+            for (j = tgt->pos[i * 2]; j < tgt->pos[i * 2 + 1]; ++j) {
+                if (write_all(params->fd, params->buffer, BLOCKSIZE) == -1) {
+                    goto pceout;
+                }
+            }
+
             // offset in bytes
             blocks[0] = tgt->pos[i * 2] * (uint64_t) BLOCKSIZE;
             // length in bytes