5 files changed, 93 insertions, 7 deletions

diff --git a/install.cpp b/install.cpp
index 33c1f5498..144a353d6 100644
--- a/install.cpp
+++ b/install.cpp
@@ -144,6 +144,7 @@ try_update_binary(const char* path, ZipArchive* zip, bool* wipe_cache) {
     close(pipefd[1]);
 
     *wipe_cache = false;
+    bool retry_update = false;
 
     char buffer[1024];
     FILE* from_child = fdopen(pipefd[0], "r");
@@ -180,6 +181,8 @@ try_update_binary(const char* path, ZipArchive* zip, bool* wipe_cache) {
             // to be able to reboot during installation (useful for
             // debugging packages that don't exit).
             ui->SetEnableReboot(true);
+        } else if (strcmp(command, "retry_update") == 0) {
+            retry_update = true;
         } else {
             LOGE("unknown command [%s]\n", command);
         }
@@ -188,6 +191,9 @@ try_update_binary(const char* path, ZipArchive* zip, bool* wipe_cache) {
 
     int status;
     waitpid(pid, &status, 0);
+    if (retry_update) {
+        return INSTALL_RETRY;
+    }
     if (!WIFEXITED(status) || WEXITSTATUS(status) != 0) {
         LOGE("Error in %s\n(Status %d)\n", path, WEXITSTATUS(status));
         return INSTALL_ERROR;
diff --git a/install.h b/install.h
index f92f061df..fd08e3c49 100644
--- a/install.h
+++ b/install.h
@@ -23,7 +23,8 @@
 extern "C" {
 #endif
 
-enum { INSTALL_SUCCESS, INSTALL_ERROR, INSTALL_CORRUPT, INSTALL_NONE, INSTALL_SKIPPED };
+enum { INSTALL_SUCCESS, INSTALL_ERROR, INSTALL_CORRUPT, INSTALL_NONE, INSTALL_SKIPPED,
+        INSTALL_RETRY };
 // Install the package specified by root_path.  If INSTALL_SUCCESS is
 // returned and *wipe_cache is true on exit, caller should wipe the
 // cache partition.
diff --git a/otafault/ota_io.cpp b/otafault/ota_io.cpp
index 2ce3dfcc2..dd805e56e 100644
--- a/otafault/ota_io.cpp
+++ b/otafault/ota_io.cpp
@@ -29,6 +29,7 @@ static std::map<intptr_t, const char*> filename_cache;
 static std::string read_fault_file_name = "";
 static std::string write_fault_file_name = "";
 static std::string fsync_fault_file_name = "";
+bool have_eio_error = false;
 
 static bool get_hit_file(const char* cached_path, std::string ffn) {
     return should_hit_cache()
@@ -85,10 +86,16 @@ size_t ota_fread(void* ptr, size_t size, size_t nitems, FILE* stream) {
                 get_hit_file(cached_path, read_fault_file_name)) {
             read_fault_file_name = "";
             errno = EIO;
+            have_eio_error = true;
             return 0;
         }
     }
-    return fread(ptr, size, nitems, stream);
+    size_t status = fread(ptr, size, nitems, stream);
+    // If I/O error occurs, set the retry-update flag.
+    if (status != nitems && errno == EIO) {
+        have_eio_error = true;
+    }
+    return status;
 }
 
 ssize_t ota_read(int fd, void* buf, size_t nbyte) {
@@ -99,10 +106,15 @@ ssize_t ota_read(int fd, void* buf, size_t nbyte) {
                 && get_hit_file(cached_path, read_fault_file_name)) {
             read_fault_file_name = "";
             errno = EIO;
+            have_eio_error = true;
             return -1;
         }
     }
-    return read(fd, buf, nbyte);
+    ssize_t status = read(fd, buf, nbyte);
+    if (status == -1 && errno == EIO) {
+        have_eio_error = true;
+    }
+    return status;
 }
 
 size_t ota_fwrite(const void* ptr, size_t size, size_t count, FILE* stream) {
@@ -113,10 +125,15 @@ size_t ota_fwrite(const void* ptr, size_t size, size_t count, FILE* stream) {
                 get_hit_file(cached_path, write_fault_file_name)) {
             write_fault_file_name = "";
             errno = EIO;
+            have_eio_error = true;
             return 0;
         }
     }
-    return fwrite(ptr, size, count, stream);
+    size_t status = fwrite(ptr, size, count, stream);
+    if (status != count && errno == EIO) {
+        have_eio_error = true;
+    }
+    return status;
 }
 
 ssize_t ota_write(int fd, const void* buf, size_t nbyte) {
@@ -127,10 +144,15 @@ ssize_t ota_write(int fd, const void* buf, size_t nbyte) {
                 get_hit_file(cached_path, write_fault_file_name)) {
             write_fault_file_name = "";
             errno = EIO;
+            have_eio_error = true;
             return -1;
         }
     }
-    return write(fd, buf, nbyte);
+    ssize_t status = write(fd, buf, nbyte);
+    if (status == -1 && errno == EIO) {
+        have_eio_error = true;
+    }
+    return status;
 }
 
 int ota_fsync(int fd) {
@@ -141,9 +163,14 @@ int ota_fsync(int fd) {
                 get_hit_file(cached_path, fsync_fault_file_name)) {
             fsync_fault_file_name = "";
             errno = EIO;
+            have_eio_error = true;
             return -1;
         }
     }
