Adding testing for optimized flash super

Adding test cases for correct formation of optimized flash super task.
We are adding an explicit pattern match for this task to be correctly
formed. Changing Optimized flash task to only remove the reboot to
userspace as a user might want to reboot back to bootloader after
flashing. We also need to change a couple functions to take a
IFastbootDriver to mock up the initialization path.

Test: fastboot_test
Bug: 297085098
Change-Id: Ic5c63bd4057ca6d64647134e5ce33fef12077fdb
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 8c607dd..81787f5 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -1285,7 +1285,7 @@
     return current_slot;
 }
 
-static int get_slot_count() {
+static int get_slot_count(fastboot::IFastBootDriver* fb) {
     std::string var;
     int count = 0;
     if (fb->GetVar("slot-count", &var) != fastboot::SUCCESS ||
@@ -1295,8 +1295,8 @@
     return count;
 }
 
-bool supports_AB() {
-    return get_slot_count() >= 2;
+bool supports_AB(fastboot::IFastBootDriver* fb) {
+    return get_slot_count(fb) >= 2;
 }
 
 // Given a current slot, this returns what the 'other' slot is.
@@ -1308,7 +1308,7 @@
 }
 
 static std::string get_other_slot(const std::string& current_slot) {
-    return get_other_slot(current_slot, get_slot_count());
+    return get_other_slot(current_slot, get_slot_count(fb));
 }
 
 static std::string get_other_slot(int count) {
@@ -1316,7 +1316,7 @@
 }
 
 static std::string get_other_slot() {
-    return get_other_slot(get_current_slot(), get_slot_count());
+    return get_other_slot(get_current_slot(), get_slot_count(fb));
 }
 
 static std::string verify_slot(const std::string& slot_name, bool allow_all) {
@@ -1325,7 +1325,7 @@
         if (allow_all) {
             return "all";
         } else {
-            int count = get_slot_count();
+            int count = get_slot_count(fb);
             if (count > 0) {
                 return "a";
             } else {
@@ -1334,7 +1334,7 @@
         }
     }
 
-    int count = get_slot_count();
+    int count = get_slot_count(fb);
     if (count == 0) die("Device does not support slots");
 
     if (slot == "other") {
@@ -1407,7 +1407,7 @@
                 slot.c_str());
         }
         if (has_slot == "yes") {
-            for (int i = 0; i < get_slot_count(); i++) {
+            for (int i = 0; i < get_slot_count(fb); i++) {
                 do_for_partition(part, std::string(1, (char)(i + 'a')), func, force_slot);
             }
         } else {
@@ -1528,7 +1528,7 @@
 // Sets slot_override as the active slot. If slot_override is blank,
 // set current slot as active instead. This clears slot-unbootable.
 static void set_active(const std::string& slot_override) {
-    if (!supports_AB()) return;
+    if (!supports_AB(fb)) return;
 
     if (slot_override != "") {
         fb->SetActive(slot_override);
@@ -1845,7 +1845,7 @@
         fp_->secondary_slot = get_other_slot();
     }
     if (fp_->secondary_slot == "") {
-        if (supports_AB()) {
+        if (supports_AB(fb)) {
             fprintf(stderr, "Warning: Could not determine slot for secondary images. Ignoring.\n");
         }
         fp_->skip_secondary = true;
diff --git a/fastboot/fastboot.h b/fastboot/fastboot.h
index 75b8d29..35deea7 100644
--- a/fastboot/fastboot.h
+++ b/fastboot/fastboot.h
@@ -29,10 +29,8 @@
 
 #include <functional>
 #include <string>
-#include "fastboot_driver.h"
 #include "fastboot_driver_interface.h"
 #include "filesystem.h"
-#include "super_flash_helper.h"
 #include "task.h"
 #include "util.h"
 
@@ -183,12 +181,12 @@
 };
 
 Result<NetworkSerial, FastbootError> ParseNetworkSerial(const std::string& serial);
-bool supports_AB();
 std::string GetPartitionName(const ImageEntry& entry, const std::string& current_slot_);
 void flash_partition_files(const std::string& partition, const std::vector<SparsePtr>& files);
 int64_t get_sparse_limit(int64_t size, const FlashingPlan* fp);
 std::vector<SparsePtr> resparse_file(sparse_file* s, int64_t max_size);
 
+bool supports_AB(fastboot::IFastBootDriver* fb);
 bool is_retrofit_device(fastboot::IFastBootDriver* fb);
 bool is_logical(const std::string& partition);
 void fb_perform_format(const std::string& partition, int skip_if_not_supported,
diff --git a/fastboot/fastboot_driver.h b/fastboot/fastboot_driver.h
index 8774ead..48b90ed 100644
--- a/fastboot/fastboot_driver.h
+++ b/fastboot/fastboot_driver.h
@@ -27,7 +27,6 @@
  */
 #pragma once
 #include <cstdlib>
-#include <deque>
 #include <functional>
 #include <limits>
 #include <string>
@@ -37,10 +36,8 @@
 #include <android-base/stringprintf.h>
 #include <android-base/unique_fd.h>
 #include <bootimg.h>
-#include <inttypes.h>
 #include <sparse/sparse.h>
 
-#include "constants.h"
 #include "fastboot_driver_interface.h"
 #include "transport.h"
 
diff --git a/fastboot/fuzzy_fastboot/main.cpp b/fastboot/fuzzy_fastboot/main.cpp
index e635937..95778e1 100644
--- a/fastboot/fuzzy_fastboot/main.cpp
+++ b/fastboot/fuzzy_fastboot/main.cpp
@@ -50,6 +50,7 @@
 #include <gtest/gtest.h>
 #include <sparse/sparse.h>
 
+#include "constants.h"
 #include "fastboot_driver.h"
 #include "usb.h"
 
@@ -929,8 +930,7 @@
 
     ASSERT_TRUE(UsbStillAvailible()) << USB_PORT_GONE;
     std::string resp;
-    EXPECT_EQ(fb->GetVar("product", &resp), SUCCESS)
-                << "Device is unresponsive to getvar command";
+    EXPECT_EQ(fb->GetVar("product", &resp), SUCCESS) << "Device is unresponsive to getvar command";
 }
 
 TEST_F(Fuzz, CommandTooLarge) {
@@ -986,11 +986,10 @@
 TEST_F(Fuzz, SparseZeroBlkSize) {
     // handcrafted malform sparse file with zero as block size
     const std::vector<char> buf = {
-        '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00', '\x1c', '\x00', '\x0c', '\x00',
-        '\x00', '\x00', '\x00', '\x00', '\x01', '\x00', '\x00', '\x00', '\x01', '\x00', '\x00', '\x00',
-        '\x00', '\x00', '\x00', '\x00', '\xc2', '\xca', '\x00', '\x00', '\x01', '\x00', '\x00', '\x00',
-        '\x10', '\x00', '\x00', '\x00', '\x11', '\x22', '\x33', '\x44'
-    };
+            '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00', '\x1c', '\x00', '\x0c',
+            '\x00', '\x00', '\x00', '\x00', '\x00', '\x01', '\x00', '\x00', '\x00', '\x01', '\x00',
+            '\x00', '\x00', '\x00', '\x00', '\x00', '\x00', '\xc2', '\xca', '\x00', '\x00', '\x01',
+            '\x00', '\x00', '\x00', '\x10', '\x00', '\x00', '\x00', '\x11', '\x22', '\x33', '\x44'};
 
     ASSERT_EQ(DownloadCommand(buf.size()), SUCCESS) << "Device rejected download command";
     ASSERT_EQ(SendBuffer(buf), SUCCESS) << "Downloading payload failed";
@@ -1005,13 +1004,10 @@
 TEST_F(Fuzz, SparseVeryLargeBlkSize) {
     // handcrafted sparse file with block size of ~4GB and divisible 4
     const std::vector<char> buf = {
-        '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00',
-        '\x1c', '\x00', '\x0c', '\x00', '\xF0', '\xFF', '\xFF', '\xFF',
-        '\x01', '\x00', '\x00', '\x00', '\x01', '\x00', '\x00', '\x00',
-        '\x00', '\x00', '\x00', '\x00', '\xc3', '\xca', '\x00', '\x00',
-        '\x01', '\x00', '\x00', '\x00', '\x0c', '\x00', '\x00', '\x00',
-        '\x11', '\x22', '\x33', '\x44'
-    };
+            '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00', '\x1c', '\x00', '\x0c',
+            '\x00', '\xF0', '\xFF', '\xFF', '\xFF', '\x01', '\x00', '\x00', '\x00', '\x01', '\x00',
+            '\x00', '\x00', '\x00', '\x00', '\x00', '\x00', '\xc3', '\xca', '\x00', '\x00', '\x01',
+            '\x00', '\x00', '\x00', '\x0c', '\x00', '\x00', '\x00', '\x11', '\x22', '\x33', '\x44'};
 
     ASSERT_EQ(DownloadCommand(buf.size()), SUCCESS) << "Device rejected download command";
     ASSERT_EQ(SendBuffer(buf), SUCCESS) << "Downloading payload failed";
@@ -1022,11 +1018,10 @@
 TEST_F(Fuzz, SparseTrimmed) {
     // handcrafted malform sparse file which is trimmed
     const std::vector<char> buf = {
-        '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00', '\x1c', '\x00', '\x0c', '\x00',
-        '\x00', '\x10', '\x00', '\x00', '\x00', '\x00', '\x08', '\x00', '\x01', '\x00', '\x00', '\x00',
-        '\x00', '\x00', '\x00', '\x00', '\xc1', '\xca', '\x00', '\x00', '\x01', '\x00', '\x00', '\x00',
-        '\x00', '\x00', '\x00', '\x80', '\x11', '\x22', '\x33', '\x44'
-    };
+            '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00', '\x1c', '\x00', '\x0c',
+            '\x00', '\x00', '\x10', '\x00', '\x00', '\x00', '\x00', '\x08', '\x00', '\x01', '\x00',
+            '\x00', '\x00', '\x00', '\x00', '\x00', '\x00', '\xc1', '\xca', '\x00', '\x00', '\x01',
+            '\x00', '\x00', '\x00', '\x00', '\x00', '\x00', '\x80', '\x11', '\x22', '\x33', '\x44'};
 
     ASSERT_EQ(DownloadCommand(buf.size()), SUCCESS) << "Device rejected download command";
     ASSERT_EQ(SendBuffer(buf), SUCCESS) << "Downloading payload failed";
@@ -1041,11 +1036,10 @@
 TEST_F(Fuzz, SparseInvalidChurk) {
     // handcrafted malform sparse file with invalid churk
     const std::vector<char> buf = {
-        '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00', '\x1c', '\x00', '\x0c', '\x00',
-        '\x00', '\x10', '\x00', '\x00', '\x00', '\x00', '\x08', '\x00', '\x01', '\x00', '\x00', '\x00',
-        '\x00', '\x00', '\x00', '\x00', '\xc1', '\xca', '\x00', '\x00', '\x01', '\x00', '\x00', '\x00',
-        '\x10', '\x00', '\x00', '\x00', '\x11', '\x22', '\x33', '\x44'
-    };
+            '\x3a', '\xff', '\x26', '\xed', '\x01', '\x00', '\x00', '\x00', '\x1c', '\x00', '\x0c',
+            '\x00', '\x00', '\x10', '\x00', '\x00', '\x00', '\x00', '\x08', '\x00', '\x01', '\x00',
+            '\x00', '\x00', '\x00', '\x00', '\x00', '\x00', '\xc1', '\xca', '\x00', '\x00', '\x01',
+            '\x00', '\x00', '\x00', '\x10', '\x00', '\x00', '\x00', '\x11', '\x22', '\x33', '\x44'};
 
     ASSERT_EQ(DownloadCommand(buf.size()), SUCCESS) << "Device rejected download command";
     ASSERT_EQ(SendBuffer(buf), SUCCESS) << "Downloading payload failed";
@@ -1895,7 +1889,8 @@
     if (!fastboot::FastBootTest::IsFastbootOverTcp()) {
         printf("<Waiting for Device>\n");
         const auto matcher = [](usb_ifc_info* info) -> int {
-            return fastboot::FastBootTest::MatchFastboot(info, fastboot::FastBootTest::device_serial);
+            return fastboot::FastBootTest::MatchFastboot(info,
+                                                         fastboot::FastBootTest::device_serial);
         };
         Transport* transport = nullptr;
         while (!transport) {
diff --git a/fastboot/task.cpp b/fastboot/task.cpp
index f0eed0c..f13dd55 100644
--- a/fastboot/task.cpp
+++ b/fastboot/task.cpp
@@ -15,8 +15,7 @@
 //
 #include "task.h"
 
-#include <cstddef>
-#include <iostream>
+#include "fastboot_driver.h"
 
 #include <android-base/logging.h>
 #include <android-base/parseint.h>
@@ -130,6 +129,7 @@
     // Send the data to the device.
     flash_partition_files(super_name_, files);
 }
+
 std::string OptimizedFlashSuperTask::ToString() const {
     return "optimized-flash-super";
 }
@@ -165,7 +165,7 @@
         LOG(INFO) << "super optimization is disabled";
         return nullptr;
     }
-    if (!supports_AB()) {
+    if (!supports_AB(fp->fb)) {
         LOG(VERBOSE) << "Cannot optimize flashing super on non-AB device";
         return nullptr;
     }
@@ -218,17 +218,21 @@
 
     auto s = helper->GetSparseLayout();
     if (!s) return nullptr;
-    // Remove images that we already flashed, just in case we have non-dynamic OS images.
+
+    // Remove tasks that are concatenated into this optimized task
     auto remove_if_callback = [&](const auto& task) -> bool {
         if (auto flash_task = task->AsFlashTask()) {
             return helper->WillFlash(flash_task->GetPartitionAndSlot());
         } else if (auto update_super_task = task->AsUpdateSuperTask()) {
             return true;
         } else if (auto reboot_task = task->AsRebootTask()) {
-            return true;
+            if (reboot_task->GetTarget() == "fastboot") {
+                return true;
+            }
         }
         return false;
     };
+
     tasks.erase(std::remove_if(tasks.begin(), tasks.end(), remove_if_callback), tasks.end());
 
     return std::make_unique<OptimizedFlashSuperTask>(super_name, std::move(helper), std::move(s),
diff --git a/fastboot/task.h b/fastboot/task.h
index 6ebe381..a98c874 100644
--- a/fastboot/task.h
+++ b/fastboot/task.h
@@ -15,10 +15,8 @@
 //
 #pragma once
 
-#include <sstream>
 #include <string>
 
-#include "fastboot_driver.h"
 #include "super_flash_helper.h"
 #include "util.h"
 
@@ -29,6 +27,7 @@
 class FlashTask;
 class RebootTask;
 class UpdateSuperTask;
+class OptimizedFlashSuperTask;
 class WipeTask;
 class ResizeTask;
 class Task {
@@ -40,6 +39,7 @@
     virtual FlashTask* AsFlashTask() { return nullptr; }
     virtual RebootTask* AsRebootTask() { return nullptr; }
     virtual UpdateSuperTask* AsUpdateSuperTask() { return nullptr; }
+    virtual OptimizedFlashSuperTask* AsOptimizedFlashSuperTask() { return nullptr; }
     virtual WipeTask* AsWipeTask() { return nullptr; }
     virtual ResizeTask* AsResizeTask() { return nullptr; }
 
@@ -86,13 +86,13 @@
   public:
     OptimizedFlashSuperTask(const std::string& super_name, std::unique_ptr<SuperFlashHelper> helper,
                             SparsePtr sparse_layout, uint64_t super_size, const FlashingPlan* fp);
+    virtual OptimizedFlashSuperTask* AsOptimizedFlashSuperTask() override { return this; }
 
     static std::unique_ptr<OptimizedFlashSuperTask> Initialize(
             const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>& tasks);
     static bool CanOptimize(const ImageSource* source,
                             const std::vector<std::unique_ptr<Task>>& tasks);
 
-    using ImageEntry = std::pair<const Image*, std::string>;
     void Run() override;
     std::string ToString() const override;
 
diff --git a/fastboot/task_test.cpp b/fastboot/task_test.cpp
index 1ba3f4a..42be3cb 100644
--- a/fastboot/task_test.cpp
+++ b/fastboot/task_test.cpp
@@ -19,11 +19,10 @@
 #include "fastboot_driver_mock.h"
 
 #include <gtest/gtest.h>
-#include <fstream>
 #include <iostream>
 #include <memory>
-#include <unordered_map>
 #include "android-base/strings.h"
+#include "gmock/gmock.h"
 
 using android::base::Split;
 using testing::_;
@@ -235,3 +234,111 @@
             << "size of fastboot-info task list: " << fastboot_info_tasks.size()
             << " size of hardcoded task list: " << hardcoded_tasks.size();
 }
+
+TEST_F(ParseTest, CanOptimizeTest) {
+    if (!get_android_product_out()) {
+        GTEST_SKIP();
+    }
+
+    LocalImageSource s;
+    fp->source = &s;
+    fp->sparse_limit = std::numeric_limits<int64_t>::max();
+
+    fastboot::MockFastbootDriver fb;
+    fp->fb = &fb;
+    fp->should_optimize_flash_super = false;
+    fp->should_use_fastboot_info = true;
+
+    std::vector<std::pair<std::vector<std::string>, bool>> patternmatchtest = {
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "reboot fastboot",
+              "update-super", "flash product", "flash system", "flash system_ext", "flash odm",
+              "if-wipe erase userdata"},
+             true},
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "reboot fastboot",
+              "update-super", "flash product", "flash system", "flash system_ext", "flash odm",
+              "if-wipe erase userdata"},
+             true},
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "reboot fastboot",
+              "flash product", "flash system", "flash system_ext", "flash odm",
+              "if-wipe erase userdata"},
+             false},
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "update-super", "flash product",
+              "flash system", "flash system_ext", "flash odm", "if-wipe erase userdata"},
+             false},
+    };
+
+    auto remove_if_callback = [&](const auto& task) -> bool { return !!task->AsResizeTask(); };
+
+    for (auto& test : patternmatchtest) {
+        std::vector<std::unique_ptr<Task>> tasks = ParseFastbootInfo(fp.get(), test.first);
+        tasks.erase(std::remove_if(tasks.begin(), tasks.end(), remove_if_callback), tasks.end());
+        ASSERT_EQ(OptimizedFlashSuperTask::CanOptimize(fp->source, tasks), test.second);
+    }
+}
+// Note: this test is exclusively testing that optimized flash super pattern matches a given task
+// list and is able to optimized based on a correct sequence of tasks
+TEST_F(ParseTest, OptimizedFlashSuperPatternMatchTest) {
+    if (!get_android_product_out()) {
+        GTEST_SKIP();
+    }
+
+    LocalImageSource s;
+    fp->source = &s;
+    fp->sparse_limit = std::numeric_limits<int64_t>::max();
+
+    fastboot::MockFastbootDriver fb;
+    fp->fb = &fb;
+    fp->should_optimize_flash_super = true;
+    fp->should_use_fastboot_info = true;
+
+    ON_CALL(fb, GetVar("super-partition-name", _, _))
+            .WillByDefault(testing::Return(fastboot::BAD_ARG));
+
+    ON_CALL(fb, GetVar("slot-count", _, _))
+            .WillByDefault(testing::DoAll(testing::SetArgPointee<1>("2"),
+                                          testing::Return(fastboot::SUCCESS)));
+
+    ON_CALL(fb, GetVar("partition-size:super", _, _))
+            .WillByDefault(testing::DoAll(testing::SetArgPointee<1>("1000"),
+                                          testing::Return(fastboot::SUCCESS)));
+
+    std::vector<std::pair<std::vector<std::string>, bool>> patternmatchtest = {
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "reboot fastboot",
+              "update-super", "flash product", "flash system", "flash system_ext", "flash odm",
+              "if-wipe erase userdata"},
+             true},
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "reboot fastboot",
+              "update-super", "flash product", "flash system", "flash system_ext", "flash odm",
+              "if-wipe erase userdata"},
+             true},
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "reboot fastboot",
+              "flash product", "flash system", "flash system_ext", "flash odm",
+              "if-wipe erase userdata"},
+             false},
+            {{"flash boot", "flash init_boot", "flash vendor_boot", "update-super", "flash product",
+              "flash system", "flash system_ext", "flash odm", "if-wipe erase userdata"},
+             false},
+    };
+
+    for (auto& test : patternmatchtest) {
+        std::vector<std::unique_ptr<Task>> tasks = ParseFastbootInfo(fp.get(), test.first);
+        // Check to make sure we have an optimized flash super task && no more dynamic partition
+        // flashing tasks
+        auto&& IsOptimized = [](const FlashingPlan* fp,
+                                const std::vector<std::unique_ptr<Task>>& tasks) {
+            bool contains_optimized_task = false;
+            for (auto& task : tasks) {
+                if (auto optimized_task = task->AsOptimizedFlashSuperTask()) {
+                    contains_optimized_task = true;
+                }
+                if (auto flash_task = task->AsFlashTask()) {
+                    if (FlashTask::IsDynamicParitition(fp->source, flash_task)) {
+                        return false;
+                    }
+                }
+            }
+            return contains_optimized_task;
+        };
+        ASSERT_EQ(IsOptimized(fp.get(), tasks), test.second);
+    }
+}