Merge "task_profiles_test: Skip this test if cgroups is read-only"
diff --git a/fastboot/Android.bp b/fastboot/Android.bp
index 3b786e8..7794c4b 100644
--- a/fastboot/Android.bp
+++ b/fastboot/Android.bp
@@ -381,6 +381,7 @@
         "socket_mock.cpp",
         "socket_test.cpp",
         "super_flash_helper_test.cpp",
+        "task_test.cpp",
         "tcp_test.cpp",
         "udp_test.cpp",
     ],
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 2887402..cdcd036 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -115,19 +115,6 @@
 
 fastboot::FastBootDriver* fb = nullptr;
 
-enum fb_buffer_type {
-    FB_BUFFER_FD,
-    FB_BUFFER_SPARSE,
-};
-
-struct fastboot_buffer {
-    enum fb_buffer_type type;
-    std::vector<SparsePtr> files;
-    int64_t sz;
-    unique_fd fd;
-    int64_t image_size;
-};
-
 static std::vector<Image> images = {
         // clang-format off
     { "boot",     "boot.img",         "boot.sig",     "boot",     false, ImageType::BootCritical },
@@ -1284,7 +1271,7 @@
     }
 }
 
-static std::string get_current_slot() {
+std::string get_current_slot() {
     std::string current_slot;
     if (fb->GetVar("current-slot", &current_slot) != fastboot::SUCCESS) return "";
     if (current_slot[0] == '_') current_slot.erase(0, 1);
@@ -1560,7 +1547,7 @@
     }
 }
 
