1 files changed, 18 insertions, 0 deletions

diff --git a/recovery.cpp b/recovery.cpp
index e852ef87a..8f2183d30 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' },
   { "shutdown_after", no_argument, NULL, 'p' },
   { NULL, 0, NULL, 0 },
 };
@@ -77,6 +78,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.
@@ -173,6 +175,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", (int)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 'p': shutdown_after = true; break;
         case '?':
             LOGE("Invalid command argument\n");
@@ -970,6 +981,7 @@ 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();
@@ -977,6 +989,12 @@ main(int argc, char **argv) {
 
     ui->SetLocale(locale);
     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->SetBackground(RecoveryUI::NONE);
     if (show_text) ui->ShowText(true);