Add partial extraction capability to ota_extractor

Test: th
Bug: 221875271
Change-Id: I48adddbcbfce3a437e0a716e9f5afd6ecf969f50

1 files changed, 26 insertions, 3 deletions

diff --git a/aosp/ota_extractor.cc b/aosp/ota_extractor.cc
index f200fc81..aa0d8230 100644
--- a/aosp/ota_extractor.cc
+++ b/aosp/ota_extractor.cc
@@ -18,11 +18,13 @@
 
 #include <cstdint>
 #include <cstdio>
+#include <iterator>
 #include <memory>
 
 #include <sys/mman.h>
 #include <sys/stat.h>
 
+#include <android-base/strings.h>
 #include <base/files/file_path.h>
 #include <gflags/gflags.h>
 #include <unistd.h>
@@ -41,6 +43,10 @@ DEFINE_int64(payload_offset,
              0,
              "Offset to start of payload.bin. Useful if payload path actually "
              "points to a .zip file containing payload.bin");
+DEFINE_string(partitions,
+              "",
+              "Comma separated list of partitions to extract, leave empty for "
+              "extracting all partitions");
 
 using chromeos_update_engine::DeltaArchiveManifest;
 using chromeos_update_engine::PayloadMetadata;
@@ -51,7 +57,8 @@ bool ExtractImagesFromOTA(const DeltaArchiveManifest& manifest,
                           const PayloadMetadata& metadata,
                           int payload_fd,
                           size_t payload_offset,
-                          std::string_view output_dir) {
+                          std::string_view output_dir,
+                          const std::set<std::string>& partitions) {
   InstallOperationExecutor executor(manifest.block_size());
   const size_t data_begin = metadata.GetMetadataSize() +
                             metadata.GetMetadataSignatureSize() +
@@ -60,6 +67,10 @@ bool ExtractImagesFromOTA(const DeltaArchiveManifest& manifest,
       base::StringPiece(output_dir.data(), output_dir.size()));
   std::vector<unsigned char> blob;
   for (const auto& partition : manifest.partitions()) {
+    if (!partitions.empty() &&
+        partitions.count(partition.partition_name()) == 0) {
+      continue;
+    }
     LOG(INFO) << "Extracting partition " << partition.partition_name()
               << " size: " << partition.new_partition_info().size();
     const auto output_path =
@@ -98,9 +109,20 @@ bool ExtractImagesFromOTA(const DeltaArchiveManifest& manifest,
 int main(int argc, char* argv[]) {
   gflags::ParseCommandLineFlags(&argc, &argv, true);
   xz_crc32_init();
+  auto tokens = android::base::Tokenize(FLAGS_partitions, ",");
+  const std::set<std::string> partitions(
+      std::make_move_iterator(tokens.begin()),
+      std::make_move_iterator(tokens.end()));
+  if (FLAGS_payload.empty()) {
+    LOG(ERROR) << "--payload <payload path> is required";
+    return 1;
+  }
+  if (!partitions.empty()) {
+    LOG(INFO) << "Extracting " << android::base::Join(partitions, ", ");
+  }
   int payload_fd = open(FLAGS_payload.c_str(), O_RDONLY | O_CLOEXEC);
   if (payload_fd < 0) {
-    PLOG(ERROR) << "Failed to open payload file:";
+    PLOG(ERROR) << "Failed to open payload file";
     return 1;
   }
   chromeos_update_engine::ScopedFdCloser closer{&payload_fd};
@@ -143,5 +165,6 @@ int main(int argc, char* argv[]) {
                                payload_metadata,
                                payload_fd,
                                FLAGS_payload_offset,
-                               FLAGS_output_dir);
+                               FLAGS_output_dir,
+                               partitions);
 }