-    return fsync(fd);
+    int status = fsync(fd);
+    if (status == -1 && errno == EIO) {
+        have_eio_error = true;
+    }
+    return status;
 }
 
diff --git a/recovery.cpp b/recovery.cpp
index 3ce218548..90742d879 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -36,6 +36,7 @@
 
 #include <adb.h>
 #include <android-base/file.h>
+#include <android-base/parseint.h>
 #include <android-base/stringprintf.h>
 #include <cutils/android_reboot.h>
 #include <cutils/properties.h>
@@ -60,6 +61,7 @@ struct selabel_handle *sehandle;
 static const struct option OPTIONS[] = {
   { "send_intent", required_argument, NULL, 'i' },
   { "update_package", required_argument, NULL, 'u' },
+  { "retry_count", required_argument, NULL, 'n' },
   { "wipe_data", no_argument, NULL, 'w' },
   { "wipe_cache", no_argument, NULL, 'c' },
   { "show_text", no_argument, NULL, 't' },
@@ -86,6 +88,7 @@ static const char *TEMPORARY_INSTALL_FILE = "/tmp/last_install";
 static const char *LAST_KMSG_FILE = "/cache/recovery/last_kmsg";
 static const char *LAST_LOG_FILE = "/cache/recovery/last_log";
 static const int KEEP_LOG_COUNT = 10;
+static const int EIO_RETRY_COUNT = 2;
 static const int BATTERY_READ_TIMEOUT_IN_SEC = 10;
 // GmsCore enters recovery mode to install package when having enough battery
 // percentage. Normally, the threshold is 40% without charger and 20% with charger.
@@ -1120,6 +1123,29 @@ static bool is_battery_ok() {
     }
 }
 
+static void set_retry_bootloader_message(int retry_count, int argc, char** argv) {
+    struct bootloader_message boot {};
+    strlcpy(boot.command, "boot-recovery", sizeof(boot.command));
+    strlcpy(boot.recovery, "recovery\n", sizeof(boot.recovery));
+
+    for (int i = 1; i < argc; ++i) {
+        if (strstr(argv[i], "retry_count") == nullptr) {
+            strlcat(boot.recovery, argv[i], sizeof(boot.recovery));
+            strlcat(boot.recovery, "\n", sizeof(boot.recovery));
+        }
+    }
+
+    // Initialize counter to 1 if it's not in BCB, otherwise increment it by 1.
+    if (retry_count == 0) {
+        strlcat(boot.recovery, "--retry_count=1\n", sizeof(boot.recovery));
+    } else {
+        char buffer[20];
+        snprintf(buffer, sizeof(buffer), "--retry_count=%d\n", retry_count+1);
+        strlcat(boot.recovery, buffer, sizeof(boot.recovery));
+    }
+    set_bootloader_message(&boot);
+}
+
 int main(int argc, char **argv) {
     // If this binary is started with the single argument "--adbd",
     // instead of being the normal recovery binary, it turns into kind
@@ -1153,11 +1179,13 @@ int main(int argc, char **argv) {
     bool sideload_auto_reboot = false;
     bool just_exit = false;
     bool shutdown_after = false;
+    int retry_count = 0;
 
     int arg;
     while ((arg = getopt_long(argc, argv, "", OPTIONS, NULL)) != -1) {
         switch (arg) {
         case 'i': send_intent = optarg; break;
+        case 'n': android::base::ParseInt(optarg, &retry_count, 0); break;
         case 'u': update_package = optarg; break;
         case 'w': should_wipe_data = true; break;
         case 'c': should_wipe_cache = true; break;
@@ -1262,7 +1290,24 @@ int main(int argc, char **argv) {
             }
             if (status != INSTALL_SUCCESS) {
                 ui->Print("Installation aborted.\n");
-
+                // When I/O error happens, reboot and retry installation EIO_RETRY_COUNT
+                // times before we abandon this OTA update.
+                if (status == INSTALL_RETRY && retry_count < EIO_RETRY_COUNT) {
+                    copy_logs();
+                    set_retry_bootloader_message(retry_count, argc, argv);
+                    // Print retry count on screen.
+                    ui->Print("Retry attempt %d\n", retry_count);
+
+                    // Reboot and retry the update
+                    int ret = property_set(ANDROID_RB_PROPERTY, "reboot,recovery");
+                    if (ret < 0) {
+                        ui->Print("Reboot failed\n");
+                    } else {
+                        while (true) {
+                            pause();
+                        }
+                    }
+                }
                 // If this is an eng or userdebug build, then automatically
                 // turn the text display on if the script fails so the error
                 // message is visible.
diff --git a/updater/updater.cpp b/updater/updater.cpp
index efb4a8cec..1693fa1db 100644
--- a/updater/updater.cpp
+++ b/updater/updater.cpp
@@ -36,6 +36,8 @@
 // (Note it's "updateR-script", not the older "update-script".)
 #define SCRIPT_NAME "META-INF/com/google/android/updater-script"
 
+extern bool have_eio_error;
+
 struct selabel_handle *sehandle;
 
 int main(int argc, char** argv) {
@@ -141,6 +143,11 @@ int main(int argc, char** argv) {
     state.errmsg = NULL;
 
     char* result = Evaluate(&state, root);
+
+    if (have_eio_error) {
+        fprintf(cmd_pipe, "retry_update\n");
+    }
+
     if (result == NULL) {
         if (state.errmsg == NULL) {
             printf("script aborted (no error message)\n");