Fixing optimization logic

Adding a hard pattern check for optimized task formation. To keep
behavior consistent, we will remove the old initialization path and add
resize tasks after attempting optimization.

Test: fastboot_test
Bug: 297085098
Change-Id: Ie0e656af9be7abdd130290fe547ffbf385ce75d6

3 files changed, 93 insertions, 122 deletions

diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 71a228ea1..8c607dd83 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -1678,7 +1678,7 @@ bool AddResizeTasks(const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>*
     }
     for (size_t i = 0; i < tasks->size(); i++) {
         if (auto flash_task = tasks->at(i)->AsFlashTask()) {
-            if (should_flash_in_userspace(*metadata.get(), flash_task->GetPartitionAndSlot())) {
+            if (FlashTask::IsDynamicParitition(fp->source, flash_task)) {
                 if (!loc) {
                     loc = i;
                 }
@@ -1760,25 +1760,15 @@ std::vector<std::unique_ptr<Task>> ParseFastbootInfo(const FlashingPlan* fp,
         }
         tasks.emplace_back(std::move(task));
     }
-    if (auto flash_super_task = OptimizedFlashSuperTask::InitializeFromTasks(fp, tasks)) {
-        auto it = tasks.begin();
-        for (size_t i = 0; i < tasks.size(); i++) {
-            if (auto flash_task = tasks[i]->AsFlashTask()) {
-                if (should_flash_in_userspace(flash_task->GetPartitionAndSlot())) {
-                    break;
-                }
-            }
-            if (auto wipe_task = tasks[i]->AsWipeTask()) {
-                break;
-            }
-            it++;
-        }
-        tasks.insert(it, std::move(flash_super_task));
+
+    if (auto flash_super_task = OptimizedFlashSuperTask::Initialize(fp, tasks)) {
+        tasks.emplace_back(std::move(flash_super_task));
     } else {
         if (!AddResizeTasks(fp, &tasks)) {
             LOG(WARNING) << "Failed to add resize tasks";
-        };
+        }
     }
+
     return tasks;
 }
 
@@ -1885,30 +1875,35 @@ std::vector<std::unique_ptr<Task>> FlashAllTool::CollectTasksFromImageList() {
     // or in bootloader fastboot.
     std::vector<std::unique_ptr<Task>> tasks;
     AddFlashTasks(boot_images_, tasks);
-    if (auto flash_super_task = OptimizedFlashSuperTask::Initialize(fp_, os_images_)) {
+
+    // Sync the super partition. This will reboot to userspace fastboot if needed.
+    tasks.emplace_back(std::make_unique<UpdateSuperTask>(fp_));
+    for (const auto& [image, slot] : os_images_) {
+        // Retrofit devices have two super partitions, named super_a and super_b.
+        // On these devices, secondary slots must be flashed as physical
+        // partitions (otherwise they would not mount on first boot). To enforce
+        // this, we delete any logical partitions for the "other" slot.
+        if (is_retrofit_device(fp_->fb)) {
+            std::string partition_name = image->part_name + "_"s + slot;
+            if (image->IsSecondary() && should_flash_in_userspace(partition_name)) {
+                fp_->fb->DeletePartition(partition_name);
+            }
+            tasks.emplace_back(std::make_unique<DeleteTask>(fp_, partition_name));
+        }
+    }
+
+    AddFlashTasks(os_images_, tasks);
+
+    if (auto flash_super_task = OptimizedFlashSuperTask::Initialize(fp_, tasks)) {
         tasks.emplace_back(std::move(flash_super_task));
     } else {
-        // Sync the super partition. This will reboot to userspace fastboot if needed.
-        tasks.emplace_back(std::make_unique<UpdateSuperTask>(fp_));
         // Resize any logical partition to 0, so each partition is reset to 0
         // extents, and will achieve more optimal allocation.
-        for (const auto& [image, slot] : os_images_) {
-            // Retrofit devices have two super partitions, named super_a and super_b.
-            // On these devices, secondary slots must be flashed as physical
-            // partitions (otherwise they would not mount on first boot). To enforce
-            // this, we delete any logical partitions for the "other" slot.
-            if (is_retrofit_device(fp_->fb)) {
-                std::string partition_name = image->part_name + "_"s + slot;
-                if (image->IsSecondary() && should_flash_in_userspace(partition_name)) {
-                    fp_->fb->DeletePartition(partition_name);
-                }
-                tasks.emplace_back(std::make_unique<DeleteTask>(fp_, partition_name));
-            }
-            tasks.emplace_back(std::make_unique<ResizeTask>(fp_, image->part_name, "0", slot));
+        if (!AddResizeTasks(fp_, &tasks)) {
+            LOG(WARNING) << "Failed to add resize tasks";
         }
     }
 
-    AddFlashTasks(os_images_, tasks);
     return tasks;
 }
 
diff --git a/fastboot/task.cpp b/fastboot/task.cpp
index 146064cc5..f0eed0cd1 100644
--- a/fastboot/task.cpp
+++ b/fastboot/task.cpp
@@ -15,6 +15,7 @@
 //
 #include "task.h"
 
+#include <cstddef>
 #include <iostream>
 
 #include <android-base/logging.h>
@@ -30,6 +31,15 @@ FlashTask::FlashTask(const std::string& slot, const std::string& pname, const st
                      const bool apply_vbmeta, const FlashingPlan* fp)
     : pname_(pname), fname_(fname), slot_(slot), apply_vbmeta_(apply_vbmeta), fp_(fp) {}
 
+bool FlashTask::IsDynamicParitition(const ImageSource* source, const FlashTask* task) {
+    std::vector<char> contents;
+    if (!source->ReadFile("super_empty.img", &contents)) {
+        return false;
+    }
+    auto metadata = android::fs_mgr::ReadFromImageBlob(contents.data(), contents.size());
+    return should_flash_in_userspace(*metadata.get(), task->GetPartitionAndSlot());
+}
+
 void FlashTask::Run() {
     auto flash = [&](const std::string& partition) {
         if (should_flash_in_userspace(partition) && !is_userspace_fastboot() && !fp_->force_flash) {
@@ -46,7 +56,7 @@ void FlashTask::Run() {
     do_for_partitions(pname_, slot_, flash, true);
 }
 
-std::string FlashTask::ToString() {
+std::string FlashTask::ToString() const {
     std::string apply_vbmeta_string = "";
     if (apply_vbmeta_) {
         apply_vbmeta_string = " --apply_vbmeta";
@@ -54,7 +64,7 @@ std::string FlashTask::ToString() {
     return "flash" + apply_vbmeta_string + " " + pname_ + " " + fname_;
 }
 
-std::string FlashTask::GetPartitionAndSlot() {
+std::string FlashTask::GetPartitionAndSlot() const {
     auto slot = slot_;
     if (slot.empty()) {
         slot = get_current_slot();
@@ -92,7 +102,7 @@ void RebootTask::Run() {
     }
 }
 
-std::string RebootTask::ToString() {
+std::string RebootTask::ToString() const {
     return "reboot " + reboot_target_;
 }
 
@@ -120,79 +130,36 @@ void OptimizedFlashSuperTask::Run() {
     // Send the data to the device.
     flash_partition_files(super_name_, files);
 }
-std::string OptimizedFlashSuperTask::ToString() {
+std::string OptimizedFlashSuperTask::ToString() const {
     return "optimized-flash-super";
 }
 
-std::unique_ptr<OptimizedFlashSuperTask> OptimizedFlashSuperTask::Initialize(
-        const FlashingPlan* fp, std::vector<ImageEntry>& os_images) {
-    if (!fp->should_optimize_flash_super) {
-        LOG(INFO) << "super optimization is disabled";
-        return nullptr;
-    }
-    if (!supports_AB()) {
-        LOG(VERBOSE) << "Cannot optimize flashing super on non-AB device";
-        return nullptr;
-    }
-    if (fp->slot_override == "all") {
-        LOG(VERBOSE) << "Cannot optimize flashing super for all slots";
-        return nullptr;
-    }
-
-    // Does this device use dynamic partitions at all?
-    unique_fd fd = fp->source->OpenFile("super_empty.img");
-
-    if (fd < 0) {
-        LOG(VERBOSE) << "could not open super_empty.img";
-        return nullptr;
-    }
-
-    std::string super_name;
-    // Try to find whether there is a super partition.
-    if (fp->fb->GetVar("super-partition-name", &super_name) != fastboot::SUCCESS) {
-        super_name = "super";
-    }
-
-    uint64_t partition_size;
-    std::string partition_size_str;
-    if (fp->fb->GetVar("partition-size:" + super_name, &partition_size_str) != fastboot::SUCCESS) {
-        LOG(VERBOSE) << "Cannot optimize super flashing: could not determine super partition";
-        return nullptr;
-    }
-    partition_size_str = fb_fix_numeric_var(partition_size_str);
-    if (!android::base::ParseUint(partition_size_str, &partition_size)) {
-        LOG(VERBOSE) << "Could not parse " << super_name << " size: " << partition_size_str;
-        return nullptr;
-    }
-
-    std::unique_ptr<SuperFlashHelper> helper = std::make_unique<SuperFlashHelper>(*fp->source);
-    if (!helper->Open(fd)) {
-        return nullptr;
-    }
-
-    for (const auto& entry : os_images) {
-        auto partition = GetPartitionName(entry, fp->current_slot);
-        auto image = entry.first;
-
-        if (!helper->AddPartition(partition, image->img_name, image->optional_if_no_image)) {
-            return nullptr;
+// This looks for a block within tasks that has the following pattern [reboot fastboot,
+// update-super, $LIST_OF_DYNAMIC_FLASH_TASKS] and returns true if this is found.Theoretically
+// this check is just a pattern match and could break if fastboot-info has a bunch of junk commands
+// but all devices should pretty much follow this pattern
+bool OptimizedFlashSuperTask::CanOptimize(const ImageSource* source,
+                                          const std::vector<std::unique_ptr<Task>>& tasks) {
+    for (size_t i = 0; i < tasks.size(); i++) {
+        auto reboot_task = tasks[i]->AsRebootTask();
+        if (!reboot_task || reboot_task->GetTarget() != "fastboot") {
+            continue;
+        }
+        // The check for i >= tasks.size() - 2 is because we are peeking two tasks ahead. We need to
+        // check for an update-super && flash {dynamic_partition}
+        if (i >= tasks.size() - 2 || !tasks[i + 1]->AsUpdateSuperTask()) {
+            continue;
+        }
+        auto flash_task = tasks[i + 2]->AsFlashTask();
+        if (!FlashTask::IsDynamicParitition(source, flash_task)) {
+            continue;
         }
+        return true;
     }
-
-    auto s = helper->GetSparseLayout();
-    if (!s) return nullptr;
-
-    // Remove images that we already flashed, just in case we have non-dynamic OS images.
-    auto remove_if_callback = [&](const ImageEntry& entry) -> bool {
-        return helper->WillFlash(GetPartitionName(entry, fp->current_slot));
-    };
-    os_images.erase(std::remove_if(os_images.begin(), os_images.end(), remove_if_callback),
-                    os_images.end());
-    return std::make_unique<OptimizedFlashSuperTask>(super_name, std::move(helper), std::move(s),
-                                                     partition_size, fp);
+    return false;
 }
 
-std::unique_ptr<OptimizedFlashSuperTask> OptimizedFlashSuperTask::InitializeFromTasks(
+std::unique_ptr<OptimizedFlashSuperTask> OptimizedFlashSuperTask::Initialize(
         const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>& tasks) {
     if (!fp->should_optimize_flash_super) {
         LOG(INFO) << "super optimization is disabled";
@@ -206,6 +173,9 @@ std::unique_ptr<OptimizedFlashSuperTask> OptimizedFlashSuperTask::InitializeFrom
         LOG(VERBOSE) << "Cannot optimize flashing super for all slots";
         return nullptr;
     }
+    if (!CanOptimize(fp->source, tasks)) {
+        return nullptr;
+    }
 
     // Does this device use dynamic partitions at all?
     unique_fd fd = fp->source->OpenFile("super_empty.img");
@@ -288,7 +258,7 @@ void UpdateSuperTask::Run() {
     }
     fp_->fb->RawCommand(command, "Updating super partition");
 }
-std::string UpdateSuperTask::ToString() {
+std::string UpdateSuperTask::ToString() const {
     return "update-super";
 }
 
@@ -305,7 +275,7 @@ void ResizeTask::Run() {
     do_for_partitions(pname_, slot_, resize_partition, false);
 }
 
-std::string ResizeTask::ToString() {
+std::string ResizeTask::ToString() const {
     return "resize " + pname_;
 }
 
@@ -315,7 +285,7 @@ void DeleteTask::Run() {
     fp_->fb->DeletePartition(pname_);
 }
 
-std::string DeleteTask::ToString() {
+std::string DeleteTask::ToString() const {
     return "delete " + pname_;
 }
 
@@ -335,6 +305,6 @@ void WipeTask::Run() {
     fb_perform_format(pname_, 1, partition_type, "", fp_->fs_options, fp_);
 }
 
-std::string WipeTask::ToString() {
+std::string WipeTask::ToString() const {
     return "erase " + pname_;
 }
diff --git a/fastboot/task.h b/fastboot/task.h
index f7c8801f9..6ebe38123 100644
--- a/fastboot/task.h
+++ b/fastboot/task.h
@@ -30,17 +30,18 @@ class FlashTask;
 class RebootTask;
 class UpdateSuperTask;
 class WipeTask;
-
+class ResizeTask;
 class Task {
   public:
     Task() = default;
     virtual void Run() = 0;
-    virtual std::string ToString() = 0;
+    virtual std::string ToString() const = 0;
 
     virtual FlashTask* AsFlashTask() { return nullptr; }
     virtual RebootTask* AsRebootTask() { return nullptr; }
     virtual UpdateSuperTask* AsUpdateSuperTask() { return nullptr; }
     virtual WipeTask* AsWipeTask() { return nullptr; }
+    virtual ResizeTask* AsResizeTask() { return nullptr; }
 
     virtual ~Task() = default;
 };
@@ -51,12 +52,13 @@ class FlashTask : public Task {
               const bool apply_vbmeta, const FlashingPlan* fp);
     virtual FlashTask* AsFlashTask() override { return this; }
 
+    static bool IsDynamicParitition(const ImageSource* source, const FlashTask* task);
     void Run() override;
-    std::string ToString() override;
-    std::string GetPartition() { return pname_; }
-    std::string GetImageName() { return fname_; }
-    std::string GetSlot() { return slot_; }
-    std::string GetPartitionAndSlot();
+    std::string ToString() const override;
+    std::string GetPartition() const { return pname_; }
+    std::string GetImageName() const { return fname_; }
+    std::string GetSlot() const { return slot_; }
+    std::string GetPartitionAndSlot() const;
 
   private:
     const std::string pname_;
@@ -72,7 +74,8 @@ class RebootTask : public Task {
     RebootTask(const FlashingPlan* fp, const std::string& reboot_target);
     virtual RebootTask* AsRebootTask() override { return this; }
     void Run() override;
-    std::string ToString() override;
+    std::string ToString() const override;
+    std::string GetTarget() const { return reboot_target_; };
 
   private:
     const std::string reboot_target_ = "";
@@ -83,13 +86,15 @@ class OptimizedFlashSuperTask : public Task {
   public:
     OptimizedFlashSuperTask(const std::string& super_name, std::unique_ptr<SuperFlashHelper> helper,
                             SparsePtr sparse_layout, uint64_t super_size, const FlashingPlan* fp);
-    static std::unique_ptr<OptimizedFlashSuperTask> Initialize(const FlashingPlan* fp,
-                                                               std::vector<ImageEntry>& os_images);
-    static std::unique_ptr<OptimizedFlashSuperTask> InitializeFromTasks(
+
+    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() override;
+    std::string ToString() const override;
 
   private:
     const std::string super_name_;
@@ -105,7 +110,7 @@ class UpdateSuperTask : public Task {
     virtual UpdateSuperTask* AsUpdateSuperTask() override { return this; }
 
     void Run() override;
-    std::string ToString() override;
+    std::string ToString() const override;
 
   private:
     const FlashingPlan* fp_;
@@ -116,7 +121,8 @@ class ResizeTask : public Task {
     ResizeTask(const FlashingPlan* fp, const std::string& pname, const std::string& size,
                const std::string& slot);
     void Run() override;
-    std::string ToString() override;
+    std::string ToString() const override;
+    virtual ResizeTask* AsResizeTask() override { return this; }
 
   private:
     const FlashingPlan* fp_;
@@ -129,7 +135,7 @@ class DeleteTask : public Task {
   public:
     DeleteTask(const FlashingPlan* fp, const std::string& pname);
     void Run() override;
-    std::string ToString() override;
+    std::string ToString() const override;
 
   private:
     const FlashingPlan* fp_;
@@ -141,7 +147,7 @@ class WipeTask : public Task {
     WipeTask(const FlashingPlan* fp, const std::string& pname);
     virtual WipeTask* AsWipeTask() override { return this; }
     void Run() override;
-    std::string ToString() override;
+    std::string ToString() const override;
 
   private:
     const FlashingPlan* fp_;