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

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_;
+};