12 files changed, 328 insertions, 27 deletions

diff --git a/Android.mk b/Android.mk
index 7b2843d17..b25c1f07a 100644
--- a/Android.mk
+++ b/Android.mk
@@ -28,7 +28,32 @@ recovery_common_cflags := \
     -Werror \
     -DRECOVERY_API_VERSION=$(RECOVERY_API_VERSION)
 
-# librecovery_ui (static library)
+# librecovery_ui_ext (shared library)
+# ===================================
+include $(CLEAR_VARS)
+
+LOCAL_MODULE := librecovery_ui_ext
+
+# LOCAL_MODULE_PATH for shared libraries is unsupported in multiarch builds.
+LOCAL_MULTILIB := first
+
+ifeq ($(TARGET_IS_64_BIT),true)
+LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/system/lib64
+else
+LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/system/lib
+endif
+
+LOCAL_WHOLE_STATIC_LIBRARIES := \
+    $(TARGET_RECOVERY_UI_LIB)
+
+LOCAL_SHARED_LIBRARIES := \
+    libbase \
+    liblog \
+    librecovery_ui
+
+include $(BUILD_SHARED_LIBRARY)
+
+# librecovery_ui (shared library)
 # ===============================
 include $(CLEAR_VARS)
 LOCAL_SRC_FILES := \
@@ -40,13 +65,50 @@ LOCAL_SRC_FILES := \
 
 LOCAL_MODULE := librecovery_ui
 
+LOCAL_CFLAGS := $(recovery_common_cflags)
+
+LOCAL_MULTILIB := first
+
+ifeq ($(TARGET_IS_64_BIT),true)
+LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/system/lib64
+else
+LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/system/lib
+endif
+
 LOCAL_STATIC_LIBRARIES := \
     libminui \
     libotautil \
-    libbase
+
+LOCAL_SHARED_LIBRARIES := \
+    libbase \
+    libpng \
+    libz \
+
+include $(BUILD_SHARED_LIBRARY)
+
+# librecovery_ui (static library)
+# ===============================
+include $(CLEAR_VARS)
+LOCAL_SRC_FILES := \
+    device.cpp \
+    screen_ui.cpp \
+    ui.cpp \
+    vr_ui.cpp \
+    wear_ui.cpp
+
+LOCAL_MODULE := librecovery_ui
 
 LOCAL_CFLAGS := $(recovery_common_cflags)
 
+LOCAL_STATIC_LIBRARIES := \
+    libminui \
+    libotautil \
+
+LOCAL_SHARED_LIBRARIES := \
+    libbase \
+    libpng \
+    libz \
+
 include $(BUILD_STATIC_LIBRARY)
 
 # Health HAL dependency
@@ -63,11 +125,9 @@ health_hal_static_libraries := \
     libbatterymonitor
 
 librecovery_static_libraries := \
-    $(TARGET_RECOVERY_UI_LIB) \
     libbootloader_message \
     libfusesideload \
     libminadbd \
-    librecovery_ui \
     libminui \
     libverifier \
     libotautil \
@@ -125,8 +185,6 @@ LOCAL_SRC_FILES := \
 
 LOCAL_MODULE := recovery
 
-LOCAL_FORCE_STATIC_EXECUTABLE := true
-
 LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/system/bin
 
 # Cannot link with LLD: undefined symbol: UsbNoPermissionsLongHelpText
@@ -137,8 +195,12 @@ LOCAL_CFLAGS := $(recovery_common_cflags)
 
 LOCAL_STATIC_LIBRARIES := \
     librecovery \
+    librecovery_ui_default \
     $(librecovery_static_libraries)
 
+LOCAL_SHARED_LIBRARIES := \
+    librecovery_ui \
+
 LOCAL_HAL_STATIC_LIBRARIES := libhealthd
 
 LOCAL_REQUIRED_MODULES := \
@@ -167,6 +229,17 @@ LOCAL_REQUIRED_MODULES += \
     recovery-refresh
 endif
 
+LOCAL_REQUIRED_MODULES += \
+    librecovery_ui_ext
+
+# TODO(b/110380063): Explicitly install the following shared libraries to recovery, until `recovery`
+# module is built with Soong (with `recovery: true` flag).
+LOCAL_REQUIRED_MODULES += \
+    libbase.recovery \
+    liblog.recovery \
+    libpng.recovery \
+    libz.recovery \
+
 include $(BUILD_EXECUTABLE)
 
 include \
