15 files changed, 222 insertions, 131 deletions

diff --git a/minui/graphics.cpp b/minui/graphics.cpp
index 3b386015a..c1aab412d 100644
--- a/minui/graphics.cpp
+++ b/minui/graphics.cpp
@@ -356,6 +356,10 @@ int gr_init() {
 
   gr_flip();
   gr_flip();
+  if (!gr_draw) {
+    printf("gr_init: gr_draw becomes nullptr after gr_flip\n");
+    return -1;
+  }
 
   gr_rotate(DEFAULT_ROTATION);
 
diff --git a/minui/graphics_drm.cpp b/minui/graphics_drm.cpp
index e7d4b38ef..4c98507f6 100644
--- a/minui/graphics_drm.cpp
+++ b/minui/graphics_drm.cpp
@@ -45,15 +45,17 @@ void MinuiBackendDrm::DrmDisableCrtc(int drm_fd, drmModeCrtc* crtc) {
   }
 }
 
-void MinuiBackendDrm::DrmEnableCrtc(int drm_fd, drmModeCrtc* crtc, GRSurfaceDrm* surface) {
-  int32_t ret = drmModeSetCrtc(drm_fd, crtc->crtc_id, surface->fb_id, 0, 0,  // x,y
-                               &main_monitor_connector->connector_id,
-                               1,  // connector_count
-                               &main_monitor_crtc->mode);
+int MinuiBackendDrm::DrmEnableCrtc(int drm_fd, drmModeCrtc* crtc, GRSurfaceDrm* surface) {
+  int ret = drmModeSetCrtc(drm_fd, crtc->crtc_id, surface->fb_id, 0, 0,  // x,y
+                           &main_monitor_connector->connector_id,
+                           1,  // connector_count
+                           &main_monitor_crtc->mode);
 
   if (ret) {
     printf("drmModeSetCrtc failed ret=%d\n", ret);
   }
+
+  return ret;
 }
 
 void MinuiBackendDrm::Blank(bool blank) {
@@ -368,7 +370,10 @@ GRSurface* MinuiBackendDrm::Init() {
 
   current_buffer = 0;
 
-  DrmEnableCrtc(drm_fd, main_monitor_crtc, GRSurfaceDrms[1]);
+  // We will likely encounter errors in the backend functions (i.e. Flip) if EnableCrtc fails.
+  if (DrmEnableCrtc(drm_fd, main_monitor_crtc, GRSurfaceDrms[1]) != 0) {
+    return nullptr;
+  }
 
   return GRSurfaceDrms[0];
 }
diff --git a/minui/graphics_drm.h b/minui/graphics_drm.h
index de9621205..756625b03 100644
--- a/minui/graphics_drm.h
+++ b/minui/graphics_drm.h
@@ -42,7 +42,7 @@ class MinuiBackendDrm : public MinuiBackend {
 
  private:
   void DrmDisableCrtc(int drm_fd, drmModeCrtc* crtc);
-  void DrmEnableCrtc(int drm_fd, drmModeCrtc* crtc, GRSurfaceDrm* surface);
+  int DrmEnableCrtc(int drm_fd, drmModeCrtc* crtc, GRSurfaceDrm* surface);
   GRSurfaceDrm* DrmCreateSurface(int width, int height);
   void DrmDestroySurface(GRSurfaceDrm* surface);
   void DisableNonMainCrtcs(int fd, drmModeRes* resources, drmModeCrtc* main_crtc);
diff --git a/screen_ui.cpp b/screen_ui.cpp
index b9aba807d..b4ef054ce 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -173,7 +173,9 @@ ScreenRecoveryUI::ScreenRecoveryUI(bool scrollable_menu)
 
 ScreenRecoveryUI::~ScreenRecoveryUI() {
   progress_thread_stopped_ = true;
-  progress_thread_.join();
+  if (progress_thread_.joinable()) {
+    progress_thread_.join();
+  }
 }
 
 GRSurface* ScreenRecoveryUI::GetCurrentFrame() const {
diff --git a/tests/unit/rangeset_test.cpp b/tests/unit/rangeset_test.cpp
index 7ae193e18..fc72f2f6d 100644
--- a/tests/unit/rangeset_test.cpp
+++ b/tests/unit/rangeset_test.cpp
@@ -209,6 +209,7 @@ TEST(RangeSetTest, GetBlockNumber) {
   ASSERT_EQ(static_cast<size_t>(6), rs.GetBlockNumber(5));
   ASSERT_EQ(static_cast<size_t>(9), rs.GetBlockNumber(8));
 
+  ::testing::FLAGS_gtest_death_test_style = "threadsafe";
   // Out of bound.
   ASSERT_EXIT(rs.GetBlockNumber(9), ::testing::KilledBySignal(SIGABRT), "");
 }
@@ -284,6 +285,8 @@ TEST(SortedRangeSetTest, file_range) {
 
   ASSERT_EQ(static_cast<size_t>(10), rs.GetOffsetInRangeSet(4106));
   ASSERT_EQ(static_cast<size_t>(40970), rs.GetOffsetInRangeSet(4096 * 16 + 10));
+
+  ::testing::FLAGS_gtest_death_test_style = "threadsafe";
   // block#10 not in range.
   ASSERT_EXIT(rs.GetOffsetInRangeSet(40970), ::testing::KilledBySignal(SIGABRT), "");
 }
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index 25623074c..a3dd2add9 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -293,6 +293,11 @@ TEST_F(ScreenRecoveryUITest, Init) {
   ASSERT_FALSE(ui_->WasTextEverVisible());
 }
 
+TEST_F(ScreenRecoveryUITest, dtor_NotCallingInit) {
+  ui_.reset();
+  ASSERT_FALSE(ui_);
+}
+
 TEST_F(ScreenRecoveryUITest, ShowText) {
   ASSERT_TRUE(ui_->Init(kTestLocale));
   ASSERT_FALSE(ui_->IsTextVisible());
@@ -408,5 +413,7 @@ TEST_F(ScreenRecoveryUITest, LoadAnimation_MissingAnimation) {
   ASSERT_TRUE(ui_->Init(kTestLocale));
   TemporaryDir resource_dir;
   Paths::Get().set_resource_dir(resource_dir.path);
+
+  ::testing::FLAGS_gtest_death_test_style = "threadsafe";
   ASSERT_EXIT(ui_->RunLoadAnimation(), ::testing::KilledBySignal(SIGABRT), "");
 }
diff --git a/ui.cpp b/ui.cpp
index 51d7f129c..6c91d01b8 100644
--- a/ui.cpp
+++ b/ui.cpp
@@ -78,7 +78,9 @@ RecoveryUI::RecoveryUI()
 RecoveryUI::~RecoveryUI() {
   ev_exit();
   input_thread_stopped_ = true;
-  input_thread_.join();
+  if (input_thread_.joinable()) {
+    input_thread_.join();
+  }
 }
 
 void RecoveryUI::OnKeyDetected(int key_code) {
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index bdb64636b..f2811bccf 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -197,8 +197,8 @@ static int read_all(int fd, uint8_t* data, size_t size) {
     return 0;
 }
 
-static int read_all(int fd, std::vector<uint8_t>& buffer, size_t size) {
-    return read_all(fd, buffer.data(), size);
+static int read_all(int fd, std::vector<uint8_t>* buffer, size_t size) {
+  return read_all(fd, buffer->data(), size);
 }
 
 static int write_all(int fd, const uint8_t* data, size_t size) {
@@ -244,11 +244,10 @@ static bool check_lseek(int fd, off64_t offset, int whence) {
     return true;
 }
 
-static void allocate(size_t size, std::vector<uint8_t>& buffer) {
-    // if the buffer's big enough, reuse it.
-    if (size <= buffer.size()) return;
-
-    buffer.resize(size);
+static void allocate(size_t size, std::vector<uint8_t>* buffer) {
+  // If the buffer's big enough, reuse it.
+  if (size <= buffer->size()) return;
+  buffer->resize(size);
 }
 
 /**
@@ -502,7 +501,7 @@ static void* unzip_new_data(void* cookie) {
   return nullptr;
 }
 
-static int ReadBlocks(const RangeSet& src, std::vector<uint8_t>& buffer, int fd) {
+static int ReadBlocks(const RangeSet& src, std::vector<uint8_t>* buffer, int fd) {
   size_t p = 0;
   for (const auto& range : src) {
     if (!check_lseek(fd, static_cast<off64_t>(range.first) * BLOCKSIZE, SEEK_SET)) {
@@ -510,7 +509,7 @@ static int ReadBlocks(const RangeSet& src, std::vector<uint8_t>& buffer, int fd)
     }
 
     size_t size = (range.second - range.first) * BLOCKSIZE;
-    if (read_all(fd, buffer.data() + p, size) == -1) {
+    if (read_all(fd, buffer->data() + p, size) == -1) {
       return -1;
     }
 
@@ -662,31 +661,31 @@ static void PrintHashForMissingStashedBlocks(const std::string& id, int fd) {
   LOG(INFO) << "print hash in hex for source blocks in missing stash: " << id;
   const RangeSet& src = stash_map[id];
   std::vector<uint8_t> buffer(src.blocks() * BLOCKSIZE);
-  if (ReadBlocks(src, buffer, fd) == -1) {
-      LOG(ERROR) << "failed to read source blocks for stash: " << id;
-      return;
+  if (ReadBlocks(src, &buffer, fd) == -1) {
+    LOG(ERROR) << "failed to read source blocks for stash: " << id;
+    return;
   }
   PrintHashForCorruptedStashedBlocks(id, buffer, src);
 }
 
 static int VerifyBlocks(const std::string& expected, const std::vector<uint8_t>& buffer,
-        const size_t blocks, bool printerror) {
-    uint8_t digest[SHA_DIGEST_LENGTH];
-    const uint8_t* data = buffer.data();
+                        const size_t blocks, bool printerror) {
+  uint8_t digest[SHA_DIGEST_LENGTH];
+  const uint8_t* data = buffer.data();
 
-    SHA1(data, blocks * BLOCKSIZE, digest);
+  SHA1(data, blocks * BLOCKSIZE, digest);
 
-    std::string hexdigest = print_sha1(digest);
+  std::string hexdigest = print_sha1(digest);
 
-    if (hexdigest != expected) {
-        if (printerror) {
-            LOG(ERROR) << "failed to verify blocks (expected " << expected << ", read "
-                       << hexdigest << ")";
-        }
-        return -1;
+  if (hexdigest != expected) {
+    if (printerror) {
+      LOG(ERROR) << "failed to verify blocks (expected " << expected << ", read " << hexdigest
+                 << ")";
     }
+    return -1;
+  }
 
-    return 0;
+  return 0;
 }
 
 static std::string GetStashFileName(const std::string& base, const std::string& id,
@@ -751,8 +750,8 @@ static void DeleteStash(const std::string& base) {
   }
 }
 
-static int LoadStash(CommandParameters& params, const std::string& id, bool verify,
-                     std::vector<uint8_t>& buffer, bool printnoent) {
+static int LoadStash(const 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.
   if (!params.canwrite) {
@@ -764,9 +763,9 @@ static int LoadStash(CommandParameters& params, const std::string& id, bool veri
         LOG(ERROR) << "failed to read source blocks in stash map.";
         return -1;
       }
-      if (VerifyBlocks(id, buffer, src.blocks(), true) != 0) {
+      if (VerifyBlocks(id, *buffer, src.blocks(), true) != 0) {
         LOG(ERROR) << "failed to verify loaded source blocks in stash map.";
-        PrintHashForCorruptedStashedBlocks(id, buffer, src);
+        PrintHashForCorruptedStashedBlocks(id, *buffer, src);
         return -1;
       }
       return 0;
@@ -804,14 +803,14 @@ static int LoadStash(CommandParameters& params, const std::string& id, bool veri
   }
 
   size_t blocks = sb.st_size / BLOCKSIZE;
-  if (verify && VerifyBlocks(id, buffer, blocks, true) != 0) {
+  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
                  << " when executing command: " << params.cmdname;
     } else {
       const RangeSet& src = stash_map[id];
-      PrintHashForCorruptedStashedBlocks(id, buffer, src);
+      PrintHashForCorruptedStashedBlocks(id, *buffer, src);
     }
     DeleteFile(fn);
     return -1;
@@ -821,71 +820,71 @@ static int LoadStash(CommandParameters& params, const std::string& id, bool veri
 }
 
 static int WriteStash(const std::string& base, const std::string& id, int blocks,
-                      std::vector<uint8_t>& buffer, bool checkspace, bool* exists) {
-    if (base.empty()) {
-        return -1;
-    }
-
-    if (checkspace && CacheSizeCheck(blocks * BLOCKSIZE) != 0) {
-        LOG(ERROR) << "not enough space to write stash";
-        return -1;
-    }
+                      const std::vector<uint8_t>& buffer, bool checkspace, bool* exists) {
+  if (base.empty()) {
+    return -1;
+  }
 
-    std::string fn = GetStashFileName(base, id, ".partial");
-    std::string cn = GetStashFileName(base, id, "");
+  if (checkspace && CacheSizeCheck(blocks * BLOCKSIZE) != 0) {
+    LOG(ERROR) << "not enough space to write stash";
+    return -1;
+  }
 
-    if (exists) {
-        struct stat sb;
-        int res = stat(cn.c_str(), &sb);
+  std::string fn = GetStashFileName(base, id, ".partial");
+  std::string cn = GetStashFileName(base, id, "");
 
-        if (res == 0) {
-            // The file already exists and since the name is the hash of the contents,
-            // it's safe to assume the contents are identical (accidental hash collisions
-            // are unlikely)
-            LOG(INFO) << " skipping " << blocks << " existing blocks in " << cn;
-            *exists = true;
-            return 0;
-        }
-
-        *exists = false;
+  if (exists) {
+    struct stat sb;
+    int res = stat(cn.c_str(), &sb);
+
+    if (res == 0) {
+      // The file already exists and since the name is the hash of the contents,
+      // it's safe to assume the contents are identical (accidental hash collisions
+      // are unlikely)
+      LOG(INFO) << " skipping " << blocks << " existing blocks in " << cn;
+      *exists = true;
+      return 0;
     }
 
-    LOG(INFO) << " writing " << blocks << " blocks to " << cn;
+    *exists = false;
+  }
 
-    android::base::unique_fd fd(
-        TEMP_FAILURE_RETRY(ota_open(fn.c_str(), O_WRONLY | O_CREAT | O_TRUNC, STASH_FILE_MODE)));
-    if (fd == -1) {
-        PLOG(ERROR) << "failed to create \"" << fn << "\"";
-        return -1;
-    }
+  LOG(INFO) << " writing " << blocks << " blocks to " << cn;
 
-    if (fchown(fd, AID_SYSTEM, AID_SYSTEM) != 0) {  // system user
-        PLOG(ERROR) << "failed to chown \"" << fn << "\"";
-        return -1;
-    }
+  android::base::unique_fd fd(
+      TEMP_FAILURE_RETRY(ota_open(fn.c_str(), O_WRONLY | O_CREAT | O_TRUNC, STASH_FILE_MODE)));
+  if (fd == -1) {
+    PLOG(ERROR) << "failed to create \"" << fn << "\"";
+    return -1;
+  }
 
-    if (write_all(fd, buffer, blocks * BLOCKSIZE) == -1) {
-        return -1;
-    }
+  if (fchown(fd, AID_SYSTEM, AID_SYSTEM) != 0) {  // system user
+    PLOG(ERROR) << "failed to chown \"" << fn << "\"";
+    return -1;
+  }
 
-    if (ota_fsync(fd) == -1) {
-        failure_type = kFsyncFailure;
-        PLOG(ERROR) << "fsync \"" << fn << "\" failed";
-        return -1;
-    }
+  if (write_all(fd, buffer, blocks * BLOCKSIZE) == -1) {
+    return -1;
+  }
 
-    if (rename(fn.c_str(), cn.c_str()) == -1) {
-        PLOG(ERROR) << "rename(\"" << fn << "\", \"" << cn << "\") failed";
-        return -1;
-    }
+  if (ota_fsync(fd) == -1) {
+    failure_type = kFsyncFailure;
+    PLOG(ERROR) << "fsync \"" << fn << "\" failed";
+    return -1;
+  }
 
-    std::string dname = GetStashFileName(base, "", "");
-    if (!FsyncDir(dname)) {
-      failure_type = kFsyncFailure;
-      return -1;
-    }
+  if (rename(fn.c_str(), cn.c_str()) == -1) {
+    PLOG(ERROR) << "rename(\"" << fn << "\", \"" << cn << "\") failed";
+    return -1;
+  }
 
-    return 0;
+  std::string dname = GetStashFileName(base, "", "");
+  if (!FsyncDir(dname)) {
+    failure_type = kFsyncFailure;
+    return -1;
+  }
+
+  return 0;
 }
 
 // Creates a directory for storing stash files and checks if the /cache partition
@@ -1014,7 +1013,7 @@ static int LoadSourceBlocks(CommandParameters& params, const RangeSet& tgt, size
     return -1;
   }
 
-  allocate(*src_blocks * BLOCKSIZE, params.buffer);
+  allocate(*src_blocks * BLOCKSIZE, &params.buffer);
 
   // "-" or <src_range> [<src_loc>]
   if (params.tokens[params.cpos] == "-") {
@@ -1025,7 +1024,7 @@ static int LoadSourceBlocks(CommandParameters& params, const RangeSet& tgt, size
     CHECK(static_cast<bool>(src));
     *overlap = src.Overlaps(tgt);
 
-    if (ReadBlocks(src, params.buffer, params.fd) == -1) {
+    if (ReadBlocks(src, &params.buffer, params.fd) == -1) {
       return -1;
     }
 
@@ -1050,7 +1049,7 @@ static int LoadSourceBlocks(CommandParameters& params, const RangeSet& tgt, size
     }
 
     std::vector<uint8_t> stash;
-    if (LoadStash(params, tokens[0], false, 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];
@@ -1091,7 +1090,7 @@ static int LoadSourceBlocks(CommandParameters& params, const RangeSet& tgt, size
  *
  * If the return value is 0, source blocks have expected content and the command can be performed.
  */
-static int LoadSrcTgtVersion3(CommandParameters& params, RangeSet& tgt, size_t* src_blocks,
+static int LoadSrcTgtVersion3(CommandParameters& params, RangeSet* tgt, size_t* src_blocks,
                               bool onehash, bool* overlap) {
   CHECK(src_blocks != nullptr);
   CHECK(overlap != nullptr);
@@ -1122,21 +1121,21 @@ static int LoadSrcTgtVersion3(CommandParameters& params, RangeSet& tgt, size_t*
   }
 
   // <tgt_range>
-  tgt = RangeSet::Parse(params.tokens[params.cpos++]);
-  CHECK(static_cast<bool>(tgt));
+  *tgt = RangeSet::Parse(params.tokens[params.cpos++]);
+  CHECK(static_cast<bool>(*tgt));
 
-  std::vector<uint8_t> tgtbuffer(tgt.blocks() * BLOCKSIZE);
-  if (ReadBlocks(tgt, tgtbuffer, params.fd) == -1) {
+  std::vector<uint8_t> tgtbuffer(tgt->blocks() * BLOCKSIZE);
+  if (ReadBlocks(*tgt, &tgtbuffer, params.fd) == -1) {
     return -1;
   }
 
   // Return now if target blocks already have expected content.
-  if (VerifyBlocks(tgthash, tgtbuffer, tgt.blocks(), false) == 0) {
+  if (VerifyBlocks(tgthash, tgtbuffer, tgt->blocks(), false) == 0) {
     return 1;
   }
 
   // Load source blocks.
-  if (LoadSourceBlocks(params, tgt, src_blocks, overlap) == -1) {
+  if (LoadSourceBlocks(params, *tgt, src_blocks, overlap) == -1) {
     return -1;
   }
 
@@ -1165,7 +1164,7 @@ static int LoadSrcTgtVersion3(CommandParameters& params, RangeSet& tgt, size_t*
     return 0;
   }
 
-  if (*overlap && LoadStash(params, srchash, true, 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.
@@ -1185,7 +1184,7 @@ static int PerformCommandMove(CommandParameters& params) {
   size_t blocks = 0;
   bool overlap = false;
   RangeSet tgt;
-  int status = LoadSrcTgtVersion3(params, tgt, &blocks, true, &overlap);
+  int status = LoadSrcTgtVersion3(params, &tgt, &blocks, true, &overlap);
 
   if (status == -1) {
     LOG(ERROR) << "failed to read blocks for move";
@@ -1231,7 +1230,7 @@ static int PerformCommandStash(CommandParameters& params) {
   }
 
   const std::string& id = params.tokens[params.cpos++];
-  if (LoadStash(params, id, true, 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;
@@ -1241,8 +1240,8 @@ static int PerformCommandStash(CommandParameters& params) {
   CHECK(static_cast<bool>(src));
 
   size_t blocks = src.blocks();
-  allocate(blocks * BLOCKSIZE, params.buffer);
-  if (ReadBlocks(src, params.buffer, params.fd) == -1) {
+  allocate(blocks * BLOCKSIZE, &params.buffer);
+  if (ReadBlocks(src, &params.buffer, params.fd) == -1) {
     return -1;
   }
   stash_map[id] = src;
@@ -1296,7 +1295,7 @@ static int PerformCommandZero(CommandParameters& params) {
 
   LOG(INFO) << "  zeroing " << tgt.blocks() << " blocks";
 
-  allocate(BLOCKSIZE, params.buffer);
+  allocate(BLOCKSIZE, &params.buffer);
   memset(params.buffer.data(), 0, BLOCKSIZE);
 
   if (params.canwrite) {
@@ -1384,7 +1383,7 @@ static int PerformCommandDiff(CommandParameters& params) {
   RangeSet tgt;
   size_t blocks = 0;
   bool overlap = false;
-  int status = LoadSrcTgtVersion3(params, tgt, &blocks, false, &overlap);
+  int status = LoadSrcTgtVersion3(params, &tgt, &blocks, false, &overlap);
 
   if (status == -1) {
     LOG(ERROR) << "failed to read blocks for diff";
@@ -1975,7 +1974,7 @@ Value* RangeSha1Fn(const char* name, State* state, const std::vector<std::unique
     }
 
     for (size_t j = range.first; j < range.second; ++j) {
-      if (read_all(fd, buffer, BLOCKSIZE) == -1) {
+      if (read_all(fd, &buffer, BLOCKSIZE) == -1) {
         ErrorAbort(state, kFreadFailure, "failed to read %s: %s", blockdev_filename->data.c_str(),
                    strerror(errno));
         return StringValue("");
@@ -2025,7 +2024,7 @@ Value* CheckFirstBlockFn(const char* name, State* state,
   RangeSet blk0(std::vector<Range>{ Range{ 0, 1 } });
   std::vector<uint8_t> block0_buffer(BLOCKSIZE);
 
-  if (ReadBlocks(blk0, block0_buffer, fd) == -1) {
+  if (ReadBlocks(blk0, &block0_buffer, fd) == -1) {
     ErrorAbort(state, kFreadFailure, "failed to read %s: %s", arg_filename->data.c_str(),
                strerror(errno));
     return StringValue("");
diff --git a/updater_sample/README.md b/updater_sample/README.md
index 8ec43d3c6..11b55eb91 100644
--- a/updater_sample/README.md
+++ b/updater_sample/README.md
@@ -140,6 +140,9 @@ 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.
 
+`key_value_pair_headers` argument also accepts properties other than HTTP Headers.
+List of allowed properties can be found in `system/update_engine/common/constants.cc`.
+
 ### UpdateEngine#cancel
 
 Cancel the ongoing update. The update could be running or suspended, but it
@@ -181,9 +184,8 @@ Called whenever an update attempt is completed.
 - [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`
+- [x] Add smart update completion detection using onStatusUpdate
+- [x] Add pause/resume demo
 - [ ] Verify system partition checksum for package
 - [?] Add non-A/B updates demo
 
diff --git a/updater_sample/res/layout/activity_main.xml b/updater_sample/res/layout/activity_main.xml
index 7cde42cec..b560827d8 100644
--- a/updater_sample/res/layout/activity_main.xml
+++ b/updater_sample/res/layout/activity_main.xml
@@ -197,6 +197,29 @@
                     android:text="Reset" />
             </LinearLayout>
 
+            <LinearLayout
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:layout_marginTop="12dp"
+                android:orientation="horizontal">
+
+                <Button
+                    android:id="@+id/buttonSuspend"
+                    android:layout_width="0dp"
+                    android:layout_height="wrap_content"
+                    android:layout_weight="1"
+                    android:onClick="onSuspendClick"
+                    android:text="Suspend" />
+
+                <Button
+                    android:id="@+id/buttonResume"
+                    android:layout_width="0dp"
+                    android:layout_height="wrap_content"
+                    android:layout_weight="1"
+                    android:onClick="onResumeClick"
+                    android:text="Resume" />
+            </LinearLayout>
+
             <TextView
                 android:id="@+id/textViewUpdateInfo"
                 android:layout_width="wrap_content"
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
index e4c09346b..a9783e70a 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -198,6 +198,24 @@ public class UpdateManager {
     }
 
     /**
+     * Suspend running update.
+     */
+    public synchronized void suspend() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "suspend invoked");
+        setUpdaterState(UpdaterState.PAUSED);
+        mUpdateEngine.cancel();
+    }
+
+    /**
+     * Resume suspended update.
+     */
+    public synchronized void resume() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "resume invoked");
+        setUpdaterState(UpdaterState.RUNNING);
+        updateEngineReApplyPayload();
+    }
+
+    /**
      * Updates {@link this.mState} and if state is changed,
      * it also notifies {@link this.mOnStateChangeCallback}.
      */
@@ -237,10 +255,6 @@ public class UpdateManager {
     /**
      * Requests update engine to stop any ongoing update. If an update has been applied,
      * leave it as is.
-     *
-     * <p>Sometimes it's possible that the
-     * update engine would throw an error when the method is called, and the only way to
-     * handle it is to catch the exception.</p>
      */
     public synchronized void cancelRunningUpdate() throws UpdaterState.InvalidTransitionException {
         Log.d(TAG, "cancelRunningUpdate invoked");
@@ -250,10 +264,6 @@ public class UpdateManager {
 
     /**
      * Resets update engine to IDLE state. If an update has been applied it reverts it.
-     *
-     * <p>Sometimes it's possible that the
-     * update engine would throw an error when the method is called, and the only way to
-     * handle it is to catch the exception.</p>
      */
     public synchronized void resetUpdate() throws UpdaterState.InvalidTransitionException {
         Log.d(TAG, "resetUpdate invoked");
@@ -506,7 +516,7 @@ public class UpdateManager {
             synchronizeUpdaterStateWithUpdateEngineStatus();
         }
 
-        getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(progress));
+        getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(mProgress.get()));
 
         if (previousStatus != status) {
             getOnEngineStatusUpdateCallback().ifPresent(callback -> callback.accept(status));
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
index 573d336e9..4eb0b68c7 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
@@ -52,12 +52,12 @@ public class UpdaterState {
      */
     private static final ImmutableMap<Integer, ImmutableSet<Integer>> TRANSITIONS =
             ImmutableMap.<Integer, ImmutableSet<Integer>>builder()
-                    .put(IDLE, ImmutableSet.of(ERROR, RUNNING))
+                    .put(IDLE, ImmutableSet.of(IDLE, ERROR, RUNNING))
+                    .put(ERROR, ImmutableSet.of(IDLE))
                     .put(RUNNING, ImmutableSet.of(
-                            ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED))
+                            IDLE, ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED))
                     .put(PAUSED, ImmutableSet.of(ERROR, RUNNING, IDLE))
-                    .put(SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR, IDLE))
-                    .put(ERROR, ImmutableSet.of(IDLE))
+                    .put(SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR, REBOOT_REQUIRED, IDLE))
                     .put(REBOOT_REQUIRED, ImmutableSet.of(IDLE))
                     .build();
 
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 6c71cb6f4..fc9fddd70 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -55,6 +55,8 @@ public class MainActivity extends Activity {
     private Button mButtonApplyConfig;
     private Button mButtonStop;
     private Button mButtonReset;
+    private Button mButtonSuspend;
+    private Button mButtonResume;
     private ProgressBar mProgressBar;
     private TextView mTextViewUpdaterState;
     private TextView mTextViewEngineStatus;
@@ -79,6 +81,8 @@ public class MainActivity extends Activity {
         this.mButtonApplyConfig = findViewById(R.id.buttonApplyConfig);
         this.mButtonStop = findViewById(R.id.buttonStop);
         this.mButtonReset = findViewById(R.id.buttonReset);
+        this.mButtonSuspend = findViewById(R.id.buttonSuspend);
+        this.mButtonResume = findViewById(R.id.buttonResume);
         this.mProgressBar = findViewById(R.id.progressBar);
         this.mTextViewUpdaterState = findViewById(R.id.textViewUpdaterState);
         this.mTextViewEngineStatus = findViewById(R.id.textViewEngineStatus);
@@ -209,9 +213,34 @@ public class MainActivity extends Activity {
     }
 
     /**
+     * suspend button clicked
+     */
+    public void onSuspendClick(View view) {
+        try {
+            mUpdateManager.suspend();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to suspend running update", e);
+        }
+    }
+
+    /**
+     * resume button clicked
+     */
+    public void onResumeClick(View view) {
+        try {
+            uiResetWidgets();
+            uiResetEngineText();
+            mUpdateManager.resume();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to resume running update", e);
+        }
+    }
+
+    /**
      * switch slot button clicked
      */
     public void onSwitchSlotClick(View view) {
+        uiResetWidgets();
         mUpdateManager.setSwitchSlotOnReboot();
     }
 
@@ -289,6 +318,8 @@ public class MainActivity extends Activity {
         mButtonApplyConfig.setEnabled(false);
         mButtonStop.setEnabled(false);
         mButtonReset.setEnabled(false);
+        mButtonSuspend.setEnabled(false);
+        mButtonResume.setEnabled(false);
         mProgressBar.setEnabled(false);
         mProgressBar.setVisibility(ProgressBar.INVISIBLE);
         mButtonSwitchSlot.setEnabled(false);
@@ -303,6 +334,7 @@ public class MainActivity extends Activity {
 
     private void uiStateIdle() {
         uiResetWidgets();
+        mButtonReset.setEnabled(true);
         mSpinnerConfigs.setEnabled(true);
         mButtonReload.setEnabled(true);
         mButtonApplyConfig.setEnabled(true);
@@ -314,6 +346,7 @@ public class MainActivity extends Activity {
         mProgressBar.setEnabled(true);
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
         mButtonStop.setEnabled(true);
+        mButtonSuspend.setEnabled(true);
     }
 
     private void uiStatePaused() {
@@ -321,6 +354,7 @@ public class MainActivity extends Activity {
         mButtonReset.setEnabled(true);
         mProgressBar.setEnabled(true);
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
+        mButtonResume.setEnabled(true);
     }
 
     private void uiStateSlotSwitchRequired() {
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java
index 3ba84c116..03086930e 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/util/PayloadSpecsTest.java
@@ -65,7 +65,7 @@ public class PayloadSpecsTest {
         mTargetContext = InstrumentationRegistry.getTargetContext();
         mTestContext = InstrumentationRegistry.getContext();
 
-        mTestDir = mTargetContext.getFilesDir();
+        mTestDir = mTargetContext.getCacheDir();
         mPayloadSpecs = new PayloadSpecs();
     }
 
diff --git a/updater_sample/tools/gen_update_config.py b/updater_sample/tools/gen_update_config.py
index 7fb64f7fc..b43e49df8 100755
--- a/updater_sample/tools/gen_update_config.py
+++ b/updater_sample/tools/gen_update_config.py
@@ -131,10 +131,10 @@ def main():  # pylint: disable=missing-docstring
                         choices=ab_install_type_choices,
                         help='A/B update installation type')
     parser.add_argument('--ab_force_switch_slot',
-                        type=bool,
                         default=False,
-                        help='if set true device will boot to a new slot, otherwise user manually '
-                             'switches slot on the screen')
+                        action='store_true',
+                        help='if set device will boot to a new slot, otherwise user '
+                              'manually switches slot on the screen')
     parser.add_argument('package',
                         type=str,
                         help='OTA package zip file')