-std::string GetPartitionName(const ImageEntry& entry, std::string& current_slot) {
+std::string GetPartitionName(const ImageEntry& entry, const std::string& current_slot) {
     auto slot = entry.second;
     if (slot.empty()) {
         slot = current_slot;
@@ -1574,23 +1561,215 @@
     return entry.first->part_name + "_" + slot;
 }
 
-class FlashAllTool {
-  public:
-    FlashAllTool(FlashingPlan* fp);
+std::unique_ptr<FlashTask> ParseFlashCommand(const FlashingPlan* fp,
+                                             const std::vector<std::string>& parts) {
+    bool apply_vbmeta = false;
+    std::string slot = fp->slot_override;
+    std::string partition;
+    std::string img_name;
+    for (auto& part : parts) {
+        if (part == "--apply-vbmeta") {
+            apply_vbmeta = true;
+        } else if (part == "--slot-other") {
+            slot = fp->secondary_slot;
+        } else if (partition.empty()) {
+            partition = part;
+        } else if (img_name.empty()) {
+            img_name = part;
+        } else {
+            LOG(ERROR) << "unknown argument" << part
+                       << " in fastboot-info.txt. parts: " << android::base::Join(parts, " ");
+            return nullptr;
+        }
+    }
+    if (partition.empty()) {
+        LOG(ERROR) << "partition name not found when parsing fastboot-info.txt. parts: "
+                   << android::base::Join(parts, " ");
+        return nullptr;
+    }
+    if (img_name.empty()) {
+        img_name = partition + ".img";
+    }
+    return std::make_unique<FlashTask>(slot, partition, img_name, apply_vbmeta);
+}
 
-    void Flash();
+std::unique_ptr<RebootTask> ParseRebootCommand(const FlashingPlan* fp,
+                                               const std::vector<std::string>& parts) {
+    if (parts.empty()) return std::make_unique<RebootTask>(fp);
+    if (parts.size() > 1) {
+        LOG(ERROR) << "unknown arguments in reboot {target} in fastboot-info.txt: "
+                   << android::base::Join(parts, " ");
+        return nullptr;
+    }
+    return std::make_unique<RebootTask>(fp, parts[0]);
+}
 
-  private:
-    void CheckRequirements();
-    void DetermineSlot();
-    void CollectImages();
-    void FlashImages(const std::vector<std::pair<const Image*, std::string>>& images);
-    void FlashImage(const Image& image, const std::string& slot, fastboot_buffer* buf);
+std::unique_ptr<WipeTask> ParseWipeCommand(const FlashingPlan* fp,
+                                           const std::vector<std::string>& parts) {
+    if (parts.size() != 1) {
+        LOG(ERROR) << "unknown arguments in erase {partition} in fastboot-info.txt: "
+                   << android::base::Join(parts, " ");
+        return nullptr;
+    }
+    return std::make_unique<WipeTask>(fp, parts[0]);
+}
 
-    std::vector<ImageEntry> boot_images_;
-    std::vector<ImageEntry> os_images_;
-    FlashingPlan* fp_;
-};
+std::unique_ptr<Task> ParseFastbootInfoLine(const FlashingPlan* fp,
+                                            const std::vector<std::string>& command) {
+    if (command.size() == 0) {
+        return nullptr;
+    }
+    std::unique_ptr<Task> task;
+
+    if (command[0] == "flash") {
+        task = ParseFlashCommand(fp, std::vector<std::string>{command.begin() + 1, command.end()});
+    } else if (command[0] == "reboot") {
+        task = ParseRebootCommand(fp, std::vector<std::string>{command.begin() + 1, command.end()});
+    } else if (command[0] == "update-super" && command.size() == 1) {
+        task = std::make_unique<UpdateSuperTask>(fp);
+    } else if (command[0] == "erase" && command.size() == 2) {
+        task = ParseWipeCommand(fp, std::vector<std::string>{command.begin() + 1, command.end()});
+    }
+    if (!task) {
+        LOG(ERROR) << "unknown command parsing fastboot-info.txt line: "
+                   << android::base::Join(command, " ");
+    }
+    return task;
+}
+
+void AddResizeTasks(const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>* tasks) {
+    // expands "resize-partitions" into individual commands : resize {os_partition_1}, resize
+    // {os_partition_2}, etc.
+    std::vector<std::unique_ptr<Task>> resize_tasks;
+    std::optional<size_t> loc;
+    for (size_t i = 0; i < tasks->size(); i++) {
+        if (auto flash_task = tasks->at(i)->AsFlashTask()) {
+            if (should_flash_in_userspace(flash_task->GetPartitionAndSlot())) {
+                if (!loc) {
+                    loc = i;
+                }
+                resize_tasks.emplace_back(std::make_unique<ResizeTask>(
+                        fp, flash_task->GetPartition(), "0", fp->slot_override));
+            }
+        }
+    }
+    // if no logical partitions (although should never happen since system will always need to be
+    // flashed)
+    if (!loc) {
+        return;
+    }
+    tasks->insert(tasks->begin() + loc.value(), std::make_move_iterator(resize_tasks.begin()),
+                  std::make_move_iterator(resize_tasks.end()));
+    return;
+}
+
+static bool IsNumber(const std::string& s) {
+    bool period = false;
+    for (size_t i = 0; i < s.length(); i++) {
+        if (!isdigit(s[i])) {
+            if (!period && s[i] == '.' && i != 0 && i != s.length() - 1) {
+                period = true;
+            } else {
+                return false;
+            }
+        }
+    }
+    return true;
+}
+
+static bool IsIgnore(const std::vector<std::string>& command) {
+    if (command[0][0] == '#') {
+        return true;
+    }
+    return false;
+}
+
+bool CheckFastbootInfoRequirements(const std::vector<std::string>& command) {
+    if (command.size() != 2) {
+        LOG(ERROR) << "unknown characters in version info in fastboot-info.txt -> "
+                   << android::base::Join(command, " ");
+        return false;
+    }
+    if (command[0] != "version") {
+        LOG(ERROR) << "unknown characters in version info in fastboot-info.txt -> "
+                   << android::base::Join(command, " ");
+        return false;
+    }
+
+    if (!IsNumber(command[1])) {
+        LOG(ERROR) << "version number contains non-numeric values in fastboot-info.txt -> "
+                   << android::base::Join(command, " ");
+        return false;
+    }
+
+    LOG(VERBOSE) << "Checking 'fastboot-info.txt version'";
+    if (command[1] < PLATFORM_TOOLS_VERSION) {
+        return true;
+    }
+    LOG(ERROR) << "fasboot-info.txt version: " << command[1]
+               << " not compatible with host tool version --> " << PLATFORM_TOOLS_VERSION;
+    return false;
+}
+
+std::vector<std::unique_ptr<Task>> ParseFastbootInfo(const FlashingPlan* fp,
+                                                     const std::vector<std::string>& file) {
+    std::vector<std::unique_ptr<Task>> tasks;
+    // Get os_partitions that need to be resized
+    for (auto& text : file) {
+        std::vector<std::string> command = android::base::Tokenize(text, " ");
+        if (IsIgnore(command)) {
+            continue;
+        }
+        if (command.size() > 1 && command[0] == "version") {
+            if (!CheckFastbootInfoRequirements(command)) {
+                return {};
+            }
+            continue;
+        } else if (command.size() >= 2 && command[0] == "if-wipe") {
+            if (!fp->wants_wipe) {
+                continue;
+            }
+            command.erase(command.begin());
+        }
+        auto task = ParseFastbootInfoLine(fp, command);
+        if (!task) {
+            LOG(ERROR) << "Error when parsing fastboot-info.txt, falling back on Hardcoded list: "
+                       << text;
+            return {};
+        }
+        tasks.emplace_back(std::move(task));
+    }
+    if (auto flash_super_task = FlashSuperLayoutTask::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));
+    } else {
+        AddResizeTasks(fp, &tasks);
+    }
+    return tasks;
+}
+
+std::vector<std::unique_ptr<Task>> ParseFastbootInfo(const FlashingPlan* fp, std::ifstream& fs) {
+    if (!fs || fs.eof()) return {};
+
+    std::string text;
+    std::vector<std::string> file;
+    // Get os_partitions that need to be resized
+    while (std::getline(fs, text)) {
+        file.emplace_back(text);
+    }
+    return ParseFastbootInfo(fp, file);
+}
 
 FlashAllTool::FlashAllTool(FlashingPlan* fp) : fp_(fp) {}
 
