1 files changed, 64 insertions, 35 deletions

diff --git a/tests/component/uncrypt_test.cpp b/tests/component/uncrypt_test.cpp
index 3925236a5..55baca2e3 100644
--- a/tests/component/uncrypt_test.cpp
+++ b/tests/component/uncrypt_test.cpp
@@ -20,6 +20,7 @@
 #include <sys/un.h>
 #include <unistd.h>
 
+#include <algorithm>
 #include <string>
 
 #include <android-base/file.h>
@@ -38,43 +39,49 @@ static const std::string INIT_SVC_CLEAR_BCB = "init.svc.clear-bcb";
 static const std::string INIT_SVC_UNCRYPT = "init.svc.uncrypt";
 static constexpr int SOCKET_CONNECTION_MAX_RETRY = 30;
 
+static void StopService() {
+  ASSERT_TRUE(android::base::SetProperty("ctl.stop", "setup-bcb"));
+  ASSERT_TRUE(android::base::SetProperty("ctl.stop", "clear-bcb"));
+  ASSERT_TRUE(android::base::SetProperty("ctl.stop", "uncrypt"));
+
+  bool success = false;
+  for (int retry = 0; retry < SOCKET_CONNECTION_MAX_RETRY; retry++) {
+    std::string setup_bcb = android::base::GetProperty(INIT_SVC_SETUP_BCB, "");
+    std::string clear_bcb = android::base::GetProperty(INIT_SVC_CLEAR_BCB, "");
+    std::string uncrypt = android::base::GetProperty(INIT_SVC_UNCRYPT, "");
+    GTEST_LOG_(INFO) << "setup-bcb: [" << setup_bcb << "] clear-bcb: [" << clear_bcb
+                     << "] uncrypt: [" << uncrypt << "]";
+    if (setup_bcb != "running" && clear_bcb != "running" && uncrypt != "running") {
+      success = true;
+      break;
+    }
+    sleep(1);
+  }
+
+  ASSERT_TRUE(success) << "uncrypt service is not available.";
+}
+
 class UncryptTest : public ::testing::Test {
  protected:
   UncryptTest() : has_misc(true) {}
 
-  virtual void SetUp() override {
-    ASSERT_TRUE(android::base::SetProperty("ctl.stop", "setup-bcb"));
-    ASSERT_TRUE(android::base::SetProperty("ctl.stop", "clear-bcb"));
-    ASSERT_TRUE(android::base::SetProperty("ctl.stop", "uncrypt"));
-
-    bool success = false;
-    for (int retry = 0; retry < SOCKET_CONNECTION_MAX_RETRY; retry++) {
-      std::string setup_bcb = android::base::GetProperty(INIT_SVC_SETUP_BCB, "");
-      std::string clear_bcb = android::base::GetProperty(INIT_SVC_CLEAR_BCB, "");
-      std::string uncrypt = android::base::GetProperty(INIT_SVC_UNCRYPT, "");
-      LOG(INFO) << "setup-bcb: [" << setup_bcb << "] clear-bcb: [" << clear_bcb << "] uncrypt: ["
-                << uncrypt << "]";
-      if (setup_bcb != "running" && clear_bcb != "running" && uncrypt != "running") {
-        success = true;
-        break;
-      }
-      sleep(1);
-    }
-
-    ASSERT_TRUE(success) << "uncrypt service is not available.";
-
+  void SetUp() override {
     std::string err;
     has_misc = !get_bootloader_message_blk_device(&err).empty();
   }
 
-  void SetupOrClearBcb(bool isSetup, const std::string& message,
-                       const std::string& message_in_bcb) const {
-    if (!has_misc) {
-      GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
-      return;
+  void TearDown() override {
+    // Clear the BCB.
+    if (has_misc) {
+      std::string err;
+      ASSERT_TRUE(clear_bootloader_message(&err)) << "Failed to clear BCB: " << err;
     }
+  }
 
-    // Trigger the setup-bcb service.
+  void SetupOrClearBcb(bool isSetup, const std::string& message,
+                       const std::string& message_in_bcb) const {
+    // Restart the setup-bcb service.
+    StopService();
     ASSERT_TRUE(android::base::SetProperty("ctl.start", isSetup ? "setup-bcb" : "clear-bcb"));
 
     // Test tends to be flaky if proceeding immediately ("Transport endpoint is not connected").
@@ -144,27 +151,49 @@ class UncryptTest : public ::testing::Test {
     }
   }
 
+  void VerifyBootloaderMessage(const std::string& expected) {
+    std::string err;
+    bootloader_message boot;
+    ASSERT_TRUE(read_bootloader_message(&boot, &err)) << "Failed to read BCB: " << err;
+
+    // Check that we have all the expected bytes.
+    ASSERT_EQ(expected, std::string(reinterpret_cast<const char*>(&boot), sizeof(boot)));
+  }
+
   bool has_misc;
 };
 
 TEST_F(UncryptTest, setup_bcb) {
+  if (!has_misc) {
+    GTEST_LOG_(INFO) << "Test skipped due to no /misc partition found on the device.";
+    return;
+  }
+
+  std::string random_data;
+  random_data.reserve(sizeof(bootloader_message));
+  generate_n(back_inserter(random_data), sizeof(bootloader_message), []() { return rand() % 128; });
+
+  bootloader_message boot;
+  memcpy(&boot, random_data.c_str(), random_data.size());
+
+  std::string err;
+  ASSERT_TRUE(write_bootloader_message(boot, &err)) << "Failed to write BCB: " << err;
+  VerifyBootloaderMessage(random_data);
+
+  ASSERT_TRUE(clear_bootloader_message(&err)) << "Failed to clear BCB: " << err;
+  VerifyBootloaderMessage(std::string(sizeof(bootloader_message), '\0'));
+
   std::string message = "--update_message=abc value";
   std::string message_in_bcb = "recovery\n--update_message=abc value\n";
   SetupOrClearBcb(true, message, message_in_bcb);
-}
 
-TEST_F(UncryptTest, clear_bcb) {
   SetupOrClearBcb(false, "", "");
-}
 
-TEST_F(UncryptTest, setup_bcb_wipe_ab) {
   TemporaryFile wipe_package;
   ASSERT_TRUE(android::base::WriteStringToFile(std::string(345, 'a'), wipe_package.path));
 
   // It's expected to store a wipe package in /misc, with the package size passed to recovery.
-  std::string message =
-      "--wipe_ab\n--wipe_package="s + wipe_package.path + "\n--reason=wipePackage"s;
-  std::string message_in_bcb =
-      "recovery\n--wipe_ab\n--wipe_package_size=345\n--reason=wipePackage\n";
+  message = "--wipe_ab\n--wipe_package="s + wipe_package.path + "\n--reason=wipePackage"s;
+  message_in_bcb = "recovery\n--wipe_ab\n--wipe_package_size=345\n--reason=wipePackage\n";
   SetupOrClearBcb(true, message, message_in_bcb);
 }