add the functions for multi-stage packages to updater

In order to support multi-stage recovery packages, we add the
set_stage() and get_stage() functions, which store a short string
somewhere it can be accessed across invocations of recovery.  We also
add reboot_now() which updater can invoke to immediately reboot the
device, without doing normal recovery cleanup.  (It can also choose
whether to boot off the boot or recovery partition.)

If the stage string is of the form "#/#", recovery's UI will be
augmented with a simple indicator of what stage you're in, so it
doesn't look like a reboot loop.

Change-Id: I62f7ff0bc802b549c9bcf3cc154a6bad99f94603

8 files changed, 181 insertions, 9 deletions

diff --git a/bootloader.h b/bootloader.h
index 712aa1a..c2895dd 100644
--- a/bootloader.h
+++ b/bootloader.h
@@ -38,11 +38,24 @@ extern "C" {
  * The recovery field is only written by linux and used
  * for the system to send a message to recovery or the
  * other way around.
+ *
+ * The stage field is written by packages which restart themselves
+ * multiple times, so that the UI can reflect which invocation of the
+ * package it is.  If the value is of the format "#/#" (eg, "1/3"),
+ * the UI will add a simple indicator of that status.
  */
 struct bootloader_message {
     char command[32];
     char status[32];
-    char recovery[1024];
+    char recovery[768];
+
+    // The 'recovery' field used to be 1024 bytes.  It has only ever
+    // been used to store the recovery command line, so 768 bytes
+    // should be plenty.  We carve off the last 256 bytes to store the
+    // stage string (for multistage packages) and possible future
+    // expansion.
+    char stage[32];
+    char reserved[224];
 };
 
 /* Read and write the bootloader command from the "misc" partition.
diff --git a/recovery.cpp b/recovery.cpp
index d803cad..43cd9da 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -56,6 +56,7 @@ static const struct option OPTIONS[] = {
   { "show_text", no_argument, NULL, 't' },
   { "just_exit", no_argument, NULL, 'x' },
   { "locale", required_argument, NULL, 'l' },
+  { "stages", required_argument, NULL, 'g' },
   { NULL, 0, NULL, 0 },
 };
 
@@ -76,6 +77,7 @@ static const char *SIDELOAD_TEMP_DIR = "/tmp/sideload";
 RecoveryUI* ui = NULL;
 char* locale = NULL;
 char recovery_version[PROPERTY_VALUE_MAX+1];
+char* stage = NULL;
 
 /*
  * The recovery tool communicates with the main system through /cache files.
@@ -172,6 +174,7 @@ get_args(int *argc, char ***argv) {
     struct bootloader_message boot;
     memset(&boot, 0, sizeof(boot));
     get_bootloader_message(&boot);  // this may fail, leaving a zeroed structure
+    stage = strndup(boot.stage, sizeof(boot.stage));
 
     if (boot.command[0] != 0 && boot.command[0] != 255) {
         LOGI("Boot command: %.*s\n", sizeof(boot.command), boot.command);
@@ -959,6 +962,14 @@ main(int argc, char **argv) {
         case 't': show_text = 1; break;
         case 'x': just_exit = true; break;
         case 'l': locale = optarg; break;
+        case 'g': {
+            if (stage == NULL || *stage == '\0') {
+                char buffer[20] = "1/";
+                strncat(buffer, optarg, sizeof(buffer)-3);
+                stage = strdup(buffer);
+            }
+            break;
+        }
         case '?':
             LOGE("Invalid command argument\n");
             continue;
@@ -969,12 +980,19 @@ main(int argc, char **argv) {
         load_locale_from_cache();
     }
     printf("locale is [%s]\n", locale);
+    printf("stage is [%s]\n", stage, stage);
 
     Device* device = make_device();
     ui = device->GetUI();
     gCurrentUI = ui;
 
     ui->Init();
+
+    int st_cur, st_max;
+    if (stage != NULL && sscanf(stage, "%d/%d", &st_cur, &st_max) == 2) {
+        ui->SetStage(st_cur, st_max);
+    }
+
     ui->SetLocale(locale);
     ui->SetBackground(RecoveryUI::NONE);
     if (show_text) ui->ShowText(true);
diff --git a/res/images/stage_empty.png b/res/images/stage_empty.png
new file mode 100644
index 0000000..251ec19
--- /dev/null
+++ b/res/images/stage_empty.png
diff --git a/res/images/stage_fill.png b/res/images/stage_fill.png
new file mode 100644
index 0000000..1ab79e8
--- /dev/null
+++ b/res/images/stage_fill.png
diff --git a/screen_ui.cpp b/screen_ui.cpp
index 8376341..27d0a24 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -81,7 +81,9 @@ ScreenRecoveryUI::ScreenRecoveryUI() :
     install_overlay_offset_x(13),
     install_overlay_offset_y(190),
     overlay_offset_x(-1),
-    overlay_offset_y(-1) {
+    overlay_offset_y(-1),
+    stage(-1),
+    max_stage(-1) {
 
     for (int i = 0; i < 5; i++)
         backgroundIcon[i] = NULL;
@@ -120,14 +122,28 @@ void ScreenRecoveryUI::draw_background_locked(Icon icon)
         int iconHeight = gr_get_height(surface);
         int textWidth = gr_get_width(text_surface);
         int textHeight = gr_get_height(text_surface);
+        int stageHeight = gr_get_height(stageMarkerEmpty);
+
+        int sh = (max_stage >= 0) ? stageHeight : 0;
 
         int iconX = (gr_fb_width() - iconWidth) / 2;
-        int iconY = (gr_fb_height() - (iconHeight+textHeight+40)) / 2;
+        int iconY = (gr_fb_height() - (iconHeight+textHeight+40+sh)) / 2;
 
         int textX = (gr_fb_width() - textWidth) / 2;
-        int textY = ((gr_fb_height() - (iconHeight+textHeight+40)) / 2) + iconHeight + 40;
+        int textY = ((gr_fb_height() - (iconHeight+textHeight+40+sh)) / 2) + iconHeight + 40;
 
         gr_blit(surface, 0, 0, iconWidth, iconHeight, iconX, iconY);
+        if (stageHeight > 0) {
+            int sw = gr_get_width(stageMarkerEmpty);
+            int x = (gr_fb_width() - max_stage * gr_get_width(stageMarkerEmpty)) / 2;
+            int y = iconY + iconHeight + 20;
+            for (int i = 0; i < max_stage; ++i) {
+                gr_blit((i < stage) ? stageMarkerFill : stageMarkerEmpty,
+                        0, 0, sw, stageHeight, x, y);
+                x += sw;
+            }
+        }
+
         if (icon == INSTALLING_UPDATE || icon == ERASING) {
             draw_install_overlay_locked(installingFrame);
         }
@@ -383,6 +399,8 @@ void ScreenRecoveryUI::Init()
 
     LoadBitmap("progress_empty", &progressBarEmpty);
     LoadBitmap("progress_fill", &progressBarFill);
+    LoadBitmap("stage_empty", &stageMarkerEmpty);
+    LoadBitmap("stage_fill", &stageMarkerFill);
 
     LoadLocalizedBitmap("installing_text", &backgroundText[INSTALLING_UPDATE]);
     LoadLocalizedBitmap("erasing_text", &backgroundText[ERASING]);
@@ -453,7 +471,10 @@ void ScreenRecoveryUI::SetBackground(Icon icon)
         gr_surface text = backgroundText[icon];
         overlay_offset_x = install_overlay_offset_x + (gr_fb_width() - gr_get_width(bg)) / 2;
         overlay_offset_y = install_overlay_offset_y +
-            (gr_fb_height() - (gr_get_height(bg) + gr_get_height(text) + 40)) / 2;
+            (gr_fb_height() - (gr_get_height(bg) +
+                               gr_get_height(text) +
+                               40 +
+                               ((max_stage >= 0) ? gr_get_height(stageMarkerEmpty) : 0))) / 2;
     }
 
     currentIcon = icon;
@@ -505,6 +526,13 @@ void ScreenRecoveryUI::SetProgress(float fraction)
     pthread_mutex_unlock(&updateMutex);
 }
 
+void ScreenRecoveryUI::SetStage(int current, int max) {
+    pthread_mutex_lock(&updateMutex);
+    stage = current;
+    max_stage = max;
+    pthread_mutex_unlock(&updateMutex);
+}
+
 void ScreenRecoveryUI::Print(const char *fmt, ...)
 {
     char buf[256];
diff --git a/screen_ui.h b/screen_ui.h
index fc35d95..5c4366d 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -39,6 +39,8 @@ class ScreenRecoveryUI : public RecoveryUI {
     void ShowProgress(float portion, float seconds);
     void SetProgress(float fraction);
 
+    void SetStage(int current, int max);
+
     // text log
     void ShowText(bool visible);
     bool IsTextVisible();
@@ -58,9 +60,6 @@ class ScreenRecoveryUI : public RecoveryUI {
     enum UIElement { HEADER, MENU, MENU_SEL_BG, MENU_SEL_FG, LOG, TEXT_FILL };
     virtual void SetColor(UIElement e);
 
-  protected:
-    int install_overlay_offset_x, install_overlay_offset_y;
-
   private:
     Icon currentIcon;
     int installingFrame;
@@ -73,6 +72,8 @@ class ScreenRecoveryUI : public RecoveryUI {
     gr_surface *progressBarIndeterminate;
     gr_surface progressBarEmpty;
     gr_surface progressBarFill;
+    gr_surface stageMarkerEmpty;
+    gr_surface stageMarkerFill;
 
     ProgressType progressBarType;
 
@@ -102,8 +103,13 @@ class ScreenRecoveryUI : public RecoveryUI {
     int animation_fps;
     int indeterminate_frames;
     int installing_frames;
+  protected:
+    int install_overlay_offset_x, install_overlay_offset_y;
+  private:
     int overlay_offset_x, overlay_offset_y;
 
+    int stage, max_stage;
+
     void draw_install_overlay_locked(int frame);
     void draw_background_locked(Icon icon);
     void draw_progress_locked();
diff --git a/ui.h b/ui.h
index 6c8987a..0757260 100644
--- a/ui.h
+++ b/ui.h
@@ -30,6 +30,8 @@ class RecoveryUI {
 
     // Initialize the object; called before anything else.
     virtual void Init();
+    // Show a stage indicator.  Call immediately after Init().
+    virtual void SetStage(int current, int max) { }
 
     // After calling Init(), you can tell the UI what locale it is operating in.
     virtual void SetLocale(const char* locale) { }
diff --git a/updater/install.c b/updater/install.c
index 9fb1b63..9f299f6 100644
--- a/updater/install.c
+++ b/updater/install.c
@@ -34,6 +34,9 @@
 #include <linux/xattr.h>
 #include <inttypes.h>
 
+#include "bootloader.h"
+#include "applypatch/applypatch.h"
+#include "cutils/android_reboot.h"
 #include "cutils/misc.h"
 #include "cutils/properties.h"
 #include "edify/expr.h"
@@ -42,7 +45,6 @@
 #include "mtdutils/mounts.h"
 #include "mtdutils/mtdutils.h"
 #include "updater.h"
-#include "applypatch/applypatch.h"
 
 #ifdef USE_EXT4
 #include "make_ext4fs.h"
@@ -1419,6 +1421,105 @@ Value* ReadFileFn(const char* name, State* state, int argc, Expr* argv[]) {
     return v;
 }
 
+// Immediately reboot the device.  Recovery is not finished normally,
+// so if you reboot into recovery it will re-start applying the
+// current package (because nothing has cleared the copy of the
+// arguments stored in the BCB).
+//
+// The argument is the partition name passed to the android reboot
+// property.  It can be "recovery" to boot from the recovery
+// partition, or "" (empty string) to boot from the regular boot
+// partition.
+Value* RebootNowFn(const char* name, State* state, int argc, Expr* argv[]) {
+    if (argc != 2) {
+        return ErrorAbort(state, "%s() expects 2 args, got %d", name, argc);
+    }
+
+    char* filename;
+    char* property;
+    if (ReadArgs(state, argv, 2, &filename, &property) < 0) return NULL;
+
+    char buffer[80];
+
+    // zero out the 'command' field of the bootloader message.
+    memset(buffer, 0, sizeof(((struct bootloader_message*)0)->command));
+    FILE* f = fopen(filename, "r+b");
+    fseek(f, offsetof(struct bootloader_message, command), SEEK_SET);
+    fwrite(buffer, sizeof(((struct bootloader_message*)0)->command), 1, f);
+    fclose(f);
+    free(filename);
+
+    strcpy(buffer, "reboot,");
+    if (property != NULL) {
+        strncat(buffer, property, sizeof(buffer)-10);
+    }
+
+    property_set(ANDROID_RB_PROPERTY, buffer);
+
+    sleep(5);
+    free(property);
+    ErrorAbort(state, "%s() failed to reboot", name);
+    return NULL;
+}
+
+// Store a string value somewhere that future invocations of recovery
+// can access it.  This value is called the "stage" and can be used to
+// drive packages that need to do reboots in the middle of
+// installation and keep track of where they are in the multi-stage
+// install.
+//
+// The first argument is the block device for the misc partition
+// ("/misc" in the fstab), which is where this value is stored.  The
+// second argument is the string to store; it should not exceed 31
+// bytes.
+Value* SetStageFn(const char* name, State* state, int argc, Expr* argv[]) {
+    if (argc != 2) {
+        return ErrorAbort(state, "%s() expects 2 args, got %d", name, argc);
+    }
+
+    char* filename;
+    char* stagestr;
+    if (ReadArgs(state, argv, 2, &filename, &stagestr) < 0) return NULL;
+
+    // Store this value in the misc partition, immediately after the
+    // bootloader message that the main recovery uses to save its
+    // arguments in case of the device restarting midway through
+    // package installation.
+    FILE* f = fopen(filename, "r+b");
+    fseek(f, offsetof(struct bootloader_message, stage), SEEK_SET);
+    int to_write = strlen(stagestr)+1;
+    int max_size = sizeof(((struct bootloader_message*)0)->stage);
+    if (to_write > max_size) {
+        to_write = max_size;
+        stagestr[max_size-1] = 0;
+    }
+    fwrite(stagestr, to_write, 1, f);
+    fclose(f);
+
+    free(stagestr);
+    return StringValue(filename);
+}
+
+// Return the value most recently saved with SetStageFn.  The argument
+// is the block device for the misc partition.
+Value* GetStageFn(const char* name, State* state, int argc, Expr* argv[]) {
+    if (argc != 2) {
+        return ErrorAbort(state, "%s() expects 1 arg, got %d", name, argc);
+    }
+
+    char* filename;
+    if (ReadArgs(state, argv, 1, &filename) < 0) return NULL;
+
+    char buffer[sizeof(((struct bootloader_message*)0)->stage)];
+    FILE* f = fopen(filename, "rb");
+    fseek(f, offsetof(struct bootloader_message, stage), SEEK_SET);
+    fread(buffer, sizeof(buffer), 1, f);
+    fclose(f);
+    buffer[sizeof(buffer)-1] = '\0';
+
+    return StringValue(strdup(buffer));
+}
+
 void RegisterInstallFunctions() {
     RegisterFunction("mount", MountFn);
     RegisterFunction("is_mounted", IsMountedFn);
@@ -1466,4 +1567,8 @@ void RegisterInstallFunctions() {
     RegisterFunction("ui_print", UIPrintFn);
 
     RegisterFunction("run_program", RunProgramFn);
+
+    RegisterFunction("reboot_now", RebootNowFn);
+    RegisterFunction("get_stage", GetStageFn);
+    RegisterFunction("set_stage", SetStageFn);
 }