fastboot: Implement helper commands for Virtual A/B.

This introduces two new commands to the fastboot protocol:

  - getvar snapshot-update-status - Return "none", "snapshotted", or
    "merging" depending on the current status set by the boot control
    HAL.
  - snapshot-update [cancel] - Cancel any pending snapshot-based updates
    via the boot control HAL. After this, the HAL should return
    MergeStatus::CANCELLED and "update-merge-status" should be "none".
    If no argument is specified, the snapshot-update-status is returned
    via an INFO response.

Bootloaders are expected to implement this in a manner consistent with
the boot control HAL.

Fastboot-based tooling should expect wipes of userdata to fail when
update-merge-status returns "merging". Thus, the force flag now cancel
any pending snapshots.

Bug: 139154945
Test: fastboot getvar snapshot-update-status
      fastboot snapshot-update cancel
      fastboot snapshot-update

Change-Id: Idc423fe7656b212e929e64eb0e6b85b453e0e8dc
diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index f452a64..02a887e 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -122,6 +122,7 @@
 
     shared_libs: [
         "android.hardware.boot@1.0",
+        "android.hardware.boot@1.1",
         "android.hardware.fastboot@1.0",
         "android.hardware.health@2.0",
         "libadbd",
diff --git a/fastboot/constants.h b/fastboot/constants.h
index 8a72627..7fba67c 100644
--- a/fastboot/constants.h
+++ b/fastboot/constants.h
@@ -34,6 +34,7 @@
 #define FB_CMD_UPDATE_SUPER "update-super"
 #define FB_CMD_OEM "oem"
 #define FB_CMD_GSI "gsi"
+#define FB_CMD_SNAPSHOT_UPDATE "snapshot-update"
 
 #define RESPONSE_OKAY "OKAY"
 #define RESPONSE_FAIL "FAIL"
@@ -66,3 +67,4 @@
 #define FB_VAR_BATTERY_VOLTAGE "battery-voltage"
 #define FB_VAR_BATTERY_SOC_OK "battery-soc-ok"
 #define FB_VAR_SUPER_PARTITION_NAME "super-partition-name"
+#define FB_VAR_SNAPSHOT_UPDATE_STATUS "snapshot-update-status"
diff --git a/fastboot/device/commands.cpp b/fastboot/device/commands.cpp
index 4c77c75..dfd5690 100644
--- a/fastboot/device/commands.cpp
+++ b/fastboot/device/commands.cpp
@@ -25,6 +25,7 @@
 #include <android-base/stringprintf.h>
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
+#include <android/hardware/boot/1.1/IBootControl.h>
 #include <cutils/android_reboot.h>
 #include <ext4_utils/wipe.h>
 #include <fs_mgr.h>
@@ -44,8 +45,10 @@
 using ::android::hardware::boot::V1_0::BoolResult;
 using ::android::hardware::boot::V1_0::CommandResult;
 using ::android::hardware::boot::V1_0::Slot;
+using ::android::hardware::boot::V1_1::MergeStatus;
 using ::android::hardware::fastboot::V1_0::Result;
 using ::android::hardware::fastboot::V1_0::Status;
+using IBootControl1_1 = ::android::hardware::boot::V1_1::IBootControl;
 
 struct VariableHandlers {
     // Callback to retrieve the value of a single variable.
@@ -101,7 +104,8 @@
             {FB_VAR_BATTERY_VOLTAGE, {GetBatteryVoltage, nullptr}},
             {FB_VAR_BATTERY_SOC_OK, {GetBatterySoCOk, nullptr}},
             {FB_VAR_HW_REVISION, {GetHardwareRevision, nullptr}},
-            {FB_VAR_SUPER_PARTITION_NAME, {GetSuperPartitionName, nullptr}}};
+            {FB_VAR_SUPER_PARTITION_NAME, {GetSuperPartitionName, nullptr}},
+            {FB_VAR_SNAPSHOT_UPDATE_STATUS, {GetSnapshotUpdateStatus, nullptr}}};
 
     if (args.size() < 2) {
         return device->WriteFail("Missing argument");
@@ -547,3 +551,40 @@
     }
     return device->WriteStatus(FastbootResult::OKAY, "Success");
 }
