[paver] support ZBI A/B/R partitions on EFI

Test: manual (zedboot -> zedboot -> pave -> observe lsblk)
Test: manual (install-disk-image wipe -> observe lsblk)
Test: manual (zedboot -> zedboot -> pave -> observe lsblk)
Bug: ZX-2220 #comment adding paver support
Change-Id: I0a34407cc741cdababa7b841a8759bb691dacd8a
diff --git a/system/uapp/disk-pave/device-partitioner.cpp b/system/uapp/disk-pave/device-partitioner.cpp
index 83a07c4..5b74ed8 100644
--- a/system/uapp/disk-pave/device-partitioner.cpp
+++ b/system/uapp/disk-pave/device-partitioner.cpp
@@ -29,8 +29,14 @@
 
 namespace {
 
-bool KernelFilterCallback(const gpt_partition_t& part, fbl::StringPiece partition_name) {
-    const uint8_t kern_type[GPT_GUID_LEN] = GUID_CROS_KERNEL_VALUE;
+constexpr char kEfiName[] = "EFI Gigaboot";
+constexpr char kGptDriverName[] = "/boot/driver/gpt.so";
+constexpr char kFvmPartitionName[] = "fvm";
+constexpr char kZirconAName[] = "ZIRCON-A";
+constexpr char kZirconBName[] = "ZIRCON-B";
+constexpr char kZirconRName[] = "ZIRCON-R";
+
+bool KernelFilterCallback(const gpt_partition_t& part, const uint8_t kern_type[GPT_GUID_LEN], fbl::StringPiece partition_name) {
     char cstring_name[GPT_NAME_LEN];
     utf16_to_cstring(cstring_name, reinterpret_cast<const uint16_t*>(part.name), GPT_NAME_LEN);
     return memcmp(part.type, kern_type, GPT_GUID_LEN) == 0 &&
@@ -42,13 +48,45 @@
     return memcmp(part.type, partition_type, GPT_GUID_LEN) == 0;
 }
 
+bool IsGigabootPartition(const gpt_partition_t& part) {
+    const uint8_t efi_type[GPT_GUID_LEN] = GUID_EFI_VALUE;
+    char cstring_name[GPT_NAME_LEN];
+    utf16_to_cstring(cstring_name, reinterpret_cast<const uint16_t*>(part.name), GPT_NAME_LEN);
+    // Disk-paved EFI: Identified by "EFI Gigaboot" label.
+    const bool gigaboot_efi = strncmp(cstring_name, kEfiName, strlen(kEfiName)) == 0;
+    return memcmp(part.type, efi_type, GPT_GUID_LEN) == 0 && gigaboot_efi;
+}
+
+bool WipeFilterCallback(const gpt_partition_t& part) {
+    const uint8_t efi_removable_types[][GPT_GUID_LEN] = {
+        GUID_FVM_VALUE,
+        GUID_INSTALL_VALUE,
+        GUID_SYSTEM_VALUE,
+        GUID_BLOB_VALUE,
+        GUID_DATA_VALUE,
+        GUID_ZIRCON_A_VALUE,
+        GUID_ZIRCON_B_VALUE,
+        GUID_ZIRCON_R_VALUE,
+    };
+
+    if (IsGigabootPartition(part)) {
+        return true;
+    }
+
+    for (const auto &type : efi_removable_types) {
+        if (memcmp(part.type, type, GPT_GUID_LEN) == 0) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 constexpr size_t ReservedHeaderBlocks(size_t blk_size) {
     constexpr size_t kReservedEntryBlocks = (16 * 1024);
     return (kReservedEntryBlocks + 2 * blk_size) / blk_size;
 };
 
-constexpr char kGptDriverName[] = "/boot/driver/gpt.so";
-constexpr char kFvmPartitionName[] = "fvm";
 
 // Helper function to auto-deduce type.
 template <typename T>
@@ -561,12 +599,6 @@
     return ZX_OK;
 }
 
-// Name used by previous Fuchsia Installer.
-constexpr char kOldEfiName[] = "EFI";
-
-// Name used for EFI partitions added by paver.
-constexpr char kEfiName[] = "EFI Gigaboot";
-
 zx_status_t EfiDevicePartitioner::AddPartition(Partition partition_type, fbl::unique_fd* out_fd) {
     const char* name;
     uint8_t type[GPT_GUID_LEN];
@@ -577,10 +609,31 @@
     case Partition::kEfi: {
         const uint8_t efi_type[GPT_GUID_LEN] = GUID_EFI_VALUE;
         memcpy(type, efi_type, GPT_GUID_LEN);
-        minimum_size_bytes = 1LU * (1 << 30);
+        minimum_size_bytes = 20LU * (1 << 20);
         name = kEfiName;
         break;
     }
+    case Partition::kZirconA: {
+        const uint8_t zircon_a_type[GPT_GUID_LEN] = GUID_ZIRCON_A_VALUE;
+        memcpy(type, zircon_a_type, GPT_GUID_LEN);
+        minimum_size_bytes = 16LU * (1 << 20);
+        name = kZirconAName;
+        break;
+    }
+    case Partition::kZirconB: {
+        const uint8_t zircon_b_type[GPT_GUID_LEN] = GUID_ZIRCON_B_VALUE;
+        memcpy(type, zircon_b_type, GPT_GUID_LEN);
+        minimum_size_bytes = 16LU * (1 << 20);
+        name = kZirconBName;
+        break;
+    }
+    case Partition::kZirconR: {
+        const uint8_t zircon_r_type[GPT_GUID_LEN] = GUID_ZIRCON_R_VALUE;
+        memcpy(type, zircon_r_type, GPT_GUID_LEN);
+        minimum_size_bytes = 24LU * (1 << 20);
+        name = kZirconRName;
+        break;
+    }
     case Partition::kFuchsiaVolumeManager: {
         const uint8_t fvm_type[GPT_GUID_LEN] = GUID_FVM_VALUE;
         memcpy(type, fvm_type, GPT_GUID_LEN);
@@ -597,34 +650,30 @@
                               optional_reserve_bytes, out_fd);
 }
 
-bool EfiDevicePartitioner::FilterZirconPartition(const block_info_t& info,
-                                                 const gpt_partition_t& part) {
-    const uint8_t efi_type[GPT_GUID_LEN] = GUID_EFI_VALUE;
-    char cstring_name[GPT_NAME_LEN];
-    utf16_to_cstring(cstring_name, reinterpret_cast<const uint16_t*>(part.name), GPT_NAME_LEN);
-    // Old EFI: Installed by the legacy Fuchsia installer, identified by
-    // large size and "EFI" label.
-    constexpr unsigned int k512MB = (1LU << 29);
-    const bool old_efi = strncmp(cstring_name, kOldEfiName, strlen(kOldEfiName)) == 0 &&
-                         ((part.last - part.first + 1) * info.block_size) > k512MB;
-    // Disk-paved EFI: Identified by "EFI Gigaboot" label.
-    const bool new_efi = strncmp(cstring_name, kEfiName, strlen(kEfiName)) == 0;
-    return memcmp(part.type, efi_type, GPT_GUID_LEN) == 0 && (old_efi || new_efi);
-}
-
 zx_status_t EfiDevicePartitioner::FindPartition(Partition partition_type,
                                                 fbl::unique_fd* out_fd) const {
-    block_info_t info;
-    zx_status_t status;
-    if ((status = gpt_->GetBlockInfo(&info)) != ZX_OK) {
-        ERROR("Unable to get block info\n");
-        return ZX_ERR_IO;
-    }
-
     switch (partition_type) {
     case Partition::kEfi: {
-        const auto filter = [&info](const gpt_partition_t& part) {
-            return FilterZirconPartition(info, part);
+        return gpt_->FindPartition(IsGigabootPartition, out_fd);
+    }
+    case Partition::kZirconA: {
+        const auto filter = [](const gpt_partition_t& part) {
+            const uint8_t guid[GPT_GUID_LEN] = GUID_ZIRCON_A_VALUE;
+            return KernelFilterCallback(part, guid, kZirconAName);
+        };
+        return gpt_->FindPartition(filter, out_fd);
+    }
+    case Partition::kZirconB: {
+        const auto filter = [](const gpt_partition_t& part) {
+            const uint8_t guid[GPT_GUID_LEN] = GUID_ZIRCON_B_VALUE;
+            return KernelFilterCallback(part, guid, kZirconBName);
+        };
+        return gpt_->FindPartition(filter, out_fd);
+    }
+    case Partition::kZirconR: {
+        const auto filter = [](const gpt_partition_t& part) {
+            const uint8_t guid[GPT_GUID_LEN] = GUID_ZIRCON_R_VALUE;
+            return KernelFilterCallback(part, guid, kZirconRName);
         };
         return gpt_->FindPartition(filter, out_fd);
     }
@@ -637,67 +686,8 @@
     }
 }
 
-zx_status_t EfiDevicePartitioner::WipePartitions(const fbl::Vector<Partition>& partitions) {
-    const uint8_t fvm_type[GPT_GUID_LEN] = GUID_FVM_VALUE;
-    const uint8_t install_type[GPT_GUID_LEN] = GUID_INSTALL_VALUE;
-    const uint8_t system_type[GPT_GUID_LEN] = GUID_SYSTEM_VALUE;
-    const uint8_t blob_type[GPT_GUID_LEN] = GUID_BLOB_VALUE;
-    const uint8_t data_type[GPT_GUID_LEN] = GUID_DATA_VALUE;
-
-    block_info_t info;
-    zx_status_t status;
-    if ((status = gpt_->GetBlockInfo(&info)) != ZX_OK) {
-        ERROR("Unable to get block info\n");
-        return ZX_ERR_IO;
-    }
-
-    fbl::Vector<const uint8_t*> partition_list;
-    bool efi = false;
-    for (const Partition& partition_type : partitions) {
-        switch (partition_type) {
-        case Partition::kEfi: {
-            // Special case.
-            efi = true;
-            break;
-        }
-        case Partition::kKernelC:
-            break;
-        case Partition::kFuchsiaVolumeManager:
-            partition_list.push_back(fvm_type);
-            break;
-        case Partition::kInstallType:
-            partition_list.push_back(install_type);
-            break;
-        case Partition::kSystem:
-            partition_list.push_back(system_type);
-            break;
-        case Partition::kBlob:
-            partition_list.push_back(blob_type);
-            break;
-        case Partition::kData:
-            partition_list.push_back(data_type);
-            break;
-        default:
-            return ZX_ERR_NOT_SUPPORTED;
-        }
-    }
-
-    // Early return if nothing to wipe.
-    if (partition_list.is_empty() && !efi) {
-        return ZX_OK;
-    }
-
-    const auto filter = [&info, &partition_list, efi](const gpt_partition_t& part) {
-        for (const auto& type : partition_list) {
-            if (memcmp(part.type, type, GPT_GUID_LEN) == 0)
-                return true;
-        }
-        if (efi) {
-            return FilterZirconPartition(info, part);
-        }
-        return false;
-    };
-    return gpt_->WipePartitions(filter);
+zx_status_t EfiDevicePartitioner::WipePartitions() {
+    return gpt_->WipePartitions(WipeFilterCallback);
 }
 
 zx_status_t EfiDevicePartitioner::GetBlockSize(const fbl::unique_fd& device_fd,
@@ -743,10 +733,6 @@
     return ZX_OK;
 }
 
-constexpr char kZirconAName[] = "ZIRCON-A";
-// TODO(raggi): near future - constexpr char kZirconBName[] = "ZIRCON-B";
-// TODO(raggi): near future - constexpr char kZirconRName[] = "ZIRCON-R";
-
 zx_status_t CrosDevicePartitioner::AddPartition(Partition partition_type,
                                                 fbl::unique_fd* out_fd) {
     const char* name;
@@ -782,7 +768,8 @@
     switch (partition_type) {
     case Partition::kKernelC: {
         const auto filter = [](const gpt_partition_t& part) {
-            return KernelFilterCallback(part, kZirconAName);
+            const uint8_t guid[GPT_GUID_LEN] = GUID_CROS_KERNEL_VALUE;
+            return KernelFilterCallback(part, guid, kZirconAName);
         };
         return gpt_->FindPartition(filter, out_fd);
     }
@@ -825,7 +812,8 @@
     }
 
     const auto filter_zircona = [](const gpt_partition_t& part) {
-        return KernelFilterCallback(part, kZirconAName);
+        const uint8_t guid[GPT_GUID_LEN] = GUID_CROS_KERNEL_VALUE;
+        return KernelFilterCallback(part, guid, kZirconAName);
     };
     zx_status_t status;
     gpt_partition_t* partition;
@@ -860,49 +848,8 @@
     return ZX_OK;
 }
 
-zx_status_t CrosDevicePartitioner::WipePartitions(const fbl::Vector<Partition>& partitions) {
-    const uint8_t fvm_type[GPT_GUID_LEN] = GUID_FVM_VALUE;
-    const uint8_t install_type[GPT_GUID_LEN] = GUID_INSTALL_VALUE;
-    const uint8_t system_type[GPT_GUID_LEN] = GUID_SYSTEM_VALUE;
-    const uint8_t blob_type[GPT_GUID_LEN] = GUID_BLOB_VALUE;
-    const uint8_t data_type[GPT_GUID_LEN] = GUID_DATA_VALUE;
-
-    // TODO(raggi): add logic here to cleanup the kernc, rootc, and a/b/r partitions.
-
-    fbl::Vector<const uint8_t*> partition_list;
-    for (const auto& partition_type : partitions) {
-        switch (partition_type) {
-        case Partition::kEfi:
-            continue;
-        case Partition::kFuchsiaVolumeManager:
-            partition_list.push_back(fvm_type);
-            break;
-        case Partition::kInstallType:
-            partition_list.push_back(install_type);
-            break;
-        case Partition::kSystem:
-            partition_list.push_back(system_type);
-            break;
-        case Partition::kBlob:
-            partition_list.push_back(blob_type);
-            break;
-        case Partition::kData:
-            partition_list.push_back(data_type);
-            break;
-        default:
-            return ZX_ERR_NOT_SUPPORTED;
-        }
-    }
-
-    auto filter = [&](const gpt_partition_t& part) {
-        for (const auto& type : partition_list) {
-            if (memcmp(part.type, type, GPT_GUID_LEN) == 0) {
-                return true;
-            }
-        }
-        return false;
-    };
-    return gpt_->WipePartitions(filter);
+zx_status_t CrosDevicePartitioner::WipePartitions() {
+    return gpt_->WipePartitions(WipeFilterCallback);
 }
 
 zx_status_t CrosDevicePartitioner::GetBlockSize(const fbl::unique_fd& device_fd,
@@ -961,22 +908,13 @@
     return OpenBlockPartition(nullptr, type, ZX_SEC(5), out_fd);
 }
 
-zx_status_t FixedDevicePartitioner::WipePartitions(const fbl::Vector<Partition>& partitions) {
+zx_status_t FixedDevicePartitioner::WipePartitions() {
     const uint8_t fvm_type[GPT_GUID_LEN] = GUID_FVM_VALUE;
     zx_status_t status;
-    for (const Partition& partition_type : partitions) {
-        switch (partition_type) {
-        case Partition::kFuchsiaVolumeManager:
-            if ((status = WipeBlockPartition(nullptr, fvm_type)) != ZX_OK) {
-                ERROR("Failed to wipe FVM.\n");
-            } else {
-                LOG("Wiped FVM successfully.\n");
-            }
-            break;
-        default:
-            // All non-FVM partitions are currently ignored on FixedDevices.
-            continue;
-        }
+    if ((status = WipeBlockPartition(nullptr, fvm_type)) != ZX_OK) {
+        ERROR("Failed to wipe FVM.\n");
+    } else {
+        LOG("Wiped FVM successfully.\n");
     }
     LOG("Immediate reboot strongly recommended\n");
     return ZX_OK;
@@ -1047,22 +985,13 @@
     return OpenSkipBlockPartition(type, ZX_SEC(5), out_fd);
 }
 
-zx_status_t SkipBlockDevicePartitioner::WipePartitions(const fbl::Vector<Partition>& partitions) {
+zx_status_t SkipBlockDevicePartitioner::WipePartitions() {
     const uint8_t fvm_type[GPT_GUID_LEN] = GUID_FVM_VALUE;
     zx_status_t status;
-    for (const Partition& partition_type : partitions) {
-        switch (partition_type) {
-        case Partition::kFuchsiaVolumeManager:
-            if ((status = WipeBlockPartition(nullptr, fvm_type)) != ZX_OK) {
-                ERROR("Failed to wipe FVM.\n");
-            } else {
-                LOG("Wiped FVM successfully.\n");
-            }
-            break;
-        default:
-            // All non-FVM partitions are currently ignored on SkipBlockDevices.
-            continue;
-        }
+    if ((status = WipeBlockPartition(nullptr, fvm_type)) != ZX_OK) {
+        ERROR("Failed to wipe FVM.\n");
+    } else {
+        LOG("Wiped FVM successfully.\n");
     }
     LOG("Immediate reboot strongly recommended\n");
     return ZX_OK;
diff --git a/system/uapp/disk-pave/device-partitioner.h b/system/uapp/disk-pave/device-partitioner.h
index cbcacc9..6a05fb6 100644
--- a/system/uapp/disk-pave/device-partitioner.h
+++ b/system/uapp/disk-pave/device-partitioner.h
@@ -17,6 +17,7 @@
 namespace paver {
 
 enum class Partition {
+    kUnknown,
     kBootloader,
     kKernelC,
     kEfi,
@@ -63,8 +64,8 @@
     // written.
     virtual zx_status_t FinalizePartition(Partition partition_type) = 0;
 
-    // Wipes partition list specified.
-    virtual zx_status_t WipePartitions(const fbl::Vector<Partition>& partitions) = 0;
+    // Wipes Fuchsia specific partitions.
+    virtual zx_status_t WipePartitions() = 0;
 
     // Returns block size in bytes for specified device.
     virtual zx_status_t GetBlockSize(const fbl::unique_fd& device_fd,
@@ -147,7 +148,7 @@
 
     zx_status_t FinalizePartition(Partition unused) override { return ZX_OK; }
 
-    zx_status_t WipePartitions(const fbl::Vector<Partition>& partitions) override;
+    zx_status_t WipePartitions() override;
 
     zx_status_t GetBlockSize(const fbl::unique_fd& device_fd, uint32_t* block_size) const override;
 
@@ -155,8 +156,6 @@
     EfiDevicePartitioner(fbl::unique_ptr<GptDevicePartitioner> gpt)
         : gpt_(fbl::move(gpt)) {}
 
-    static bool FilterZirconPartition(const block_info_t& info, const gpt_partition_t& part);
-
     fbl::unique_ptr<GptDevicePartitioner> gpt_;
 };
 
@@ -175,7 +174,7 @@
 
     zx_status_t FinalizePartition(Partition unused) override;
 
-    zx_status_t WipePartitions(const fbl::Vector<Partition>& partitions) override;
+    zx_status_t WipePartitions() override;
 
     zx_status_t GetBlockSize(const fbl::unique_fd& device_fd, uint32_t* block_size) const override;
 
@@ -206,7 +205,7 @@
 
     zx_status_t FinalizePartition(Partition unused) override { return ZX_OK; }
 
-    zx_status_t WipePartitions(const fbl::Vector<Partition>& partitions) override;
+    zx_status_t WipePartitions() override;
 
     zx_status_t GetBlockSize(const fbl::unique_fd& device_fd, uint32_t* block_size) const override;
 
@@ -235,7 +234,7 @@
 
     zx_status_t FinalizePartition(Partition unused) override { return ZX_OK; }
 
-    zx_status_t WipePartitions(const fbl::Vector<Partition>& partitions) override;
+    zx_status_t WipePartitions() override;
 
     zx_status_t GetBlockSize(const fbl::unique_fd& device_fd, uint32_t* block_size) const override;
 
diff --git a/system/uapp/disk-pave/pave-lib.cpp b/system/uapp/disk-pave/pave-lib.cpp
index 7d846fc..fa36581 100644
--- a/system/uapp/disk-pave/pave-lib.cpp
+++ b/system/uapp/disk-pave/pave-lib.cpp
@@ -42,6 +42,27 @@
 namespace paver {
 namespace {
 
+static Partition PartitionType(const Command command) {
+    switch(command) {
+    case Command::kInstallBootloader:
+        return Partition::kBootloader;
+    case Command::kInstallEfi:
+        return Partition::kEfi;
+    case Command::kInstallKernc:
+        return Partition::kKernelC;
+    case Command::kInstallZirconA:
+        return Partition::kZirconA;
+    case Command::kInstallZirconB:
+        return Partition::kZirconB;
+    case Command::kInstallZirconR:
+        return Partition::kZirconR;
+    case Command::kInstallFvm:
+        return Partition::kFuchsiaVolumeManager;
+    default:
+        return Partition::kUnknown;
+    }
+}
+
 // The number of additional slices a partition will need to become
 // zxcrypt'd.
 //
@@ -825,36 +846,6 @@
 
 } // namespace
 
-zx_status_t FvmPave(fbl::unique_ptr<DevicePartitioner> device_partitioner,
-                    fbl::unique_fd payload_fd) {
-    LOG("Paving FVM\n");
-
-    fbl::unique_fd partition_fd;
-    zx_status_t status;
-    status = device_partitioner->FindPartition(Partition::kFuchsiaVolumeManager, &partition_fd);
-    if (status == ZX_OK) {
-        LOG("FVM partition already exists\n");
-    } else if (status != ZX_ERR_NOT_FOUND) {
-        ERROR("Failure finding FVM partition: %s\n", zx_status_get_string(status));
-        return status;
-    } else {
-        LOG("Could not find FVM; attempting to add it: %s\n", zx_status_get_string(status));
-        status = device_partitioner->AddPartition(Partition::kFuchsiaVolumeManager, &partition_fd);
-        if (status != ZX_OK) {
-            ERROR("Failure creating partition: %s\n", zx_status_get_string(status));
-            return status;
-        }
-    }
-
-    LOG("Streaming partitions...\n");
-    if ((status = FvmStreamPartitions(fbl::move(partition_fd), fbl::move(payload_fd))) != ZX_OK) {
-        ERROR("Failed to stream partitions: %s\n", zx_status_get_string(status));
-        return status;
-    }
-    LOG("Completed successfully\n");
-    return ZX_OK;
-}
-
 zx_status_t PartitionPave(fbl::unique_ptr<DevicePartitioner> partitioner,
                           fbl::unique_fd payload_fd, Partition partition_type, Arch arch) {
     LOG("Paving partition.\n");
@@ -874,6 +865,16 @@
         LOG("Partition already exists\n");
     }
 
+    if (partition_type == Partition::kFuchsiaVolumeManager) {
+        LOG("Streaming partitions...\n");
+        if ((status = FvmStreamPartitions(fbl::move(partition_fd), fbl::move(payload_fd))) != ZX_OK) {
+            ERROR("Failed to stream partitions: %s\n", zx_status_get_string(status));
+            return status;
+        }
+        LOG("Completed successfully\n");
+        return ZX_OK;
+    }
+
     uint32_t block_size_bytes;
     if ((status = partitioner->GetBlockSize(partition_fd, &block_size_bytes)) != ZX_OK) {
         ERROR("Couldn't get partition block size\n");
@@ -918,17 +919,6 @@
     return ZX_OK;
 }
 
-zx_status_t FvmClean(fbl::unique_ptr<DevicePartitioner> partitioner) {
-    const fbl::Vector<Partition> partition_list = {
-        Partition::kEfi,
-        Partition::kFuchsiaVolumeManager,
-        Partition::kSystem,
-        Partition::kBlob,
-        Partition::kData,
-    };
-    return partitioner->WipePartitions(partition_list);
-}
-
 void Drain(fbl::unique_fd fd) {
     char buf[8192];
     while (read(fd.get(), &buf, sizeof(buf)) > 0)
@@ -944,70 +934,46 @@
     const bool is_cros_device = device_partitioner->IsCros();
 
     switch (flags.cmd) {
+    case Command::kWipe:
+        return device_partitioner->WipePartitions();
+    case Command::kInstallFvm:
+        break;
     case Command::kInstallBootloader:
         if (flags.arch == Arch::X64 && !flags.force) {
             LOG("SKIPPING BOOTLOADER install on x64 device, pass --force if desired.\n");
             Drain(fbl::move(flags.payload_fd));
             return ZX_OK;
         }
-        return PartitionPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd),
-                             Partition::kBootloader, flags.arch);
-
+        break;
     case Command::kInstallEfi:
         if ((is_cros_device || flags.arch == Arch::ARM64) && !flags.force) {
             LOG("SKIPPING EFI install on ARM64/CROS device, pass --force if desired.\n");
             Drain(fbl::move(flags.payload_fd));
             return ZX_OK;
         }
-        return PartitionPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd),
-                             Partition::kEfi, flags.arch);
-
+        break;
     case Command::kInstallKernc:
         if (!is_cros_device && !flags.force) {
             LOG("SKIPPING KERNC install on non-CROS device, pass --force if desired.\n");
             Drain(fbl::move(flags.payload_fd));
             return ZX_OK;
         }
-        return PartitionPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd),
-                             Partition::kKernelC, flags.arch);
-
+        break;
     case Command::kInstallZirconA:
-        // TODO(ZX-2220): At some point x64 devices will want to pave A/B/R partitions as well.
-        if (flags.arch == Arch::X64 && !flags.force) {
-            LOG("SKIPPING ZIRCON-A install on x64 device, pass --force if desired.\n");
-            Drain(fbl::move(flags.payload_fd));
-            return ZX_OK;
-        }
-        return PartitionPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd),
-                             Partition::kZirconA, flags.arch);
-
     case Command::kInstallZirconB:
-        if (flags.arch == Arch::X64 && !flags.force) {
-            LOG("SKIPPING ZIRCON-B install on x64 device, pass --force if desired.\n");
-            Drain(fbl::move(flags.payload_fd));
-            return ZX_OK;
-        }
-        return PartitionPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd),
-                             Partition::kZirconB, flags.arch);
-
     case Command::kInstallZirconR:
-        if (flags.arch == Arch::X64 && !flags.force) {
-            LOG("SKIPPING ZIRCON-R install on x64 device, pass --force if desired.\n");
+        if (is_cros_device && !flags.force) {
+            LOG("SKIPPING Zircon-{A/B/R} install on CROS device, pass --force if desired.\n");
             Drain(fbl::move(flags.payload_fd));
             return ZX_OK;
         }
-        return PartitionPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd),
-                             Partition::kZirconR, flags.arch);
-    case Command::kInstallFvm:
-        return FvmPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd));
-
-    case Command::kWipe:
-        return FvmClean(fbl::move(device_partitioner));
-
+        break;
     default:
         ERROR("Unsupported command.");
         return ZX_ERR_NOT_SUPPORTED;
     }
+    return PartitionPave(fbl::move(device_partitioner), fbl::move(flags.payload_fd),
+                            PartitionType(flags.cmd), flags.arch);
 }
 
 } //  namespace paver