diff --git a/device.h b/device.h
index cbecc4371..a6ad62788 100644
--- a/device.h
+++ b/device.h
@@ -119,8 +119,12 @@ class Device {
   std::unique_ptr<RecoveryUI> ui_;
 };
 
+// Disable name mangling, as this function will be loaded via dlsym(3).
+extern "C" {
+
 // The device-specific library must define this function (or the default one will be used, if there
 // is no device-specific library). It returns the Device object that recovery should use.
 Device* make_device();
+}
 
 #endif  // _DEVICE_H
diff --git a/etc/init.rc b/etc/init.rc
index 8e18438b5..3821eb6a8 100644
--- a/etc/init.rc
+++ b/etc/init.rc
@@ -77,7 +77,7 @@ on late-init
     trigger early-boot
     trigger boot
 
-service ueventd /sbin/ueventd
+service ueventd /system/bin/ueventd
     critical
     seclabel u:r:ueventd:s0
 
diff --git a/otautil/include/otautil/error_code.h b/otautil/include/otautil/error_code.h
index b0ff42d8d..0f6c9f85f 100644
--- a/otautil/include/otautil/error_code.h
+++ b/otautil/include/otautil/error_code.h
@@ -48,6 +48,7 @@ enum CauseCode : int {
   kRebootFailure,
   kPackageExtractFileFailure,
   kPatchApplicationFailure,
+  kHashTreeComputationFailure,
   kVendorFailure = 200
 };
 
diff --git a/recovery_main.cpp b/recovery_main.cpp
index c79d7d8d8..9a9890de0 100644
--- a/recovery_main.cpp
+++ b/recovery_main.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <dlfcn.h>
 #include <errno.h>
 #include <fcntl.h>
 #include <getopt.h>
