From 997ce15d5e50e1e8ff5b8f7ac3a8d45f4519ea47 Mon Sep 17 00:00:00 2001
From: Daniel Zheng <zhengdaniel@google.com>
Date: Tue, 30 May 2023 09:24:46 -0700
Subject: [PATCH] Update add resizetasks to work in update

Old should_flash_in_userspace doesn't work if $ANDROID_PRODUCT_OUT is
not set. Also adding in a check to see if resizetasks were added
correctly

Test: fastboot update update.zip without $ANDROID_PRODUCT_OUT and
removing the flashsuperlayout code.
Bug: 283330320

Change-Id: Ib72f6a1cf07745daf70fffae3d1a6b8352e3f79c
---
 fastboot/fastboot.cpp | 20 +++++++++++++++-----
 fastboot/fastboot.h   |  2 +-
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index fcad0d915..3a135bc46 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -1646,14 +1646,22 @@ std::unique_ptr<Task> ParseFastbootInfoLine(const FlashingPlan* fp,
     return task;
 }
 
-void AddResizeTasks(const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>* tasks) {
+bool 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;
+    std::vector<char> contents;
+    if (!fp->source->ReadFile("super_empty.img", &contents)) {
+        return false;
+    }
+    auto metadata = android::fs_mgr::ReadFromImageBlob(contents.data(), contents.size());
+    if (!metadata) {
+        return false;
+    }
     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 (should_flash_in_userspace(*metadata.get(), flash_task->GetPartitionAndSlot())) {
                 if (!loc) {
                     loc = i;
                 }
@@ -1665,11 +1673,11 @@ void AddResizeTasks(const FlashingPlan* fp, std::vector<std::unique_ptr<Task>>*
     // if no logical partitions (although should never happen since system will always need to be
     // flashed)
     if (!loc) {
-        return;
+        return false;
     }
     tasks->insert(tasks->begin() + loc.value(), std::make_move_iterator(resize_tasks.begin()),
                   std::make_move_iterator(resize_tasks.end()));
-    return;
+    return true;
 }
 
 static bool IsIgnore(const std::vector<std::string>& command) {
@@ -1750,7 +1758,9 @@ std::vector<std::unique_ptr<Task>> ParseFastbootInfo(const FlashingPlan* fp,
         }
         tasks.insert(it, std::move(flash_super_task));
     } else {
-        AddResizeTasks(fp, &tasks);
+        if (!AddResizeTasks(fp, &tasks)) {
+            LOG(WARNING) << "Failed to add resize tasks";
+        };
     }
     return tasks;
 }
diff --git a/fastboot/fastboot.h b/fastboot/fastboot.h
index 0a1f1eb66..abdf63695 100644
--- a/fastboot/fastboot.h
+++ b/fastboot/fastboot.h
@@ -144,7 +144,7 @@ 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);
+bool 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);