@@ -1611,38 +1790,23 @@
 
     CancelSnapshotIfNeeded();
 
-    // First flash boot partitions. We allow this to happen either in userspace
-    // or in bootloader fastboot.
-    FlashImages(boot_images_);
-
-    std::vector<std::unique_ptr<Task>> tasks;
-
-    if (auto flash_super_task = FlashSuperLayoutTask::Initialize(fp_, os_images_)) {
-        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()) {
-                std::string partition_name = image->part_name + "_"s + slot;
-                if (image->IsSecondary() && is_logical(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));
-        }
+    std::string path = find_item_given_name("fastboot-info.txt");
+    std::ifstream stream(path);
+    std::vector<std::unique_ptr<Task>> tasks = ParseFastbootInfo(fp_, stream);
+    if (tasks.empty()) {
+        LOG(VERBOSE) << "Flashing from hardcoded images. fastboot-info.txt is empty or does not "
+                        "exist";
+        HardcodedFlash();
+        return;
     }
+    LOG(VERBOSE) << "Flashing from fastboot-info.txt";
     for (auto& task : tasks) {
         task->Run();
     }
-    FlashImages(os_images_);
+    if (fp_->wants_wipe) {
+        // avoid adding duplicate wipe tasks in fastboot main code.
+        fp_->wants_wipe = false;
+    }
 }
 
 void FlashAllTool::CheckRequirements() {
@@ -1693,6 +1857,42 @@
     }
 }
 
