Snap for 4816034 from f5f2d8cc13a2a7136f2e6c7eee6975350d72eab7 to qt-release

Change-Id: I46fbcd79edf81eff6a6d21df859a17fa621cee1c
diff --git a/Android.mk b/Android.mk
index efd7462..2992f06 100644
--- a/Android.mk
+++ b/Android.mk
@@ -145,6 +145,7 @@
 
 LOCAL_SRC_FILES := \
     adb_install.cpp \
+    fsck_unshare_blocks.cpp \
     fuse_sdcard_provider.cpp \
     install.cpp \
     recovery.cpp \
@@ -205,6 +206,13 @@
 endif
 endif
 
+# e2fsck is needed for adb remount -R.
+ifeq ($(BOARD_EXT4_SHARE_DUP_BLOCKS),true)
+ifneq (,$(filter userdebug eng,$(TARGET_BUILD_VARIANT)))
+LOCAL_REQUIRED_MODULES += e2fsck_static
+endif
+endif
+
 ifeq ($(BOARD_CACHEIMAGE_PARTITION_SIZE),)
 LOCAL_REQUIRED_MODULES += \
     recovery-persist \
@@ -217,6 +225,5 @@
     $(LOCAL_PATH)/boot_control/Android.mk \
     $(LOCAL_PATH)/minui/Android.mk \
     $(LOCAL_PATH)/tests/Android.mk \
-    $(LOCAL_PATH)/tools/Android.mk \
     $(LOCAL_PATH)/updater/Android.mk \
     $(LOCAL_PATH)/updater_sample/Android.mk \
diff --git a/fsck_unshare_blocks.cpp b/fsck_unshare_blocks.cpp
new file mode 100644
index 0000000..a100368
--- /dev/null
+++ b/fsck_unshare_blocks.cpp
@@ -0,0 +1,163 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "fsck_unshare_blocks.h"
+
+#include <errno.h>
+#include <fcntl.h>
+#include <spawn.h>
+#include <string.h>
+#include <sys/mount.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <android-base/logging.h>
+#include <android-base/properties.h>
+#include <android-base/unique_fd.h>
+#include <fstab/fstab.h>
+
+#include "roots.h"
+
+static constexpr const char* SYSTEM_E2FSCK_BIN = "/system/bin/e2fsck_static";
+static constexpr const char* TMP_E2FSCK_BIN = "/tmp/e2fsck.bin";
+
+static bool copy_file(const char* source, const char* dest) {
+  android::base::unique_fd source_fd(open(source, O_RDONLY));
+  if (source_fd < 0) {
+    PLOG(ERROR) << "open %s failed" << source;
+    return false;
+  }
+
+  android::base::unique_fd dest_fd(open(dest, O_CREAT | O_WRONLY, S_IRWXU));
+  if (dest_fd < 0) {
+    PLOG(ERROR) << "open %s failed" << dest;
+    return false;
+  }
+
+  for (;;) {
+    char buf[4096];
+    ssize_t rv = read(source_fd, buf, sizeof(buf));
+    if (rv < 0) {
+      PLOG(ERROR) << "read failed";
+      return false;
+    }
+    if (rv == 0) {
+      break;
+    }
+    if (write(dest_fd, buf, rv) != rv) {
+      PLOG(ERROR) << "write failed";
+      return false;
+    }
+  }
+  return true;
+}
+
+static bool run_e2fsck(const std::string& partition) {
+  Volume* volume = volume_for_mount_point(partition);
+  if (!volume) {
+    LOG(INFO) << "No fstab entry for " << partition << ", skipping.";
+    return true;
+  }
+
+  LOG(INFO) << "Running e2fsck on device " << volume->blk_device;
+
+  std::vector<std::string> args = { TMP_E2FSCK_BIN, "-p", "-E", "unshare_blocks",
+                                    volume->blk_device };
+  std::vector<char*> argv(args.size());
+  std::transform(args.cbegin(), args.cend(), argv.begin(),
+                 [](const std::string& arg) { return const_cast<char*>(arg.c_str()); });
+  argv.push_back(nullptr);
+
+  pid_t child;
+  char* env[] = { nullptr };
+  if (posix_spawn(&child, argv[0], nullptr, nullptr, argv.data(), env)) {
+    PLOG(ERROR) << "posix_spawn failed";
+    return false;
+  }
+
+  int status = 0;
+  int ret = TEMP_FAILURE_RETRY(waitpid(child, &status, 0));
+  if (ret < 0) {
+    PLOG(ERROR) << "waitpid failed";
+    return false;
+  }
+  if (!WIFEXITED(status)) {
+    LOG(ERROR) << "e2fsck exited abnormally: " << status;
+    return false;
+  }
+  int return_code = WEXITSTATUS(status);
+  if (return_code >= 8) {
+    LOG(ERROR) << "e2fsck could not unshare blocks: " << return_code;
+    return false;
+  }
+
+  LOG(INFO) << "Successfully unshared blocks on " << partition;
+  return true;
+}
+
+static const char* get_system_root() {
+  if (android::base::GetBoolProperty("ro.build.system_root_image", false)) {
+    return "/system_root";
+  } else {
+    return "/system";
+  }
+}
+
+bool do_fsck_unshare_blocks() {
+  // List of partitions we will try to e2fsck -E unshare_blocks.
+  std::vector<std::string> partitions = { "/odm", "/oem", "/product", "/vendor" };
+
+  // Temporarily mount system so we can copy e2fsck_static.
+  bool mounted = false;
+  if (android::base::GetBoolProperty("ro.build.system_root_image", false)) {
+    mounted = ensure_path_mounted_at("/", "/system_root") != -1;
+    partitions.push_back("/");
+  } else {
+    mounted = ensure_path_mounted("/system") != -1;
+    partitions.push_back("/system");
+  }
+  if (!mounted) {
+    LOG(ERROR) << "Failed to mount system image.";
+    return false;
+  }
+  if (!copy_file(SYSTEM_E2FSCK_BIN, TMP_E2FSCK_BIN)) {
+    LOG(ERROR) << "Could not copy e2fsck to /tmp.";
+    return false;
+  }
+  if (umount(get_system_root()) < 0) {
+    PLOG(ERROR) << "umount failed";
+    return false;
+  }
+
+  bool ok = true;
+  for (const auto& partition : partitions) {
+    ok &= run_e2fsck(partition);
+  }
+
+  if (ok) {
+    LOG(INFO) << "Finished running e2fsck.";
+  } else {
+    LOG(ERROR) << "Finished running e2fsck, but not all partitions succceeded.";
+  }
+  return ok;
+}
diff --git a/fsck_unshare_blocks.h b/fsck_unshare_blocks.h
new file mode 100644
index 0000000..9de8ef9
--- /dev/null
+++ b/fsck_unshare_blocks.h
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef _FILESYSTEM_CMDS_H
+#define _FILESYSTEM_CMDS_H
+
+bool do_fsck_unshare_blocks();
+
+#endif  // _FILESYSTEM_CMDS_H
diff --git a/recovery.cpp b/recovery.cpp
index 69b1499..8f39679 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -55,6 +55,7 @@
 #include "adb_install.h"
 #include "common.h"
 #include "device.h"
+#include "fsck_unshare_blocks.h"
 #include "fuse_sdcard_provider.h"
 #include "fuse_sideload.h"
 #include "install.h"
