9 files changed, 419 insertions, 340 deletions

diff --git a/tests/unit/updater_test.cpp b/tests/unit/updater_test.cpp
index a0a7b66ab..4a8d1e6ff 100644
--- a/tests/unit/updater_test.cpp
+++ b/tests/unit/updater_test.cpp
@@ -57,16 +57,14 @@ using namespace std::string_literals;
 
 using PackageEntries = std::unordered_map<std::string, std::string>;
 
-struct selabel_handle* sehandle = nullptr;
-
 static void expect(const char* expected, const std::string& expr_str, CauseCode cause_code,
-                   UpdaterInfo* info = nullptr) {
+                   Updater* updater = nullptr) {
   std::unique_ptr<Expr> e;
   int error_count = 0;
   ASSERT_EQ(0, ParseString(expr_str, &e, &error_count));
   ASSERT_EQ(0, error_count);
 
-  State state(expr_str, info);
+  State state(expr_str, updater);
 
   std::string result;
   bool status = Evaluate(&state, e, &result);
@@ -102,38 +100,6 @@ static void BuildUpdatePackage(const PackageEntries& entries, int fd) {
   ASSERT_EQ(0, fclose(zip_file_ptr));
 }
 
-static void RunBlockImageUpdate(bool is_verify, const PackageEntries& entries,
-                                const std::string& image_file, const std::string& result,
-                                CauseCode cause_code = kNoCause) {
-  CHECK(entries.find("transfer_list") != entries.end());
-
-  // Build the update package.
-  TemporaryFile zip_file;
-  BuildUpdatePackage(entries, zip_file.release());
-
-  MemMapping map;
-  ASSERT_TRUE(map.MapFile(zip_file.path));
-  ZipArchiveHandle handle;
-  ASSERT_EQ(0, OpenArchiveFromMemory(map.addr, map.length, zip_file.path, &handle));
-
-  // Set up the handler, command_pipe, patch offset & length.
-  UpdaterInfo updater_info;
-  updater_info.package_zip = handle;
-  TemporaryFile temp_pipe;
-  updater_info.cmd_pipe = fdopen(temp_pipe.release(), "wbe");
-  updater_info.package_zip_addr = map.addr;
-  updater_info.package_zip_len = map.length;
-
-  std::string new_data = entries.find("new_data.br") != entries.end() ? "new_data.br" : "new_data";
-  std::string script = is_verify ? "block_image_verify" : "block_image_update";
-  script += R"((")" + image_file + R"(", package_extract_file("transfer_list"), ")" + new_data +
-            R"(", "patch_data"))";
-  expect(result.c_str(), script, cause_code, &updater_info);
-
-  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
-  CloseArchive(handle);
-}
-
 static std::string GetSha1(std::string_view content) {
   uint8_t digest[SHA_DIGEST_LENGTH];
   SHA1(reinterpret_cast<const uint8_t*>(content.data()), content.size(), digest);
@@ -159,29 +125,24 @@ static Value* BlobToString(const char* name, State* state,
   return args[0].release();
 }
 
-class UpdaterTest : public ::testing::Test {
+class UpdaterTestBase {
  protected:
-  void SetUp() override {
+  void SetUp() {
     RegisterBuiltins();
     RegisterInstallFunctions();
     RegisterBlockImageFunctions();
 
-    RegisterFunction("blob_to_string", BlobToString);
-
     // Each test is run in a separate process (isolated mode). Shared temporary files won't cause
     // conflicts.
     Paths::Get().set_cache_temp_source(temp_saved_source_.path);
     Paths::Get().set_last_command_file(temp_last_command_.path);
     Paths::Get().set_stash_directory_base(temp_stash_base_.path);
 
-    // Enable a special command "abort" to simulate interruption.
-    Command::abort_allowed_ = true;
-
     last_command_file_ = temp_last_command_.path;
     image_file_ = image_temp_file_.path;
   }
 
-  void TearDown() override {
+  void TearDown() {
     // Clean up the last_command_file if any.
     ASSERT_TRUE(android::base::RemoveFileIfExists(last_command_file_));
 
@@ -191,16 +152,80 @@ class UpdaterTest : public ::testing::Test {
     ASSERT_TRUE(android::base::RemoveFileIfExists(updated_marker));
   }
 
+  void RunBlockImageUpdate(bool is_verify, PackageEntries entries, const std::string& image_file,
+                           const std::string& result, CauseCode cause_code = kNoCause) {
+    CHECK(entries.find("transfer_list") != entries.end());
+    std::string new_data =
+        entries.find("new_data.br") != entries.end() ? "new_data.br" : "new_data";
+    std::string script = is_verify ? "block_image_verify" : "block_image_update";
+    script += R"((")" + image_file + R"(", package_extract_file("transfer_list"), ")" + new_data +
+              R"(", "patch_data"))";
+    entries.emplace(Updater::SCRIPT_NAME, script);
+
+    // Build the update package.
+    TemporaryFile zip_file;
+    BuildUpdatePackage(entries, zip_file.release());
+
+    // Set up the handler, command_pipe, patch offset & length.
+    TemporaryFile temp_pipe;
+    ASSERT_TRUE(updater_.Init(temp_pipe.release(), zip_file.path, false, nullptr));
+    ASSERT_TRUE(updater_.RunUpdate());
+    ASSERT_EQ(result, updater_.result());
+
+    // Parse the cause code written to the command pipe.
+    int received_cause_code = kNoCause;
+    std::string pipe_content;
+    ASSERT_TRUE(android::base::ReadFileToString(temp_pipe.path, &pipe_content));
+    auto lines = android::base::Split(pipe_content, "\n");
+    for (std::string_view line : lines) {
+      if (android::base::ConsumePrefix(&line, "log cause: ")) {
+        ASSERT_TRUE(android::base::ParseInt(line.data(), &received_cause_code));
+      }
+    }
+    ASSERT_EQ(cause_code, received_cause_code);
+  }
+
   TemporaryFile temp_saved_source_;
   TemporaryDir temp_stash_base_;
   std::string last_command_file_;
   std::string image_file_;
 
+  Updater updater_;
+
  private:
   TemporaryFile temp_last_command_;
   TemporaryFile image_temp_file_;
 };
 
+class UpdaterTest : public UpdaterTestBase, public ::testing::Test {
+ protected:
+  void SetUp() override {
+    UpdaterTestBase::SetUp();
+
+    RegisterFunction("blob_to_string", BlobToString);
+    // Enable a special command "abort" to simulate interruption.
+    Command::abort_allowed_ = true;
+  }
+
+  void TearDown() override {
+    UpdaterTestBase::TearDown();
+  }
+
+  void SetUpdaterCmdPipe(int fd) {
+    FILE* cmd_pipe = fdopen(fd, "w");
+    ASSERT_NE(nullptr, cmd_pipe);
+    updater_.cmd_pipe_.reset(cmd_pipe);
+  }
+
+  void SetUpdaterOtaPackageHandle(ZipArchiveHandle handle) {
+    updater_.package_handle_ = handle;
+  }
+
+  void FlushUpdaterCommandPipe() const {
+    fflush(updater_.cmd_pipe_.get());
+  }
+};
+
 TEST_F(UpdaterTest, getprop) {
     expect(android::base::GetProperty("ro.product.device", "").c_str(),
            "getprop(\"ro.product.device\")",
@@ -317,13 +342,12 @@ TEST_F(UpdaterTest, package_extract_file) {
   ASSERT_EQ(0, OpenArchive(zip_path.c_str(), &handle));
 
   // Need to set up the ziphandle.
-  UpdaterInfo updater_info;
-  updater_info.package_zip = handle;
+  SetUpdaterOtaPackageHandle(handle);
 
   // Two-argument version.
   TemporaryFile temp_file1;
   std::string script("package_extract_file(\"a.txt\", \"" + std::string(temp_file1.path) + "\")");
-  expect("t", script, kNoCause, &updater_info);
+  expect("t", script, kNoCause, &updater_);
 
   // Verify the extracted entry.
   std::string data;
@@ -332,32 +356,30 @@ TEST_F(UpdaterTest, package_extract_file) {
 
   // Now extract another entry to the same location, which should overwrite.
   script = "package_extract_file(\"b.txt\", \"" + std::string(temp_file1.path) + "\")";
-  expect("t", script, kNoCause, &updater_info);
+  expect("t", script, kNoCause, &updater_);
 
   ASSERT_TRUE(android::base::ReadFileToString(temp_file1.path, &data));
   ASSERT_EQ(kBTxtContents, data);
 
   // Missing zip entry. The two-argument version doesn't abort.
   script = "package_extract_file(\"doesntexist\", \"" + std::string(temp_file1.path) + "\")";
-  expect("", script, kNoCause, &updater_info);
+  expect("", script, kNoCause, &updater_);
 
   // Extract to /dev/full should fail.
   script = "package_extract_file(\"a.txt\", \"/dev/full\")";
-  expect("", script, kNoCause, &updater_info);
+  expect("", script, kNoCause, &updater_);
 
   // One-argument version. package_extract_file() gives a VAL_BLOB, which needs to be converted to
   // VAL_STRING for equality test.
   script = "blob_to_string(package_extract_file(\"a.txt\")) == \"" + kATxtContents + "\"";
-  expect("t", script, kNoCause, &updater_info);
+  expect("t", script, kNoCause, &updater_);
 
   script = "blob_to_string(package_extract_file(\"b.txt\")) == \"" + kBTxtContents + "\"";
-  expect("t", script, kNoCause, &updater_info);
+  expect("t", script, kNoCause, &updater_);
 
   // Missing entry. The one-argument version aborts the evaluation.
   script = "package_extract_file(\"doesntexist\")";
-  expect(nullptr, script, kPackageExtractFileFailure, &updater_info);
-
-  CloseArchive(handle);
+  expect(nullptr, script, kPackageExtractFileFailure, &updater_);
 }
 
 TEST_F(UpdaterTest, read_file) {
@@ -563,17 +585,15 @@ TEST_F(UpdaterTest, set_progress) {
   expect(nullptr, "set_progress(\".3.5\")", kArgsParsingFailure);
 
   TemporaryFile tf;
-  UpdaterInfo updater_info;
-  updater_info.cmd_pipe = fdopen(tf.release(), "w");
-  expect(".52", "set_progress(\".52\")", kNoCause, &updater_info);
-  fflush(updater_info.cmd_pipe);
+  SetUpdaterCmdPipe(tf.release());
+  expect(".52", "set_progress(\".52\")", kNoCause, &updater_);
+  FlushUpdaterCommandPipe();
 
   std::string cmd;
   ASSERT_TRUE(android::base::ReadFileToString(tf.path, &cmd));
   ASSERT_EQ(android::base::StringPrintf("set_progress %f\n", .52), cmd);
   // recovery-updater protocol expects 2 tokens ("set_progress <frac>").
   ASSERT_EQ(2U, android::base::Split(cmd, " ").size());
-  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
 }
 
 TEST_F(UpdaterTest, show_progress) {
@@ -588,17 +608,15 @@ TEST_F(UpdaterTest, show_progress) {
   expect(nullptr, "show_progress(\".3\", \"5a\")", kArgsParsingFailure);
 
   TemporaryFile tf;
-  UpdaterInfo updater_info;
-  updater_info.cmd_pipe = fdopen(tf.release(), "w");
-  expect(".52", "show_progress(\".52\", \"10\")", kNoCause, &updater_info);
-  fflush(updater_info.cmd_pipe);
+  SetUpdaterCmdPipe(tf.release());
+  expect(".52", "show_progress(\".52\", \"10\")", kNoCause, &updater_);
+  FlushUpdaterCommandPipe();
 
   std::string cmd;
   ASSERT_TRUE(android::base::ReadFileToString(tf.path, &cmd));
   ASSERT_EQ(android::base::StringPrintf("progress %f %d\n", .52, 10), cmd);
   // recovery-updater protocol expects 3 tokens ("progress <frac> <secs>").
   ASSERT_EQ(3U, android::base::Split(cmd, " ").size());
-  ASSERT_EQ(0, fclose(updater_info.cmd_pipe));
 }
 
 TEST_F(UpdaterTest, block_image_update_parsing_error) {
@@ -993,44 +1011,20 @@ TEST_F(UpdaterTest, last_command_verify) {
   ASSERT_EQ(-1, access(last_command_file_.c_str(), R_OK));
 }
 
-class ResumableUpdaterTest : public testing::TestWithParam<size_t> {
+class ResumableUpdaterTest : public UpdaterTestBase, public testing::TestWithParam<size_t> {
  protected:
   void SetUp() override {
-    RegisterBuiltins();
-    RegisterInstallFunctions();
-    RegisterBlockImageFunctions();
-
-    Paths::Get().set_cache_temp_source(temp_saved_source_.path);
-    Paths::Get().set_last_command_file(temp_last_command_.path);
-    Paths::Get().set_stash_directory_base(temp_stash_base_.path);
-
+    UpdaterTestBase::SetUp();
     // Enable a special command "abort" to simulate interruption.
     Command::abort_allowed_ = true;
-
     index_ = GetParam();
-    image_file_ = image_temp_file_.path;
-    last_command_file_ = temp_last_command_.path;
   }
 
   void TearDown() override {
-    // Clean up the last_command_file if any.
-    ASSERT_TRUE(android::base::RemoveFileIfExists(last_command_file_));
-
-    // Clear partition updated marker if any.
-    std::string updated_marker{ temp_stash_base_.path };
-    updated_marker += "/" + GetSha1(image_temp_file_.path) + ".UPDATED";
-    ASSERT_TRUE(android::base::RemoveFileIfExists(updated_marker));
+    UpdaterTestBase::TearDown();
   }
 
-  TemporaryFile temp_saved_source_;
-  TemporaryDir temp_stash_base_;
-  std::string last_command_file_;
-  std::string image_file_;
   size_t index_;
-
- private:
-  TemporaryFile temp_last_command_;
-  TemporaryFile image_temp_file_;
 };
 
 static std::string g_source_image;
diff --git a/updater/Android.bp b/updater/Android.bp
index b80cdb3a0..daf7e3277 100644
--- a/updater/Android.bp
+++ b/updater/Android.bp
@@ -70,6 +70,7 @@ cc_library_static {
         "commands.cpp",
         "dynamic_partitions.cpp",
         "install.cpp",
+        "updater.cpp",
     ],
 
     include_dirs: [
diff --git a/updater/Android.mk b/updater/Android.mk
index c7a6ba989..0178239e0 100644
--- a/updater/Android.mk
+++ b/updater/Android.mk
@@ -59,7 +59,7 @@ include $(CLEAR_VARS)
 LOCAL_MODULE := updater
 
 LOCAL_SRC_FILES := \
-    updater.cpp
+    updater_main.cpp
 
 LOCAL_C_INCLUDES := \
     $(LOCAL_PATH)/include
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index b008c28b4..3089865c7 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -42,6 +42,7 @@
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <android-base/parseint.h>
+#include <android-base/stringprintf.h>
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 #include <applypatch/applypatch.h>
@@ -1668,15 +1669,9 @@ static Value* PerformBlockImageUpdate(const char* name, State* state,
     return StringValue("");
   }
 
-  UpdaterInfo* ui = static_cast<UpdaterInfo*>(state->cookie);
-  if (ui == nullptr) {
-    return StringValue("");
-  }
-
-  FILE* cmd_pipe = ui->cmd_pipe;
-  ZipArchiveHandle za = ui->package_zip;
-
-  if (cmd_pipe == nullptr || za == nullptr) {
+  auto updater = static_cast<Updater*>(state->cookie);
+  ZipArchiveHandle za = updater->package_handle();
+  if (za == nullptr) {
     return StringValue("");
   }
 
@@ -1686,8 +1681,8 @@ static Value* PerformBlockImageUpdate(const char* name, State* state,
     LOG(ERROR) << name << "(): no file \"" << patch_data_fn->data << "\" in package";
     return StringValue("");
   }
+  params.patch_start = updater->GetMappedPackageAddress() + patch_entry.offset;
 
-  params.patch_start = ui->package_zip_addr + patch_entry.offset;
   std::string_view new_data(new_data_fn->data);
   ZipEntry new_entry;
   if (FindEntry(za, new_data, &new_entry) != 0) {
@@ -1887,8 +1882,10 @@ static Value* PerformBlockImageUpdate(const char* name, State* state,
         LOG(WARNING) << "Failed to update the last command file.";
       }
 
-      fprintf(cmd_pipe, "set_progress %.4f\n", static_cast<double>(params.written) / total_blocks);
-      fflush(cmd_pipe);
+      updater->WriteToCommandPipe(
+          android::base::StringPrintf("set_progress %.4f",
+                                      static_cast<double>(params.written) / total_blocks),
+          true);
     }
   }
 
@@ -1915,11 +1912,13 @@ pbiudone:
 
       const char* partition = strrchr(blockdev_filename->data.c_str(), '/');
       if (partition != nullptr && *(partition + 1) != 0) {
-        fprintf(cmd_pipe, "log bytes_written_%s: %" PRIu64 "\n", partition + 1,
-                static_cast<uint64_t>(params.written) * BLOCKSIZE);
-        fprintf(cmd_pipe, "log bytes_stashed_%s: %" PRIu64 "\n", partition + 1,
-                static_cast<uint64_t>(params.stashed) * BLOCKSIZE);
-        fflush(cmd_pipe);
+        updater->WriteToCommandPipe(
+            android::base::StringPrintf("log bytes_written_%s: %" PRIu64, partition + 1,
+                                        static_cast<uint64_t>(params.written) * BLOCKSIZE));
+        updater->WriteToCommandPipe(
+            android::base::StringPrintf("log bytes_stashed_%s: %" PRIu64, partition + 1,
+                                        static_cast<uint64_t>(params.stashed) * BLOCKSIZE),
+            true);
       }
       // Delete stash only after successfully completing the update, as it may contain blocks needed
       // to complete the update later.
@@ -2172,8 +2171,11 @@ Value* CheckFirstBlockFn(const char* name, State* state,
   uint16_t mount_count = *reinterpret_cast<uint16_t*>(&block0_buffer[0x400 + 0x34]);
 
   if (mount_count > 0) {
-    uiPrintf(state, "Device was remounted R/W %" PRIu16 " times", mount_count);
-    uiPrintf(state, "Last remount happened on %s", ctime(&mount_time));
+    auto updater = static_cast<Updater*>(state->cookie);
+    updater->UiPrint(
+        android::base::StringPrintf("Device was remounted R/W %" PRIu16 " times", mount_count));
+    updater->UiPrint(
+        android::base::StringPrintf("Last remount happened on %s", ctime(&mount_time)));
   }
 
   return StringValue("t");
diff --git a/updater/include/updater/install.h b/updater/include/updater/install.h
index 8d6ca4728..9fe203149 100644
--- a/updater/include/updater/install.h
+++ b/updater/include/updater/install.h
@@ -14,15 +14,6 @@
  * limitations under the License.
  */
 
-#ifndef _UPDATER_INSTALL_H_
-#define _UPDATER_INSTALL_H_
-
-struct State;
+#pragma once
 
 void RegisterInstallFunctions();
-
-// uiPrintf function prints msg to screen as well as logs
-void uiPrintf(State* _Nonnull state, const char* _Nonnull format, ...)
-    __attribute__((__format__(printf, 2, 3)));
-
-#endif
diff --git a/updater/include/updater/updater.h b/updater/include/updater/updater.h
index f4a2fe874..d5468292b 100644
--- a/updater/include/updater/updater.h
+++ b/updater/include/updater/updater.h
@@ -14,22 +14,75 @@
  * limitations under the License.
  */
 
-#ifndef _UPDATER_UPDATER_H_
-#define _UPDATER_UPDATER_H_
+#pragma once
 
+#include <stdint.h>
 #include <stdio.h>
-#include <ziparchive/zip_archive.h>
 
-typedef struct {
-    FILE* cmd_pipe;
-    ZipArchiveHandle package_zip;
-    int version;
+#include <memory>
+#include <string>
+
+#include <ziparchive/zip_archive.h>
 
-    uint8_t* package_zip_addr;
-    size_t package_zip_len;
-} UpdaterInfo;
+#include "edify/expr.h"
+#include "otautil/error_code.h"
+#include "otautil/sysutil.h"
 
 struct selabel_handle;
-extern struct selabel_handle *sehandle;
 
-#endif
+class Updater {
+ public:
+  ~Updater();
+
+  // Memory-maps the OTA package and opens it as a zip file. Also sets up the command pipe and
+  // selabel handle. TODO(xunchang) implement a run time environment class and move sehandle there.
+  bool Init(int fd, const std::string& package_filename, bool is_retry,
+            struct selabel_handle* sehandle);
+
+  // Parses and evaluates the updater-script in the OTA package. Reports the error code if the
+  // evaluation fails.
+  bool RunUpdate();
+
+  // Writes the message to command pipe, adds a new line in the end.
+  void WriteToCommandPipe(const std::string& message, bool flush = false) const;
+
+  // Sends over the message to recovery to print it on the screen.
+  void UiPrint(const std::string& message) const;
+
+  ZipArchiveHandle package_handle() const {
+    return package_handle_;
+  }
+  struct selabel_handle* sehandle() const {
+    return sehandle_;
+  }
+  std::string result() const {
+    return result_;
+  }
+
+  uint8_t* GetMappedPackageAddress() const {
+    return mapped_package_.addr;
+  }
+
+ private:
+  friend class UpdaterTestBase;
+  friend class UpdaterTest;
+  // Where in the package we expect to find the edify script to execute.
+  // (Note it's "updateR-script", not the older "update-script".)
+  static constexpr const char* SCRIPT_NAME = "META-INF/com/google/android/updater-script";
+
+  // Reads the entry |name| in the zip archive and put the result in |content|.
+  bool ReadEntryToString(ZipArchiveHandle za, const std::string& entry_name, std::string* content);
+
+  // Parses the error code embedded in state->errmsg; and reports the error code and cause code.
+  void ParseAndReportErrorCode(State* state);
+
+  MemMapping mapped_package_;
+  ZipArchiveHandle package_handle_{ nullptr };
+  std::string updater_script_;
+
+  bool is_retry_{ false };
+  std::unique_ptr<FILE, decltype(&fclose)> cmd_pipe_{ nullptr, fclose };
+  struct selabel_handle* sehandle_{ nullptr };
+
+  std::string result_;
+};
diff --git a/updater/install.cpp b/updater/install.cpp
index c30f63960..b4d88403c 100644
--- a/updater/install.cpp
+++ b/updater/install.cpp
@@ -64,36 +64,6 @@
 #include "otautil/sysutil.h"
 #include "updater/updater.h"
 
-// Send over the buffer to recovery though the command pipe.
-static void uiPrint(State* state, const std::string& buffer) {
-  UpdaterInfo* ui = static_cast<UpdaterInfo*>(state->cookie);
-
-  // "line1\nline2\n" will be split into 3 tokens: "line1", "line2" and "".
-  // So skip sending empty strings to UI.
-  std::vector<std::string> lines = android::base::Split(buffer, "\n");
-  for (auto& line : lines) {
-    if (!line.empty()) {
-      fprintf(ui->cmd_pipe, "ui_print %s\n", line.c_str());
-    }
-  }
-
-  // On the updater side, we need to dump the contents to stderr (which has
-  // been redirected to the log file). Because the recovery will only print
-  // the contents to screen when processing pipe command ui_print.
-  LOG(INFO) << buffer;
-}
-
-void uiPrintf(State* _Nonnull state, const char* _Nonnull format, ...) {
-  std::string error_msg;
-
-  va_list ap;
-  va_start(ap, format);
-  android::base::StringAppendV(&error_msg, format, ap);
-  va_end(ap);
-
-  uiPrint(state, error_msg);
-}
-
 // This is the updater side handler for ui_print() in edify script. Contents will be sent over to
 // the recovery side for on-screen display.
 Value* UIPrintFn(const char* name, State* state, const std::vector<std::unique_ptr<Expr>>& argv) {
@@ -103,7 +73,7 @@ Value* UIPrintFn(const char* name, State* state, const std::vector<std::unique_p
   }
 
   std::string buffer = android::base::Join(args, "");
-  uiPrint(state, buffer);
+  static_cast<Updater*>(state->cookie)->UiPrint(buffer);
   return StringValue(buffer);
 }
 
@@ -129,7 +99,7 @@ Value* PackageExtractFileFn(const char* name, State* state,
     const std::string& zip_path = args[0];
     const std::string& dest_path = args[1];
 
-    ZipArchiveHandle za = static_cast<UpdaterInfo*>(state->cookie)->package_zip;
+    ZipArchiveHandle za = static_cast<Updater*>(state->cookie)->package_handle();
     ZipEntry entry;
     if (FindEntry(za, zip_path, &entry) != 0) {
       LOG(ERROR) << name << ": no " << zip_path << " in package";
@@ -172,7 +142,7 @@ Value* PackageExtractFileFn(const char* name, State* state,
     }
     const std::string& zip_path = args[0];
 
-    ZipArchiveHandle za = static_cast<UpdaterInfo*>(state->cookie)->package_zip;
+    ZipArchiveHandle za = static_cast<Updater*>(state->cookie)->package_handle();
     ZipEntry entry;
     if (FindEntry(za, zip_path, &entry) != 0) {
       return ErrorAbort(state, kPackageExtractFileFailure, "%s(): no %s in package", name,
@@ -311,11 +281,11 @@ Value* MountFn(const char* name, State* state, const std::vector<std::unique_ptr
                       name);
   }
 
+  auto updater = static_cast<Updater*>(state->cookie);
   {
     char* secontext = nullptr;
-
-    if (sehandle) {
-      selabel_lookup(sehandle, &secontext, mount_point.c_str(), 0755);
+    if (updater->sehandle()) {
+      selabel_lookup(updater->sehandle(), &secontext, mount_point.c_str(), 0755);
       setfscreatecon(secontext);
     }
 
@@ -329,8 +299,9 @@ Value* MountFn(const char* name, State* state, const std::vector<std::unique_ptr
 
   if (mount(location.c_str(), mount_point.c_str(), fs_type.c_str(),
             MS_NOATIME | MS_NODEV | MS_NODIRATIME, mount_options.c_str()) < 0) {
-    uiPrintf(state, "%s: Failed to mount %s at %s: %s", name, location.c_str(), mount_point.c_str(),
-             strerror(errno));
+    updater->UiPrint(android::base::StringPrintf("%s: Failed to mount %s at %s: %s", name,
+                                                 location.c_str(), mount_point.c_str(),
+                                                 strerror(errno)));
     return StringValue("");
   }
 
@@ -376,15 +347,18 @@ Value* UnmountFn(const char* name, State* state, const std::vector<std::unique_p
                       "mount_point argument to unmount() can't be empty");
   }
 
+  auto updater = static_cast<Updater*>(state->cookie);
   scan_mounted_volumes();
   MountedVolume* vol = find_mounted_volume_by_mount_point(mount_point.c_str());
   if (vol == nullptr) {
-    uiPrintf(state, "Failed to unmount %s: No such volume", mount_point.c_str());
+    updater->UiPrint(
+        android::base::StringPrintf("Failed to unmount %s: No such volume", mount_point.c_str()));
     return nullptr;
   } else {
     int ret = unmount_mounted_volume(vol);
     if (ret != 0) {
-      uiPrintf(state, "Failed to unmount %s: %s", mount_point.c_str(), strerror(errno));
+      updater->UiPrint(android::base::StringPrintf("Failed to unmount %s: %s", mount_point.c_str(),
+                                                   strerror(errno)));
     }
   }
 
@@ -529,8 +503,8 @@ Value* ShowProgressFn(const char* name, State* state,
                       sec_str.c_str());
   }
 
-  UpdaterInfo* ui = static_cast<UpdaterInfo*>(state->cookie);
-  fprintf(ui->cmd_pipe, "progress %f %d\n", frac, sec);
+  auto updater = static_cast<Updater*>(state->cookie);
+  updater->WriteToCommandPipe(android::base::StringPrintf("progress %f %d", frac, sec));
 
   return StringValue(frac_str);
 }
@@ -553,8 +527,8 @@ Value* SetProgressFn(const char* name, State* state,
                       frac_str.c_str());
   }
 
-  UpdaterInfo* ui = static_cast<UpdaterInfo*>(state->cookie);
-  fprintf(ui->cmd_pipe, "set_progress %f\n", frac);
+  auto updater = static_cast<Updater*>(state->cookie);
+  updater->WriteToCommandPipe(android::base::StringPrintf("set_progress %f", frac));
 
   return StringValue(frac_str);
 }
@@ -653,7 +627,8 @@ Value* WipeCacheFn(const char* name, State* state, const std::vector<std::unique
     return ErrorAbort(state, kArgsParsingFailure, "%s() expects no args, got %zu", name,
                       argv.size());
   }
-  fprintf(static_cast<UpdaterInfo*>(state->cookie)->cmd_pipe, "wipe_cache\n");
+
+  static_cast<Updater*>(state->cookie)->WriteToCommandPipe("wipe_cache");
   return StringValue("t");
 }
 
@@ -881,8 +856,7 @@ Value* EnableRebootFn(const char* name, State* state, const std::vector<std::uni
     return ErrorAbort(state, kArgsParsingFailure, "%s() expects no args, got %zu", name,
                       argv.size());
   }
-  UpdaterInfo* ui = static_cast<UpdaterInfo*>(state->cookie);
-  fprintf(ui->cmd_pipe, "enable_reboot\n");
+  static_cast<Updater*>(state->cookie)->WriteToCommandPipe("enable_reboot");
   return StringValue("t");
 }
 
diff --git a/updater/updater.cpp b/updater/updater.cpp
index a020699ca..e0679fb0c 100644
--- a/updater/updater.cpp
+++ b/updater/updater.cpp
@@ -16,8 +16,6 @@
 
 #include "updater/updater.h"
 
-#include <stdio.h>
-#include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
 
@@ -25,197 +23,155 @@
 
 #include <android-base/logging.h>
 #include <android-base/strings.h>
-#include <selinux/android.h>
-#include <selinux/label.h>
-#include <selinux/selinux.h>
-#include <ziparchive/zip_archive.h>
-
-#include "edify/expr.h"
-#include "otautil/dirutil.h"
-#include "otautil/error_code.h"
-#include "otautil/sysutil.h"
-#include "updater/blockimg.h"
-#include "updater/dynamic_partitions.h"
-#include "updater/install.h"
-
-// Generated by the makefile, this function defines the
-// RegisterDeviceExtensions() function, which calls all the
-// registration functions for device-specific extensions.
-#include "register.inc"
-
-// Where in the package we expect to find the edify script to execute.
-// (Note it's "updateR-script", not the older "update-script".)
-static constexpr const char* SCRIPT_NAME = "META-INF/com/google/android/updater-script";
-
-struct selabel_handle *sehandle;
-
-static void UpdaterLogger(android::base::LogId /* id */, android::base::LogSeverity /* severity */,
-                          const char* /* tag */, const char* /* file */, unsigned int /* line */,
-                          const char* message) {
-  fprintf(stdout, "%s\n", message);
-}
-
-int main(int argc, char** argv) {
-  // Various things log information to stdout or stderr more or less
-  // at random (though we've tried to standardize on stdout).  The
-  // log file makes more sense if buffering is turned off so things
-  // appear in the right order.
-  setbuf(stdout, nullptr);
-  setbuf(stderr, nullptr);
-
-  // We don't have logcat yet under recovery. Update logs will always be written to stdout
-  // (which is redirected to recovery.log).
-  android::base::InitLogging(argv, &UpdaterLogger);
-
-  if (argc != 4 && argc != 5) {
-    LOG(ERROR) << "unexpected number of arguments: " << argc;
-    return 1;
-  }
 
-  char* version = argv[1];
-  if ((version[0] != '1' && version[0] != '2' && version[0] != '3') || version[1] != '\0') {
-    // We support version 1, 2, or 3.
-    LOG(ERROR) << "wrong updater binary API; expected 1, 2, or 3; got " << argv[1];
-    return 2;
+Updater::~Updater() {
+  if (package_handle_) {
+    CloseArchive(package_handle_);
   }
+}
 
+bool Updater::Init(int fd, const std::string& package_filename, bool is_retry,
+                   struct selabel_handle* sehandle) {
   // Set up the pipe for sending commands back to the parent process.
+  cmd_pipe_.reset(fdopen(fd, "wb"));
+  if (!cmd_pipe_) {
+    LOG(ERROR) << "Failed to open the command pipe";
+    return false;
+  }
 
-  int fd = atoi(argv[2]);
-  FILE* cmd_pipe = fdopen(fd, "wb");
-  setlinebuf(cmd_pipe);
-
-  // Extract the script from the package.
+  setlinebuf(cmd_pipe_.get());
 
-  const char* package_filename = argv[3];
-  MemMapping map;
-  if (!map.MapFile(package_filename)) {
-    LOG(ERROR) << "failed to map package " << argv[3];
-    return 3;
+  if (!mapped_package_.MapFile(package_filename)) {
+    LOG(ERROR) << "failed to map package " << package_filename;
+    return false;
   }
-  ZipArchiveHandle za;
-  int open_err = OpenArchiveFromMemory(map.addr, map.length, argv[3], &za);
-  if (open_err != 0) {
-    LOG(ERROR) << "failed to open package " << argv[3] << ": " << ErrorCodeString(open_err);
-    CloseArchive(za);
-    return 3;
+  if (int open_err = OpenArchiveFromMemory(mapped_package_.addr, mapped_package_.length,
+                                           package_filename.c_str(), &package_handle_);
+      open_err != 0) {
+    LOG(ERROR) << "failed to open package " << package_filename << ": "
+               << ErrorCodeString(open_err);
+    return false;
   }
-
-  ZipEntry script_entry;
-  int find_err = FindEntry(za, SCRIPT_NAME, &script_entry);
-  if (find_err != 0) {
-    LOG(ERROR) << "failed to find " << SCRIPT_NAME << " in " << package_filename << ": "
-               << ErrorCodeString(find_err);
-    CloseArchive(za);
-    return 4;
+  if (!ReadEntryToString(package_handle_, SCRIPT_NAME, &updater_script_)) {
+    return false;
   }
 
-  std::string script;
-  script.resize(script_entry.uncompressed_length);
-  int extract_err = ExtractToMemory(za, &script_entry, reinterpret_cast<uint8_t*>(&script[0]),
-                                    script_entry.uncompressed_length);
-  if (extract_err != 0) {
-    LOG(ERROR) << "failed to read script from package: " << ErrorCodeString(extract_err);
-    CloseArchive(za);
-    return 5;
-  }
-
-  // Configure edify's functions.
+  is_retry_ = is_retry;
 
-  RegisterBuiltins();
-  RegisterInstallFunctions();
-  RegisterBlockImageFunctions();
-  RegisterDynamicPartitionsFunctions();
-  RegisterDeviceExtensions();
+  sehandle_ = sehandle;
+  if (!sehandle_) {
+    fprintf(cmd_pipe_.get(), "ui_print Warning: No file_contexts\n");
+  }
+  return true;
+}
 
+bool Updater::RunUpdate() {
   // Parse the script.
-
   std::unique_ptr<Expr> root;
   int error_count = 0;
-  int error = ParseString(script, &root, &error_count);
+  int error = ParseString(updater_script_, &root, &error_count);
   if (error != 0 || error_count > 0) {
     LOG(ERROR) << error_count << " parse errors";
-    CloseArchive(za);
-    return 6;
-  }
-
-  sehandle = selinux_android_file_context_handle();
-  selinux_android_set_sehandle(sehandle);
-
-  if (!sehandle) {
-    fprintf(cmd_pipe, "ui_print Warning: No file_contexts\n");
+    return false;
   }
 
   // Evaluate the parsed script.
+  State state(updater_script_, this);
+  state.is_retry = is_retry_;
+
+  bool status = Evaluate(&state, root, &result_);
+  if (status) {
+    fprintf(cmd_pipe_.get(), "ui_print script succeeded: result was [%s]\n", result_.c_str());
+    // Even though the script doesn't abort, still log the cause code if result is empty.
+    if (result_.empty() && state.cause_code != kNoCause) {
+      fprintf(cmd_pipe_.get(), "log cause: %d\n", state.cause_code);
+    }
+    return true;
+  }
 
-  UpdaterInfo updater_info;
-  updater_info.cmd_pipe = cmd_pipe;
-  updater_info.package_zip = za;
-  updater_info.version = atoi(version);
-  updater_info.package_zip_addr = map.addr;
-  updater_info.package_zip_len = map.length;
+  ParseAndReportErrorCode(&state);
+  return false;
+}
 
-  State state(script, &updater_info);
+void Updater::WriteToCommandPipe(const std::string& message, bool flush) const {
+  fprintf(cmd_pipe_.get(), "%s\n", message.c_str());
+  if (flush) {
+    fflush(cmd_pipe_.get());
+  }
+}
 
-  if (argc == 5) {
-    if (strcmp(argv[4], "retry") == 0) {
-      state.is_retry = true;
-    } else {
-      printf("unexpected argument: %s", argv[4]);
+void Updater::UiPrint(const std::string& message) const {
+  // "line1\nline2\n" will be split into 3 tokens: "line1", "line2" and "".
+  // so skip sending empty strings to ui.
+  std::vector<std::string> lines = android::base::Split(message, "\n");
+  for (const auto& line : lines) {
+    if (!line.empty()) {
+      fprintf(cmd_pipe_.get(), "ui_print %s\n", line.c_str());
     }
   }
 
-  std::string result;
-  bool status = Evaluate(&state, root, &result);
-
-  if (!status) {
-    if (state.errmsg.empty()) {
-      LOG(ERROR) << "script aborted (no error message)";
-      fprintf(cmd_pipe, "ui_print script aborted (no error message)\n");
-    } else {
-      LOG(ERROR) << "script aborted: " << state.errmsg;
-      const std::vector<std::string> lines = android::base::Split(state.errmsg, "\n");
-      for (const std::string& line : lines) {
-        // Parse the error code in abort message.
-        // Example: "E30: This package is for bullhead devices."
-        if (!line.empty() && line[0] == 'E') {
-          if (sscanf(line.c_str(), "E%d: ", &state.error_code) != 1) {
-            LOG(ERROR) << "Failed to parse error code: [" << line << "]";
-          }
+  // on the updater side, we need to dump the contents to stderr (which has
+  // been redirected to the log file). because the recovery will only print
+  // the contents to screen when processing pipe command ui_print.
+  LOG(INFO) << message;
+}
+
+void Updater::ParseAndReportErrorCode(State* state) {
+  CHECK(state);
+  if (state->errmsg.empty()) {
+    LOG(ERROR) << "script aborted (no error message)";
+    fprintf(cmd_pipe_.get(), "ui_print script aborted (no error message)\n");
+  } else {
+    LOG(ERROR) << "script aborted: " << state->errmsg;
+    const std::vector<std::string> lines = android::base::Split(state->errmsg, "\n");
+    for (const std::string& line : lines) {
+      // Parse the error code in abort message.
+      // Example: "E30: This package is for bullhead devices."
+      if (!line.empty() && line[0] == 'E') {
+        if (sscanf(line.c_str(), "E%d: ", &state->error_code) != 1) {
+          LOG(ERROR) << "Failed to parse error code: [" << line << "]";
         }
-        fprintf(cmd_pipe, "ui_print %s\n", line.c_str());
       }
+      fprintf(cmd_pipe_.get(), "ui_print %s\n", line.c_str());
     }
+  }
 
-    // Installation has been aborted. Set the error code to kScriptExecutionFailure unless
-    // a more specific code has been set in errmsg.
-    if (state.error_code == kNoError) {
-      state.error_code = kScriptExecutionFailure;
-    }
-    fprintf(cmd_pipe, "log error: %d\n", state.error_code);
-    // Cause code should provide additional information about the abort.
-    if (state.cause_code != kNoCause) {
-      fprintf(cmd_pipe, "log cause: %d\n", state.cause_code);
-      if (state.cause_code == kPatchApplicationFailure) {
-        LOG(INFO) << "Patch application failed, retry update.";
-        fprintf(cmd_pipe, "retry_update\n");
-      } else if (state.cause_code == kEioFailure) {
-        LOG(INFO) << "Update failed due to EIO, retry update.";
-        fprintf(cmd_pipe, "retry_update\n");
-      }
+  // Installation has been aborted. Set the error code to kScriptExecutionFailure unless
+  // a more specific code has been set in errmsg.
+  if (state->error_code == kNoError) {
+    state->error_code = kScriptExecutionFailure;
+  }
+  fprintf(cmd_pipe_.get(), "log error: %d\n", state->error_code);
+  // Cause code should provide additional information about the abort.
+  if (state->cause_code != kNoCause) {
+    fprintf(cmd_pipe_.get(), "log cause: %d\n", state->cause_code);
+    if (state->cause_code == kPatchApplicationFailure) {
+      LOG(INFO) << "Patch application failed, retry update.";
+      fprintf(cmd_pipe_.get(), "retry_update\n");
+    } else if (state->cause_code == kEioFailure) {
+      LOG(INFO) << "Update failed due to EIO, retry update.";
+      fprintf(cmd_pipe_.get(), "retry_update\n");
     }
+  }
+}
 
-    if (updater_info.package_zip) {
-      CloseArchive(updater_info.package_zip);
-    }
-    return 7;
-  } else {
-    fprintf(cmd_pipe, "ui_print script succeeded: result was [%s]\n", result.c_str());
+bool Updater::ReadEntryToString(ZipArchiveHandle za, const std::string& entry_name,
+                                std::string* content) {
+  ZipEntry entry;
+  int find_err = FindEntry(za, entry_name, &entry);
+  if (find_err != 0) {
+    LOG(ERROR) << "failed to find " << entry_name
+               << " in the package: " << ErrorCodeString(find_err);
+    return false;
   }
 
-  if (updater_info.package_zip) {
-    CloseArchive(updater_info.package_zip);
+  content->resize(entry.uncompressed_length);
+  int extract_err = ExtractToMemory(za, &entry, reinterpret_cast<uint8_t*>(&content->at(0)),
+                                    entry.uncompressed_length);
+  if (extract_err != 0) {
+    LOG(ERROR) << "failed to read script from package: " << ErrorCodeString(extract_err);
+    return false;
   }
 
-  return 0;
+  return true;
 }
diff --git a/updater/updater_main.cpp b/updater/updater_main.cpp
new file mode 100644
index 000000000..dd22c137d
--- /dev/null
+++ b/updater/updater_main.cpp
@@ -0,0 +1,108 @@
+/*
+ * Copyright (C) 2019 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 <stdio.h>
+#include <string.h>
+#include <unistd.h>
+
+#include <string>
+
+#include <android-base/logging.h>
+#include <android-base/parseint.h>
+#include <selinux/android.h>
+#include <selinux/label.h>
+#include <selinux/selinux.h>
+
+#include "edify/expr.h"
+#include "updater/blockimg.h"
+#include "updater/dynamic_partitions.h"
+#include "updater/install.h"
+#include "updater/updater.h"
+
+// Generated by the makefile, this function defines the
+// RegisterDeviceExtensions() function, which calls all the
+// registration functions for device-specific extensions.
+#include "register.inc"
+
+static void UpdaterLogger(android::base::LogId /* id */, android::base::LogSeverity /* severity */,
+                          const char* /* tag */, const char* /* file */, unsigned int /* line */,
+                          const char* message) {
+  fprintf(stdout, "%s\n", message);
+}
+
+int main(int argc, char** argv) {
+  // Various things log information to stdout or stderr more or less
+  // at random (though we've tried to standardize on stdout).  The
+  // log file makes more sense if buffering is turned off so things
+  // appear in the right order.
+  setbuf(stdout, nullptr);
+  setbuf(stderr, nullptr);
+
+  // We don't have logcat yet under recovery. Update logs will always be written to stdout
+  // (which is redirected to recovery.log).
+  android::base::InitLogging(argv, &UpdaterLogger);
+
+  if (argc != 4 && argc != 5) {
+    LOG(ERROR) << "unexpected number of arguments: " << argc;
+    return 1;
+  }
+
+  char* version = argv[1];
+  if ((version[0] != '1' && version[0] != '2' && version[0] != '3') || version[1] != '\0') {
+    // We support version 1, 2, or 3.
+    LOG(ERROR) << "wrong updater binary API; expected 1, 2, or 3; got " << argv[1];
+    return 1;
+  }
+
+  int fd;
+  if (!android::base::ParseInt(argv[2], &fd)) {
+    LOG(ERROR) << "Failed to parse fd in " << argv[2];
+    return 1;
+  }
+
+  std::string package_name = argv[3];
+
+  bool is_retry = false;
+  if (argc == 5) {
+    if (strcmp(argv[4], "retry") == 0) {
+      is_retry = true;
+    } else {
+      LOG(ERROR) << "unexpected argument: " << argv[4];
+      return 1;
+    }
+  }
+
+  // Configure edify's functions.
+  RegisterBuiltins();
+  RegisterInstallFunctions();
+  RegisterBlockImageFunctions();
+  RegisterDynamicPartitionsFunctions();
+  RegisterDeviceExtensions();
+
+  auto sehandle = selinux_android_file_context_handle();
+  selinux_android_set_sehandle(sehandle);
+
+  Updater updater;
+  if (!updater.Init(fd, package_name, is_retry, sehandle)) {
+    return 1;
+  }
+
+  if (!updater.RunUpdate()) {
+    return 1;
+  }
+
+  return 0;
+}
\ No newline at end of file