From 0feef55859ed7ca3bee283579e275d094cb4a2ef Mon Sep 17 00:00:00 2001
From: Florian Mayer <fmayer@google.com>
Date: Mon, 30 Oct 2023 17:40:56 +0000
Subject: [PATCH] Add --keep_memtag_mode for --wipe_data

This allows us to erase the system image without resetting the MTE
state. This is useful for TestHarness looking to re-use an MTE enabled
device without incurring an extra reboot to re-enable after reboot.

Bug: 300694575
Change-Id: Ie1ea6891361e561189b9390b97b0b4a4c3a6d7e8
---
 install/include/install/wipe_data.h |  2 +-
 install/wipe_data.cpp               | 16 ++++++++++------
 recovery.cpp                        |  6 +++++-
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/install/include/install/wipe_data.h b/install/include/install/wipe_data.h
index 42cad871..255d9b17 100644
--- a/install/include/install/wipe_data.h
+++ b/install/include/install/wipe_data.h
@@ -27,4 +27,4 @@ struct selabel_handle;
 bool WipeCache(RecoveryUI* ui, const std::function<bool()>& confirm);
 
 // Returns true on success.
-bool WipeData(Device* device);
+bool WipeData(Device* device, bool keep_memtag_mode = false);
diff --git a/install/wipe_data.cpp b/install/wipe_data.cpp
index c65e6f48..7aff6226 100644
--- a/install/wipe_data.cpp
+++ b/install/wipe_data.cpp
@@ -79,7 +79,7 @@ bool WipeCache(RecoveryUI* ui, const std::function<bool()>& confirm_func) {
   return success;
 }
 
-bool WipeData(Device* device) {
+bool WipeData(Device* device, bool keep_memtag_mode) {
   RecoveryUI* ui = device->GetUI();
   ui->Print("\n-- Wiping data...\n");
   ui->SetBackground(RecoveryUI::ERASING);
@@ -101,11 +101,15 @@ bool WipeData(Device* device) {
       success &= EraseVolume(METADATA_ROOT, ui);
     }
   }
-  ui->Print("Resetting memtag message...\n");
-  std::string err;
-  if (!WriteMiscMemtagMessage({}, &err)) {
-    ui->Print("Failed to reset memtag message: %s\n", err.c_str());
-    success = false;
+  if (keep_memtag_mode) {
+    ui->Print("NOT resetting memtag message as per request...\n");
+  } else {
+    ui->Print("Resetting memtag message...\n");
+    std::string err;
+    if (!WriteMiscMemtagMessage({}, &err)) {
+      ui->Print("Failed to reset memtag message: %s\n", err.c_str());
+      success = false;
+    }
   }
   if (success) {
     success &= device->PostWipeData();
diff --git a/recovery.cpp b/recovery.cpp
index 4d390195..1a6a7d6c 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -608,6 +608,7 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
     { "wipe_ab", no_argument, nullptr, 0 },
     { "wipe_cache", no_argument, nullptr, 0 },
     { "wipe_data", no_argument, nullptr, 0 },
+    { "keep_memtag_mode", no_argument, nullptr, 0 },
     { "wipe_package_size", required_argument, nullptr, 0 },
     { nullptr, 0, nullptr, 0 },
   };
@@ -616,6 +617,7 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
   bool install_with_fuse = false;  // memory map the update package by default.
   bool should_wipe_data = false;
   bool should_prompt_and_wipe_data = false;
+  bool should_keep_memtag_mode = false;
   bool should_wipe_cache = false;
   bool should_wipe_ab = false;
   size_t wipe_package_size = 0;
@@ -675,6 +677,8 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
           should_wipe_data = true;
         } else if (option == "wipe_package_size") {
           android::base::ParseUint(optarg, &wipe_package_size);
+        } else if (option == "keep_memtag_mode") {
+          should_keep_memtag_mode = true;
         }
         break;
       }
@@ -793,7 +797,7 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
   } else if (should_wipe_data) {
     save_current_log = true;
     CHECK(device->GetReason().has_value());
-    if (!WipeData(device)) {
+    if (!WipeData(device, should_keep_memtag_mode)) {
       status = INSTALL_ERROR;
     }
   } else if (should_prompt_and_wipe_data) {