6 files changed, 193 insertions, 91 deletions

diff --git a/bootloader_message/bootloader_message.cpp b/bootloader_message/bootloader_message.cpp
index f6f8005f6..9a5671843 100644
--- a/bootloader_message/bootloader_message.cpp
+++ b/bootloader_message/bootloader_message.cpp
@@ -80,26 +80,23 @@ static bool wait_for_device(const std::string& blk_device, std::string* err) {
   return ret == 0;
 }
 
-static bool read_misc_partition(void* p, size_t size, size_t offset, std::string* err) {
-  std::string misc_blk_device = get_misc_blk_device(err);
-  if (misc_blk_device.empty()) {
-    return false;
-  }
+static bool read_misc_partition(void* p, size_t size, const std::string& misc_blk_device,
+                                size_t offset, std::string* err) {
   if (!wait_for_device(misc_blk_device, err)) {
     return false;
   }
   android::base::unique_fd fd(open(misc_blk_device.c_str(), O_RDONLY));
-  if (fd.get() == -1) {
+  if (fd == -1) {
     *err = android::base::StringPrintf("failed to open %s: %s", misc_blk_device.c_str(),
                                        strerror(errno));
     return false;
   }
-  if (lseek(fd.get(), static_cast<off_t>(offset), SEEK_SET) != static_cast<off_t>(offset)) {
+  if (lseek(fd, static_cast<off_t>(offset), SEEK_SET) != static_cast<off_t>(offset)) {
     *err = android::base::StringPrintf("failed to lseek %s: %s", misc_blk_device.c_str(),
                                        strerror(errno));
     return false;
   }
-  if (!android::base::ReadFully(fd.get(), p, size)) {
+  if (!android::base::ReadFully(fd, p, size)) {
     *err = android::base::StringPrintf("failed to read %s: %s", misc_blk_device.c_str(),
                                        strerror(errno));
     return false;
@@ -107,29 +104,25 @@ static bool read_misc_partition(void* p, size_t size, size_t offset, std::string
   return true;
 }
 
-static bool write_misc_partition(const void* p, size_t size, size_t offset, std::string* err) {
-  std::string misc_blk_device = get_misc_blk_device(err);
-  if (misc_blk_device.empty()) {
-    return false;
-  }
-  android::base::unique_fd fd(open(misc_blk_device.c_str(), O_WRONLY | O_SYNC));
-  if (fd.get() == -1) {
+static bool write_misc_partition(const void* p, size_t size, const std::string& misc_blk_device,
+                                 size_t offset, std::string* err) {
+  android::base::unique_fd fd(open(misc_blk_device.c_str(), O_WRONLY));
+  if (fd == -1) {
     *err = android::base::StringPrintf("failed to open %s: %s", misc_blk_device.c_str(),
                                        strerror(errno));
     return false;
   }
-  if (lseek(fd.get(), static_cast<off_t>(offset), SEEK_SET) != static_cast<off_t>(offset)) {
+  if (lseek(fd, static_cast<off_t>(offset), SEEK_SET) != static_cast<off_t>(offset)) {
     *err = android::base::StringPrintf("failed to lseek %s: %s", misc_blk_device.c_str(),
                                        strerror(errno));
     return false;
   }
-  if (!android::base::WriteFully(fd.get(), p, size)) {
+  if (!android::base::WriteFully(fd, p, size)) {
     *err = android::base::StringPrintf("failed to write %s: %s", misc_blk_device.c_str(),
                                        strerror(errno));
     return false;
   }
-  // TODO: O_SYNC and fsync duplicates each other?
-  if (fsync(fd.get()) == -1) {
+  if (fsync(fd) == -1) {
     *err = android::base::StringPrintf("failed to fsync %s: %s", misc_blk_device.c_str(),
                                        strerror(errno));
     return false;
@@ -137,12 +130,32 @@ static bool write_misc_partition(const void* p, size_t size, size_t offset, std:
   return true;
 }
 
+bool read_bootloader_message_from(bootloader_message* boot, const std::string& misc_blk_device,
+                                  std::string* err) {
+  return read_misc_partition(boot, sizeof(*boot), misc_blk_device,
+                             BOOTLOADER_MESSAGE_OFFSET_IN_MISC, err);
+}
+
 bool read_bootloader_message(bootloader_message* boot, std::string* err) {
-  return read_misc_partition(boot, sizeof(*boot), BOOTLOADER_MESSAGE_OFFSET_IN_MISC, err);
+  std::string misc_blk_device = get_misc_blk_device(err);
+  if (misc_blk_device.empty()) {
+    return false;
+  }
+  return read_bootloader_message_from(boot, misc_blk_device, err);
+}
+
+bool write_bootloader_message_to(const bootloader_message& boot, const std::string& misc_blk_device,
+                                 std::string* err) {
+  return write_misc_partition(&boot, sizeof(boot), misc_blk_device,
+                              BOOTLOADER_MESSAGE_OFFSET_IN_MISC, err);
 }
 
 bool write_bootloader_message(const bootloader_message& boot, std::string* err) {
-  return write_misc_partition(&boot, sizeof(boot), BOOTLOADER_MESSAGE_OFFSET_IN_MISC, err);
+  std::string misc_blk_device = get_misc_blk_device(err);
+  if (misc_blk_device.empty()) {
+    return false;
+  }
+  return write_bootloader_message_to(boot, misc_blk_device, err);
 }
 
 bool clear_bootloader_message(std::string* err) {
@@ -177,12 +190,21 @@ bool write_reboot_bootloader(std::string* err) {
 }
 
 bool read_wipe_package(std::string* package_data, size_t size, std::string* err) {
+  std::string misc_blk_device = get_misc_blk_device(err);
+  if (misc_blk_device.empty()) {
+    return false;
+  }
   package_data->resize(size);
-  return read_misc_partition(&(*package_data)[0], size, WIPE_PACKAGE_OFFSET_IN_MISC, err);
+  return read_misc_partition(&(*package_data)[0], size, misc_blk_device,
+                             WIPE_PACKAGE_OFFSET_IN_MISC, err);
 }
 
 bool write_wipe_package(const std::string& package_data, std::string* err) {
-  return write_misc_partition(package_data.data(), package_data.size(),
+  std::string misc_blk_device = get_misc_blk_device(err);
+  if (misc_blk_device.empty()) {
+    return false;
+  }
+  return write_misc_partition(package_data.data(), package_data.size(), misc_blk_device,
                               WIPE_PACKAGE_OFFSET_IN_MISC, err);
 }
 
diff --git a/bootloader_message/include/bootloader_message/bootloader_message.h b/bootloader_message/include/bootloader_message/bootloader_message.h
index 5a5dd8793..e45f42487 100644
--- a/bootloader_message/include/bootloader_message/bootloader_message.h
+++ b/bootloader_message/include/bootloader_message/bootloader_message.h
@@ -178,15 +178,33 @@ static_assert(sizeof(struct bootloader_control) ==
 #include <string>
 #include <vector>
 
+// Read bootloader message into boot. Error message will be set in err.
 bool read_bootloader_message(bootloader_message* boot, std::string* err);
+
+// Read bootloader message from the specified misc device into boot.
+bool read_bootloader_message_from(bootloader_message* boot, const std::string& misc_blk_device,
+                                  std::string* err);
+
+// Write bootloader message to BCB.
 bool write_bootloader_message(const bootloader_message& boot, std::string* err);
+
+// Write bootloader message to the specified BCB device.
+bool write_bootloader_message_to(const bootloader_message& boot,
+                                 const std::string& misc_blk_device, std::string* err);
+
+// Write bootloader message (boots into recovery with the options) to BCB.
 bool write_bootloader_message(const std::vector<std::string>& options, std::string* err);
+
+// Clear BCB.
 bool clear_bootloader_message(std::string* err);
 
 // Writes the reboot-bootloader reboot reason to the bootloader_message.
 bool write_reboot_bootloader(std::string* err);
 
+// Read the wipe package from BCB (from offset WIPE_PACKAGE_OFFSET_IN_MISC).
 bool read_wipe_package(std::string* package_data, size_t size, std::string* err);
+
+// Write the wipe package into BCB (to offset WIPE_PACKAGE_OFFSET_IN_MISC).
 bool write_wipe_package(const std::string& package_data, std::string* err);
 
 #else
diff --git a/tests/Android.mk b/tests/Android.mk
index fdc947028..5f6a7ce0c 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -80,10 +80,12 @@ LOCAL_STATIC_LIBRARIES := \
     libedify \
     libotafault \
     libupdater \
+    libbootloader_message \
     libverifier \
     libminui \
     libotautil \
     libmounts \
+    libfs_mgr \
     liblog \
     libselinux \
     libext4_utils_static \
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index cd285729d..f31f1f82a 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -23,6 +23,7 @@
 #include <android-base/file.h>
 #include <android-base/properties.h>
 #include <android-base/test_utils.h>
+#include <bootloader_message/bootloader_message.h>
 #include <gtest/gtest.h>
 #include <ziparchive/zip_archive.h>
 
@@ -451,3 +452,61 @@ TEST_F(UpdaterTest, write_value) {
   script = "write_value(\"value\", \"/proc/0/file1\")";
   expect("", script.c_str(), kNoCause);
 }
+
+TEST_F(UpdaterTest, get_stage) {
+  // get_stage() expects one argument.
+  expect(nullptr, "get_stage()", kArgsParsingFailure);
+  expect(nullptr, "get_stage(\"arg1\", \"arg2\")", kArgsParsingFailure);
+  expect(nullptr, "get_stage(\"arg1\", \"arg2\", \"arg3\")", kArgsParsingFailure);
+
+  // Set up a local file as BCB.
+  TemporaryFile tf;
+  std::string temp_file(tf.path);
+  bootloader_message boot;
+  strlcpy(boot.stage, "2/3", sizeof(boot.stage));
+  std::string err;
+  ASSERT_TRUE(write_bootloader_message_to(boot, temp_file, &err));
+
+  // Can read the stage value.
+  std::string script("get_stage(\"" + temp_file + "\")");
+  expect("2/3", script.c_str(), kNoCause);
+
+  // Bad BCB path.
+  script = "get_stage(\"doesntexist\")";
+  expect("", script.c_str(), kNoCause);
+}
+
+TEST_F(UpdaterTest, set_stage) {
+  // set_stage() expects two arguments.
+  expect(nullptr, "set_stage()", kArgsParsingFailure);
+  expect(nullptr, "set_stage(\"arg1\")", kArgsParsingFailure);
+  expect(nullptr, "set_stage(\"arg1\", \"arg2\", \"arg3\")", kArgsParsingFailure);
+
+  // Set up a local file as BCB.
+  TemporaryFile tf;
+  std::string temp_file(tf.path);
+  bootloader_message boot;
+  strlcpy(boot.command, "command", sizeof(boot.command));
+  strlcpy(boot.stage, "2/3", sizeof(boot.stage));
+  std::string err;
+  ASSERT_TRUE(write_bootloader_message_to(boot, temp_file, &err));
+
+  // Write with set_stage().
+  std::string script("set_stage(\"" + temp_file + "\", \"1/3\")");
+  expect(tf.path, script.c_str(), kNoCause);
+
+  // Verify.
+  bootloader_message boot_verify;
+  ASSERT_TRUE(read_bootloader_message_from(&boot_verify, temp_file, &err));
+
+  // Stage should be updated, with command part untouched.
+  ASSERT_STREQ("1/3", boot_verify.stage);
+  ASSERT_STREQ(boot.command, boot_verify.command);
+
+  // Bad BCB path.
+  script = "set_stage(\"doesntexist\", \"1/3\")";
+  expect("", script.c_str(), kNoCause);
+
+  script = "set_stage(\"/dev/full\", \"1/3\")";
+  expect("", script.c_str(), kNoCause);
+}
diff --git a/updater/Android.mk b/updater/Android.mk
index 3c1d0d41f..5d328a3fa 100644
--- a/updater/Android.mk
+++ b/updater/Android.mk
@@ -27,12 +27,14 @@ updater_common_static_libraries := \
     libedify \
     libziparchive \
     libotautil \
+    libbootloader_message \
     libutils \
     libmounts \
     libotafault \
     libext4_utils_static \
     libfec \
     libfec_rs \
+    libfs_mgr \
     liblog \
     libselinux \
     libsparse_static \
diff --git a/updater/install.cpp b/updater/install.cpp
index 6c110732a..b9bc19e2c 100644
--- a/updater/install.cpp
+++ b/updater/install.cpp
@@ -1196,31 +1196,35 @@ Value* WriteValueFn(const char* name, State* state, int argc, Expr* argv[]) {
 // partition, or "" (empty string) to boot from the regular boot
 // partition.
 Value* RebootNowFn(const char* name, State* state, int argc, Expr* argv[]) {
-    if (argc != 2) {
-        return ErrorAbort(state, kArgsParsingFailure, "%s() expects 2 args, got %d", name, argc);
-    }
+  if (argc != 2) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() expects 2 args, got %d", name, argc);
+  }
 
-    std::vector<std::string> args;
-    if (!ReadArgs(state, 2, argv, &args)) {
-        return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse the argument(s)", name);
-    }
-    const std::string& filename = args[0];
-    const std::string& property = args[1];
-
-    // zero out the 'command' field of the bootloader message.
-    char buffer[80];
-    memset(buffer, 0, sizeof(((struct bootloader_message*)0)->command));
-    FILE* f = ota_fopen(filename.c_str(), "r+b");
-    fseek(f, offsetof(struct bootloader_message, command), SEEK_SET);
-    ota_fwrite(buffer, sizeof(((struct bootloader_message*)0)->command), 1, f);
-    ota_fclose(f);
+  std::vector<std::string> args;
+  if (!ReadArgs(state, 2, argv, &args)) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s(): Failed to parse the argument(s)", name);
+  }
+  const std::string& filename = args[0];
+  const std::string& property = args[1];
+
+  // Zero out the 'command' field of the bootloader message. Leave the rest intact.
+  bootloader_message boot;
+  std::string err;
+  if (!read_bootloader_message_from(&boot, filename, &err)) {
+    printf("%s(): Failed to read from \"%s\": %s", name, filename.c_str(), err.c_str());
+    return StringValue("");
+  }
+  memset(boot.command, 0, sizeof(boot.command));
+  if (!write_bootloader_message_to(boot, filename, &err)) {
+    printf("%s(): Failed to write to \"%s\": %s", name, filename.c_str(), err.c_str());
+    return StringValue("");
+  }
 
-    std::string reboot_cmd = "reboot,";
-    reboot_cmd += property;
-    android::base::SetProperty(ANDROID_RB_PROPERTY, reboot_cmd);
+  const std::string reboot_cmd = "reboot," + property;
+  android::base::SetProperty(ANDROID_RB_PROPERTY, reboot_cmd);
 
-    sleep(5);
-    return ErrorAbort(state, kRebootFailure, "%s() failed to reboot", name);
+  sleep(5);
+  return ErrorAbort(state, kRebootFailure, "%s() failed to reboot", name);
 }
 
 // Store a string value somewhere that future invocations of recovery
@@ -1234,62 +1238,57 @@ Value* RebootNowFn(const char* name, State* state, int argc, Expr* argv[]) {
 // second argument is the string to store; it should not exceed 31
 // bytes.
 Value* SetStageFn(const char* name, State* state, int argc, Expr* argv[]) {
-    if (argc != 2) {
-        return ErrorAbort(state, kArgsParsingFailure, "%s() expects 2 args, got %d", name, argc);
-    }
+  if (argc != 2) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() expects 2 args, got %d", name, argc);
+  }
 
-    std::vector<std::string> args;
-    if (!ReadArgs(state, 2, argv, &args)) {
-        return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse the argument(s)", name);
-    }
-    const std::string& filename = args[0];
-    std::string& stagestr = args[1];
-
-    // Store this value in the misc partition, immediately after the
-    // bootloader message that the main recovery uses to save its
-    // arguments in case of the device restarting midway through
-    // package installation.
-    FILE* f = ota_fopen(filename.c_str(), "r+b");
-    fseek(f, offsetof(struct bootloader_message, stage), SEEK_SET);
-    size_t to_write = stagestr.size();
-    size_t max_size = sizeof(((struct bootloader_message*)0)->stage);
-    if (to_write > max_size) {
-        to_write = max_size;
-        stagestr = stagestr.substr(0, max_size-1);
-    }
-    size_t status = ota_fwrite(stagestr.c_str(), to_write, 1, f);
-    ota_fclose(f);
+  std::vector<std::string> args;
+  if (!ReadArgs(state, 2, argv, &args)) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse the argument(s)", name);
+  }
+  const std::string& filename = args[0];
+  const std::string& stagestr = args[1];
+
+  // Store this value in the misc partition, immediately after the
+  // bootloader message that the main recovery uses to save its
+  // arguments in case of the device restarting midway through
+  // package installation.
+  bootloader_message boot;
+  std::string err;
+  if (!read_bootloader_message_from(&boot, filename, &err)) {
+    printf("%s(): Failed to read from \"%s\": %s", name, filename.c_str(), err.c_str());
+    return StringValue("");
+  }
+  strlcpy(boot.stage, stagestr.c_str(), sizeof(boot.stage));
+  if (!write_bootloader_message_to(boot, filename, &err)) {
+    printf("%s(): Failed to write to \"%s\": %s", name, filename.c_str(), err.c_str());
+    return StringValue("");
+  }
 
-    if (status != to_write) {
-        return StringValue("");
-    }
-    return StringValue(filename);
+  return StringValue(filename);
 }
 
 // Return the value most recently saved with SetStageFn.  The argument
 // is the block device for the misc partition.
 Value* GetStageFn(const char* name, State* state, int argc, Expr* argv[]) {
-    if (argc != 1) {
-        return ErrorAbort(state, kArgsParsingFailure, "%s() expects 1 arg, got %d", name, argc);
-    }
+  if (argc != 1) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() expects 1 arg, got %d", name, argc);
+  }
 
-    std::vector<std::string> args;
-    if (!ReadArgs(state, 1, argv, &args)) {
-        return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse the argument(s)", name);
-    }
-    const std::string& filename = args[0];
+  std::vector<std::string> args;
+  if (!ReadArgs(state, 1, argv, &args)) {
+    return ErrorAbort(state, kArgsParsingFailure, "%s() Failed to parse the argument(s)", name);
+  }
+  const std::string& filename = args[0];
 
-    char buffer[sizeof(((struct bootloader_message*)0)->stage)];
-    FILE* f = ota_fopen(filename.c_str(), "rb");
-    fseek(f, offsetof(struct bootloader_message, stage), SEEK_SET);
-    size_t status = ota_fread(buffer, sizeof(buffer), 1, f);
-    ota_fclose(f);
-    if (status != sizeof(buffer)) {
-        return StringValue("");
-    }
+  bootloader_message boot;
+  std::string err;
+  if (!read_bootloader_message_from(&boot, filename, &err)) {
+    printf("%s(): Failed to read from \"%s\": %s", name, filename.c_str(), err.c_str());
+    return StringValue("");
+  }
 
-    buffer[sizeof(buffer)-1] = '\0';
-    return StringValue(buffer);
+  return StringValue(boot.stage);
 }
 
 Value* WipeBlockDeviceFn(const char* name, State* state, int argc, Expr* argv[]) {