4 files changed, 76 insertions, 14 deletions

diff --git a/edify/expr.h b/edify/expr.h
index 5c06de846..886347991 100644
--- a/edify/expr.h
+++ b/edify/expr.h
@@ -49,6 +49,8 @@ typedef struct {
     // in addition to the error code.
     CauseCode cause_code = kNoCause;
 
+    bool is_retry = false;
+
 } State;
 
 #define VAL_STRING  1  // data will be NULL-terminated; size doesn't count null
diff --git a/install.cpp b/install.cpp
index 9d2971c8c..190d8525b 100644
--- a/install.cpp
+++ b/install.cpp
@@ -56,7 +56,7 @@ static const float DEFAULT_IMAGE_PROGRESS_FRACTION = 0.1;
 // If the package contains an update binary, extract it and run it.
 static int
 try_update_binary(const char* path, ZipArchive* zip, bool* wipe_cache,
-                  std::vector<std::string>& log_buffer)
+                  std::vector<std::string>& log_buffer, int retry_count)
 {
     const ZipEntry* binary_entry =
             mzFindZipEntry(zip, ASSUMED_UPDATE_BINARY_NAME);
@@ -130,15 +130,19 @@ try_update_binary(const char* path, ZipArchive* zip, bool* wipe_cache,
     //
     //   - the name of the package zip file.
     //
+    //   - an optional argument "retry" if this update is a retry of a failed
+    //   update attempt.
+    //
 
-    const char* args[5];
+    const char* args[6];
     args[0] = binary;
     args[1] = EXPAND(RECOVERY_API_VERSION);   // defined in Android.mk
     char temp[16];
     snprintf(temp, sizeof(temp), "%d", pipefd[1]);
     args[2] = temp;
     args[3] = path;
-    args[4] = NULL;
+    args[4] = retry_count > 0 ? "retry" : NULL;
+    args[5] = NULL;
 
     pid_t pid = fork();
     if (pid == 0) {
@@ -215,7 +219,7 @@ try_update_binary(const char* path, ZipArchive* zip, bool* wipe_cache,
 
 static int
 really_install_package(const char *path, bool* wipe_cache, bool needs_mount,
-                       std::vector<std::string>& log_buffer)
+                       std::vector<std::string>& log_buffer, int retry_count)
 {
     ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
     ui->Print("Finding update package...\n");
@@ -276,8 +280,11 @@ really_install_package(const char *path, bool* wipe_cache, bool needs_mount,
 
     // Verify and install the contents of the package.
     ui->Print("Installing update...\n");
+    if (retry_count > 0) {
+        ui->Print("Retry attempt: %d\n", retry_count);
+    }
     ui->SetEnableReboot(false);
-    int result = try_update_binary(path, &zip, wipe_cache, log_buffer);
+    int result = try_update_binary(path, &zip, wipe_cache, log_buffer, retry_count);
     ui->SetEnableReboot(true);
     ui->Print("\n");
 
@@ -306,7 +313,7 @@ install_package(const char* path, bool* wipe_cache, const char* install_file,
         LOGE("failed to set up expected mounts for install; aborting\n");
         result = INSTALL_ERROR;
     } else {
-        result = really_install_package(path, wipe_cache, needs_mount, log_buffer);
+        result = really_install_package(path, wipe_cache, needs_mount, log_buffer, retry_count);
     }
     if (install_log != nullptr) {
         fputc(result == INSTALL_SUCCESS ? '1' : '0', install_log);
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index f78342d36..7aff7fff7 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -70,6 +70,7 @@ struct RangeSet {
 };
 
 static CauseCode failure_type = kNoCause;
+static bool is_retry = false;
 static std::map<std::string, RangeSet> stash_map;
 
 static void parse_range(const std::string& range_text, RangeSet& rs) {
@@ -179,6 +180,21 @@ static int write_all(int fd, const std::vector<uint8_t>& buffer, size_t size) {
     return write_all(fd, buffer.data(), size);
 }
 
+static bool discard_blocks(int fd, off64_t offset, uint64_t size) {
+    // Don't discard blocks unless the update is a retry run.
+    if (!is_retry) {
+        return true;
+    }
+
+    uint64_t args[2] = {static_cast<uint64_t>(offset), size};
+    int status = ioctl(fd, BLKDISCARD, &args);
+    if (status == -1) {
+        fprintf(stderr, "BLKDISCARD ioctl failed: %s\n", strerror(errno));
+        return false;
+    }
+    return true;
+}
+
 static bool check_lseek(int fd, off64_t offset, int whence) {
     off64_t rc = TEMP_FAILURE_RETRY(lseek64(fd, offset, whence));
     if (rc == -1) {
@@ -238,10 +254,15 @@ static ssize_t RangeSinkWrite(const uint8_t* data, ssize_t size, void* token) {
                 rss->p_remain = (rss->tgt.pos[rss->p_block * 2 + 1] -
                                  rss->tgt.pos[rss->p_block * 2]) * BLOCKSIZE;
 
-                if (!check_lseek(rss->fd, (off64_t)rss->tgt.pos[rss->p_block*2] * BLOCKSIZE,
-                                 SEEK_SET)) {
+                off64_t offset = static_cast<off64_t>(rss->tgt.pos[rss->p_block*2]) * BLOCKSIZE;
+                if (!discard_blocks(rss->fd, offset, rss->p_remain)) {
+                    break;
+                }
+
+                if (!check_lseek(rss->fd, offset, SEEK_SET)) {
                     break;
                 }
+
             } else {
                 // we can't write any more; return how many bytes have
                 // been written so far.
@@ -347,11 +368,15 @@ static int WriteBlocks(const RangeSet& tgt, const std::vector<uint8_t>& buffer,
 
     size_t p = 0;
     for (size_t i = 0; i < tgt.count; ++i) {
-        if (!check_lseek(fd, (off64_t) tgt.pos[i * 2] * BLOCKSIZE, SEEK_SET)) {
+        off64_t offset = static_cast<off64_t>(tgt.pos[i * 2]) * BLOCKSIZE;
+        size_t size = (tgt.pos[i * 2 + 1] - tgt.pos[i * 2]) * BLOCKSIZE;
+        if (!discard_blocks(fd, offset, size)) {
             return -1;
         }
 
-        size_t size = (tgt.pos[i * 2 + 1] - tgt.pos[i * 2]) * BLOCKSIZE;
+        if (!check_lseek(fd, offset, SEEK_SET)) {
+            return -1;
+        }
 
         if (write_all(fd, data + p, size) == -1) {
             return -1;
@@ -1101,7 +1126,13 @@ static int PerformCommandZero(CommandParameters& params) {
 
     if (params.canwrite) {
         for (size_t i = 0; i < tgt.count; ++i) {
-            if (!check_lseek(params.fd, (off64_t) tgt.pos[i * 2] * BLOCKSIZE, SEEK_SET)) {
+            off64_t offset = static_cast<off64_t>(tgt.pos[i * 2]) * BLOCKSIZE;
+            size_t size = (tgt.pos[i * 2 + 1] - tgt.pos[i * 2]) * BLOCKSIZE;
+            if (!discard_blocks(params.fd, offset, size)) {
+                return -1;
+            }
+
+            if (!check_lseek(params.fd, offset, SEEK_SET)) {
                 return -1;
             }
 
@@ -1140,7 +1171,12 @@ static int PerformCommandNew(CommandParameters& params) {
         rss.p_block = 0;
         rss.p_remain = (tgt.pos[1] - tgt.pos[0]) * BLOCKSIZE;
 
-        if (!check_lseek(params.fd, (off64_t) tgt.pos[0] * BLOCKSIZE, SEEK_SET)) {
+        off64_t offset = static_cast<off64_t>(tgt.pos[0]) * BLOCKSIZE;
+        if (!discard_blocks(params.fd, offset, tgt.size * BLOCKSIZE)) {
+            return -1;
+        }
+
+        if (!check_lseek(params.fd, offset, SEEK_SET)) {
             return -1;
         }
 
@@ -1218,7 +1254,12 @@ static int PerformCommandDiff(CommandParameters& params) {
             rss.p_block = 0;
             rss.p_remain = (tgt.pos[1] - tgt.pos[0]) * BLOCKSIZE;
 
-            if (!check_lseek(params.fd, (off64_t) tgt.pos[0] * BLOCKSIZE, SEEK_SET)) {
+            off64_t offset = static_cast<off64_t>(tgt.pos[0]) * BLOCKSIZE;
+            if (!discard_blocks(params.fd, offset, rss.p_remain)) {
+                return -1;
+            }
+
+            if (!check_lseek(params.fd, offset, SEEK_SET)) {
                 return -1;
             }
 
@@ -1336,6 +1377,10 @@ static Value* PerformBlockImageUpdate(const char* name, State* state, int /* arg
     params.canwrite = !dryrun;
 
     fprintf(stderr, "performing %s\n", dryrun ? "verification" : "update");
+    if (state->is_retry) {
+        is_retry = true;
+        fprintf(stderr, "This update is a retry.\n");
+    }
 
     Value* blockdev_filename = nullptr;
     Value* transfer_list_value = nullptr;
diff --git a/updater/updater.cpp b/updater/updater.cpp
index f51a6db73..c222cee0d 100644
--- a/updater/updater.cpp
+++ b/updater/updater.cpp
@@ -51,7 +51,7 @@ int main(int argc, char** argv) {
     setbuf(stdout, NULL);
     setbuf(stderr, NULL);
 
-    if (argc != 4) {
+    if (argc != 4 && argc != 5) {
         printf("unexpected number of arguments (%d)\n", argc);
         return 1;
     }
@@ -145,6 +145,14 @@ int main(int argc, char** argv) {
     state.script = script;
     state.errmsg = NULL;
 
+    if (argc == 5) {
+        if (strcmp(argv[4], "retry") == 0) {
+            state.is_retry = true;
+        } else {
+            printf("unexpected argument: %s", argv[4]);
+        }
+    }
+
     char* result = Evaluate(&state, root);
 
     if (have_eio_error) {