+void FlashAllTool::HardcodedFlash() {
+    CollectImages();
+    // First flash boot partitions. We allow this to happen either in userspace
+    // or in bootloader fastboot.
+    FlashImages(boot_images_);
+
+    std::vector<std::unique_ptr<Task>> tasks;
+
+    if (auto flash_super_task = FlashSuperLayoutTask::Initialize(fp_, os_images_)) {
+        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()) {
+                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));
+        }
+    }
+    for (auto& i : tasks) {
+        i->Run();
+    }
+    FlashImages(os_images_);
+}
+
 void FlashAllTool::FlashImages(const std::vector<std::pair<const Image*, std::string>>& images) {
     for (const auto& [image, slot] : images) {
         fastboot_buffer buf;
diff --git a/fastboot/fastboot.h b/fastboot/fastboot.h
index 6462a4f..b3dc67d 100644
--- a/fastboot/fastboot.h
+++ b/fastboot/fastboot.h
@@ -29,7 +29,10 @@
 
 #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"
 
 #include <bootimg.h>
@@ -47,6 +50,19 @@
     unsigned ParseFsOption(const char*);
 };
 
+enum fb_buffer_type {
+    FB_BUFFER_FD,
+    FB_BUFFER_SPARSE,
+};
+
+struct fastboot_buffer {
+    enum fb_buffer_type type;
+    std::vector<SparsePtr> files;
+    int64_t sz;
+    unique_fd fd;
+    int64_t image_size;
+};
+
 enum class ImageType {
     // Must be flashed for device to boot into the kernel.
     BootCritical,
@@ -83,7 +99,27 @@
     std::string slot_override;
     std::string current_slot;
     std::string secondary_slot;
-    fastboot::FastBootDriver* fb;
+
+    fastboot::IFastBootDriver* fb;
+};
+
+class FlashAllTool {
+  public:
+    FlashAllTool(FlashingPlan* fp);
+
+    void Flash();
+
+  private:
+    void CheckRequirements();
+    void DetermineSlot();
+    void CollectImages();
+    void FlashImages(const std::vector<std::pair<const Image*, std::string>>& images);
+    void FlashImage(const Image& image, const std::string& slot, fastboot_buffer* buf);
+    void HardcodedFlash();
+
+    std::vector<ImageEntry> boot_images_;
+    std::vector<ImageEntry> os_images_;
+    FlashingPlan* fp_;
 };
 
 bool should_flash_in_userspace(const std::string& partition_name);
@@ -94,6 +130,20 @@
 std::string find_item(const std::string& item);
 void reboot_to_userspace_fastboot();
 void syntax_error(const char* fmt, ...);
+std::string get_current_slot();
+
+// Code for Parsing fastboot-info.txt
+std::unique_ptr<FlashTask> ParseFlashCommand(const FlashingPlan* fp,
+                                             const std::vector<std::string>& parts);
+std::unique_ptr<RebootTask> ParseRebootCommand(const FlashingPlan* fp,
+                                               const std::vector<std::string>& parts);
+std::unique_ptr<WipeTask> ParseWipeCommand(const FlashingPlan* fp,
+                                           const std::vector<std::string>& parts);
+std::unique_ptr<Task> ParseFastbootInfoLine(const FlashingPlan* fp,
+                                            const std::vector<std::string>& command);
+void AddResizeTasks(const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>& tasks);
+std::vector<std::unique_ptr<Task>> ParseFastbootInfo(const FlashingPlan* fp,
+                                                     const std::vector<std::string>& file);
 
 struct NetworkSerial {
     Socket::Protocol protocol;
@@ -103,7 +153,7 @@
 
 Result<NetworkSerial, FastbootError> ParseNetworkSerial(const std::string& serial);
 bool supports_AB();
-std::string GetPartitionName(const ImageEntry& entry, std::string& current_slot_);
+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);
 std::vector<SparsePtr> resparse_file(sparse_file* s, int64_t max_size);
diff --git a/fastboot/fastboot_driver.h b/fastboot/fastboot_driver.h
index 64e3c2b..3d6c7b0 100644
--- a/fastboot/fastboot_driver.h
+++ b/fastboot/fastboot_driver.h
@@ -40,21 +40,13 @@
 #include <sparse/sparse.h>
 
 #include "constants.h"
+#include "fastboot_driver_interface.h"
 #include "transport.h"
 
 class Transport;
 
 namespace fastboot {
 
-enum RetCode : int {
-    SUCCESS = 0,
-    BAD_ARG,
-    IO_ERROR,
-    BAD_DEV_RESP,
-    DEVICE_FAIL,
-    TIMEOUT,
-};
-
 struct DriverCallbacks {
     std::function<void(const std::string&)> prolog = [](const std::string&) {};
     std::function<void(int)> epilog = [](int) {};
@@ -62,7 +54,7 @@
     std::function<void(const std::string&)> text = [](const std::string&) {};
 };
 
-class FastBootDriver {
+class FastBootDriver : public IFastBootDriver {
     friend class FastBootTest;
 
   public:
@@ -77,9 +69,10 @@
     RetCode Boot(std::string* response = nullptr, std::vector<std::string>* info = nullptr);
     RetCode Continue(std::string* response = nullptr, std::vector<std::string>* info = nullptr);
     RetCode CreatePartition(const std::string& partition, const std::string& size);
-    RetCode DeletePartition(const std::string& partition);
+    RetCode DeletePartition(const std::string& partition) override;
     RetCode Download(const std::string& name, android::base::borrowed_fd fd, size_t size,
-                     std::string* response = nullptr, std::vector<std::string>* info = nullptr);
+                     std::string* response = nullptr,
+                     std::vector<std::string>* info = nullptr) override;
     RetCode Download(android::base::borrowed_fd fd, size_t size, std::string* response = nullptr,
                      std::vector<std::string>* info = nullptr);
     RetCode Download(const std::string& name, const std::vector<char>& buf,
@@ -92,16 +85,17 @@
     RetCode Download(sparse_file* s, bool use_crc = false, std::string* response = nullptr,
                      std::vector<std::string>* info = nullptr);
     RetCode Erase(const std::string& partition, std::string* response = nullptr,
-                  std::vector<std::string>* info = nullptr);
+                  std::vector<std::string>* info = nullptr) override;
     RetCode Flash(const std::string& partition, std::string* response = nullptr,
                   std::vector<std::string>* info = nullptr);
     RetCode GetVar(const std::string& key, std::string* val,
-                   std::vector<std::string>* info = nullptr);
+                   std::vector<std::string>* info = nullptr) override;
     RetCode GetVarAll(std::vector<std::string>* response);
-    RetCode Reboot(std::string* response = nullptr, std::vector<std::string>* info = nullptr);
+    RetCode Reboot(std::string* response = nullptr,
+                   std::vector<std::string>* info = nullptr) override;
     RetCode RebootTo(std::string target, std::string* response = nullptr,
-                     std::vector<std::string>* info = nullptr);
-    RetCode ResizePartition(const std::string& partition, const std::string& size);
+                     std::vector<std::string>* info = nullptr) override;
+    RetCode ResizePartition(const std::string& partition, const std::string& size) override;
     RetCode SetActive(const std::string& slot, std::string* response = nullptr,
                       std::vector<std::string>* info = nullptr);
     RetCode Upload(const std::string& outfile, std::string* response = nullptr,
@@ -115,7 +109,7 @@
     /* HIGHER LEVEL COMMANDS -- Composed of the commands above */
     RetCode FlashPartition(const std::string& partition, const std::vector<char>& data);
     RetCode FlashPartition(const std::string& partition, android::base::borrowed_fd fd,
-                           uint32_t sz);
+                           uint32_t sz) override;
     RetCode FlashPartition(const std::string& partition, sparse_file* s, uint32_t sz,
                            size_t current, size_t total);
 
@@ -127,7 +121,7 @@
     void SetInfoCallback(std::function<void(const std::string&)> info);
     static const std::string RCString(RetCode rc);
     std::string Error();
-    RetCode WaitForDisconnect();
+    RetCode WaitForDisconnect() override;
 
     // Note: set_transport will return the previous transport.
     Transport* set_transport(Transport* transport);
diff --git a/fastboot/fastboot_driver_interface.h b/fastboot/fastboot_driver_interface.h
new file mode 100644
index 0000000..795938f
--- /dev/null
+++ b/fastboot/fastboot_driver_interface.h
@@ -0,0 +1,59 @@
+//
+// Copyright (C) 2023 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+#pragma once
+
+#include <string>
+
+#include "android-base/unique_fd.h"
+
+class Transport;
+
+namespace fastboot {
+
+enum RetCode : int {
+    SUCCESS = 0,
+    BAD_ARG,
+    IO_ERROR,
+    BAD_DEV_RESP,
+    DEVICE_FAIL,
+    TIMEOUT,
+};
+
+class IFastBootDriver {
+  public:
+    RetCode virtual FlashPartition(const std::string& partition, android::base::borrowed_fd fd,
+                                   uint32_t sz) = 0;
+    RetCode virtual DeletePartition(const std::string& partition) = 0;
+    RetCode virtual WaitForDisconnect() = 0;
+    RetCode virtual Reboot(std::string* response = nullptr,
+                           std::vector<std::string>* info = nullptr) = 0;
+
+    RetCode virtual RebootTo(std::string target, std::string* response = nullptr,
+                             std::vector<std::string>* info = nullptr) = 0;
+    RetCode virtual GetVar(const std::string& key, std::string* val,
+                           std::vector<std::string>* info = nullptr) = 0;
+    RetCode virtual Download(const std::string& name, android::base::borrowed_fd fd, size_t size,
+                             std::string* response = nullptr,
+                             std::vector<std::string>* info = nullptr) = 0;
+    RetCode virtual RawCommand(const std::string& cmd, const std::string& message,
+                               std::string* response = nullptr,
+                               std::vector<std::string>* info = nullptr, int* dsize = nullptr) = 0;
+    RetCode virtual ResizePartition(const std::string& partition, const std::string& size) = 0;
+    RetCode virtual Erase(const std::string& partition, std::string* response = nullptr,
+                          std::vector<std::string>* info = nullptr) = 0;
+    virtual ~IFastBootDriver() = default;
+};
+}  // namespace fastboot
\ No newline at end of file
diff --git a/fastboot/fastboot_driver_mock.h b/fastboot/fastboot_driver_mock.h
new file mode 100644
index 0000000..62a5708
--- /dev/null
+++ b/fastboot/fastboot_driver_mock.h
@@ -0,0 +1,51 @@
+//
+// Copyright (C) 2023 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+#pragma once
+
+#include <gmock/gmock.h>
+#include "fastboot_driver_interface.h"
+
+namespace fastboot {
+
+class MockFastbootDriver : public IFastBootDriver {
+  public:
+    MOCK_METHOD(RetCode, FlashPartition,
+                (const std::string& partition, android::base::borrowed_fd fd, uint32_t sz),
+                (override));
+    MOCK_METHOD(RetCode, DeletePartition, (const std::string&), (override));
+    MOCK_METHOD(RetCode, Reboot, (std::string*, std::vector<std::string>*), (override));
+    MOCK_METHOD(RetCode, RebootTo, (std::string, std::string*, std::vector<std::string>*),
+                (override));
+
+    MOCK_METHOD(RetCode, GetVar, (const std::string&, std::string*, std::vector<std::string>*),
+                (override));
+
+    MOCK_METHOD(RetCode, Download,
+                (const std::string&, android::base::borrowed_fd, size_t, std::string*,
+                 std::vector<std::string>*),
+                (override));
+
+    MOCK_METHOD(RetCode, RawCommand,
+                (const std::string&, const std::string&, std::string*, std::vector<std::string>*,
+                 int*),
+                (override));
+    MOCK_METHOD(RetCode, ResizePartition, (const std::string&, const std::string&), (override));
+    MOCK_METHOD(RetCode, Erase, (const std::string&, std::string*, std::vector<std::string>*),
+                (override));
+    MOCK_METHOD(RetCode, WaitForDisconnect, (), (override));
+};
+
+}  // namespace fastboot
\ No newline at end of file
diff --git a/fastboot/task.cpp b/fastboot/task.cpp
index d43fd80..054c1ed 100644
--- a/fastboot/task.cpp
+++ b/fastboot/task.cpp
@@ -23,6 +23,7 @@
 #include "fastboot.h"
 #include "filesystem.h"
 #include "super_flash_helper.h"
+#include "util.h"
 
 using namespace std::string_literals;
 FlashTask::FlashTask(const std::string& _slot, const std::string& _pname, const std::string& _fname,
@@ -45,6 +46,20 @@
     do_for_partitions(pname_, slot_, flash, true);
 }
 
+std::string FlashTask::GetPartitionAndSlot() {
+    auto slot = slot_;
+    if (slot.empty()) {
+        slot = get_current_slot();
+    }
+    if (slot.empty()) {
+        return pname_;
+    }
+    if (slot == "all") {
+        LOG(FATAL) << "Cannot retrieve a singular name when using all slots";
+    }
+    return pname_ + "_" + slot;
+}
+
 RebootTask::RebootTask(const FlashingPlan* fp) : fp_(fp){};
 RebootTask::RebootTask(const FlashingPlan* fp, const std::string& reboot_target)
     : reboot_target_(reboot_target), fp_(fp){};
@@ -93,7 +108,7 @@
 }
 
 std::unique_ptr<FlashSuperLayoutTask> FlashSuperLayoutTask::Initialize(
-        FlashingPlan* fp, std::vector<ImageEntry>& os_images) {
+        const FlashingPlan* fp, std::vector<ImageEntry>& os_images) {
     if (!supports_AB()) {
         LOG(VERBOSE) << "Cannot optimize flashing super on non-AB device";
         return nullptr;
@@ -156,6 +171,77 @@
                                                   partition_size);
 }
 