diff --git a/system/uapp/disk-pave/pave-lib.h b/system/uapp/disk-pave/pave-lib.h
index 1bee15f..e64b4fe 100644
--- a/system/uapp/disk-pave/pave-lib.h
+++ b/system/uapp/disk-pave/pave-lib.h
@@ -38,26 +38,10 @@
     fbl::unique_fd payload_fd;
 };
 
-// Paves a sparse_file to the underlying disk, on top partition.
-// Expects FVM partition to not exist
-extern zx_status_t FvmPave(fbl::unique_ptr<DevicePartitioner> device_partitioner,
-                           fbl::unique_fd payload_fd);
-
 // Paves an image onto the disk.
 extern zx_status_t PartitionPave(fbl::unique_ptr<DevicePartitioner> partitioner,
                                  fbl::unique_fd payload_fd, Partition partition_type, Arch arch);
 
-// Wipes the following partitions:
-// - System
-// - Data
-// - Blob
-// - FVM
-// - EFI
-//
-// From the target, leaving it (hopefully) in a state ready for a sparse FVM
-// image to be installed.
-extern zx_status_t FvmClean(fbl::unique_ptr<DevicePartitioner> partitioner);
-
 // Reads the entire file from supplied file descriptor. This is necessary due to
 // implementation of streaming protocol which forces entire file to be
 // transfered.
