2 files changed, 34 insertions, 4 deletions

diff --git a/install.c b/install.c
index cbb35805c..ab19478cf 100644
--- a/install.c
+++ b/install.c
@@ -210,7 +210,15 @@ try_update_binary(const char *path, ZipArchive *zip) {
     //     progress bar.  The program can write single-line commands:
     //
     //        progress <frac> <secs>
-    //            fill up <frac> of the progress bar over <secs> seconds.
+    //            fill up the next <frac> part of of the progress bar
+    //            over <secs> seconds.  If <secs> is zero, use
+    //            set_progress commands to manually control the
+    //            progress of this segment of the bar
+    //
+    //        set_progress <frac>
+    //            <frac> should be between 0.0 and 1.0; sets the
+    //            progress bar within the segment defined by the most
+    //            recent progress command.
     //
     //        firmware <"hboot"|"radio"> <filename>
     //            arrange to install the contents of <filename> in the
@@ -261,6 +269,10 @@ try_update_binary(const char *path, ZipArchive *zip) {
 
             ui_show_progress(fraction * (1-VERIFICATION_PROGRESS_FRACTION),
                              seconds);
+        } else if (strcmp(command, "set_progress") == 0) {
+            char* fraction_s = strtok(NULL, " \n");
+            float fraction = strtof(fraction_s, NULL);
+            ui_set_progress(fraction);
         } else if (strcmp(command, "firmware") == 0) {
             char* type = strtok(NULL, " \n");
             char* filename = strtok(NULL, " \n");
diff --git a/updater/install.c b/updater/install.c
index 0bd09399c..e1f3c9acf 100644
--- a/updater/install.c
+++ b/updater/install.c
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-#include <stdio.h>
+#include <ctype.h>
 #include <errno.h>
 #include <stdarg.h>
+#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 #include <sys/mount.h>
@@ -265,9 +266,25 @@ char* ShowProgressFn(const char* name, State* state, int argc, Expr* argv[]) {
     UpdaterInfo* ui = (UpdaterInfo*)(state->cookie);
     fprintf(ui->cmd_pipe, "progress %f %d\n", frac, sec);
 
-    free(frac_str);
     free(sec_str);
-    return strdup("");
+    return frac_str;
+}
+
+char* SetProgressFn(const char* name, State* state, int argc, Expr* argv[]) {
+    if (argc != 1) {
+        return ErrorAbort(state, "%s() expects 1 arg, got %d", name, argc);
+    }
+    char* frac_str;
+    if (ReadArgs(state, argv, 1, &frac_str) < 0) {
+        return NULL;
+    }
+
+    double frac = strtod(frac_str, NULL);
+
+    UpdaterInfo* ui = (UpdaterInfo*)(state->cookie);
+    fprintf(ui->cmd_pipe, "set_progress %f\n", frac);
+
+    return frac_str;
 }
 
 // package_extract_dir(package_path, destination_path)
@@ -749,6 +766,7 @@ void RegisterInstallFunctions() {
     RegisterFunction("unmount", UnmountFn);
     RegisterFunction("format", FormatFn);
     RegisterFunction("show_progress", ShowProgressFn);
+    RegisterFunction("set_progress", SetProgressFn);
     RegisterFunction("delete", DeleteFn);
     RegisterFunction("delete_recursive", DeleteFn);
     RegisterFunction("package_extract_dir", PackageExtractDirFn);