1 files changed, 38 insertions, 13 deletions

diff --git a/install/install.cpp b/install/install.cpp
index bb8c3b825..811bcb021 100644
--- a/install/install.cpp
+++ b/install/install.cpp
@@ -235,30 +235,41 @@ bool CheckPackageMetadata(const std::map<std::string, std::string>& metadata, Ot
   return true;
 }
 
-bool SetUpAbUpdateCommands(const std::string& package, ZipArchiveHandle zip, int status_fd,
-                           std::vector<std::string>* cmd) {
-  CHECK(cmd != nullptr);
-
+static std::string ExtractPayloadProperties(ZipArchiveHandle zip) {
   // For A/B updates we extract the payload properties to a buffer and obtain the RAW payload offset
   // in the zip file.
   static constexpr const char* AB_OTA_PAYLOAD_PROPERTIES = "payload_properties.txt";
   ZipEntry64 properties_entry;
   if (FindEntry(zip, AB_OTA_PAYLOAD_PROPERTIES, &properties_entry) != 0) {
     LOG(ERROR) << "Failed to find " << AB_OTA_PAYLOAD_PROPERTIES;
-    return false;
+    return {};
   }
   auto properties_entry_length = properties_entry.uncompressed_length;
   if (properties_entry_length > std::numeric_limits<size_t>::max()) {
     LOG(ERROR) << "Failed to extract " << AB_OTA_PAYLOAD_PROPERTIES
                << " because's uncompressed size exceeds size of address space. "
                << properties_entry_length;
-    return false;
+    return {};
   }
-  std::vector<uint8_t> payload_properties(properties_entry_length);
+  std::string payload_properties(properties_entry_length, '\0');
   int32_t err =
-      ExtractToMemory(zip, &properties_entry, payload_properties.data(), properties_entry_length);
+      ExtractToMemory(zip, &properties_entry, reinterpret_cast<uint8_t*>(payload_properties.data()),
+                      properties_entry_length);
   if (err != 0) {
     LOG(ERROR) << "Failed to extract " << AB_OTA_PAYLOAD_PROPERTIES << ": " << ErrorCodeString(err);
+    return {};
+  }
+  return payload_properties;
+}
+
+bool SetUpAbUpdateCommands(const std::string& package, ZipArchiveHandle zip, int status_fd,
+                           std::vector<std::string>* cmd) {
+  CHECK(cmd != nullptr);
+
+  // For A/B updates we extract the payload properties to a buffer and obtain the RAW payload offset
+  // in the zip file.
+  const auto payload_properties = ExtractPayloadProperties(zip);
+  if (payload_properties.empty()) {
     return false;
   }
 
@@ -332,10 +343,20 @@ static void log_max_temperature(int* max_temperature, const std::atomic<bool>& l
   }
 }
 
+static bool PerformPowerwashIfRequired(ZipArchiveHandle zip, Device *device) {
+  const auto payload_properties = ExtractPayloadProperties(zip);
+  if (payload_properties.find("POWERWASH=1") != std::string::npos) {
+    LOG(INFO) << "Payload properties has POWERWASH=1, wiping userdata...";
+    return WipeData(device, true);
+  }
+  return true;
+}
+
 // If the package contains an update binary, extract it and run it.
 static InstallResult TryUpdateBinary(Package* package, bool* wipe_cache,
                                      std::vector<std::string>* log_buffer, int retry_count,
-                                     int* max_temperature, RecoveryUI* ui) {
+                                     int* max_temperature, Device* device) {
+  auto ui = device->GetUI();
   std::map<std::string, std::string> metadata;
   auto zip = package->GetZipArchiveHandle();
   if (!ReadMetadataFromPackage(zip, &metadata)) {
@@ -530,13 +551,15 @@ static InstallResult TryUpdateBinary(Package* package, bool* wipe_cache,
   } else {
     LOG(FATAL) << "Invalid status code " << status;
   }
+  PerformPowerwashIfRequired(zip, device);
 
   return INSTALL_SUCCESS;
 }
 
 static InstallResult VerifyAndInstallPackage(Package* package, bool* wipe_cache,
                                              std::vector<std::string>* log_buffer, int retry_count,
-                                             int* max_temperature, RecoveryUI* ui) {
+                                             int* max_temperature, Device* device) {
+  auto ui = device->GetUI();
   ui->SetBackground(RecoveryUI::INSTALLING_UPDATE);
   // Give verification half the progress bar...
   ui->SetProgressType(RecoveryUI::DETERMINATE);
@@ -554,7 +577,8 @@ static InstallResult VerifyAndInstallPackage(Package* package, bool* wipe_cache,
     ui->Print("Retry attempt: %d\n", retry_count);
   }
   ui->SetEnableReboot(false);
-  auto result = TryUpdateBinary(package, wipe_cache, log_buffer, retry_count, max_temperature, ui);
+  auto result =
+      TryUpdateBinary(package, wipe_cache, log_buffer, retry_count, max_temperature, device);
   ui->SetEnableReboot(true);
   ui->Print("\n");
 
@@ -562,7 +586,8 @@ static InstallResult VerifyAndInstallPackage(Package* package, bool* wipe_cache,
 }
 
 InstallResult InstallPackage(Package* package, const std::string_view package_id,
-                             bool should_wipe_cache, int retry_count, RecoveryUI* ui) {
+                             bool should_wipe_cache, int retry_count, Device* device) {
+  auto ui = device->GetUI();
   auto start = std::chrono::system_clock::now();
 
   int start_temperature = GetMaxValueFromThermalZone();
@@ -584,7 +609,7 @@ InstallResult InstallPackage(Package* package, const std::string_view package_id
   } else {
     bool updater_wipe_cache = false;
     result = VerifyAndInstallPackage(package, &updater_wipe_cache, &log_buffer, retry_count,
-                                     &max_temperature, ui);
+                                     &max_temperature, device);
     should_wipe_cache = should_wipe_cache || updater_wipe_cache;
   }