diff --git a/system/uapp/disk-pave/test/device-partitioner-test.cpp b/system/uapp/disk-pave/test/device-partitioner-test.cpp
index 93a7261..073612b 100644
--- a/system/uapp/disk-pave/test/device-partitioner-test.cpp
+++ b/system/uapp/disk-pave/test/device-partitioner-test.cpp
@@ -337,7 +337,7 @@
 
     fbl::unique_ptr<paver::DevicePartitioner> partitioner;
     ASSERT_EQ(paver::FixedDevicePartitioner::Initialize(&partitioner), ZX_OK);
-    ASSERT_EQ(partitioner->WipePartitions(fbl::Vector<paver::Partition>()), ZX_OK);
+    ASSERT_EQ(partitioner->WipePartitions(), ZX_OK);
 
     END_TEST;
 }
@@ -476,7 +476,7 @@
 
     fbl::unique_ptr<paver::DevicePartitioner> partitioner;
     ASSERT_EQ(paver::SkipBlockDevicePartitioner::Initialize(&partitioner), ZX_OK);
-    ASSERT_EQ(partitioner->WipePartitions(fbl::Vector<paver::Partition>()), ZX_OK);
+    ASSERT_EQ(partitioner->WipePartitions(), ZX_OK);
 
     END_TEST;
 }