@@ -989,6 +990,7 @@
                  [](const std::string& arg) { return const_cast<char*>(arg.c_str()); });
 
   static constexpr struct option OPTIONS[] = {
+    { "fsck_unshare_blocks", no_argument, nullptr, 0 },
     { "just_exit", no_argument, nullptr, 'x' },
     { "locale", required_argument, nullptr, 0 },
     { "prompt_and_wipe_data", no_argument, nullptr, 0 },
@@ -1017,6 +1019,7 @@
   bool sideload_auto_reboot = false;
   bool just_exit = false;
   bool shutdown_after = false;
+  bool fsck_unshare_blocks = false;
   int retry_count = 0;
   bool security_update = false;
   std::string locale;
@@ -1034,7 +1037,9 @@
         break;
       case 0: {
         std::string option = OPTIONS[option_index].name;
-        if (option == "locale") {
+        if (option == "fsck_unshare_blocks") {
+          fsck_unshare_blocks = true;
+        } else if (option == "locale") {
           // Handled in recovery_main.cpp
         } else if (option == "prompt_and_wipe_data") {
           should_prompt_and_wipe_data = true;
@@ -1201,6 +1206,10 @@
     if (sideload_auto_reboot) {
       ui->Print("Rebooting automatically.\n");
     }
+  } else if (fsck_unshare_blocks) {
+    if (!do_fsck_unshare_blocks()) {
+      status = INSTALL_ERROR;
+    }
   } else if (!just_exit) {
     // If this is an eng or userdebug build, automatically turn on the text display if no command
     // is specified. Note that this should be called before setting the background to avoid
diff --git a/screen_ui.cpp b/screen_ui.cpp
index f1b3878..b9aba80 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -19,8 +19,6 @@
 #include <dirent.h>
 #include <errno.h>
 #include <fcntl.h>
-#include <linux/input.h>
-#include <pthread.h>
 #include <stdarg.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -32,8 +30,10 @@
 #include <unistd.h>
 
 #include <algorithm>
+#include <chrono>
 #include <memory>
 #include <string>
+#include <thread>
 #include <unordered_map>
 #include <vector>
 
@@ -169,8 +169,12 @@
       stage(-1),
       max_stage(-1),
       locale_(""),
-      rtl_locale_(false),
-      updateMutex(PTHREAD_MUTEX_INITIALIZER) {}
+      rtl_locale_(false) {}
+
+ScreenRecoveryUI::~ScreenRecoveryUI() {
+  progress_thread_stopped_ = true;
+  progress_thread_.join();
+}
 
 GRSurface* ScreenRecoveryUI::GetCurrentFrame() const {
   if (currentIcon == INSTALLING_UPDATE || currentIcon == ERASING) {
@@ -361,7 +365,7 @@
     surfaces.emplace(name, std::unique_ptr<GRSurface, decltype(&free)>(text_image, &free));
   }
 
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   gr_color(0, 0, 0, 255);
   gr_clear();
 
@@ -393,7 +397,6 @@
   }
   // Update the whole screen.
   gr_flip();
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::CheckBackgroundTextImages() {
@@ -613,52 +616,46 @@
   gr_flip();
 }
 
-// Keeps the progress bar updated, even when the process is otherwise busy.
-void* ScreenRecoveryUI::ProgressThreadStartRoutine(void* data) {
-  reinterpret_cast<ScreenRecoveryUI*>(data)->ProgressThreadLoop();
-  return nullptr;
-}
-
 void ScreenRecoveryUI::ProgressThreadLoop() {
   double interval = 1.0 / kAnimationFps;
-  while (true) {
+  while (!progress_thread_stopped_) {
     double start = now();
-    pthread_mutex_lock(&updateMutex);
-
     bool redraw = false;
+    {
+      std::lock_guard<std::mutex> lg(updateMutex);
 
-    // update the installation animation, if active
-    // skip this if we have a text overlay (too expensive to update)
-    if ((currentIcon == INSTALLING_UPDATE || currentIcon == ERASING) && !show_text) {
-      if (!intro_done) {
-        if (current_frame == intro_frames - 1) {
-          intro_done = true;
-          current_frame = 0;
+      // update the installation animation, if active
+      // skip this if we have a text overlay (too expensive to update)
+      if ((currentIcon == INSTALLING_UPDATE || currentIcon == ERASING) && !show_text) {
+        if (!intro_done) {
+          if (current_frame == intro_frames - 1) {
+            intro_done = true;
+            current_frame = 0;
+          } else {
+            ++current_frame;
+          }
         } else {
-          ++current_frame;
+          current_frame = (current_frame + 1) % loop_frames;
         }
-      } else {
-        current_frame = (current_frame + 1) % loop_frames;
-      }
 
-      redraw = true;
-    }
-
-    // move the progress bar forward on timed intervals, if configured
-    int duration = progressScopeDuration;
-    if (progressBarType == DETERMINATE && duration > 0) {
-      double elapsed = now() - progressScopeTime;
-      float p = 1.0 * elapsed / duration;
-      if (p > 1.0) p = 1.0;
-      if (p > progress) {
-        progress = p;
         redraw = true;
       }
+
+      // move the progress bar forward on timed intervals, if configured
+      int duration = progressScopeDuration;
+      if (progressBarType == DETERMINATE && duration > 0) {
+        double elapsed = now() - progressScopeTime;
+        float p = 1.0 * elapsed / duration;
+        if (p > 1.0) p = 1.0;
+        if (p > progress) {
+          progress = p;
+          redraw = true;
+        }
+      }
+
+      if (redraw) update_progress_locked();
     }
 
-    if (redraw) update_progress_locked();
-
-    pthread_mutex_unlock(&updateMutex);
     double end = now();
     // minimum of 20ms delay between frames
     double delay = interval - (end - start);
@@ -749,7 +746,8 @@
 
   LoadAnimation();
 
-  pthread_create(&progress_thread_, nullptr, ProgressThreadStartRoutine, this);
+  // Keep the progress bar updated, even when the process is otherwise busy.
+  progress_thread_ = std::thread(&ScreenRecoveryUI::ProgressThreadLoop, this);
 
   return true;
 }
@@ -797,16 +795,14 @@
 }
 
 void ScreenRecoveryUI::SetBackground(Icon icon) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
 
   currentIcon = icon;
   update_screen_locked();
-
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::SetProgressType(ProgressType type) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (progressBarType != type) {
     progressBarType = type;
   }
@@ -814,11 +810,10 @@
   progressScopeSize = 0;
   progress = 0;
   update_progress_locked();
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::ShowProgress(float portion, float seconds) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   progressBarType = DETERMINATE;
   progressScopeStart += progressScopeSize;
   progressScopeSize = portion;
@@ -826,11 +821,10 @@
   progressScopeDuration = seconds;
   progress = 0;
   update_progress_locked();
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::SetProgress(float fraction) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (fraction < 0.0) fraction = 0.0;
   if (fraction > 1.0) fraction = 1.0;
   if (progressBarType == DETERMINATE && fraction > progress) {
@@ -842,14 +836,12 @@
       update_progress_locked();
     }
   }
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::SetStage(int current, int max) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   stage = current;
   max_stage = max;
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::PrintV(const char* fmt, bool copy_to_stdout, va_list ap) {
@@ -860,7 +852,7 @@
     fputs(str.c_str(), stdout);
   }
 
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (text_rows_ > 0 && text_cols_ > 0) {
     for (const char* ptr = str.c_str(); *ptr != '\0'; ++ptr) {
       if (*ptr == '\n' || text_col_ >= text_cols_) {
@@ -873,7 +865,6 @@
     text_[text_row_][text_col_] = '\0';
     update_screen_locked();
   }
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::Print(const char* fmt, ...) {
@@ -891,23 +882,21 @@
 }
 
 void ScreenRecoveryUI::PutChar(char ch) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (ch != '\n') text_[text_row_][text_col_++] = ch;
   if (ch == '\n' || text_col_ >= text_cols_) {
     text_col_ = 0;
     ++text_row_;
   }
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::ClearText() {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   text_col_ = 0;
   text_row_ = 0;
   for (size_t i = 0; i < text_rows_; ++i) {
     memset(text_[i], 0, text_cols_ + 1);
   }
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::ShowFile(FILE* fp) {
@@ -984,17 +973,16 @@
 
 void ScreenRecoveryUI::StartMenu(const std::vector<std::string>& headers,
                                  const std::vector<std::string>& items, size_t initial_selection) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (text_rows_ > 0 && text_cols_ > 1) {
     menu_ = std::make_unique<Menu>(scrollable_menu_, text_rows_, text_cols_ - 1, headers, items,
                                    initial_selection);
     update_screen_locked();
   }
-  pthread_mutex_unlock(&updateMutex);
 }
 
 int ScreenRecoveryUI::SelectMenu(int sel) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (menu_) {
     int old_sel = menu_->selection();
     sel = menu_->Select(sel);
@@ -1003,17 +991,15 @@
       update_screen_locked();
     }
   }
-  pthread_mutex_unlock(&updateMutex);
   return sel;
 }
 
 void ScreenRecoveryUI::EndMenu() {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (menu_) {
     menu_.reset();
     update_screen_locked();
   }
-  pthread_mutex_unlock(&updateMutex);
 }
 
 size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
@@ -1065,31 +1051,27 @@
 }
 
 bool ScreenRecoveryUI::IsTextVisible() {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   int visible = show_text;
-  pthread_mutex_unlock(&updateMutex);
   return visible;
 }
 
 bool ScreenRecoveryUI::WasTextEverVisible() {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   int ever_visible = show_text_ever;
-  pthread_mutex_unlock(&updateMutex);
   return ever_visible;
 }
 
 void ScreenRecoveryUI::ShowText(bool visible) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   show_text = visible;
   if (show_text) show_text_ever = true;
   update_screen_locked();
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::Redraw() {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   update_screen_locked();
-  pthread_mutex_unlock(&updateMutex);
 }
 
 void ScreenRecoveryUI::KeyLongPress(int) {
diff --git a/screen_ui.h b/screen_ui.h
index c90a2cd..b76d470 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -17,12 +17,13 @@
 #ifndef RECOVERY_SCREEN_UI_H
 #define RECOVERY_SCREEN_UI_H
 
-#include <pthread.h>
 #include <stdio.h>
 
+#include <atomic>
 #include <functional>
 #include <memory>
 #include <string>
+#include <thread>
 #include <vector>
 
 #include "ui.h"
@@ -112,6 +113,7 @@
 
   ScreenRecoveryUI();
   explicit ScreenRecoveryUI(bool scrollable_menu);
+  ~ScreenRecoveryUI() override;
 
   bool Init(const std::string& locale) override;
   std::string GetLocale() const override;
@@ -189,7 +191,6 @@
   GRSurface* GetCurrentFrame() const;
   GRSurface* GetCurrentText() const;
 
-  static void* ProgressThreadStartRoutine(void* data);
   void ProgressThreadLoop();
 
   virtual void ShowFile(FILE*);
@@ -275,7 +276,8 @@
   // An alternate text screen, swapped with 'text_' when we're viewing a log file.
   char** file_viewer_text_;
 
-  pthread_t progress_thread_;
+  std::thread progress_thread_;
+  std::atomic<bool> progress_thread_stopped_{ false };
 
   // Number of intro frames and loop frames in the animation.
   size_t intro_frames;
@@ -293,7 +295,7 @@
   std::string locale_;
   bool rtl_locale_;
 
-  pthread_mutex_t updateMutex;
+  std::mutex updateMutex;
 
  private:
   void SetLocale(const std::string&);
diff --git a/tests/Android.mk b/tests/Android.mk
index efe46b8..cee94dc 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -37,6 +37,7 @@
 
 LOCAL_SRC_FILES := \
     unit/asn1_decoder_test.cpp \
+    unit/commands_test.cpp \
     unit/dirutil_test.cpp \
     unit/locale_test.cpp \
     unit/rangeset_test.cpp \
diff --git a/tests/unit/commands_test.cpp b/tests/unit/commands_test.cpp
new file mode 100644
index 0000000..18aa471
--- /dev/null
+++ b/tests/unit/commands_test.cpp
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <string>
+
+#include <gtest/gtest.h>
+
+#include "private/commands.h"
+
+TEST(CommandsTest, ParseType) {
+  ASSERT_EQ(Command::Type::ZERO, Command::ParseType("zero"));
+  ASSERT_EQ(Command::Type::NEW, Command::ParseType("new"));
+  ASSERT_EQ(Command::Type::ERASE, Command::ParseType("erase"));
+  ASSERT_EQ(Command::Type::MOVE, Command::ParseType("move"));
+  ASSERT_EQ(Command::Type::BSDIFF, Command::ParseType("bsdiff"));
+  ASSERT_EQ(Command::Type::IMGDIFF, Command::ParseType("imgdiff"));
+  ASSERT_EQ(Command::Type::STASH, Command::ParseType("stash"));
+  ASSERT_EQ(Command::Type::FREE, Command::ParseType("free"));
+}
+
+TEST(CommandsTest, ParseType_InvalidCommand) {
+  ASSERT_EQ(Command::Type::LAST, Command::ParseType("foo"));
+  ASSERT_EQ(Command::Type::LAST, Command::ParseType("bar"));
+}
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index 269222f..2562307 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -279,8 +279,6 @@
     testdata_dir_ = from_testdata_base("");
     Paths::Get().set_resource_dir(testdata_dir_);
     res_set_resource_dir(testdata_dir_);
-
-    ASSERT_TRUE(ui_->Init(kTestLocale));
   }
 
   std::unique_ptr<TestableScreenRecoveryUI> ui_;
@@ -288,6 +286,7 @@
 };
 
 TEST_F(ScreenRecoveryUITest, Init) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   ASSERT_EQ(kTestLocale, ui_->GetLocale());
   ASSERT_FALSE(ui_->GetRtlLocale());
   ASSERT_FALSE(ui_->IsTextVisible());
@@ -295,6 +294,7 @@
 }
 
 TEST_F(ScreenRecoveryUITest, ShowText) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   ASSERT_FALSE(ui_->IsTextVisible());
   ui_->ShowText(true);
   ASSERT_TRUE(ui_->IsTextVisible());
@@ -308,12 +308,15 @@
 TEST_F(ScreenRecoveryUITest, RtlLocale) {
   ASSERT_TRUE(ui_->Init(kTestRtlLocale));
   ASSERT_TRUE(ui_->GetRtlLocale());
+}
 
+TEST_F(ScreenRecoveryUITest, RtlLocaleWithSuffix) {
   ASSERT_TRUE(ui_->Init(kTestRtlLocaleWithSuffix));
   ASSERT_TRUE(ui_->GetRtlLocale());
 }
 
 TEST_F(ScreenRecoveryUITest, ShowMenu) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   ui_->SetKeyBuffer({
       KeyCode::UP,
       KeyCode::DOWN,
@@ -339,6 +342,7 @@
 }
 
 TEST_F(ScreenRecoveryUITest, ShowMenu_NotMenuOnly) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   ui_->SetKeyBuffer({
       KeyCode::MAGIC,
   });
@@ -349,6 +353,7 @@
 }
 
 TEST_F(ScreenRecoveryUITest, ShowMenu_TimedOut) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   ui_->SetKeyBuffer({
       KeyCode::TIMEOUT,
   });
@@ -356,6 +361,7 @@
 }
 
 TEST_F(ScreenRecoveryUITest, ShowMenu_TimedOut_TextWasEverVisible) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   ui_->ShowText(true);
   ui_->ShowText(false);
   ASSERT_TRUE(ui_->WasTextEverVisible());
@@ -371,6 +377,7 @@
 }
 
 TEST_F(ScreenRecoveryUITest, LoadAnimation) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   // Make a few copies of loop00000.png from testdata.
   std::string image_data;
   ASSERT_TRUE(android::base::ReadFileToString(testdata_dir_ + "/loop00000.png", &image_data));
@@ -398,6 +405,7 @@
 }
 
 TEST_F(ScreenRecoveryUITest, LoadAnimation_MissingAnimation) {
+  ASSERT_TRUE(ui_->Init(kTestLocale));
   TemporaryDir resource_dir;
   Paths::Get().set_resource_dir(resource_dir.path);
   ASSERT_EXIT(ui_->RunLoadAnimation(), ::testing::KilledBySignal(SIGABRT), "");
diff --git a/tools/Android.mk b/tools/Android.mk
deleted file mode 100644
index 6571161..0000000
--- a/tools/Android.mk
+++ /dev/null
@@ -1 +0,0 @@
-include $(all-subdir-makefiles)
diff --git a/tools/dumpkey/Android.bp b/tools/dumpkey/Android.bp
new file mode 100644
index 0000000..eb45e31
--- /dev/null
+++ b/tools/dumpkey/Android.bp
@@ -0,0 +1,27 @@
+// Copyright (C) 2018 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+java_library_host {
+    name: "dumpkey",
+
+    manifest: "DumpPublicKey.mf",
+
+    srcs: [
+        "DumpPublicKey.java",
+    ],
+
+    static_libs: [
+        "bouncycastle-host",
+    ],
+}
diff --git a/tools/dumpkey/Android.mk b/tools/dumpkey/Android.mk
deleted file mode 100644
index 3154914..0000000
--- a/tools/dumpkey/Android.mk
+++ /dev/null
@@ -1,22 +0,0 @@
-# Copyright (C) 2008 The Android Open Source Project
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-LOCAL_PATH := $(call my-dir)
-
-include $(CLEAR_VARS)
-LOCAL_MODULE := dumpkey
-LOCAL_SRC_FILES := DumpPublicKey.java
-LOCAL_JAR_MANIFEST := DumpPublicKey.mf
-LOCAL_STATIC_JAVA_LIBRARIES := bouncycastle-host
-include $(BUILD_HOST_JAVA_LIBRARY)
diff --git a/tools/recovery_l10n/Android.bp b/tools/recovery_l10n/Android.bp
new file mode 100644
index 0000000..d0a6d4b
--- /dev/null
+++ b/tools/recovery_l10n/Android.bp
@@ -0,0 +1,23 @@
+// Copyright (C) 2018 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+android_app {
+    name: "RecoveryLocalizer",
+
+    sdk_version: "current",
+
+    srcs: [
+        "src/**/*.java",
+    ],
+}
diff --git a/tools/recovery_l10n/Android.mk b/tools/recovery_l10n/Android.mk
deleted file mode 100644
index 7197c5c..0000000
--- a/tools/recovery_l10n/Android.mk
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright 2012 Google Inc. All Rights Reserved.
-
-LOCAL_PATH := $(call my-dir)
-
-include $(CLEAR_VARS)
-
-LOCAL_PACKAGE_NAME := RecoveryLocalizer
-LOCAL_SDK_VERSION := current
-LOCAL_MODULE_TAGS := optional
-
-LOCAL_SRC_FILES := $(call all-java-files-under, src)
-
-include $(BUILD_PACKAGE)
diff --git a/ui.cpp b/ui.cpp
index 8983d76..51d7f12 100644
--- a/ui.cpp
+++ b/ui.cpp
@@ -18,31 +18,30 @@
 
 #include <errno.h>
 #include <fcntl.h>
-#include <linux/input.h>
-#include <pthread.h>
-#include <stdarg.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#include <sys/stat.h>
 #include <sys/time.h>
 #include <sys/types.h>
 #include <time.h>
 #include <unistd.h>
 
+#include <chrono>
 #include <functional>
 #include <string>
+#include <thread>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/parseint.h>
 #include <android-base/strings.h>
-#include <minui/minui.h>
 
-#include "device.h"
+#include "minui/minui.h"
 #include "otautil/sysutil.h"
 #include "roots.h"
 
+using namespace std::chrono_literals;
+
 static constexpr int UI_WAIT_KEY_TIMEOUT_SEC = 120;
 static constexpr const char* BRIGHTNESS_FILE = "/sys/class/leds/lcd-backlight/brightness";
 static constexpr const char* MAX_BRIGHTNESS_FILE = "/sys/class/leds/lcd-backlight/max_brightness";
@@ -73,11 +72,15 @@
       touch_slot_(0),
       is_bootreason_recovery_ui_(false),
       screensaver_state_(ScreensaverState::DISABLED) {
-  pthread_mutex_init(&key_queue_mutex, nullptr);
-  pthread_cond_init(&key_queue_cond, nullptr);
   memset(key_pressed, 0, sizeof(key_pressed));
 }
 
+RecoveryUI::~RecoveryUI() {
+  ev_exit();
+  input_thread_stopped_ = true;
+  input_thread_.join();
+}
+
 void RecoveryUI::OnKeyDetected(int key_code) {
   if (key_code == KEY_POWER) {
     has_power_key = true;
@@ -90,16 +93,6 @@
   }
 }
 
-// Reads input events, handles special hot keys, and adds to the key queue.
-static void* InputThreadLoop(void*) {
-  while (true) {
-    if (!ev_wait(-1)) {
-      ev_dispatch();
-    }
-  }
-  return nullptr;
-}
-
 bool RecoveryUI::InitScreensaver() {
   // Disabled.
   if (brightness_normal_ == 0 || brightness_dimmed_ > brightness_normal_) {
@@ -166,7 +159,15 @@
     LOG(INFO) << "Screensaver disabled";
   }
 
-  pthread_create(&input_thread_, nullptr, InputThreadLoop, nullptr);
+  // Create a separate thread that handles input events.
+  input_thread_ = std::thread([this]() {
+    while (!this->input_thread_stopped_) {
+      if (!ev_wait(500)) {
+        ev_dispatch();
+      }
+    }
+  });
+
   return true;
 }
 
@@ -323,46 +324,38 @@
   return 0;
 }
 