+
+bool SnapshotUpdateHandler(FastbootDevice* device, const std::vector<std::string>& args) {
+    // Note that we use the HAL rather than mounting /metadata, since we want
+    // our results to match the bootloader.
+    auto hal = device->boot_control_hal();
+    if (!hal) return device->WriteFail("Not supported");
+
+    android::sp<IBootControl1_1> hal11 = IBootControl1_1::castFrom(hal);
+    if (!hal11) return device->WriteFail("Not supported");
+
+    // If no arguments, return the same thing as a getvar. Note that we get the
+    // HAL first so we can return "not supported" before we return the less
+    // specific error message below.
+    if (args.size() < 2 || args[1].empty()) {
+        std::string message;
+        if (!GetSnapshotUpdateStatus(device, {}, &message)) {
+            return device->WriteFail("Could not determine update status");
+        }
+        device->WriteInfo(message);
+        return device->WriteOkay("");
+    }
+
+    if (args.size() != 2 || args[1] != "cancel") {
+        return device->WriteFail("Invalid arguments");
+    }
+
+    MergeStatus status = hal11->getSnapshotMergeStatus();
+    switch (status) {
+        case MergeStatus::SNAPSHOTTED:
+        case MergeStatus::MERGING:
+            hal11->setSnapshotMergeStatus(MergeStatus::CANCELLED);
+            break;
+        default:
+            break;
+    }
+    return device->WriteStatus(FastbootResult::OKAY, "Success");
+}
diff --git a/fastboot/device/commands.h b/fastboot/device/commands.h
index afd6d08..0dc5769 100644
--- a/fastboot/device/commands.h
+++ b/fastboot/device/commands.h
@@ -49,3 +49,4 @@
 bool UpdateSuperHandler(FastbootDevice* device, const std::vector<std::string>& args);
 bool OemCmdHandler(FastbootDevice* device, const std::vector<std::string>& args);
 bool GsiHandler(FastbootDevice* device, const std::vector<std::string>& args);
+bool SnapshotUpdateHandler(FastbootDevice* device, const std::vector<std::string>& args);
diff --git a/fastboot/device/fastboot_device.cpp b/fastboot/device/fastboot_device.cpp
index 56fafab..d3c2bda 100644
--- a/fastboot/device/fastboot_device.cpp
+++ b/fastboot/device/fastboot_device.cpp
@@ -54,6 +54,7 @@
               {FB_CMD_UPDATE_SUPER, UpdateSuperHandler},
               {FB_CMD_OEM, OemCmdHandler},
               {FB_CMD_GSI, GsiHandler},
+              {FB_CMD_SNAPSHOT_UPDATE, SnapshotUpdateHandler},
       }),
       transport_(std::make_unique<ClientUsbTransport>()),
       boot_control_hal_(IBootControl::getService()),
diff --git a/fastboot/device/variables.cpp b/fastboot/device/variables.cpp
index 130a3cf..6e613d6 100644
--- a/fastboot/device/variables.cpp
+++ b/fastboot/device/variables.cpp
@@ -23,6 +23,7 @@
 #include <android-base/properties.h>
 #include <android-base/stringprintf.h>
 #include <android-base/strings.h>
+#include <android/hardware/boot/1.1/IBootControl.h>
 #include <ext4_utils/ext4_utils.h>
 #include <fs_mgr.h>
 #include <healthhalutils/HealthHalUtils.h>
@@ -34,9 +35,11 @@
 
 using ::android::hardware::boot::V1_0::BoolResult;
 using ::android::hardware::boot::V1_0::Slot;
+using ::android::hardware::boot::V1_1::MergeStatus;
 using ::android::hardware::fastboot::V1_0::FileSystemType;
 using ::android::hardware::fastboot::V1_0::Result;
 using ::android::hardware::fastboot::V1_0::Status;
+using IBootControl1_1 = ::android::hardware::boot::V1_1::IBootControl;
 using namespace android::fs_mgr;
 
 constexpr char kFastbootProtocolVersion[] = "0.4";
@@ -424,3 +427,34 @@
     *message = fs_mgr_get_super_partition_name(slot_number);
     return true;
 }
+
+bool GetSnapshotUpdateStatus(FastbootDevice* device, const std::vector<std::string>& /* args */,
+                             std::string* message) {
+    // Note that we use the HAL rather than mounting /metadata, since we want
+    // our results to match the bootloader.
+    auto hal = device->boot_control_hal();
+    if (!hal) {
+        *message = "not supported";
+        return false;
+    }
+
+    android::sp<IBootControl1_1> hal11 = IBootControl1_1::castFrom(hal);
+    if (!hal11) {
+        *message = "not supported";
+        return false;
+    }
+
+    MergeStatus status = hal11->getSnapshotMergeStatus();
+    switch (status) {
+        case MergeStatus::SNAPSHOTTED:
+            *message = "snapshotted";
+            break;
+        case MergeStatus::MERGING:
+            *message = "merging";
+            break;
+        default:
+            *message = "none";
+            break;
+    }
+    return true;
+}
diff --git a/fastboot/device/variables.h b/fastboot/device/variables.h
index 015a4c5..4dec10f 100644
--- a/fastboot/device/variables.h
+++ b/fastboot/device/variables.h
@@ -61,6 +61,8 @@
                      std::string* message);
 bool GetSuperPartitionName(FastbootDevice* device, const std::vector<std::string>& args,
                            std::string* message);
