Merge "Adding testing for optimized flash super" into main
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index d7e3db6..21df729 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -1286,7 +1286,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 ||
@@ -1296,8 +1296,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.
@@ -1309,7 +1309,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) {
@@ -1317,7 +1317,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) {
@@ -1326,7 +1326,7 @@
         if (allow_all) {
             return "all";
         } else {
-            int count = get_slot_count();
+            int count = get_slot_count(fb);
             if (count > 0) {
                 return "a";
             } else {
@@ -1335,7 +1335,7 @@
         }
     }
 
-    int count = get_slot_count();
+    int count = get_slot_count(fb);
     if (count == 0) die("Device does not support slots");
 
     if (slot == "other") {
@@ -1408,7 +1408,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 {
@@ -1529,7 +1529,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);
@@ -1844,7 +1844,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 30298cb..49cebd9 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 <memory>
@@ -38,10 +37,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 43aaec9..79f3939 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"
 
@@ -928,8 +929,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) {
@@ -985,11 +985,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";
@@ -1004,13 +1003,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";
@@ -1021,11 +1017,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";
@@ -1040,11 +1035,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";
@@ -1894,7 +1888,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);
         };
         std::unique_ptr<Transport> transport;
         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);
+    }
+}