+std::unique_ptr<FlashSuperLayoutTask> FlashSuperLayoutTask::InitializeFromTasks(
+        const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>& tasks) {
+    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& task : tasks) {
+        if (auto flash_task = task->AsFlashTask()) {
+            if (should_flash_in_userspace(flash_task->GetPartitionAndSlot())) {
+                auto partition = flash_task->GetPartitionAndSlot();
+                if (!helper->AddPartition(partition, flash_task->GetImageName(), false)) {
+                    return nullptr;
+                }
+            }
+        }
+    }
+
+    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 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;
+        }
+        return false;
+    };
+    tasks.erase(std::remove_if(tasks.begin(), tasks.end(), remove_if_callback), tasks.end());
+
+    return std::make_unique<FlashSuperLayoutTask>(super_name, std::move(helper), std::move(s),
+                                                  partition_size);
+}
+
 UpdateSuperTask::UpdateSuperTask(const FlashingPlan* fp) : fp_(fp) {}
 
 void UpdateSuperTask::Run() {
diff --git a/fastboot/task.h b/fastboot/task.h
index 801a0f6..34e3e92 100644
--- a/fastboot/task.h
+++ b/fastboot/task.h
@@ -26,10 +26,20 @@
 struct Image;
 using ImageEntry = std::pair<const Image*, std::string>;
 
+class FlashTask;
+class RebootTask;
+class UpdateSuperTask;
+class WipeTask;
+
 class Task {
   public:
     Task() = default;
     virtual void Run() = 0;
+    virtual FlashTask* AsFlashTask() { return nullptr; }
+    virtual RebootTask* AsRebootTask() { return nullptr; }
+    virtual UpdateSuperTask* AsUpdateSuperTask() { return nullptr; }
+    virtual WipeTask* AsWipeTask() { return nullptr; }
+
     virtual ~Task() = default;
 };
 
@@ -37,7 +47,12 @@
   public:
     FlashTask(const std::string& slot, const std::string& pname, const std::string& fname,
               const bool apply_vbmeta);
+    virtual FlashTask* AsFlashTask() override { return this; }
 
+    std::string GetPartition() { return pname_; }
+    std::string GetImageName() { return fname_; }
+    std::string GetPartitionAndSlot();
+    std::string GetSlot() { return slot_; }
     void Run() override;
 
   private:
@@ -51,6 +66,7 @@
   public:
     RebootTask(const FlashingPlan* fp);
     RebootTask(const FlashingPlan* fp, const std::string& reboot_target);
+    virtual RebootTask* AsRebootTask() override { return this; }
     void Run() override;
 
   private:
@@ -62,8 +78,10 @@
   public:
     FlashSuperLayoutTask(const std::string& super_name, std::unique_ptr<SuperFlashHelper> helper,
                          SparsePtr sparse_layout, uint64_t super_size);
-    static std::unique_ptr<FlashSuperLayoutTask> Initialize(FlashingPlan* fp,
+    static std::unique_ptr<FlashSuperLayoutTask> Initialize(const FlashingPlan* fp,
                                                             std::vector<ImageEntry>& os_images);
+    static std::unique_ptr<FlashSuperLayoutTask> InitializeFromTasks(
+            const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>& tasks);
     using ImageEntry = std::pair<const Image*, std::string>;
     void Run() override;
 
@@ -77,6 +95,8 @@
 class UpdateSuperTask : public Task {
   public:
     UpdateSuperTask(const FlashingPlan* fp);
+    virtual UpdateSuperTask* AsUpdateSuperTask() override { return this; }
+
     void Run() override;
 
   private:
@@ -109,6 +129,8 @@
 class WipeTask : public Task {
   public:
     WipeTask(const FlashingPlan* fp, const std::string& pname);
+    virtual WipeTask* AsWipeTask() override { return this; }
+
     void Run() override;
 
   private:
diff --git a/fastboot/task_test.cpp b/fastboot/task_test.cpp
new file mode 100644
index 0000000..400e27f
--- /dev/null
+++ b/fastboot/task_test.cpp
@@ -0,0 +1,82 @@
+//
+// Copyright (C) 2023 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include "task.h"
+#include "fastboot.h"
+
+#include <gtest/gtest.h>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <unordered_map>
+#include "android-base/strings.h"
+using android::base::Split;
+
+class ParseTest : public ::testing ::Test {
+  protected:
+    void SetUp() override {
+        fp = std::make_unique<FlashingPlan>();
+        fp->slot_override = "b";
+        fp->secondary_slot = "a";
+        fp->wants_wipe = false;
+    }
+    void TearDown() override {}
+
+    std::unique_ptr<FlashingPlan> fp;
+
+  private:
+};
+
+static std::vector<std::unique_ptr<Task>> collectTasks(FlashingPlan* fp,
+                                                       const std::vector<std::string>& commands) {
+    std::vector<std::vector<std::string>> vec_commands;
+    for (auto& command : commands) {
+        vec_commands.emplace_back(android::base::Split(command, " "));
+    }
+    std::vector<std::unique_ptr<Task>> tasks;
+    for (auto& command : vec_commands) {
+        tasks.emplace_back(ParseFastbootInfoLine(fp, command));
+    }
+    return tasks;
+}
+
+TEST_F(ParseTest, CORRECT_FlASH_TASK_FORMED) {
+    std::vector<std::string> commands = {"flash dtbo", "flash --slot-other system system_other.img",
+                                         "flash system", "flash --apply-vbmeta vbmeta"};
+
+    std::vector<std::unique_ptr<Task>> tasks = collectTasks(fp.get(), commands);
+
+    std::vector<std::vector<std::string>> expected_values{
+            {"dtbo", "dtbo_b", "b", "dtbo.img"},
+            {"system", "system_a", "a", "system_other.img"},
+            {"system", "system_b", "b", "system.img"},
+            {"vbmeta", "vbmeta_b", "b", "vbmeta.img"}
+
+    };
+
+    for (auto& task : tasks) {
+        ASSERT_TRUE(task != nullptr);
+    }
+
+    for (size_t i = 0; i < tasks.size(); i++) {
+        auto task = tasks[i]->AsFlashTask();
+        ASSERT_TRUE(task != nullptr);
+        ASSERT_EQ(task->GetPartition(), expected_values[i][0]);
+        ASSERT_EQ(task->GetPartitionAndSlot(), expected_values[i][1]);
+        ASSERT_EQ(task->GetSlot(), expected_values[i][2]);
+        ASSERT_EQ(task->GetImageName(), expected_values[i][3]);
+    }
+}