From 20bdf899b283d6e5f97a8100c076630e3bc94bc2 Mon Sep 17 00:00:00 2001
From: Hridya Valsaraju <hridya@google.com>
Date: Wed, 10 Oct 2018 11:02:19 -0700
Subject: [PATCH] set_active command should update current slot information.

Boot control HAL implementations will return the current boot
slot as the active slot. If a set_active command is issued on another
slot, it should be considered as the new active slot and subsequent
flashes should flash the same.

Test: fastboot set_active "b", fastboot getvar current-slot
Bug: 78793464
Change-Id: Ida3817670de8e74a7d7ae2a905e7ac1756c6bdf1
---
 fastboot/device/commands.cpp        | 8 +++++++-
 fastboot/device/fastboot_device.cpp | 8 +++++++-
 fastboot/device/fastboot_device.h   | 3 +++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/fastboot/device/commands.cpp b/fastboot/device/commands.cpp
index 3f663ef8e..863a1e6fd 100644
--- a/fastboot/device/commands.cpp
+++ b/fastboot/device/commands.cpp
@@ -221,7 +221,13 @@ bool SetActiveHandler(FastbootDevice* device, const std::vector<std::string>& ar
     CommandResult ret;
     auto cb = [&ret](CommandResult result) { ret = result; };
     auto result = boot_control_hal->setActiveBootSlot(slot, cb);
-    if (result.isOk() && ret.success) return device->WriteStatus(FastbootResult::OKAY, "");
+    if (result.isOk() && ret.success) {
+        // Save as slot suffix to match the suffix format as returned from
+        // the boot control HAL.
+        auto current_slot = "_" + args[1];
+        device->set_active_slot(current_slot);
+        return device->WriteStatus(FastbootResult::OKAY, "");
+    }
     return device->WriteStatus(FastbootResult::FAIL, "Unable to set slot");
 }
 
diff --git a/fastboot/device/fastboot_device.cpp b/fastboot/device/fastboot_device.cpp
index b843c05d5..6cb48928d 100644
--- a/fastboot/device/fastboot_device.cpp
+++ b/fastboot/device/fastboot_device.cpp
@@ -57,7 +57,8 @@ FastbootDevice::FastbootDevice()
       transport_(std::make_unique<ClientUsbTransport>()),
       boot_control_hal_(IBootControl::getService()),
       health_hal_(get_health_service()),
-      fastboot_hal_(IFastboot::getService()) {}
+      fastboot_hal_(IFastboot::getService()),
+      active_slot_("") {}
 
 FastbootDevice::~FastbootDevice() {
     CloseDevice();
@@ -68,6 +69,11 @@ void FastbootDevice::CloseDevice() {
 }
 
 std::string FastbootDevice::GetCurrentSlot() {
+    // Check if a set_active ccommand was issued earlier since the boot control HAL
+    // returns the slot that is currently booted into.
+    if (!active_slot_.empty()) {
+        return active_slot_;
+    }
     // Non-A/B devices must not have boot control HALs.
     if (!boot_control_hal_) {
         return "";
diff --git a/fastboot/device/fastboot_device.h b/fastboot/device/fastboot_device.h
index 2eb7177b8..091aadfd3 100644
--- a/fastboot/device/fastboot_device.h
+++ b/fastboot/device/fastboot_device.h
@@ -56,6 +56,8 @@ class FastbootDevice {
     }
     android::sp<android::hardware::health::V2_0::IHealth> health_hal() { return health_hal_; }
 
+    void set_active_slot(const std::string& active_slot) { active_slot_ = active_slot; }
+
   private:
     const std::unordered_map<std::string, CommandHandler> kCommandMap;
 
@@ -64,4 +66,5 @@ class FastbootDevice {
     android::sp<android::hardware::health::V2_0::IHealth> health_hal_;
     android::sp<android::hardware::fastboot::V1_0::IFastboot> fastboot_hal_;
     std::vector<char> download_data_;
+    std::string active_slot_;
 };