From 94f602dd05d9d5481a0c5582a7269fbbdf88631f Mon Sep 17 00:00:00 2001
From: Robin Liu <renfeng@google.com>
Date: Thu, 26 Oct 2023 11:18:35 -0400
Subject: [PATCH] Add option to cleanup device

---
 capture_service/android_application.cc |  2 +-
 capture_service/device_mgr.cc          | 26 +++++++++++++++++++++++++
 capture_service/device_mgr.h           |  1 +
 capture_service/dive_client_cli.cc     | 27 ++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/capture_service/android_application.cc b/capture_service/android_application.cc
index 32916043c..5e6a26faa 100644
--- a/capture_service/android_application.cc
+++ b/capture_service/android_application.cc
@@ -129,7 +129,7 @@ std::string AndroidApplication::GetMainActivity()
 void AndroidApplication::Start()
 {
     m_dev.Adb().Run("shell input keyevent KEYCODE_WAKEUP");
-    m_dev.Adb().Run(absl::StrFormat("shell am start %s/%s", m_package, m_main_activity));
+    m_dev.Adb().Run(absl::StrFormat("shell am start -S -W %s/%s", m_package, m_main_activity));
     m_started = true;
 }
 
diff --git a/capture_service/device_mgr.cc b/capture_service/device_mgr.cc
index b394f1cfc..ff3fa885b 100644
--- a/capture_service/device_mgr.cc
+++ b/capture_service/device_mgr.cc
@@ -163,6 +163,32 @@ std::vector<std::string> DeviceManager::ListDevice() const
 
     return dev_list;
 }
+void DeviceManager::Cleanup(const std::string &serial, const std::string &package)
+{
+    AdbSession adb(serial);
+    // Remove installed libs and libraries on device.
+    adb.Run("root");
+    adb.Run("wait-for-device");
+    adb.Run("remount");
+
+    adb.Run(absl::StrFormat("shell rm %s/%s", kTargetPath, kWrapLibName), true);
+    adb.Run(absl::StrFormat("shell rm %s/%s", kTargetPath, kVkLayerLibName), true);
+    adb.Run(absl::StrFormat("shell rm %s/%s", kTargetPath, kXrLayerLibName), true);
+    adb.Run(absl::StrFormat("shell rm -r %s", kManifestFilePath), true);
+    adb.Run(absl::StrFormat("forward --remove tcp:%d", kPort), true);
+
+    adb.Run("shell settings delete global enable_gpu_debug_layers");
+    adb.Run("shell settings delete global gpu_debug_app");
+    adb.Run("shell settings delete global gpu_debug_layers");
+    adb.Run("shell settings delete global gpu_debug_layer_app");
+    adb.Run("shell settings delete global gpu_debug_layers_gles");
+
+    // If package specified, remove package related settings.
+    if (!package.empty())
+    {
+        adb.Run(absl::StrFormat("shell setprop wrap.%s \\\"\\\"", package));
+    }
+}
 
 void AndroidDevice::RetrieveTraceFile(const std::string &trace_file_path,
                                       const std::string &save_path)
diff --git a/capture_service/device_mgr.h b/capture_service/device_mgr.h
index 0d988ef67..400073e7f 100644
--- a/capture_service/device_mgr.h
+++ b/capture_service/device_mgr.h
@@ -76,6 +76,7 @@ class DeviceManager
     }
 
     AndroidDevice *GetDevice() const { return m_device.get(); }
+    void           Cleanup(const std::string &serial, const std::string &package);
 
 private:
     std::unique_ptr<AndroidDevice> m_device;
diff --git a/capture_service/dive_client_cli.cc b/capture_service/dive_client_cli.cc
index 42c96db9a..db01a37b5 100644
--- a/capture_service/dive_client_cli.cc
+++ b/capture_service/dive_client_cli.cc
@@ -39,6 +39,7 @@ enum class Command
     kListPackage,
     kRunPackage,
     kRunAndCapture,
+    kCleanup,
 };
 
 bool AbslParseFlag(absl::string_view text, Command* command, std::string* error)
@@ -63,6 +64,11 @@ bool AbslParseFlag(absl::string_view text, Command* command, std::string* error)
         *command = Command::kRunAndCapture;
         return true;
     }
+    if (text == "cleanup")
+    {
+        *command = Command::kCleanup;
+        return true;
+    }
     if (text.empty())
     {
         *command = Command::kNone;
@@ -81,6 +87,7 @@ std::string AbslUnparseFlag(Command command)
     case Command::kListPackage: return "list_package";
     case Command::kRunPackage: return "run";
     case Command::kRunAndCapture: return "capture";
+    case Command::kCleanup: return "cleanup";
 
     default: return absl::StrCat(command);
     }
@@ -178,6 +185,21 @@ bool run_and_capture(Dive::DeviceManager& mgr, const std::string& package)
     return true;
 }
 
+bool clean_up_app_and_device(Dive::DeviceManager& mgr, const std::string& package)
+{
+    std::string serial = absl::GetFlag(FLAGS_device);
+
+    if (serial.empty() || package.empty())
+    {
+        std::cout <<"Please run with `--device [serial]` and `--package [package]` options." << std::endl;
+        print_usage();
+        return false;
+    }
+
+    mgr.Cleanup(serial, package);
+    return true;
+}
+
 int main(int argc, char** argv)
 {
     absl::SetProgramUsageMessage("Run app with --help for more details");
@@ -220,6 +242,11 @@ int main(int argc, char** argv)
         run_and_capture(mgr, package);
         break;
     }
+    case Command::kCleanup:
+    {
+        clean_up_app_and_device(mgr, package);
+        break;
+    }
     case Command::kNone:
     {
         print_usage();