+bool GetSnapshotUpdateStatus(FastbootDevice* device, const std::vector<std::string>& args,
+                             std::string* message);
 
 // Helpers for getvar all.
 std::vector<std::vector<std::string>> GetAllPartitionArgsWithSlot(FastbootDevice* device);
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 2fe3b1a..7ce7c7c 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -396,6 +396,9 @@
             " gsi wipe|disable           Wipe or disable a GSI installation (fastbootd only).\n"
             " wipe-super [SUPER_EMPTY]   Wipe the super partition. This will reset it to\n"
             "                            contain an empty set of default dynamic partitions.\n"
+            " snapshot-update cancel     On devices that support snapshot-based updates, cancel\n"
+            "                            an in-progress update. This may make the device\n"
+            "                            unbootable until it is reflashed.\n"
             "\n"
             "boot image:\n"
             " boot KERNEL [RAMDISK [SECOND]]\n"
@@ -1216,6 +1219,14 @@
     target_sparse_limit = -1;
 }
 
+static void CancelSnapshotIfNeeded() {
+    std::string merge_status = "none";
+    if (fb->GetVar(FB_VAR_SNAPSHOT_UPDATE_STATUS, &merge_status) == fastboot::SUCCESS &&
+        merge_status != "none") {
+        fb->SnapshotUpdateCommand("Cancel");
+    }
+}
+
 class ImageSource {
   public:
     virtual bool ReadFile(const std::string& name, std::vector<char>* out) const = 0;
@@ -1268,6 +1279,8 @@
     DetermineSecondarySlot();
     CollectImages();
 
+    CancelSnapshotIfNeeded();
+
     // First flash boot partitions. We allow this to happen either in userspace
     // or in bootloader fastboot.
     FlashImages(boot_images_);
@@ -2071,12 +2084,24 @@
                 image = next_arg(&args);
             }
             do_wipe_super(image, slot_override);
+        } else if (command == "snapshot-update") {
+            std::string arg;
+            if (!args.empty()) {
+                arg = next_arg(&args);
+            }
+            if (!arg.empty() && arg != "cancel") {
+                syntax_error("expected: snapshot-update [cancel]");
+            }
+            fb->SnapshotUpdateCommand(arg);
         } else {
             syntax_error("unknown command %s", command.c_str());
         }
     }
 
     if (wants_wipe) {
+        if (force_flash) {
+            CancelSnapshotIfNeeded();
+        }
         std::vector<std::string> partitions = { "userdata", "cache", "metadata" };
         for (const auto& partition : partitions) {
             std::string partition_type;
diff --git a/fastboot/fastboot_driver.cpp b/fastboot/fastboot_driver.cpp
index b897182..6a5ad20 100644
--- a/fastboot/fastboot_driver.cpp
+++ b/fastboot/fastboot_driver.cpp
@@ -122,6 +122,12 @@
                       response, info);
 }
 
+RetCode FastBootDriver::SnapshotUpdateCommand(const std::string& command, std::string* response,
+                                              std::vector<std::string>* info) {
+    std::string raw = FB_CMD_SNAPSHOT_UPDATE ":" + command;
+    return RawCommand(raw, response, info);
+}
+
 RetCode FastBootDriver::FlashPartition(const std::string& partition,
                                        const std::vector<char>& data) {
     RetCode ret;
diff --git a/fastboot/fastboot_driver.h b/fastboot/fastboot_driver.h
index af02637..7265632 100644
--- a/fastboot/fastboot_driver.h
+++ b/fastboot/fastboot_driver.h
@@ -104,6 +104,8 @@
                       std::vector<std::string>* info = nullptr);
     RetCode Upload(const std::string& outfile, std::string* response = nullptr,
                    std::vector<std::string>* info = nullptr);
+    RetCode SnapshotUpdateCommand(const std::string& command, std::string* response = nullptr,
+                                  std::vector<std::string>* info = nullptr);
 
     /* HIGHER LEVEL COMMANDS -- Composed of the commands above */
     RetCode FlashPartition(const std::string& partition, const std::vector<char>& data);