fastboot: Avoid reboots to userspace when using flashall/update.

Reboots to fastbootd (userspace fastboot) take a long time, particularly
due to the orange AVB screen and the likelihood of devices having uart
enabled. For "flashall", there is rarely a need to actually go into
userspace, because all of super is getting thrown away. We can just
flash super in the bootloader.

In the past we didn't do this because computing super.img is expensive -
both in terms of time (due to reading dependent images) and in terms of
space (it's easily over 5GB).

But we don't actually need to fully compute super.img. We can build a
sparse_file containing the metadata/headers, with additional references
to each image file containing partition data. Liblp provides the API to
do that, and here, we simply need to translate the layout to libsparse.

On Pixel, this reduces flashall time by around 35-50 seconds, or around
20% of total time, depending on whether uart is in use.

There are some caveats, in which case we'll fall back to normal
fastbootd. This does not work on non-A/B devices, on retrofit dynamic
partition devices (Pixel 3), and in some other edge-casey scenarios. If
it fails, -v will add logging information about why.

Bue: 266982466
Test: fastboot flashall on Pixel 5+
Change-Id: Ie040da597d739faa7f834202184cec8f8e412076

fastboot/Android.bp

@@ -286,6 +286,7 @@ cc_library_host_static {
         "fastboot.cpp",
         "fs.cpp",
         "socket.cpp",
+        "super_flash_helper.cpp",
         "tcp.cpp",
         "udp.cpp",
         "util.cpp",

fastboot/fastboot.cpp

@@ -63,6 +63,7 @@
 #include <build/version.h>
 #include <libavb/libavb.h>
 #include <liblp/liblp.h>
+#include <liblp/super_layout_builder.h>
 #include <platform_tools_version.h>
 #include <sparse/sparse.h>
 #include <ziparchive/zip_archive.h>
@@ -72,6 +73,7 @@
 #include "diagnose_usb.h"
 #include "fastboot_driver.h"
 #include "fs.h"
+#include "super_flash_helper.h"
 #include "tcp.h"
 #include "transport.h"
 #include "udp.h"
@@ -1405,20 +1407,30 @@ class FlashAllTool {
 
   private:
     void CheckRequirements();
-    void DetermineSecondarySlot();
+    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 UpdateSuperPartition();
+    bool OptimizedFlashSuper();
+
+    // If the image uses the default slot, or the user specified "all", then
+    // the paired string will be empty. If the image requests a specific slot
+    // (for example, system_other) it is specified instead.
+    using ImageEntry = std::pair<const Image*, std::string>;
+
+    std::string GetPartitionName(const ImageEntry& entry);
 
     const ImageSource& source_;
     std::string slot_override_;
     bool skip_secondary_;
     bool wipe_;
     bool force_flash_;
+    std::string current_slot_;
     std::string secondary_slot_;
-    std::vector<std::pair<const Image*, std::string>> boot_images_;
-    std::vector<std::pair<const Image*, std::string>> os_images_;
+
+    std::vector<ImageEntry> boot_images_;
+    std::vector<ImageEntry> os_images_;
 };
 
 FlashAllTool::FlashAllTool(const ImageSource& source, const std::string& slot_override,
@@ -1441,7 +1453,7 @@ void FlashAllTool::Flash() {
         set_active(slot_override_);
     }
 
-    DetermineSecondarySlot();
+    DetermineSlot();
     CollectImages();
 
     CancelSnapshotIfNeeded();
@@ -1450,24 +1462,92 @@ void FlashAllTool::Flash() {
     // or in bootloader fastboot.
     FlashImages(boot_images_);
 
-    // Sync the super partition. This will reboot to userspace fastboot if needed.
-    UpdateSuperPartition();
+    if (!OptimizedFlashSuper()) {
+        // Sync the super partition. This will reboot to userspace fastboot if needed.
+        UpdateSuperPartition();
 
-    // 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_) {
-        auto resize_partition = [](const std::string& partition) -> void {
-            if (is_logical(partition)) {
-                fb->ResizePartition(partition, "0");
-            }
-        };
-        do_for_partitions(image->part_name, slot, resize_partition, false);
+        // 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_) {
+            auto resize_partition = [](const std::string& partition) -> void {
+                if (is_logical(partition)) {
+                    fb->ResizePartition(partition, "0");
+                }
+            };
+            do_for_partitions(image->part_name, slot, resize_partition, false);
+        }
     }
 
     // Flash OS images, resizing logical partitions as needed.
     FlashImages(os_images_);
 }
 
+bool FlashAllTool::OptimizedFlashSuper() {
+    if (!supports_AB()) {
+        LOG(VERBOSE) << "Cannot optimize flashing super on non-AB device";
+        return false;
+    }
+    if (slot_override_ == "all") {
+        LOG(VERBOSE) << "Cannot optimize flashing super for all slots";
+        return false;
+    }
+
+    // Does this device use dynamic partitions at all?
+    unique_fd fd = source_.OpenFile("super_empty.img");
+    if (fd < 0) {
+        LOG(VERBOSE) << "could not open super_empty.img";
+        return false;
+    }
+
+    // Try to find whether there is a super partition.
+    std::string super_name;
+    if (fb->GetVar("super-partition-name", &super_name) != fastboot::SUCCESS) {
+        super_name = "super";
+    }
+    std::string partition_size_str;
+    if (fb->GetVar("partition-size:" + super_name, &partition_size_str) != fastboot::SUCCESS) {
+        LOG(VERBOSE) << "Cannot optimize super flashing: could not determine super partition";
+        return false;
+    }
+
+    SuperFlashHelper helper(source_);
+    if (!helper.Open(fd)) {
+        return false;
+    }
+
+    for (const auto& entry : os_images_) {
+        auto partition = GetPartitionName(entry);
+        auto image = entry.first;
+
+        if (!helper.AddPartition(partition, image->img_name, image->optional_if_no_image)) {
+            return false;
+        }
+    }
+
+    auto s = helper.GetSparseLayout();
+    if (!s) {
+        return false;
+    }
+
+    std::vector<SparsePtr> files;
+    if (int limit = get_sparse_limit(sparse_file_len(s.get(), false, false))) {
+        files = resparse_file(s.get(), limit);
+    } else {
+        files.emplace_back(std::move(s));
+    }
+
+    // Send the data to the device.
+    flash_partition_files(super_name, files);
+
+    // Remove images that we already flashed, just in case we have non-dynamic OS images.
+    auto remove_if_callback = [&, this](const ImageEntry& entry) -> bool {
+        return helper.WillFlash(GetPartitionName(entry));
+    };
+    os_images_.erase(std::remove_if(os_images_.begin(), os_images_.end(), remove_if_callback),
+                     os_images_.end());
+    return true;
+}
+
 void FlashAllTool::CheckRequirements() {
     std::vector<char> contents;
     if (!source_.ReadFile("android-info.txt", &contents)) {
@@ -1476,7 +1556,13 @@ void FlashAllTool::CheckRequirements() {
     ::CheckRequirements({contents.data(), contents.size()}, force_flash_);
 }
 
-void FlashAllTool::DetermineSecondarySlot() {
+void FlashAllTool::DetermineSlot() {
+    if (slot_override_.empty()) {
+        current_slot_ = get_current_slot();
+    } else {
+        current_slot_ = slot_override_;
+    }
+
     if (skip_secondary_) {
         return;
     }
@@ -1575,6 +1661,20 @@ void FlashAllTool::UpdateSuperPartition() {
     }
 }
 
+std::string FlashAllTool::GetPartitionName(const ImageEntry& entry) {
+    auto slot = entry.second;
+    if (slot.empty()) {
+        slot = current_slot_;
+    }
+    if (slot.empty()) {
+        return entry.first->part_name;
+    }
+    if (slot == "all") {
+        LOG(FATAL) << "Cannot retrieve a singular name when using all slots";
+    }
+    return entry.first->part_name + "_" + slot;
+}
+
 class ZipImageSource final : public ImageSource {
   public:
     explicit ZipImageSource(ZipArchiveHandle zip) : zip_(zip) {}

fastboot/super_flash_helper.cpp

@@ -0,0 +1,125 @@
+//
+// 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 "super_flash_helper.h"
+
+#include <android-base/logging.h>
+
+#include "util.h"
+
+using android::base::borrowed_fd;
+using android::base::unique_fd;
+using android::fs_mgr::SuperImageExtent;
+
+SuperFlashHelper::SuperFlashHelper(const ImageSource& source) : source_(source) {}
+
+bool SuperFlashHelper::Open(borrowed_fd fd) {
+    if (!builder_.Open(fd)) {
+        LOG(VERBOSE) << "device does not support optimized super flashing";
+        return false;
+    }
+
+    base_metadata_ = builder_.Export();
+    return !!base_metadata_;
+}
+
+bool SuperFlashHelper::IncludeInSuper(const std::string& partition) {
+    return should_flash_in_userspace(*base_metadata_.get(), partition);
+}
+
+bool SuperFlashHelper::AddPartition(const std::string& partition, const std::string& image_name,
+                                    bool optional) {
+    if (!IncludeInSuper(partition)) {
+        return true;
+    }
+    auto iter = image_fds_.find(image_name);
+    if (iter == image_fds_.end()) {
+        unique_fd fd = source_.OpenFile(image_name);
+        if (fd < 0) {
+            if (!optional) {
+                LOG(VERBOSE) << "could not find partition image: " << image_name;
+                return false;
+            }
+            return true;
+        }
+        if (is_sparse_file(fd)) {
+            LOG(VERBOSE) << "cannot optimize dynamic partitions with sparse images";
+            return false;
+        }
+        iter = image_fds_.emplace(image_name, std::move(fd)).first;
+    }
+
+    if (!builder_.AddPartition(partition, image_name, get_file_size(iter->second))) {
+        return false;
+    }
+
+    will_flash_.emplace(partition);
+    return true;
+}
+
+SparsePtr SuperFlashHelper::GetSparseLayout() {
+    // Cache extents since the sparse ptr depends on data pointers.
+    if (extents_.empty()) {
+        extents_ = builder_.GetImageLayout();
+        if (extents_.empty()) {
+            LOG(VERBOSE) << "device does not support optimized super flashing";
+            return {nullptr, nullptr};
+        }
+    }
+
+    unsigned int block_size = base_metadata_->geometry.logical_block_size;
+    int64_t flashed_size = extents_.back().offset + extents_.back().size;
+    SparsePtr s(sparse_file_new(block_size, flashed_size), sparse_file_destroy);
+
+    for (const auto& extent : extents_) {
+        if (extent.offset / block_size > UINT_MAX) {
+            // Super image is too big to send via sparse files (>8TB).
+            LOG(VERBOSE) << "super image is too big to flash";
+            return {nullptr, nullptr};
+        }
+        unsigned int block = extent.offset / block_size;
+
+        int rv = 0;
+        switch (extent.type) {
+            case SuperImageExtent::Type::DONTCARE:
+                break;
+            case SuperImageExtent::Type::ZERO:
+                rv = sparse_file_add_fill(s.get(), 0, extent.size, block);
+                break;
+            case SuperImageExtent::Type::DATA:
+                rv = sparse_file_add_data(s.get(), extent.blob->data(), extent.size, block);
+                break;
+            case SuperImageExtent::Type::PARTITION: {
+                auto iter = image_fds_.find(extent.image_name);
+                if (iter == image_fds_.end()) {
+                    LOG(FATAL) << "image added but not found: " << extent.image_name;
+                    return {nullptr, nullptr};
+                }
+                rv = sparse_file_add_fd(s.get(), iter->second.get(), extent.image_offset,
+                                        extent.size, block);
+                break;
+            }
+            default:
+                LOG(VERBOSE) << "unrecognized extent type in super image layout";
+                return {nullptr, nullptr};
+        }
+        if (rv) {
+            LOG(VERBOSE) << "sparse failure building super image layout";
+            return {nullptr, nullptr};
+        }
+    }
+    return s;
+}

fastboot/super_flash_helper.h

@@ -0,0 +1,56 @@
+//
+// 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 <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+#include <android-base/unique_fd.h>
+#include <liblp/liblp.h>
+#include <liblp/super_layout_builder.h>
+
+#include "util.h"
+
+class SuperFlashHelper final {
+  public:
+    explicit SuperFlashHelper(const ImageSource& source);
+
+    bool Open(android::base::borrowed_fd fd);
+    bool IncludeInSuper(const std::string& partition);
+    bool AddPartition(const std::string& partition, const std::string& image_name, bool optional);
+
+    // Note: the SparsePtr if non-null should not outlive SuperFlashHelper, since
+    // it depends on open fds and data pointers.
+    SparsePtr GetSparseLayout();
+
+    bool WillFlash(const std::string& partition) const {
+        return will_flash_.find(partition) != will_flash_.end();
+    }
+
+  private:
+    const ImageSource& source_;
+    android::fs_mgr::SuperLayoutBuilder builder_;
+    std::unique_ptr<android::fs_mgr::LpMetadata> base_metadata_;
+    std::vector<android::fs_mgr::SuperImageExtent> extents_;
+
+    // Cache open image fds. This keeps them alive while we flash the sparse
+    // file.
+    std::unordered_map<std::string, android::base::unique_fd> image_fds_;
+    std::unordered_set<std::string> will_flash_;
+};