@@ -329,7 +330,32 @@ int main(int argc, char** argv) {
 
   printf("locale is [%s]\n", locale.c_str());
 
-  Device* device = make_device();
+  static constexpr const char* kDefaultLibRecoveryUIExt = "librecovery_ui_ext.so";
+  // Intentionally not calling dlclose(3) to avoid potential gotchas (e.g. `make_device` may have
+  // handed out pointers to code or static [or thread-local] data and doesn't collect them all back
+  // in on dlclose).
+  void* librecovery_ui_ext = dlopen(kDefaultLibRecoveryUIExt, RTLD_NOW);
+
+  using MakeDeviceType = decltype(&make_device);
+  MakeDeviceType make_device_func = nullptr;
+  if (librecovery_ui_ext == nullptr) {
+    printf("Failed to dlopen %s: %s\n", kDefaultLibRecoveryUIExt, dlerror());
+  } else {
+    reinterpret_cast<void*&>(make_device_func) = dlsym(librecovery_ui_ext, "make_device");
+    if (make_device_func == nullptr) {
+      printf("Failed to dlsym make_device: %s\n", dlerror());
+    }
+  }
+
+  Device* device;
+  if (make_device_func == nullptr) {
+    printf("Falling back to the default make_device() instead\n");
+    device = make_device();
+  } else {
+    printf("Loading make_device from %s\n", kDefaultLibRecoveryUIExt);
+    device = (*make_device_func)();
+  }
+
   if (android::base::GetBoolProperty("ro.boot.quiescent", false)) {
     printf("Quiescent recovery mode.\n");
     device->ResetUI(new StubRecoveryUI());
diff --git a/tests/Android.mk b/tests/Android.mk
index 93286ea0d..3d3e63e7e 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -137,6 +137,7 @@ libupdater_static_libraries := \
     libext4_utils \
     libfec \
     libfec_rs \
+    libverity_tree \
     libfs_mgr \
     libgtest_prod \
     liblog \
@@ -167,10 +168,10 @@ health_hal_static_libraries := \
 
 librecovery_static_libraries := \
     librecovery \
-    $(TARGET_RECOVERY_UI_LIB) \
     libbootloader_message \
     libfusesideload \
     libminadbd \
+    librecovery_ui_default \
     librecovery_ui \
     libminui \
     libverifier \
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index 9fcf17f13..248b469b0 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -37,6 +37,7 @@
 #include <brotli/encode.h>
 #include <bsdiff/bsdiff.h>
 #include <gtest/gtest.h>
+#include <verity/hash_tree_builder.h>
 #include <ziparchive/zip_archive.h>
 #include <ziparchive/zip_writer.h>
 
@@ -389,6 +390,86 @@ TEST_F(UpdaterTest, read_file) {
   expect("", script, kNoCause);
 }
 
+TEST_F(UpdaterTest, compute_hash_tree_smoke) {
+  std::string data;
+  for (unsigned char i = 0; i < 128; i++) {
+    data += std::string(4096, i);
+  }
+  // Appends an additional block for verity data.
+  data += std::string(4096, 0);
+  ASSERT_EQ(129 * 4096, data.size());
+  ASSERT_TRUE(android::base::WriteStringToFile(data, image_file_));
+
+  std::string salt = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7";
+  std::string expected_root_hash =
+      "7e0a8d8747f54384014ab996f5b2dc4eb7ff00c630eede7134c9e3f05c0dd8ca";
+  // hash_tree_ranges, source_ranges, hash_algorithm, salt_hex, root_hash
+  std::vector<std::string> tokens{ "compute_hash_tree", "2,128,129", "2,0,128", "sha256", salt,
+                                   expected_root_hash };
+  std::string hash_tree_command = android::base::Join(tokens, " ");
+
+  std::vector<std::string> transfer_list{
+    "4", "2", "0", "2", hash_tree_command,
+  };
+
+  PackageEntries entries{
+    { "new_data", "" },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, "\n") },
+  };
+
+  RunBlockImageUpdate(false, entries, image_file_, "t");
+
+  std::string updated;
+  ASSERT_TRUE(android::base::ReadFileToString(image_file_, &updated));
+  ASSERT_EQ(129 * 4096, updated.size());
+  ASSERT_EQ(data.substr(0, 128 * 4096), updated.substr(0, 128 * 4096));
+
+  // Computes the SHA256 of the salt + hash_tree_data and expects the result to match with the
+  // root_hash.
+  std::vector<unsigned char> salt_bytes;
+  ASSERT_TRUE(HashTreeBuilder::ParseBytesArrayFromString(salt, &salt_bytes));
+  std::vector<unsigned char> hash_tree = std::move(salt_bytes);
+  hash_tree.insert(hash_tree.end(), updated.begin() + 128 * 4096, updated.end());
+
+  std::vector<unsigned char> digest(SHA256_DIGEST_LENGTH);
+  SHA256(hash_tree.data(), hash_tree.size(), digest.data());
+  ASSERT_EQ(expected_root_hash, HashTreeBuilder::BytesArrayToString(digest));
+}
+
+TEST_F(UpdaterTest, compute_hash_tree_root_mismatch) {
+  std::string data;
+  for (size_t i = 0; i < 128; i++) {
+    data += std::string(4096, i);
+  }
+  // Appends an additional block for verity data.
+  data += std::string(4096, 0);
+  ASSERT_EQ(129 * 4096, data.size());
+  // Corrupts one bit
+  data[4096] = 'A';
+  ASSERT_TRUE(android::base::WriteStringToFile(data, image_file_));
+
+  std::string salt = "aee087a5be3b982978c923f566a94613496b417f2af592639bc80d141e34dfe7";
+  std::string expected_root_hash =
+      "7e0a8d8747f54384014ab996f5b2dc4eb7ff00c630eede7134c9e3f05c0dd8ca";
+  // hash_tree_ranges, source_ranges, hash_algorithm, salt_hex, root_hash
+  std::vector<std::string> tokens{ "compute_hash_tree", "2,128,129", "2,0,128", "sha256", salt,
+                                   expected_root_hash };
+  std::string hash_tree_command = android::base::Join(tokens, " ");
+
+  std::vector<std::string> transfer_list{
+    "4", "2", "0", "2", hash_tree_command,
+  };
+
+  PackageEntries entries{
+    { "new_data", "" },
+    { "patch_data", "" },
+    { "transfer_list", android::base::Join(transfer_list, "\n") },
+  };
+
+  RunBlockImageUpdate(false, entries, image_file_, "", kHashTreeComputationFailure);
+}
+
 TEST_F(UpdaterTest, write_value) {
   // write_value() expects two arguments.
   expect(nullptr, "write_value()", kArgsParsingFailure);
diff --git a/tests/unit/commands_test.cpp b/tests/unit/commands_test.cpp
index 3daa58f33..9679a9e73 100644
--- a/tests/unit/commands_test.cpp
+++ b/tests/unit/commands_test.cpp
@@ -30,6 +30,7 @@ TEST(CommandsTest, ParseType) {
   ASSERT_EQ(Command::Type::IMGDIFF, Command::ParseType("imgdiff"));
   ASSERT_EQ(Command::Type::STASH, Command::ParseType("stash"));
   ASSERT_EQ(Command::Type::FREE, Command::ParseType("free"));
+  ASSERT_EQ(Command::Type::COMPUTE_HASH_TREE, Command::ParseType("compute_hash_tree"));
 }
 
 TEST(CommandsTest, ParseType_InvalidCommand) {
diff --git a/updater/Android.mk b/updater/Android.mk
index ac9aecb04..78d0bd451 100644
--- a/updater/Android.mk
+++ b/updater/Android.mk
@@ -34,6 +34,7 @@ updater_common_static_libraries := \
     libext4_utils \
     libfec \
     libfec_rs \
+    libverity_tree \
     libfs_mgr \
     libgtest_prod \
     liblog \
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index 2a2ab19a3..96b2d9feb 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -49,6 +49,7 @@
 #include <fec/io.h>
 #include <openssl/sha.h>
 #include <private/android_filesystem_config.h>
+#include <verity/hash_tree_builder.h>
 #include <ziparchive/zip_archive.h>
 
 #include "edify/expr.h"
@@ -1495,6 +1496,105 @@ static int PerformCommandAbort(CommandParameters&) {
   return -1;
 }
 
+// Computes the hash_tree bytes based on the parameters, checks if the root hash of the tree
+// matches the expected hash and writes the result to the specified range on the block_device.
+// Hash_tree computation arguments:
+//   hash_tree_ranges
+//   source_ranges
+//   hash_algorithm
+//   salt_hex
+//   root_hash
+static int PerformCommandComputeHashTree(CommandParameters& params) {
+  if (params.cpos + 5 != params.tokens.size()) {
+    LOG(ERROR) << "Invaild arguments count in hash computation " << params.cmdline;
+    return -1;
+  }
+
+  // Expects the hash_tree data to be contiguous.
+  RangeSet hash_tree_ranges = RangeSet::Parse(params.tokens[params.cpos++]);
+  if (!hash_tree_ranges || hash_tree_ranges.size() != 1) {
+    LOG(ERROR) << "Invalid hash tree ranges in " << params.cmdline;
+    return -1;
+  }
+
+  RangeSet source_ranges = RangeSet::Parse(params.tokens[params.cpos++]);
+  if (!source_ranges) {
+    LOG(ERROR) << "Invalid source ranges in " << params.cmdline;
+    return -1;
+  }
+
+  auto hash_function = HashTreeBuilder::HashFunction(params.tokens[params.cpos++]);
+  if (hash_function == nullptr) {
+    LOG(ERROR) << "Invalid hash algorithm in " << params.cmdline;
+    return -1;
+  }
+
+  std::vector<unsigned char> salt;
+  std::string salt_hex = params.tokens[params.cpos++];
+  if (salt_hex.empty() || !HashTreeBuilder::ParseBytesArrayFromString(salt_hex, &salt)) {
+    LOG(ERROR) << "Failed to parse salt in " << params.cmdline;
+    return -1;
+  }
+
+  std::string expected_root_hash = params.tokens[params.cpos++];
+  if (expected_root_hash.empty()) {
+    LOG(ERROR) << "Invalid root hash in " << params.cmdline;
+    return -1;
+  }
+
+  // Starts the hash_tree computation.
+  HashTreeBuilder builder(BLOCKSIZE, hash_function);
+  if (!builder.Initialize(source_ranges.blocks() * BLOCKSIZE, salt)) {
+    LOG(ERROR) << "Failed to initialize hash tree computation, source " << source_ranges.ToString()
+               << ", salt " << salt_hex;
+    return -1;
+  }
+
+  // Iterates through every block in the source_ranges and updates the hash tree structure
+  // accordingly.
+  for (const auto& range : source_ranges) {
+    uint8_t buffer[BLOCKSIZE];
+    if (!check_lseek(params.fd, static_cast<off64_t>(range.first) * BLOCKSIZE, SEEK_SET)) {
+      PLOG(ERROR) << "Failed to seek to block: " << range.first;
+      return -1;
+    }
+
+    for (size_t i = range.first; i < range.second; i++) {
+      if (read_all(params.fd, buffer, BLOCKSIZE) == -1) {
+        LOG(ERROR) << "Failed to read data in " << range.first << ":" << range.second;
+        return -1;
+      }
+
+      if (!builder.Update(reinterpret_cast<unsigned char*>(buffer), BLOCKSIZE)) {
+        LOG(ERROR) << "Failed to update hash tree builder";
+        return -1;
+      }
+    }
+  }
+
+  if (!builder.BuildHashTree()) {
+    LOG(ERROR) << "Failed to build hash tree";
+    return -1;
+  }
+
+  std::string root_hash_hex = HashTreeBuilder::BytesArrayToString(builder.root_hash());
+  if (root_hash_hex != expected_root_hash) {
+    LOG(ERROR) << "Root hash of the verity hash tree doesn't match the expected value. Expected: "
+               << expected_root_hash << ", actual: " << root_hash_hex;
+    return -1;
+  }
+
+  uint64_t write_offset = static_cast<uint64_t>(hash_tree_ranges.GetBlockNumber(0)) * BLOCKSIZE;
+  if (params.canwrite && !builder.WriteHashTreeToFd(params.fd, write_offset)) {
+    LOG(ERROR) << "Failed to write hash tree to output";
+    return -1;
+  }
+
+  // TODO(xunchang) validates the written bytes
+
+  return 0;
+}
+
 using CommandFunction = std::function<int(CommandParameters&)>;
 
 using CommandMap = std::unordered_map<Command::Type, CommandFunction>;
@@ -1737,6 +1837,9 @@ static Value* PerformBlockImageUpdate(const char* name, State* state,
 
     if (performer(params) == -1) {
       LOG(ERROR) << "failed to execute command [" << line << "]";
+      if (cmd_type == Command::Type::COMPUTE_HASH_TREE && failure_type == kNoCause) {
+        failure_type = kHashTreeComputationFailure;
+      }
       goto pbiudone;
     }
 
@@ -1894,15 +1997,16 @@ Value* BlockImageVerifyFn(const char* name, State* state,
   // Commands which are not allowed are set to nullptr to skip them completely.
   const CommandMap command_map{
     // clang-format off
-    { Command::Type::ABORT,   PerformCommandAbort },
-    { 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 },
+    { Command::Type::ABORT,             PerformCommandAbort },
+    { Command::Type::BSDIFF,            PerformCommandDiff },
+    { Command::Type::COMPUTE_HASH_TREE, PerformCommandComputeHashTree },
+    { 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());
@@ -1915,15 +2019,16 @@ Value* BlockImageUpdateFn(const char* name, State* state,
                           const std::vector<std::unique_ptr<Expr>>& argv) {
   const CommandMap command_map{
     // clang-format off
-    { Command::Type::ABORT,   PerformCommandAbort },
-    { 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 },
+    { Command::Type::ABORT,             PerformCommandAbort },
+    { Command::Type::BSDIFF,            PerformCommandDiff },
+    { Command::Type::COMPUTE_HASH_TREE, PerformCommandComputeHashTree },
+    { 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());
diff --git a/updater/commands.cpp b/updater/commands.cpp
index e88149678..15a787c51 100644
--- a/updater/commands.cpp
+++ b/updater/commands.cpp
@@ -40,6 +40,8 @@ Command::Type Command::ParseType(const std::string& type_str) {
     return Type::ABORT;
   } else if (type_str == "bsdiff") {
     return Type::BSDIFF;
+  } else if (type_str == "compute_hash_tree") {
+    return Type::COMPUTE_HASH_TREE;
   } else if (type_str == "erase") {
     return Type::ERASE;
   } else if (type_str == "free") {
@@ -175,6 +177,7 @@ Command Command::Parse(const std::string& line, size_t index, std::string* err)
   SourceInfo source_info;
   StashInfo stash_info;
 
+  // TODO(xunchang) add the parse code of compute_hash_tree
   if (op == Type::ZERO || op == Type::NEW || op == Type::ERASE) {
     // zero/new/erase <rangeset>
     if (pos + 1 != tokens.size()) {
diff --git a/updater/include/private/commands.h b/updater/include/private/commands.h
index 087d7cfbf..7f9dc79f4 100644
--- a/updater/include/private/commands.h
+++ b/updater/include/private/commands.h
@@ -213,6 +213,10 @@ class PatchInfo {
 //      - Free the given stash data.
 //      - Meaningful args: StashInfo
 //
+//    compute_hash_tree <hash_tree_ranges> <source_ranges> <hash_algorithm> <salt_hex> <root_hash>
+//      - Computes the hash_tree bytes and writes the result to the specified range on the
+//        block_device.
+//
 //    abort
 //      - Abort the current update. Allowed for testing code only.
 //
@@ -221,6 +225,7 @@ class Command {
   enum class Type {
     ABORT,
     BSDIFF,
+    COMPUTE_HASH_TREE,
     ERASE,
     FREE,
     IMGDIFF,