-// Process a key-up or -down event.  A key is "registered" when it is
-// pressed and then released, with no other keypresses or releases in
-// between.  Registered keys are passed to CheckKey() to see if it
-// should trigger a visibility toggle, an immediate reboot, or be
-// queued to be processed next time the foreground thread wants a key
-// (eg, for the menu).
+// Processes a key-up or -down event. A key is "registered" when it is pressed and then released,
+// with no other keypresses or releases in between. Registered keys are passed to CheckKey() to
+// see if it should trigger a visibility toggle, an immediate reboot, or be queued to be processed
+// next time the foreground thread wants a key (eg, for the menu).
 //
-// We also keep track of which keys are currently down so that
-// CheckKey can call IsKeyPressed to see what other keys are held when
-// a key is registered.
+// We also keep track of which keys are currently down so that CheckKey() can call IsKeyPressed()
+// to see what other keys are held when a key is registered.
 //
 // updown == 1 for key down events; 0 for key up events
 void RecoveryUI::ProcessKey(int key_code, int updown) {
   bool register_key = false;
   bool long_press = false;
-  bool reboot_enabled;
 
-  pthread_mutex_lock(&key_queue_mutex);
-  key_pressed[key_code] = updown;
-  if (updown) {
-    ++key_down_count;
-    key_last_down = key_code;
-    key_long_press = false;
-    key_timer_t* info = new key_timer_t;
-    info->ui = this;
-    info->key_code = key_code;
-    info->count = key_down_count;
-    pthread_t thread;
-    pthread_create(&thread, nullptr, &RecoveryUI::time_key_helper, info);
-    pthread_detach(thread);
-  } else {
-    if (key_last_down == key_code) {
-      long_press = key_long_press;
-      register_key = true;
+  {
+    std::lock_guard<std::mutex> lg(key_queue_mutex);
+    key_pressed[key_code] = updown;
+    if (updown) {
+      ++key_down_count;
+      key_last_down = key_code;
+      key_long_press = false;
+      std::thread time_key_thread(&RecoveryUI::TimeKey, this, key_code, key_down_count);
+      time_key_thread.detach();
+    } else {
+      if (key_last_down == key_code) {
+        long_press = key_long_press;
+        register_key = true;
+      }
+      key_last_down = -1;
     }
-    key_last_down = -1;
   }
-  reboot_enabled = enable_reboot;
-  pthread_mutex_unlock(&key_queue_mutex);
 
+  bool reboot_enabled = enable_reboot;
   if (register_key) {
     switch (CheckKey(key_code, long_press)) {
       case RecoveryUI::IGNORE:
@@ -388,54 +381,40 @@
   }
 }
 
-void* RecoveryUI::time_key_helper(void* cookie) {
-  key_timer_t* info = static_cast<key_timer_t*>(cookie);
-  info->ui->time_key(info->key_code, info->count);
-  delete info;
-  return nullptr;
-}
-
-void RecoveryUI::time_key(int key_code, int count) {
-  usleep(750000);  // 750 ms == "long"
+void RecoveryUI::TimeKey(int key_code, int count) {
+  std::this_thread::sleep_for(750ms);  // 750 ms == "long"
   bool long_press = false;
-  pthread_mutex_lock(&key_queue_mutex);
-  if (key_last_down == key_code && key_down_count == count) {
-    long_press = key_long_press = true;
+  {
+    std::lock_guard<std::mutex> lg(key_queue_mutex);
+    if (key_last_down == key_code && key_down_count == count) {
+      long_press = key_long_press = true;
+    }
   }
-  pthread_mutex_unlock(&key_queue_mutex);
   if (long_press) KeyLongPress(key_code);
 }
 
 void RecoveryUI::EnqueueKey(int key_code) {
-  pthread_mutex_lock(&key_queue_mutex);
+  std::lock_guard<std::mutex> lg(key_queue_mutex);
   const int queue_max = sizeof(key_queue) / sizeof(key_queue[0]);
   if (key_queue_len < queue_max) {
     key_queue[key_queue_len++] = key_code;
-    pthread_cond_signal(&key_queue_cond);
+    key_queue_cond.notify_one();
   }
-  pthread_mutex_unlock(&key_queue_mutex);
 }
 
 int RecoveryUI::WaitKey() {
-  pthread_mutex_lock(&key_queue_mutex);
+  std::unique_lock<std::mutex> lk(key_queue_mutex);
 
   // Time out after UI_WAIT_KEY_TIMEOUT_SEC, unless a USB cable is
   // plugged in.
   do {
-    struct timeval now;
-    struct timespec timeout;
-    gettimeofday(&now, nullptr);
-    timeout.tv_sec = now.tv_sec;
-    timeout.tv_nsec = now.tv_usec * 1000;
-    timeout.tv_sec += UI_WAIT_KEY_TIMEOUT_SEC;
-
-    int rc = 0;
-    while (key_queue_len == 0 && rc != ETIMEDOUT) {
-      rc = pthread_cond_timedwait(&key_queue_cond, &key_queue_mutex, &timeout);
+    std::cv_status rc = std::cv_status::no_timeout;
+    while (key_queue_len == 0 && rc != std::cv_status::timeout) {
+      rc = key_queue_cond.wait_for(lk, std::chrono::seconds(UI_WAIT_KEY_TIMEOUT_SEC));
     }
 
     if (screensaver_state_ != ScreensaverState::DISABLED) {
-      if (rc == ETIMEDOUT) {
+      if (rc == std::cv_status::timeout) {
         // Lower the brightness level: NORMAL -> DIMMED; DIMMED -> OFF.
         if (screensaver_state_ == ScreensaverState::NORMAL) {
           if (android::base::WriteStringToFile(std::to_string(brightness_dimmed_value_),
@@ -474,7 +453,6 @@
     key = key_queue[0];
     memcpy(&key_queue[0], &key_queue[1], sizeof(int) * --key_queue_len);
   }
-  pthread_mutex_unlock(&key_queue_mutex);
   return key;
 }
 
@@ -495,16 +473,14 @@
 }
 
 bool RecoveryUI::IsKeyPressed(int key) {
-  pthread_mutex_lock(&key_queue_mutex);
+  std::lock_guard<std::mutex> lg(key_queue_mutex);
   int pressed = key_pressed[key];
-  pthread_mutex_unlock(&key_queue_mutex);
   return pressed;
 }
 
 bool RecoveryUI::IsLongPress() {
-  pthread_mutex_lock(&key_queue_mutex);
+  std::lock_guard<std::mutex> lg(key_queue_mutex);
   bool result = key_long_press;
-  pthread_mutex_unlock(&key_queue_mutex);
   return result;
 }
 
@@ -521,15 +497,15 @@
 }
 
 void RecoveryUI::FlushKeys() {
-  pthread_mutex_lock(&key_queue_mutex);
+  std::lock_guard<std::mutex> lg(key_queue_mutex);
   key_queue_len = 0;
-  pthread_mutex_unlock(&key_queue_mutex);
 }
 
 RecoveryUI::KeyAction RecoveryUI::CheckKey(int key, bool is_long_press) {
-  pthread_mutex_lock(&key_queue_mutex);
-  key_long_press = false;
-  pthread_mutex_unlock(&key_queue_mutex);
+  {
+    std::lock_guard<std::mutex> lg(key_queue_mutex);
+    key_long_press = false;
+  }
 
   // If we have power and volume up keys, that chord is the signal to toggle the text display.
   if (HasThreeButtons() || (HasPowerKey() && HasTouchScreen() && touch_screen_allowed_)) {
@@ -552,9 +528,7 @@
 
   // Press power seven times in a row to reboot.
   if (key == KEY_POWER) {
-    pthread_mutex_lock(&key_queue_mutex);
     bool reboot_enabled = enable_reboot;
-    pthread_mutex_unlock(&key_queue_mutex);
 
     if (reboot_enabled) {
       ++consecutive_power_keys;
@@ -574,7 +548,6 @@
 }
 
 void RecoveryUI::SetEnableReboot(bool enabled) {
-  pthread_mutex_lock(&key_queue_mutex);
+  std::lock_guard<std::mutex> lg(key_queue_mutex);
   enable_reboot = enabled;
-  pthread_mutex_unlock(&key_queue_mutex);
 }
diff --git a/ui.h b/ui.h
index a74b14f..32e2809 100644
--- a/ui.h
+++ b/ui.h
@@ -17,12 +17,14 @@
 #ifndef RECOVERY_UI_H
 #define RECOVERY_UI_H
 
-#include <linux/input.h>
-#include <pthread.h>
-#include <time.h>
+#include <linux/input.h>  // KEY_MAX
 
+#include <atomic>
+#include <condition_variable>
 #include <functional>
+#include <mutex>
 #include <string>
+#include <thread>
 #include <vector>
 
 // Abstract class for controlling the user interface during recovery.
@@ -51,7 +53,7 @@
 
   RecoveryUI();
 
-  virtual ~RecoveryUI() {}
+  virtual ~RecoveryUI();
 
   // Initializes the object; called before anything else. UI texts will be initialized according to
   // the given locale. Returns true on success.
@@ -172,12 +174,6 @@
     OFF
   };
 
-  struct key_timer_t {
-    RecoveryUI* ui;
-    int key_code;
-    int count;
-  };
-
   // The sensitivity when detecting a swipe.
   const int kTouchLowThreshold;
   const int kTouchHighThreshold;
@@ -186,17 +182,15 @@
   void OnTouchDetected(int dx, int dy);
   int OnInputEvent(int fd, uint32_t epevents);
   void ProcessKey(int key_code, int updown);
+  void TimeKey(int key_code, int count);
 
   bool IsUsbConnected();
 
-  static void* time_key_helper(void* cookie);
-  void time_key(int key_code, int count);
-
   bool InitScreensaver();
 
   // Key event input queue
-  pthread_mutex_t key_queue_mutex;
-  pthread_cond_t key_queue_cond;
+  std::mutex key_queue_mutex;
+  std::condition_variable key_queue_cond;
   int key_queue[256], key_queue_len;
   char key_pressed[KEY_MAX + 1];  // under key_queue_mutex
   int key_last_down;              // under key_queue_mutex
@@ -223,7 +217,8 @@
   bool touch_swiping_;
   bool is_bootreason_recovery_ui_;
 
-  pthread_t input_thread_;
+  std::thread input_thread_;
+  std::atomic<bool> input_thread_stopped_{ false };
 
   ScreensaverState screensaver_state_;
 
diff --git a/updater/Android.mk b/updater/Android.mk
index 4762664..46c56f4 100644
--- a/updater/Android.mk
+++ b/updater/Android.mk
@@ -56,6 +56,7 @@
 LOCAL_MODULE := libupdater
 
 LOCAL_SRC_FILES := \
+    commands.cpp \
     install.cpp \
     blockimg.cpp
 
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index 4a70b98..4adb974 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -57,6 +57,7 @@
 #include "otautil/paths.h"
 #include "otautil/print_sha1.h"
 #include "otautil/rangeset.h"
+#include "private/commands.h"
 #include "updater/install.h"
 #include "updater/updater.h"
 
@@ -546,8 +547,8 @@
 struct CommandParameters {
     std::vector<std::string> tokens;
     size_t cpos;
-    const char* cmdname;
-    const char* cmdline;
+    std::string cmdname;
+    std::string cmdline;
     std::string freestash;
     std::string stashbase;
     bool canwrite;
@@ -750,7 +751,7 @@
   }
 }
 
-static int LoadStash(CommandParameters& params, const std::string& id, bool verify, size_t* blocks,
+static int LoadStash(CommandParameters& params, const std::string& id, bool verify,
                      std::vector<uint8_t>& buffer, bool printnoent) {
   // In verify mode, if source range_set was saved for the given hash, check contents in the source
   // blocks first. If the check fails, search for the stashed files on /cache as usual.
@@ -772,11 +773,6 @@
     }
   }
 
-  size_t blockcount = 0;
-  if (!blocks) {
-    blocks = &blockcount;
-  }
-
   std::string fn = GetStashFileName(params.stashbase, id, "");
 
   struct stat sb;
@@ -807,9 +803,8 @@
     return -1;
   }
 
-  *blocks = sb.st_size / BLOCKSIZE;
-
-  if (verify && VerifyBlocks(id, buffer, *blocks, true) != 0) {
+  size_t blocks = sb.st_size / BLOCKSIZE;
+  if (verify && VerifyBlocks(id, buffer, blocks, true) != 0) {
     LOG(ERROR) << "unexpected contents in " << fn;
     if (stash_map.find(id) == stash_map.end()) {
       LOG(ERROR) << "failed to find source blocks number for stash " << id
@@ -1055,7 +1050,7 @@
     }
 
     std::vector<uint8_t> stash;
-    if (LoadStash(params, tokens[0], false, nullptr, stash, true) == -1) {
+    if (LoadStash(params, tokens[0], false, stash, true) == -1) {
       // These source blocks will fail verification if used later, but we
       // will let the caller decide if this is a fatal failure
       LOG(ERROR) << "failed to load stash " << tokens[0];
@@ -1170,7 +1165,7 @@
     return 0;
   }
 
-  if (*overlap && LoadStash(params, srchash, true, nullptr, params.buffer, true) == 0) {
+  if (*overlap && LoadStash(params, srchash, true, params.buffer, true) == 0) {
     // Overlapping source blocks were previously stashed, command can proceed. We are recovering
     // from an interrupted command, so we don't know if the stash can safely be deleted after this
     // command.
@@ -1236,8 +1231,7 @@
   }
 
   const std::string& id = params.tokens[params.cpos++];
-  size_t blocks = 0;
-  if (LoadStash(params, id, true, &blocks, params.buffer, false) == 0) {
+  if (LoadStash(params, id, true, params.buffer, false) == 0) {
     // Stash file already exists and has expected contents. Do not read from source again, as the
     // source may have been already overwritten during a previous attempt.
     return 0;
@@ -1246,11 +1240,11 @@
   RangeSet src = RangeSet::Parse(params.tokens[params.cpos++]);
   CHECK(static_cast<bool>(src));
 
-  allocate(src.blocks() * BLOCKSIZE, params.buffer);
+  size_t blocks = src.blocks();
+  allocate(blocks * BLOCKSIZE, params.buffer);
   if (ReadBlocks(src, params.buffer, params.fd) == -1) {
     return -1;
   }
-  blocks = src.blocks();
   stash_map[id] = src;
 
   if (VerifyBlocks(id, params.buffer, blocks, true) != 0) {
@@ -1496,23 +1490,13 @@
   return 0;
 }
 
-// Definitions for transfer list command functions
-typedef int (*CommandFunction)(CommandParameters&);
+using CommandFunction = std::function<int(CommandParameters&)>;
 
-struct Command {
-    const char* name;
-    CommandFunction f;
-};
-
-// args:
-//    - block device (or file) to modify in-place
-//    - transfer list (blob)
-//    - new data stream (filename within package.zip)
-//    - patch stream (filename within package.zip, must be uncompressed)
+using CommandMap = std::unordered_map<Command::Type, CommandFunction>;
 
 static Value* PerformBlockImageUpdate(const char* name, State* state,
                                       const std::vector<std::unique_ptr<Expr>>& argv,
-                                      const Command* commands, size_t cmdcount, bool dryrun) {
+                                      const CommandMap& command_map, bool dryrun) {
   CommandParameters params = {};
   params.canwrite = !dryrun;
 
@@ -1532,6 +1516,11 @@
     return nullptr;
   }
 
+  // args:
+  //   - block device (or file) to modify in-place
+  //   - transfer list (blob)
+  //   - new data stream (filename within package.zip)
+  //   - patch stream (filename within package.zip, must be uncompressed)
   const std::unique_ptr<Value>& blockdev_filename = args[0];
   const std::unique_ptr<Value>& transfer_list_value = args[1];
   const std::unique_ptr<Value>& new_data_fn = args[2];
@@ -1707,16 +1696,6 @@
     skip_executed_command = false;
   }
 
-  // Build a map of the available commands
-  std::unordered_map<std::string, const Command*> cmd_map;
-  for (size_t i = 0; i < cmdcount; ++i) {
-    if (cmd_map.find(commands[i].name) != cmd_map.end()) {
-      LOG(ERROR) << "Error: command [" << commands[i].name << "] already exists in the cmd map.";
-      return StringValue("");
-    }
-    cmd_map[commands[i].name] = &commands[i];
-  }
-
   int rc = -1;
 
   static constexpr size_t kTransferListHeaderLines = 4;
@@ -1728,36 +1707,35 @@
     size_t cmdindex = i - kTransferListHeaderLines;
     params.tokens = android::base::Split(line, " ");
     params.cpos = 0;
-    params.cmdname = params.tokens[params.cpos++].c_str();
-    params.cmdline = line.c_str();
+    params.cmdname = params.tokens[params.cpos++];
+    params.cmdline = line;
     params.target_verified = false;
 
-    if (cmd_map.find(params.cmdname) == cmd_map.end()) {
+    Command::Type cmd_type = Command::ParseType(params.cmdname);
+    if (cmd_type == Command::Type::LAST) {
       LOG(ERROR) << "unexpected command [" << params.cmdname << "]";
       goto pbiudone;
     }
 
-    const Command* cmd = cmd_map[params.cmdname];
+    const CommandFunction& performer = command_map.at(cmd_type);
 
     // Skip the command if we explicitly set the corresponding function pointer to nullptr, e.g.
     // "erase" during block_image_verify.
-    if (cmd->f == nullptr) {
+    if (performer == nullptr) {
       LOG(DEBUG) << "skip executing command [" << line << "]";
       continue;
     }
 
-    std::string cmdname = std::string(params.cmdname);
-
     // Skip all commands before the saved last command index when resuming an update, except for
     // "new" command. Because new commands read in the data sequentially.
     if (params.canwrite && skip_executed_command && cmdindex <= saved_last_command_index &&
-        cmdname != "new") {
+        cmd_type != Command::Type::NEW) {
       LOG(INFO) << "Skipping already executed command: " << cmdindex
                 << ", last executed command for previous update: " << saved_last_command_index;
       continue;
     }
 
-    if (cmd->f(params) == -1) {
+    if (performer(params) == -1) {
       LOG(ERROR) << "failed to execute command [" << line << "]";
       goto pbiudone;
     }
@@ -1767,7 +1745,8 @@
     // that we will resume the update from the first command in the transfer list.
     if (!params.canwrite && skip_executed_command && cmdindex <= saved_last_command_index) {
       // TODO(xunchang) check that the cmdline of the saved index is correct.
-      if ((cmdname == "move" || cmdname == "bsdiff" || cmdname == "imgdiff") &&
+      if ((cmd_type == Command::Type::MOVE || cmd_type == Command::Type::BSDIFF ||
+           cmd_type == Command::Type::IMGDIFF) &&
           !params.target_verified) {
         LOG(WARNING) << "Previously executed command " << saved_last_command_index << ": "
                      << params.cmdline << " doesn't produce expected target blocks.";
@@ -1775,6 +1754,7 @@
         DeleteLastCommandFile();
       }
     }
+
     if (params.canwrite) {
       if (ota_fsync(params.fd) == -1) {
         failure_type = kFsyncFailure;
@@ -1911,38 +1891,42 @@
  */
 Value* BlockImageVerifyFn(const char* name, State* state,
                           const std::vector<std::unique_ptr<Expr>>& argv) {
-    // Commands which are not tested are set to nullptr to skip them completely
-    const Command commands[] = {
-        { "bsdiff",     PerformCommandDiff  },
-        { "erase",      nullptr             },
-        { "free",       PerformCommandFree  },
-        { "imgdiff",    PerformCommandDiff  },
-        { "move",       PerformCommandMove  },
-        { "new",        nullptr             },
-        { "stash",      PerformCommandStash },
-        { "zero",       nullptr             }
-    };
+  // Commands which are not allowed are set to nullptr to skip them completely.
+  const CommandMap command_map{
+    // clang-format off
+    { Command::Type::BSDIFF,  PerformCommandDiff },
+    { Command::Type::ERASE,   nullptr },
+    { Command::Type::FREE,    PerformCommandFree },
+    { Command::Type::IMGDIFF, PerformCommandDiff },
+    { Command::Type::MOVE,    PerformCommandMove },
+    { Command::Type::NEW,     nullptr },
+    { Command::Type::STASH,   PerformCommandStash },
+    { Command::Type::ZERO,    nullptr },
+    // clang-format on
+  };
+  CHECK_EQ(static_cast<size_t>(Command::Type::LAST), command_map.size());
 
-    // Perform a dry run without writing to test if an update can proceed
-    return PerformBlockImageUpdate(name, state, argv, commands,
-                sizeof(commands) / sizeof(commands[0]), true);
+  // Perform a dry run without writing to test if an update can proceed.
+  return PerformBlockImageUpdate(name, state, argv, command_map, true);
 }
 
 Value* BlockImageUpdateFn(const char* name, State* state,
                           const std::vector<std::unique_ptr<Expr>>& argv) {
-    const Command commands[] = {
-        { "bsdiff",     PerformCommandDiff  },
-        { "erase",      PerformCommandErase },
-        { "free",       PerformCommandFree  },
-        { "imgdiff",    PerformCommandDiff  },
-        { "move",       PerformCommandMove  },
-        { "new",        PerformCommandNew   },
-        { "stash",      PerformCommandStash },
-        { "zero",       PerformCommandZero  }
-    };
+  const CommandMap command_map{
+    // clang-format off
+    { Command::Type::BSDIFF,  PerformCommandDiff },
+    { Command::Type::ERASE,   PerformCommandErase },
+    { Command::Type::FREE,    PerformCommandFree },
+    { Command::Type::IMGDIFF, PerformCommandDiff },
+    { Command::Type::MOVE,    PerformCommandMove },
+    { Command::Type::NEW,     PerformCommandNew },
+    { Command::Type::STASH,   PerformCommandStash },
+    { Command::Type::ZERO,    PerformCommandZero },
+    // clang-format on
+  };
+  CHECK_EQ(static_cast<size_t>(Command::Type::LAST), command_map.size());
 
-    return PerformBlockImageUpdate(name, state, argv, commands,
-                sizeof(commands) / sizeof(commands[0]), false);
+  return PerformBlockImageUpdate(name, state, argv, command_map, false);
 }
 
 Value* RangeSha1Fn(const char* name, State* state, const std::vector<std::unique_ptr<Expr>>& argv) {
diff --git a/updater/commands.cpp b/updater/commands.cpp
new file mode 100644
index 0000000..f798c6a
--- /dev/null
+++ b/updater/commands.cpp
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "private/commands.h"
+
+#include <string>
+
+#include <android-base/logging.h>
+
+Command::Type Command::ParseType(const std::string& type_str) {
+  if (type_str == "zero") {
+    return Type::ZERO;
+  } else if (type_str == "new") {
+    return Type::NEW;
+  } else if (type_str == "erase") {
+    return Type::ERASE;
+  } else if (type_str == "move") {
+    return Type::MOVE;
+  } else if (type_str == "bsdiff") {
+    return Type::BSDIFF;
+  } else if (type_str == "imgdiff") {
+    return Type::IMGDIFF;
+  } else if (type_str == "stash") {
+    return Type::STASH;
+  } else if (type_str == "free") {
+    return Type::FREE;
+  }
+  LOG(ERROR) << "Invalid type: " << type_str;
+  return Type::LAST;
+};
diff --git a/updater/include/private/commands.h b/updater/include/private/commands.h
new file mode 100644
index 0000000..b360000
--- /dev/null
+++ b/updater/include/private/commands.h
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <string>
+
+struct Command {
+  enum class Type {
+    ZERO,
+    NEW,
+    ERASE,
+    MOVE,
+    BSDIFF,
+    IMGDIFF,
+    STASH,
+    FREE,
+    LAST,  // Not a valid type.
+  };
+
+  static Type ParseType(const std::string& type_str);
+};
diff --git a/updater_sample/Android.mk b/updater_sample/Android.mk
index 056ad66..7662111 100644
--- a/updater_sample/Android.mk
+++ b/updater_sample/Android.mk
@@ -18,8 +18,8 @@
 include $(CLEAR_VARS)
 
 LOCAL_PACKAGE_NAME := SystemUpdaterSample
-LOCAL_SDK_VERSION := system_current
 LOCAL_MODULE_TAGS := samples
+LOCAL_SDK_VERSION := system_current
 
 # TODO: enable proguard and use proguard.flags file
 LOCAL_PROGUARD_ENABLED := disabled
diff --git a/updater_sample/README.md b/updater_sample/README.md
index 3f211dd..f6c63a7 100644
--- a/updater_sample/README.md
+++ b/updater_sample/README.md
@@ -65,6 +65,32 @@
 6. Push OTA packages to the device.
 
 
+## Sample App State vs UpdateEngine Status
+
+UpdateEngine provides status for different stages of update application
+process. But it lacks of proper status codes when update fails.
+
+This creates two problems:
+
+1. If sample app is unbound from update_engine (MainActivity is paused, destroyed),
+   app doesn't receive onStatusUpdate and onPayloadApplicationCompleted notifications.
+   If app binds to update_engine after update is completed,
+   only onStatusUpdate is called, but status becomes IDLE in most cases.
+   And there is no way to know if update was successful or not.
+
+2. This sample app demostrates suspend/resume using update_engins's
+   `cancel` and `applyPayload` (which picks up from where it left).
+   When `cancel` is called, status is set to `IDLE`, which doesn't allow
+   tracking suspended state properly.
+
+To solve these problems sample app implements its own separate update
+state - `UpdaterState`. To solve the first problem, sample app persists
+`UpdaterState` on a device. When app is resumed, it checks if `UpdaterState`
+matches the update_engine's status (as onStatusUpdate is guaranteed to be called).
+If they doesn't match, sample app calls `applyPayload` again with the same
+parameters, and handles update completion properly using `onPayloadApplicationCompleted`
+callback. The second problem is solved by adding `PAUSED` updater state.
+
 ## Sending HTTP headers from UpdateEngine
 
 Sometimes OTA package server might require some HTTP headers to be present,
@@ -76,6 +102,44 @@
 which HTTP headers are supported.
 
 
+## Used update_engine APIs
+
+### UpdateEngine#bind
+
+Binds given callbacks to update_engine. When update_engine successfully
+initialized, it's guaranteed to invoke callback onStatusUpdate.
+
+### UpdateEngine#applyPayload
+
+Start an update attempt to download an apply the provided `payload_url` if
+no other update is running. The extra `key_value_pair_headers` will be
+included when fetching the payload.
+
+### UpdateEngine#cancel
+
+Cancel the ongoing update. The update could be running or suspended, but it
+can't be canceled after it was done.
+
+### UpdateEngine#resetStatus
+
+Reset the already applied update back to an idle state. This method can
+only be called when no update attempt is going on, and it will reset the
+status back to idle, deleting the currently applied update if any.
+
+### Callback: onStatusUpdate
+
+Called whenever the value of `status` or `progress` changes. For
+`progress` values changes, this method will be called only if it changes significantly.
+At this time of writing this doc, delta for `progress` is `0.005`.
+
+`onStatusUpdate` is always called when app binds to update_engine,
+except when update_engine fails to initialize.
+
+### Callback: onPayloadApplicationComplete
+
+Called whenever an update attempt is completed.
+
+
 ## Development
 
 - [x] Create a UI with list of configs, current version,
@@ -90,6 +154,10 @@
 - [x] Add demo for passing HTTP headers to `UpdateEngine#applyPayload`
 - [x] [Package compatibility check](https://source.android.com/devices/architecture/vintf/match-rules)
 - [x] Deferred switch slot demo
+- [x] Add UpdateManager; extract update logic from MainActivity
+- [x] Add Sample app update state (separate from update_engine status)
+- [-] Add smart update completion detection using onStatusUpdate
+- [ ] Add pause/resume demo
 - [ ] Add demo for passing NETWORK_ID to `UpdateEngine#applyPayload`
 - [ ] Verify system partition checksum for package
 - [?] Add non-A/B updates demo
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
index c370a4e..145cc83 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -25,11 +25,13 @@
 import com.example.android.systemupdatersample.util.PayloadSpecs;
 import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
 import com.example.android.systemupdatersample.util.UpdateEngineProperties;
-import com.example.android.systemupdatersample.util.UpdaterStates;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.AtomicDouble;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -37,6 +39,8 @@
 import java.util.function.DoubleConsumer;
 import java.util.function.IntConsumer;
 
+import javax.annotation.concurrent.GuardedBy;
+
 /**
  * Manages the update flow. It has its own state (in memory), separate from
  * {@link UpdateEngine}'s state. Asynchronously interacts with the {@link UpdateEngine}.
@@ -56,22 +60,27 @@
             new AtomicInteger(UpdateEngine.UpdateStatusConstants.IDLE);
     private AtomicInteger mEngineErrorCode = new AtomicInteger(UpdateEngineErrorCodes.UNKNOWN);
     private AtomicDouble mProgress = new AtomicDouble(0);
+    private UpdaterState mUpdaterState = new UpdaterState(UpdaterState.IDLE);
 
-    private AtomicInteger mState = new AtomicInteger(UpdaterStates.IDLE);
-
-    private final UpdateManager.UpdateEngineCallbackImpl
-            mUpdateEngineCallback = new UpdateManager.UpdateEngineCallbackImpl();
-
-    private PayloadSpec mLastPayloadSpec;
     private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
 
+    @GuardedBy("mLock")
+    private UpdateData mLastUpdateData = null;
+
+    @GuardedBy("mLock")
     private IntConsumer mOnStateChangeCallback = null;
+    @GuardedBy("mLock")
     private IntConsumer mOnEngineStatusUpdateCallback = null;
+    @GuardedBy("mLock")
     private DoubleConsumer mOnProgressUpdateCallback = null;
+    @GuardedBy("mLock")
     private IntConsumer mOnEngineCompleteCallback = null;
 
     private final Object mLock = new Object();
 
+    private final UpdateManager.UpdateEngineCallbackImpl
+            mUpdateEngineCallback = new UpdateManager.UpdateEngineCallbackImpl();
+
     public UpdateManager(UpdateEngine updateEngine, PayloadSpecs payloadSpecs) {
         this.mUpdateEngine = updateEngine;
         this.mPayloadSpecs = payloadSpecs;
@@ -108,7 +117,7 @@
 
     /**
      * Sets SystemUpdaterSample app state change callback. Value of {@code state} will be one
-     * of the values from {@link UpdaterStates}.
+     * of the values from {@link UpdaterState}.
      *
      * @param onStateChangeCallback a callback with parameter {@code state}.
      */
@@ -190,8 +199,14 @@
      * it also notifies {@link this.mOnStateChangeCallback}.
      */
     private void setUpdaterState(int updaterState) {
-        int previousState = mState.get();
-        mState.set(updaterState);
+        int previousState = mUpdaterState.get();
+        try {
+            mUpdaterState.set(updaterState);
+        } catch (UpdaterState.InvalidTransitionException e) {
+            // Note: invalid state transitions should be handled properly,
+            //       but to make sample app simple, we just throw runtime exception.
+            throw new RuntimeException("Can't set state " + updaterState, e);
+        }
         if (previousState != updaterState) {
             getOnStateChangeCallback().ifPresent(callback -> callback.accept(updaterState));
         }
@@ -208,7 +223,7 @@
     public void cancelRunningUpdate() {
         try {
             mUpdateEngine.cancel();
-            setUpdaterState(UpdaterStates.IDLE);
+            setUpdaterState(UpdaterState.IDLE);
         } catch (Exception e) {
             Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
         }
@@ -224,7 +239,7 @@
     public void resetUpdate() {
         try {
             mUpdateEngine.resetStatus();
-            setUpdaterState(UpdaterStates.IDLE);
+            setUpdaterState(UpdaterState.IDLE);
         } catch (Exception e) {
             Log.w(TAG, "UpdateEngine failed to reset the update", e);
         }
@@ -238,7 +253,12 @@
      */
     public void applyUpdate(Context context, UpdateConfig config) {
         mEngineErrorCode.set(UpdateEngineErrorCodes.UNKNOWN);
-        setUpdaterState(UpdaterStates.RUNNING);
+        setUpdaterState(UpdaterState.RUNNING);
+
+        synchronized (mLock) {
+            // Cleaning up previous update data.
+            mLastUpdateData = null;
+        }
 
         if (!config.getAbConfig().getForceSwitchSlot()) {
             mManualSwitchSlotRequired.set(true);
@@ -254,33 +274,35 @@
     }
 
     private void applyAbNonStreamingUpdate(UpdateConfig config) {
-        List<String> extraProperties = prepareExtraProperties(config);
+        UpdateData.Builder builder = UpdateData.builder()
+                .setExtraProperties(prepareExtraProperties(config));
 
-        PayloadSpec payload;
         try {
-            payload = mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile());
+            builder.setPayload(mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile()));
         } catch (IOException e) {
             Log.e(TAG, "Error creating payload spec", e);
-            setUpdaterState(UpdaterStates.ERROR);
+            setUpdaterState(UpdaterState.ERROR);
             return;
         }
-        updateEngineApplyPayload(payload, extraProperties);
+        updateEngineApplyPayload(builder.build());
     }
 
     private void applyAbStreamingUpdate(Context context, UpdateConfig config) {
-        List<String> extraProperties = prepareExtraProperties(config);
+        UpdateData.Builder builder = UpdateData.builder()
+                .setExtraProperties(prepareExtraProperties(config));
 
         Log.d(TAG, "Starting PrepareStreamingService");
         PrepareStreamingService.startService(context, config, (code, payloadSpec) -> {
             if (code == PrepareStreamingService.RESULT_CODE_SUCCESS) {
-                extraProperties.add("USER_AGENT=" + HTTP_USER_AGENT);
+                builder.setPayload(payloadSpec);
+                builder.addExtraProperty("USER_AGENT=" + HTTP_USER_AGENT);
                 config.getStreamingMetadata()
                         .getAuthorization()
-                        .ifPresent(s -> extraProperties.add("AUTHORIZATION=" + s));
-                updateEngineApplyPayload(payloadSpec, extraProperties);
+                        .ifPresent(s -> builder.addExtraProperty("AUTHORIZATION=" + s));
+                updateEngineApplyPayload(builder.build());
             } else {
                 Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
-                setUpdaterState(UpdaterStates.ERROR);
+                setUpdaterState(UpdaterState.ERROR);
             }
         });
     }
@@ -305,29 +327,40 @@
      * <p>It's possible that the update engine throws a generic error, such as upon seeing invalid
      * payload properties (which come from OTA packages), or failing to set up the network
      * with the given id.</p>
-     *
-     * @param payloadSpec contains url, offset and size to {@code PAYLOAD_BINARY_FILE_NAME}
-     * @param extraProperties additional properties to pass to {@link UpdateEngine#applyPayload}
      */
-    private void updateEngineApplyPayload(PayloadSpec payloadSpec, List<String> extraProperties) {
-        mLastPayloadSpec = payloadSpec;
-
-        ArrayList<String> properties = new ArrayList<>(payloadSpec.getProperties());
-        if (extraProperties != null) {
-            properties.addAll(extraProperties);
+    private void updateEngineApplyPayload(UpdateData update) {
+        synchronized (mLock) {
+            mLastUpdateData = update;
         }
+
+        ArrayList<String> properties = new ArrayList<>(update.getPayload().getProperties());
+        properties.addAll(update.getExtraProperties());
+
         try {
             mUpdateEngine.applyPayload(
-                    payloadSpec.getUrl(),
-                    payloadSpec.getOffset(),
-                    payloadSpec.getSize(),
+                    update.getPayload().getUrl(),
+                    update.getPayload().getOffset(),
+                    update.getPayload().getSize(),
                     properties.toArray(new String[0]));
         } catch (Exception e) {
             Log.e(TAG, "UpdateEngine failed to apply the update", e);
-            setUpdaterState(UpdaterStates.ERROR);
+            setUpdaterState(UpdaterState.ERROR);
         }
     }
 
+    private void updateEngineReApplyPayload() {
+        UpdateData lastUpdate;
+        synchronized (mLock) {
+            // mLastPayloadSpec might be empty in some cases.
+            // But to make this sample app simple, we will not handle it.
+            Preconditions.checkArgument(
+                    mLastUpdateData != null,
+                    "mLastUpdateData must be present.");
+            lastUpdate = mLastUpdateData;
+        }
+        updateEngineApplyPayload(lastUpdate);
+    }
+
     /**
      * Sets the new slot that has the updated partitions as the active slot,
      * which device will boot into next time.
@@ -342,19 +375,101 @@
      */
     public void setSwitchSlotOnReboot() {
         Log.d(TAG, "setSwitchSlotOnReboot invoked");
-        List<String> extraProperties = new ArrayList<>();
+        UpdateData.Builder builder;
+        synchronized (mLock) {
+            // To make sample app simple, we don't handle it.
+            Preconditions.checkArgument(
+                    mLastUpdateData != null,
+                    "mLastUpdateData must be present.");
+            builder = mLastUpdateData.toBuilder();
+        }
         // PROPERTY_SKIP_POST_INSTALL should be passed on to skip post-installation hooks.
-        extraProperties.add(UpdateEngineProperties.PROPERTY_SKIP_POST_INSTALL);
-        // It sets property SWITCH_SLOT_ON_REBOOT=1 by default.
+        builder.setExtraProperties(
+                Collections.singletonList(UpdateEngineProperties.PROPERTY_SKIP_POST_INSTALL));
+        // UpdateEngine sets property SWITCH_SLOT_ON_REBOOT=1 by default.
         // HTTP headers are not required, UpdateEngine is not expected to stream payload.
-        updateEngineApplyPayload(mLastPayloadSpec, extraProperties);
+        updateEngineApplyPayload(builder.build());
     }
 
+    /**
+     * Verifies if mUpdaterState matches mUpdateEngineStatus.
+     * If they don't match, runs applyPayload to trigger onPayloadApplicationComplete
+     * callback, which updates mUpdaterState.
+     */
+    private void ensureCorrectUpdaterState() {
+        // When mUpdaterState is one of IDLE, PAUSED, ERROR, SLOT_SWITCH_REQUIRED
+        //    then mUpdateEngineStatus must be IDLE.
+        // When mUpdaterState is RUNNING,
+        //    then mUpdateEngineStatus must not be IDLE or UPDATED_NEED_REBOOT.
+        // When mUpdaterState is REBOOT_REQUIRED,
+        //    then mUpdateEngineStatus must be UPDATED_NEED_REBOOT.
+        int state = mUpdaterState.get();
+        int updateEngineStatus = mUpdateEngineStatus.get();
+        if (state == UpdaterState.IDLE
+                || state == UpdaterState.ERROR
+                || state == UpdaterState.PAUSED
+                || state == UpdaterState.SLOT_SWITCH_REQUIRED) {
+            ensureUpdateEngineStatusIdle(state, updateEngineStatus);
+        } else if (state == UpdaterState.RUNNING) {
+            ensureUpdateEngineStatusRunning(state, updateEngineStatus);
+        } else if (state == UpdaterState.REBOOT_REQUIRED) {
+            ensureUpdateEngineStatusReboot(state, updateEngineStatus);
+        }
+    }
+
+    private void ensureUpdateEngineStatusIdle(int state, int updateEngineStatus) {
+        if (updateEngineStatus == UpdateEngine.UpdateStatusConstants.IDLE) {
+            return;
+        }
+        // It might happen when update is started not from the sample app.
+        // To make the sample app simple, we won't handle this case.
+        throw new RuntimeException("When mUpdaterState is " + state
+                + " mUpdateEngineStatus expected to be "
+                + UpdateEngine.UpdateStatusConstants.IDLE
+                + ", but it is " + updateEngineStatus);
+    }
+
+    private void ensureUpdateEngineStatusRunning(int state, int updateEngineStatus) {
+        if (updateEngineStatus != UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
+                && updateEngineStatus != UpdateEngine.UpdateStatusConstants.IDLE) {
+            return;
+        }
+        // Re-apply latest update. It makes update_engine to invoke
+        // onPayloadApplicationComplete callback. The callback notifies
+        // if update was successful or not.
+        updateEngineReApplyPayload();
+    }
+
+    private void ensureUpdateEngineStatusReboot(int state, int updateEngineStatus) {
+        if (updateEngineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT) {
+            return;
+        }
+        // This might happen when update is installed by other means,
+        // and sample app is not aware of it. To make the sample app simple,
+        // we won't handle this case.
+        throw new RuntimeException("When mUpdaterState is " + state
+                + " mUpdateEngineStatus expected to be "
+                + UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
+                + ", but it is " + updateEngineStatus);
+    }
+
+    /**
+     * Invoked by update_engine whenever update status or progress changes.
+     * It's also guaranteed to be invoked when app binds to the update_engine, except
+     * when update_engine fails to initialize (as defined in
+     * system/update_engine/binder_service_android.cc in
+     * function BinderUpdateEngineAndroidService::bind).
+     *
+     * @param status one of {@link UpdateEngine.UpdateStatusConstants}.
+     * @param progress a number from 0.0 to 1.0.
+     */
     private void onStatusUpdate(int status, float progress) {
         int previousStatus = mUpdateEngineStatus.get();
         mUpdateEngineStatus.set(status);
         mProgress.set(progress);
 
+        ensureCorrectUpdaterState();
+
         getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(progress));
 
         if (previousStatus != status) {
@@ -367,9 +482,11 @@
         mEngineErrorCode.set(errorCode);
         if (errorCode == UpdateEngine.ErrorCodeConstants.SUCCESS
                 || errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
-            setUpdaterState(UpdaterStates.FINISHED);
+            setUpdaterState(isManualSwitchSlotRequired()
+                    ? UpdaterState.SLOT_SWITCH_REQUIRED
+                    : UpdaterState.REBOOT_REQUIRED);
         } else if (errorCode != UpdateEngineErrorCodes.USER_CANCELLED) {
-            setUpdaterState(UpdaterStates.ERROR);
+            setUpdaterState(UpdaterState.ERROR);
         }
 
         getOnEngineCompleteCallback()
@@ -377,7 +494,7 @@
     }
 
     /**
-     * Helper class to delegate {@code update_engine} callbacks to UpdateManager
+     * Helper class to delegate {@code update_engine} callback invocations to UpdateManager.
      */
     class UpdateEngineCallbackImpl extends UpdateEngineCallback {
         @Override
@@ -391,4 +508,67 @@
         }
     }
 
+    /**
+     *
+     * Contains update data - PayloadSpec and extra properties list.
+     *
+     * <p>{@code mPayload} contains url, offset and size to {@code PAYLOAD_BINARY_FILE_NAME}.
+     * {@code mExtraProperties} is a list of additional properties to pass to
+     * {@link UpdateEngine#applyPayload}.</p>
+     */
+    private static class UpdateData {
+        private final PayloadSpec mPayload;
+        private final ImmutableList<String> mExtraProperties;
+
+        public static Builder builder() {
+            return new Builder();
+        }
+
+        UpdateData(Builder builder) {
+            this.mPayload = builder.mPayload;
+            this.mExtraProperties = ImmutableList.copyOf(builder.mExtraProperties);
+        }
+
+        public PayloadSpec getPayload() {
+            return mPayload;
+        }
+
+        public ImmutableList<String> getExtraProperties() {
+            return mExtraProperties;
+        }
+
+        public Builder toBuilder() {
+            return builder()
+                    .setPayload(mPayload)
+                    .setExtraProperties(mExtraProperties);
+        }
+
+        static class Builder {
+            private PayloadSpec mPayload;
+            private List<String> mExtraProperties;
+
+            public Builder setPayload(PayloadSpec payload) {
+                this.mPayload = payload;
+                return this;
+            }
+
+            public Builder setExtraProperties(List<String> extraProperties) {
+                this.mExtraProperties = new ArrayList<>(extraProperties);
+                return this;
+            }
+
+            public Builder addExtraProperty(String property) {
+                if (this.mExtraProperties == null) {
+                    this.mExtraProperties = new ArrayList<>();
+                }
+                this.mExtraProperties.add(property);
+                return this;
+            }
+
+            public UpdateData build() {
+                return new UpdateData(this);
+            }
+        }
+    }
+
 }
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
new file mode 100644
index 0000000..36a9098
--- /dev/null
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
@@ -0,0 +1,103 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample;
+
+import android.util.SparseArray;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Controls updater state.
+ */
+public class UpdaterState {
+
+    public static final int IDLE = 0;
+    public static final int ERROR = 1;
+    public static final int RUNNING = 2;
+    public static final int PAUSED = 3;
+    public static final int SLOT_SWITCH_REQUIRED = 4;
+    public static final int REBOOT_REQUIRED = 5;
+
+    private static final SparseArray<String> STATE_MAP = new SparseArray<>();
+
+    static {
+        STATE_MAP.put(0, "IDLE");
+        STATE_MAP.put(1, "ERROR");
+        STATE_MAP.put(2, "RUNNING");
+        STATE_MAP.put(3, "PAUSED");
+        STATE_MAP.put(4, "SLOT_SWITCH_REQUIRED");
+        STATE_MAP.put(5, "REBOOT_REQUIRED");
+    }
+
+    /**
+     * Allowed state transitions. It's a map: key is a state, value is a set of states that
+     * are allowed to transition to from key.
+     */
+    private static final ImmutableMap<Integer, ImmutableSet<Integer>> TRANSITIONS =
+            ImmutableMap.of(
+                    IDLE, ImmutableSet.of(RUNNING),
+                    RUNNING, ImmutableSet.of(ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED),
+                    PAUSED, ImmutableSet.of(RUNNING),
+                    SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR)
+            );
+
+    private AtomicInteger mState;
+
+    public UpdaterState(int state) {
+        this.mState = new AtomicInteger(state);
+    }
+
+    /**
+     * Returns updater state.
+     */
+    public int get() {
+        return mState.get();
+    }
+
+    /**
+     * Sets the updater state.
+     *
+     * @throws InvalidTransitionException if transition is not allowed.
+     */
+    public void set(int newState) throws InvalidTransitionException {
+        int oldState = mState.get();
+        if (!TRANSITIONS.get(oldState).contains(newState)) {
+            throw new InvalidTransitionException(
+                    "Can't transition from " + oldState + " to " + newState);
+        }
+        mState.set(newState);
+    }
+
+    /**
+     * Converts status code to status name.
+     */
+    public static String getStateText(int state) {
+        return STATE_MAP.get(state);
+    }
+
+    /**
+     * Defines invalid state transition exception.
+     */
+    public static class InvalidTransitionException extends Exception {
+        public InvalidTransitionException(String msg) {
+            super(msg);
+        }
+    }
+}
diff --git a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
index 9983fe3..1de72c2 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -33,11 +33,11 @@
 import com.example.android.systemupdatersample.R;
 import com.example.android.systemupdatersample.UpdateConfig;
 import com.example.android.systemupdatersample.UpdateManager;
+import com.example.android.systemupdatersample.UpdaterState;
 import com.example.android.systemupdatersample.util.PayloadSpecs;
 import com.example.android.systemupdatersample.util.UpdateConfigs;
 import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
 import com.example.android.systemupdatersample.util.UpdateEngineStatuses;
-import com.example.android.systemupdatersample.util.UpdaterStates;
 
 import java.util.List;
 
@@ -108,12 +108,16 @@
     @Override
     protected void onResume() {
         super.onResume();
+        // TODO(zhomart) load saved states
+        // Binding to UpdateEngine invokes onStatusUpdate callback,
+        // persisted updater state has to be loaded and prepared beforehand.
         this.mUpdateManager.bind();
     }
 
     @Override
     protected void onPause() {
         this.mUpdateManager.unbind();
+        // TODO(zhomart) save state
         super.onPause();
     }
 
@@ -192,7 +196,7 @@
     /**
      * Invoked when SystemUpdaterSample app state changes.
      * Value of {@code state} will be one of the
-     * values from {@link UpdaterStates}.
+     * values from {@link UpdaterState}.
      */
     private void onUpdaterStateChange(int state) {
         Log.i(TAG, "onUpdaterStateChange invoked state=" + state);
@@ -233,8 +237,8 @@
         runOnUiThread(() -> {
             Log.i(TAG,
                     "Completed - errorCode="
-                    + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
-                    + " " + completionState);
+                            + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
+                            + " " + completionState);
             setUiEngineErrorCode(errorCode);
             if (errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
                 // if update was successfully applied.
@@ -323,7 +327,7 @@
      * @param state updater sample state
      */
     private void setUiUpdaterState(int state) {
-        String stateText = UpdaterStates.getStateText(state);
+        String stateText = UpdaterState.getStateText(state);
         mTextViewUpdaterState.setText(stateText + "/" + state);
     }
 
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java b/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java
deleted file mode 100644
index fc20a79..0000000
--- a/updater_sample/src/com/example/android/systemupdatersample/util/UpdaterStates.java
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.example.android.systemupdatersample.util;
-
-import android.util.SparseArray;
-
-/**
- * SystemUpdaterSample app state.
- */
-public class UpdaterStates {
-
-    public static final int IDLE = 0;
-    public static final int ERROR = 1;
-    public static final int RUNNING = 2;
-    public static final int PAUSED = 3;
-    public static final int FINISHED = 4;
-
-    private static final SparseArray<String> STATE_MAP = new SparseArray<>();
-
-    static {
-        STATE_MAP.put(0, "IDLE");
-        STATE_MAP.put(1, "ERROR");
-        STATE_MAP.put(2, "RUNNING");
-        STATE_MAP.put(3, "PAUSED");
-        STATE_MAP.put(4, "FINISHED");
-    }
-
-    /**
-     * converts status code to status name
-     */
-    public static String getStateText(int state) {
-        return STATE_MAP.get(state);
-    }
-
-    private UpdaterStates() {}
-}
diff --git a/wear_ui.cpp b/wear_ui.cpp
index f4a8399..65c4aee 100644
--- a/wear_ui.cpp
+++ b/wear_ui.cpp
@@ -16,7 +16,6 @@
 
 #include "wear_ui.h"
 
-#include <pthread.h>
 #include <string.h>
 
 #include <string>
@@ -86,11 +85,10 @@
 
 void WearRecoveryUI::StartMenu(const std::vector<std::string>& headers,
                                const std::vector<std::string>& items, size_t initial_selection) {
-  pthread_mutex_lock(&updateMutex);
+  std::lock_guard<std::mutex> lg(updateMutex);
   if (text_rows_ > 0 && text_cols_ > 0) {
     menu_ = std::make_unique<Menu>(scrollable_menu_, text_rows_ - kMenuUnusableRows - 1,
                                    text_cols_ - 1, headers, items, initial_selection);
     update_screen_locked();
   }
-  pthread_mutex_unlock(&updateMutex);
 }