Merge "fastboot: Fix flashing both slots with dynamic partitions."
am: d016c75104

Change-Id: I139b7f506aee6554d134ba30f543b880334253e7
diff --git a/fastboot/device/flashing.cpp b/fastboot/device/flashing.cpp
index 66b90bf..7c9e1d0 100644
--- a/fastboot/device/flashing.cpp
+++ b/fastboot/device/flashing.cpp
@@ -148,13 +148,40 @@
     // image.
     std::string slot_suffix = device->GetCurrentSlot();
     uint32_t slot_number = SlotNumberForSlotSuffix(slot_suffix);
-    if (wipe || !ReadMetadata(super_name, slot_number)) {
+    std::unique_ptr<LpMetadata> old_metadata = ReadMetadata(super_name, slot_number);
+    if (wipe || !old_metadata) {
         if (!FlashPartitionTable(super_name, *new_metadata.get())) {
             return device->WriteFail("Unable to flash new partition table");
         }
         return device->WriteOkay("Successfully flashed partition table");
     }
 
+    std::set<std::string> partitions_to_keep;
+    for (const auto& partition : old_metadata->partitions) {
+        // Preserve partitions in the other slot, but not the current slot.
+        std::string partition_name = GetPartitionName(partition);
+        if (!slot_suffix.empty() && GetPartitionSlotSuffix(partition_name) == slot_suffix) {
+            continue;
+        }
+        partitions_to_keep.emplace(partition_name);
+    }
+
+    // Do not preserve the scratch partition.
+    partitions_to_keep.erase("scratch");
+
+    if (!partitions_to_keep.empty()) {
+        std::unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(*new_metadata.get());
+        if (!builder->ImportPartitions(*old_metadata.get(), partitions_to_keep)) {
+            return device->WriteFail(
+                    "Old partitions are not compatible with the new super layout; wipe needed");
+        }
+
+        new_metadata = builder->Export();
+        if (!new_metadata) {
+            return device->WriteFail("Unable to build new partition table; wipe needed");
+        }
+    }
+
     // Write the new table to every metadata slot.
     bool ok = true;
     for (size_t i = 0; i < new_metadata->geometry.metadata_slot_count; i++) {
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index 3c6b1b7..e358bec 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -1200,6 +1200,15 @@
 void FlashAllTool::Flash() {
     DumpInfo();
     CheckRequirements();
+
+    // Change the slot first, so we boot into the correct recovery image when
+    // using fastbootd.
+    if (slot_override_ == "all") {
+        set_active("a");
+    } else {
+        set_active(slot_override_);
+    }
+
     DetermineSecondarySlot();
     CollectImages();
 
@@ -1223,12 +1232,6 @@
 
     // Flash OS images, resizing logical partitions as needed.
     FlashImages(os_images_);
-
-    if (slot_override_ == "all") {
-        set_active("a");
-    } else {
-        set_active(slot_override_);
-    }
 }
 
 void FlashAllTool::CheckRequirements() {
@@ -1243,7 +1246,7 @@
     if (skip_secondary_) {
         return;
     }
-    if (slot_override_ != "") {
+    if (slot_override_ != "" && slot_override_ != "all") {
         secondary_slot_ = get_other_slot(slot_override_);
     } else {
         secondary_slot_ = get_other_slot();
diff --git a/fs_mgr/liblp/builder.cpp b/fs_mgr/liblp/builder.cpp
index 3cd33b1..4007ad9 100644
--- a/fs_mgr/liblp/builder.cpp
+++ b/fs_mgr/liblp/builder.cpp
@@ -184,22 +184,26 @@
         if (!builder) {
             return false;
         }
-
-        for (size_t i = 0; i < partition.num_extents; i++) {
-            const LpMetadataExtent& extent = metadata.extents[partition.first_extent_index + i];
-            if (extent.target_type == LP_TARGET_TYPE_LINEAR) {
-                auto copy = std::make_unique<LinearExtent>(extent.num_sectors, extent.target_source,
-                                                           extent.target_data);
-                builder->AddExtent(std::move(copy));
-            } else if (extent.target_type == LP_TARGET_TYPE_ZERO) {
-                auto copy = std::make_unique<ZeroExtent>(extent.num_sectors);
-                builder->AddExtent(std::move(copy));
-            }
-        }
+        ImportExtents(builder, metadata, partition);
     }
     return true;
 }
 
+void MetadataBuilder::ImportExtents(Partition* dest, const LpMetadata& metadata,
+                                    const LpMetadataPartition& source) {
+    for (size_t i = 0; i < source.num_extents; i++) {
+        const LpMetadataExtent& extent = metadata.extents[source.first_extent_index + i];
+        if (extent.target_type == LP_TARGET_TYPE_LINEAR) {
+            auto copy = std::make_unique<LinearExtent>(extent.num_sectors, extent.target_source,
+                                                       extent.target_data);
+            dest->AddExtent(std::move(copy));
+        } else if (extent.target_type == LP_TARGET_TYPE_ZERO) {
+            auto copy = std::make_unique<ZeroExtent>(extent.num_sectors);
+            dest->AddExtent(std::move(copy));
+        }
+    }
+}
+
 static bool VerifyDeviceProperties(const BlockDeviceInfo& device_info) {
     if (device_info.logical_block_size % LP_SECTOR_SIZE != 0) {
         LERROR << "Block device " << device_info.partition_name
@@ -471,13 +475,18 @@
     return free_regions;
 }
 
-bool MetadataBuilder::GrowPartition(Partition* partition, uint64_t aligned_size) {
+bool MetadataBuilder::ValidatePartitionSizeChange(Partition* partition, uint64_t old_size,
+                                                  uint64_t new_size) {
     PartitionGroup* group = FindGroup(partition->group_name());
     CHECK(group);
 
+    if (new_size <= old_size) {
+        return true;
+    }
+
     // Figure out how much we need to allocate, and whether our group has
     // enough space remaining.
-    uint64_t space_needed = aligned_size - partition->size();
+    uint64_t space_needed = new_size - old_size;
     if (group->maximum_size() > 0) {
         uint64_t group_size = TotalSizeOfGroup(group);
         if (group_size >= group->maximum_size() ||
@@ -488,7 +497,11 @@
             return false;
         }
     }
+    return true;
+}
 
+bool MetadataBuilder::GrowPartition(Partition* partition, uint64_t aligned_size) {
+    uint64_t space_needed = aligned_size - partition->size();
     uint64_t sectors_needed = space_needed / LP_SECTOR_SIZE;
     DCHECK(sectors_needed * LP_SECTOR_SIZE == space_needed);
 
@@ -704,6 +717,10 @@
     uint64_t aligned_size = AlignTo(requested_size, geometry_.logical_block_size);
     uint64_t old_size = partition->size();
 
+    if (!ValidatePartitionSizeChange(partition, old_size, aligned_size)) {
+        return false;
+    }
+
     if (aligned_size > old_size) {
         if (!GrowPartition(partition, aligned_size)) {
             return false;
@@ -750,5 +767,74 @@
     }
 }
 
+static bool CompareBlockDevices(const LpMetadataBlockDevice& first,
+                                const LpMetadataBlockDevice& second) {
+    // Note: we don't compare alignment, since it's a performance thing and
+    // won't affect whether old extents continue to work.
+    return first.first_logical_sector == second.first_logical_sector && first.size == second.size &&
+           GetBlockDevicePartitionName(first) == GetBlockDevicePartitionName(second);
+}
+
+bool MetadataBuilder::ImportPartitions(const LpMetadata& metadata,
+                                       const std::set<std::string>& partition_names) {
+    // The block device list must be identical. We do not try to be clever and
+    // allow ordering changes or changes that don't affect partitions. This
+    // process is designed to allow the most common flashing scenarios and more
+    // complex ones should require a wipe.
+    if (metadata.block_devices.size() != block_devices_.size()) {
+        LINFO << "Block device tables does not match.";
+        return false;
+    }
+    for (size_t i = 0; i < metadata.block_devices.size(); i++) {
+        const LpMetadataBlockDevice& old_device = metadata.block_devices[i];
+        const LpMetadataBlockDevice& new_device = block_devices_[i];
+        if (!CompareBlockDevices(old_device, new_device)) {
+            LINFO << "Block device tables do not match";
+            return false;
+        }
+    }
+
+    // Import named partitions. Note that we do not attempt to merge group
+    // information here. If the device changed its group names, the old
+    // partitions will fail to merge. The same could happen if the group
+    // allocation sizes change.
+    for (const auto& partition : metadata.partitions) {
+        std::string partition_name = GetPartitionName(partition);
+        if (partition_names.find(partition_name) == partition_names.end()) {
+            continue;
+        }
+        if (!ImportPartition(metadata, partition)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+bool MetadataBuilder::ImportPartition(const LpMetadata& metadata,
+                                      const LpMetadataPartition& source) {
+    std::string partition_name = GetPartitionName(source);
+    Partition* partition = FindPartition(partition_name);
+    if (!partition) {
+        std::string group_name = GetPartitionGroupName(metadata.groups[source.group_index]);
+        partition = AddPartition(partition_name, group_name, source.attributes);
+        if (!partition) {
+            return false;
+        }
+    }
+    if (partition->size() > 0) {
+        LINFO << "Importing partition table would overwrite non-empty partition: "
+              << partition_name;
+        return false;
+    }
+
+    ImportExtents(partition, metadata, source);
+
+    if (!ValidatePartitionSizeChange(partition, 0, partition->size())) {
+        partition->RemoveExtents();
+        return false;
+    }
+    return true;
+}
+
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/builder_test.cpp b/fs_mgr/liblp/builder_test.cpp
index c27e300..35cab38 100644
--- a/fs_mgr/liblp/builder_test.cpp
+++ b/fs_mgr/liblp/builder_test.cpp
@@ -632,3 +632,64 @@
     EXPECT_EQ(metadata->extents[2].target_data, 1472);
     EXPECT_EQ(metadata->extents[2].target_source, 2);
 }
+
+TEST(liblp, ImportPartitionsOk) {
+    unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
+    ASSERT_NE(builder, nullptr);
+
+    Partition* system = builder->AddPartition("system", LP_PARTITION_ATTR_READONLY);
+    Partition* vendor = builder->AddPartition("vendor", LP_PARTITION_ATTR_READONLY);
+    ASSERT_NE(system, nullptr);
+    ASSERT_NE(vendor, nullptr);
+    EXPECT_EQ(builder->ResizePartition(system, 65536), true);
+    EXPECT_EQ(builder->ResizePartition(vendor, 32768), true);
+    EXPECT_EQ(builder->ResizePartition(system, 98304), true);
+
+    unique_ptr<LpMetadata> exported = builder->Export();
+    ASSERT_NE(exported, nullptr);
+
+    builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
+    ASSERT_NE(builder, nullptr);
+
+    ASSERT_TRUE(builder->ImportPartitions(*exported.get(), {"vendor"}));
+    EXPECT_NE(builder->FindPartition("vendor"), nullptr);
+    EXPECT_EQ(builder->FindPartition("system"), nullptr);
+
+    unique_ptr<LpMetadata> new_metadata = builder->Export();
+    ASSERT_NE(new_metadata, nullptr);
+
+    ASSERT_EQ(exported->partitions.size(), static_cast<size_t>(2));
+    ASSERT_EQ(GetPartitionName(exported->partitions[1]), "vendor");
+    ASSERT_EQ(new_metadata->partitions.size(), static_cast<size_t>(1));
+    ASSERT_EQ(GetPartitionName(new_metadata->partitions[0]), "vendor");
+
+    const LpMetadataExtent& extent_a =
+            exported->extents[exported->partitions[1].first_extent_index];
+    const LpMetadataExtent& extent_b =
+            new_metadata->extents[new_metadata->partitions[0].first_extent_index];
+    EXPECT_EQ(extent_a.num_sectors, extent_b.num_sectors);
+    EXPECT_EQ(extent_a.target_type, extent_b.target_type);
+    EXPECT_EQ(extent_a.target_data, extent_b.target_data);
+    EXPECT_EQ(extent_a.target_source, extent_b.target_source);
+}
+
+TEST(liblp, ImportPartitionsFail) {
+    unique_ptr<MetadataBuilder> builder = MetadataBuilder::New(1024 * 1024, 1024, 2);
+    ASSERT_NE(builder, nullptr);
+
+    Partition* system = builder->AddPartition("system", LP_PARTITION_ATTR_READONLY);
+    Partition* vendor = builder->AddPartition("vendor", LP_PARTITION_ATTR_READONLY);
+    ASSERT_NE(system, nullptr);
+    ASSERT_NE(vendor, nullptr);
+    EXPECT_EQ(builder->ResizePartition(system, 65536), true);
+    EXPECT_EQ(builder->ResizePartition(vendor, 32768), true);
+    EXPECT_EQ(builder->ResizePartition(system, 98304), true);
+
+    unique_ptr<LpMetadata> exported = builder->Export();
+    ASSERT_NE(exported, nullptr);
+
+    // Different device size.
+    builder = MetadataBuilder::New(1024 * 2048, 1024, 2);
+    ASSERT_NE(builder, nullptr);
+    EXPECT_FALSE(builder->ImportPartitions(*exported.get(), {"system"}));
+}
diff --git a/fs_mgr/liblp/include/liblp/builder.h b/fs_mgr/liblp/include/liblp/builder.h
index f9de106..d5e3fed 100644
--- a/fs_mgr/liblp/include/liblp/builder.h
+++ b/fs_mgr/liblp/include/liblp/builder.h
@@ -22,6 +22,7 @@
 
 #include <map>
 #include <memory>
+#include <set>
 
 #include "liblp.h"
 #include "partition_opener.h"
@@ -225,6 +226,11 @@
     bool GetBlockDeviceInfo(const std::string& partition_name, BlockDeviceInfo* info) const;
     bool UpdateBlockDeviceInfo(const std::string& partition_name, const BlockDeviceInfo& info);
 
+    // Attempt to preserve the named partitions from an older metadata. If this
+    // is not possible (for example, the block device list has changed) then
+    // false is returned.
+    bool ImportPartitions(const LpMetadata& metadata, const std::set<std::string>& partition_names);
+
   private:
     MetadataBuilder();
     MetadataBuilder(const MetadataBuilder&) = delete;
@@ -240,6 +246,10 @@
     uint64_t TotalSizeOfGroup(PartitionGroup* group) const;
     bool UpdateBlockDeviceInfo(size_t index, const BlockDeviceInfo& info);
     bool FindBlockDeviceByName(const std::string& partition_name, uint32_t* index) const;
+    bool ValidatePartitionSizeChange(Partition* partition, uint64_t old_size, uint64_t new_size);
+    void ImportExtents(Partition* dest, const LpMetadata& metadata,
+                       const LpMetadataPartition& source);
+    bool ImportPartition(const LpMetadata& metadata, const LpMetadataPartition& source);
 
     struct Interval {
         uint32_t device_index;
diff --git a/fs_mgr/liblp/include/liblp/liblp.h b/fs_mgr/liblp/include/liblp/liblp.h
index a5dac00..8723a7f 100644
--- a/fs_mgr/liblp/include/liblp/liblp.h
+++ b/fs_mgr/liblp/include/liblp/liblp.h
@@ -93,8 +93,9 @@
 // Get the list of block device names required by the given metadata.
 std::vector<std::string> GetBlockDevicePartitionNames(const LpMetadata& metadata);
 
-// Helper to return a slot number for a slot suffix.
+// Slot suffix helpers.
 uint32_t SlotNumberForSlotSuffix(const std::string& suffix);
+std::string GetPartitionSlotSuffix(const std::string& partition_name);
 
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/utility.cpp b/fs_mgr/liblp/utility.cpp
index 199d994..2bf9a7d 100644
--- a/fs_mgr/liblp/utility.cpp
+++ b/fs_mgr/liblp/utility.cpp
@@ -124,5 +124,13 @@
     return list;
 }
 
+std::string GetPartitionSlotSuffix(const std::string& partition_name) {
+    if (partition_name.size() <= 2) {
+        return "";
+    }
+    std::string suffix = partition_name.substr(partition_name.size() - 2);
+    return (suffix == "_a" || suffix == "_b") ? suffix : "";
+}
+
 }  // namespace fs_mgr
 }  // namespace android
diff --git a/fs_mgr/liblp/utility_test.cpp b/fs_mgr/liblp/utility_test.cpp
index bdf6dfd..f1f958c 100644
--- a/fs_mgr/liblp/utility_test.cpp
+++ b/fs_mgr/liblp/utility_test.cpp
@@ -60,3 +60,11 @@
     EXPECT_EQ(AlignTo(32, 32, 30), 62);
     EXPECT_EQ(AlignTo(17, 32, 30), 30);
 }
+
+TEST(liblp, GetPartitionSlotSuffix) {
+    EXPECT_EQ(GetPartitionSlotSuffix("system"), "");
+    EXPECT_EQ(GetPartitionSlotSuffix("_"), "");
+    EXPECT_EQ(GetPartitionSlotSuffix("_a"), "");
+    EXPECT_EQ(GetPartitionSlotSuffix("system_a"), "_a");
+    EXPECT_EQ(GetPartitionSlotSuffix("system_b"), "_b");
+}