Support multiple partitions in ImageConfig

ImageConfig now have a vector of partitions instead of hardcoded rootfs and
kernel.

Bug: 23420126
TEST=cros_workon_make update_engine --test

Change-Id: Id226cc04628b44f1fcbd58f03038809867bf9e40
diff --git a/delta_performer_integration_test.cc b/delta_performer_integration_test.cc
index 143d21d..cdba884 100644
--- a/delta_performer_integration_test.cc
+++ b/delta_performer_integration_test.cc
@@ -500,24 +500,28 @@
     payload_config.major_version = kChromeOSMajorPayloadVersion;
     payload_config.minor_version = minor_version;
     if (!full_rootfs) {
-      payload_config.source.rootfs.path = state->a_img;
+      payload_config.source.partitions.emplace_back(kLegacyPartitionNameRoot);
+      payload_config.source.partitions.emplace_back(kLegacyPartitionNameKernel);
+      payload_config.source.partitions.front().path = state->a_img;
       if (!full_kernel)
-        payload_config.source.kernel.path = state->old_kernel;
+        payload_config.source.partitions.back().path = state->old_kernel;
       payload_config.source.image_info = old_image_info;
       EXPECT_TRUE(payload_config.source.LoadImageSize());
-      EXPECT_TRUE(payload_config.source.rootfs.OpenFilesystem());
-      EXPECT_TRUE(payload_config.source.kernel.OpenFilesystem());
+      for (PartitionConfig& part : payload_config.source.partitions)
+        EXPECT_TRUE(part.OpenFilesystem());
     } else {
       if (payload_config.hard_chunk_size == -1)
         // Use 1 MiB chunk size for the full unittests.
         payload_config.hard_chunk_size = 1024 * 1024;
     }
-    payload_config.target.rootfs.path = state->b_img;
-    payload_config.target.kernel.path = state->new_kernel;
+    payload_config.target.partitions.emplace_back(kLegacyPartitionNameRoot);
+    payload_config.target.partitions.back().path = state->b_img;
+    payload_config.target.partitions.emplace_back(kLegacyPartitionNameKernel);
+    payload_config.target.partitions.back().path = state->new_kernel;
     payload_config.target.image_info = new_image_info;
     EXPECT_TRUE(payload_config.target.LoadImageSize());
-    EXPECT_TRUE(payload_config.target.rootfs.OpenFilesystem());
-    EXPECT_TRUE(payload_config.target.kernel.OpenFilesystem());
+    for (PartitionConfig& part : payload_config.target.partitions)
+      EXPECT_TRUE(part.OpenFilesystem());
 
     EXPECT_TRUE(payload_config.Validate());
     EXPECT_TRUE(
diff --git a/delta_performer_unittest.cc b/delta_performer_unittest.cc
index 0b509e3..839f253 100644
--- a/delta_performer_unittest.cc
+++ b/delta_performer_unittest.cc
@@ -112,15 +112,16 @@
     PayloadGenerationConfig config;
     config.major_version = kChromeOSMajorPayloadVersion;
     config.minor_version = minor_version;
-    config.target.rootfs.path = blob_path;
-    config.target.rootfs.size = blob_data.size();
-    config.target.kernel.path = blob_path;
-    config.target.kernel.size = blob_data.size();
 
     PayloadFile payload;
     EXPECT_TRUE(payload.Init(config));
 
-    payload.AddPartition(config.source.rootfs, config.target.rootfs, aops);
+    PartitionConfig old_part(kLegacyPartitionNameRoot);
+    PartitionConfig new_part(kLegacyPartitionNameRoot);
+    new_part.path = blob_path;
+    new_part.size = blob_data.size();
+
+    payload.AddPartition(old_part, new_part, aops);
 
     string payload_path;
     EXPECT_TRUE(utils::MakeTempFile("Payload-XXXXXX", &payload_path, nullptr));
diff --git a/payload_generator/delta_diff_generator.cc b/payload_generator/delta_diff_generator.cc
index 1592917..af81e62 100644
--- a/payload_generator/delta_diff_generator.cc
+++ b/payload_generator/delta_diff_generator.cc
@@ -55,109 +55,89 @@
     const string& output_path,
     const string& private_key_path,
     uint64_t* metadata_size) {
-  if (config.is_delta) {
-    LOG_IF(WARNING, config.source.rootfs.size != config.target.rootfs.size)
-        << "Old and new images have different block counts.";
-    // TODO(deymo): Our tools only support growing the filesystem size during
-    // an update. Remove this check when that's fixed. crbug.com/192136
-    LOG_IF(FATAL, config.source.rootfs.size > config.target.rootfs.size)
-        << "Shirking the rootfs size is not supported at the moment.";
-  }
-
-  // Sanity checks for the partition size.
-  LOG(INFO) << "Rootfs partition size: " << config.rootfs_partition_size;
-  LOG(INFO) << "Actual filesystem size: " << config.target.rootfs.size;
-
-  LOG(INFO) << "Block count: "
-            << config.target.rootfs.size / config.block_size;
-
-  LOG_IF(INFO, config.source.kernel.path.empty())
-      << "Will generate full kernel update.";
-
-  const string kTempFileTemplate("CrAU_temp_data.XXXXXX");
-  string temp_file_path;
-  off_t data_file_size = 0;
-
-  LOG(INFO) << "Reading files...";
 
   // Create empty payload file object.
   PayloadFile payload;
   TEST_AND_RETURN_FALSE(payload.Init(config));
 
-  vector<AnnotatedOperation> rootfs_ops;
-  vector<AnnotatedOperation> kernel_ops;
-
-  // Select payload generation strategy based on the config.
-  unique_ptr<OperationsGenerator> strategy;
-  if (config.is_delta) {
-    // We don't efficiently support deltas on squashfs. For now, we will
-    // produce full operations in that case.
-    if (utils::IsSquashfsFilesystem(config.target.rootfs.path)) {
-      LOG(INFO) << "Using generator FullUpdateGenerator() for squashfs deltas";
-      strategy.reset(new FullUpdateGenerator());
-    } else if (utils::IsExtFilesystem(config.target.rootfs.path)) {
-      // Delta update (with possibly a full kernel update).
-      if (config.minor_version == kInPlaceMinorPayloadVersion) {
-        LOG(INFO) << "Using generator InplaceGenerator()";
-        strategy.reset(new InplaceGenerator());
-      } else if (config.minor_version == kSourceMinorPayloadVersion) {
-        LOG(INFO) << "Using generator ABGenerator()";
-        strategy.reset(new ABGenerator());
-      } else {
-        LOG(ERROR) << "Unsupported minor version given for delta payload: "
-                   << config.minor_version;
-        return false;
-      }
-    } else {
-      LOG(ERROR) << "Unsupported filesystem for delta payload in "
-                 << config.target.rootfs.path;
-      return false;
-    }
-  } else {
-    // Full update.
-    LOG(INFO) << "Using generator FullUpdateGenerator()";
-    strategy.reset(new FullUpdateGenerator());
-  }
-
+  const string kTempFileTemplate("CrAU_temp_data.XXXXXX");
+  string temp_file_path;
   int data_file_fd;
   TEST_AND_RETURN_FALSE(
       utils::MakeTempFile(kTempFileTemplate, &temp_file_path, &data_file_fd));
-  unique_ptr<ScopedPathUnlinker> temp_file_unlinker(
-      new ScopedPathUnlinker(temp_file_path));
+  ScopedPathUnlinker temp_file_unlinker(temp_file_path);
   TEST_AND_RETURN_FALSE(data_file_fd >= 0);
 
   {
+    off_t data_file_size = 0;
     ScopedFdCloser data_file_fd_closer(&data_file_fd);
     BlobFileWriter blob_file(data_file_fd, &data_file_size);
-    // Generate the operations using the strategy we selected above.
-    TEST_AND_RETURN_FALSE(strategy->GenerateOperations(config,
-                                                       config.source.rootfs,
-                                                       config.target.rootfs,
-                                                       &blob_file,
-                                                       &rootfs_ops));
-    TEST_AND_RETURN_FALSE(strategy->GenerateOperations(config,
-                                                       config.source.kernel,
-                                                       config.target.kernel,
-                                                       &blob_file,
-                                                       &kernel_ops));
+    if (config.is_delta) {
+      TEST_AND_RETURN_FALSE(config.source.partitions.size() ==
+                            config.target.partitions.size());
+    }
+    PartitionConfig empty_part("");
+    for (size_t i = 0; i < config.target.partitions.size(); i++) {
+      const PartitionConfig& old_part =
+          config.is_delta ? config.source.partitions[i] : empty_part;
+      const PartitionConfig& new_part = config.target.partitions[i];
+      LOG(INFO) << "Partition name: " << new_part.name;
+      LOG(INFO) << "Partition size: " << new_part.size;
+      LOG(INFO) << "Block count: " << new_part.size / config.block_size;
+
+      // Select payload generation strategy based on the config.
+      unique_ptr<OperationsGenerator> strategy;
+      // We don't efficiently support deltas on squashfs. For now, we will
+      // produce full operations in that case.
+      if (!old_part.path.empty() &&
+          !utils::IsSquashfsFilesystem(new_part.path)) {
+        // Delta update.
+        if (utils::IsExtFilesystem(new_part.path)) {
+          LOG_IF(WARNING, old_part.size != new_part.size)
+              << "Old and new filesystems have different size.";
+          // TODO(deymo): Our tools only support growing the filesystem size
+          // during an update. Remove this check when that's fixed.
+          // crbug.com/192136
+          LOG_IF(FATAL, old_part.size > new_part.size)
+              << "Shirking the filesystem size is not supported at the moment.";
+        }
+        if (config.minor_version == kInPlaceMinorPayloadVersion) {
+          LOG(INFO) << "Using generator InplaceGenerator().";
+          strategy.reset(new InplaceGenerator());
+        } else if (config.minor_version == kSourceMinorPayloadVersion) {
+          LOG(INFO) << "Using generator ABGenerator().";
+          strategy.reset(new ABGenerator());
+        } else {
+          LOG(ERROR) << "Unsupported minor version given for delta payload: "
+                     << config.minor_version;
+          return false;
+        }
+      } else {
+        LOG(INFO) << "Using generator FullUpdateGenerator().";
+        strategy.reset(new FullUpdateGenerator());
+      }
+
+      vector<AnnotatedOperation> aops;
+      // Generate the operations using the strategy we selected above.
+      TEST_AND_RETURN_FALSE(strategy->GenerateOperations(config,
+                                                         old_part,
+                                                         new_part,
+                                                         &blob_file,
+                                                         &aops));
+
+      // Filter the no-operations. OperationsGenerators should not output this
+      // kind of operations normally, but this is an extra step to fix that if
+      // happened.
+      diff_utils::FilterNoopOperations(&aops);
+
+      TEST_AND_RETURN_FALSE(payload.AddPartition(old_part, new_part, aops));
+    }
   }
 
-  // Filter the no-operations. OperationsGenerators should not output this kind
-  // of operations normally, but this is an extra step to fix that if
-  // happened.
-  diff_utils::FilterNoopOperations(&rootfs_ops);
-  diff_utils::FilterNoopOperations(&kernel_ops);
-
+  LOG(INFO) << "Writing payload file...";
   // Write payload file to disk.
-  TEST_AND_RETURN_FALSE(payload.AddPartition(config.source.rootfs,
-                                             config.target.rootfs,
-                                             rootfs_ops));
-  TEST_AND_RETURN_FALSE(payload.AddPartition(config.source.kernel,
-                                             config.target.kernel,
-                                             kernel_ops));
   TEST_AND_RETURN_FALSE(payload.WritePayload(output_path, temp_file_path,
                                              private_key_path, metadata_size));
-  temp_file_unlinker.reset();
 
   LOG(INFO) << "All done. Successfully created delta file with "
             << "metadata size = " << *metadata_size;
diff --git a/payload_generator/generate_delta_main.cc b/payload_generator/generate_delta_main.cc
index a600c9c..327ec38 100644
--- a/payload_generator/generate_delta_main.cc
+++ b/payload_generator/generate_delta_main.cc
@@ -197,11 +197,15 @@
   LOG(INFO) << "Calculating original checksums";
   PartitionInfo kern_info, root_info;
   ImageConfig old_image;
-  old_image.kernel.path = old_kernel;
-  old_image.rootfs.path = old_rootfs;
+  old_image.partitions.emplace_back(kLegacyPartitionNameRoot);
+  old_image.partitions.back().path = old_rootfs;
+  old_image.partitions.emplace_back(kLegacyPartitionNameKernel);
+  old_image.partitions.back().path = old_kernel;
   CHECK(old_image.LoadImageSize());
-  CHECK(diff_utils::InitializePartitionInfo(old_image.kernel, &kern_info));
-  CHECK(diff_utils::InitializePartitionInfo(old_image.rootfs, &root_info));
+  CHECK(diff_utils::InitializePartitionInfo(old_image.partitions[0],
+                                            &root_info));
+  CHECK(diff_utils::InitializePartitionInfo(old_image.partitions[1],
+                                            &kern_info));
   install_plan.kernel_hash.assign(kern_info.hash().begin(),
                                   kern_info.hash().end());
   install_plan.rootfs_hash.assign(root_info.hash().begin(),
@@ -355,11 +359,21 @@
   // A payload generation was requested. Convert the flags to a
   // PayloadGenerationConfig.
   PayloadGenerationConfig payload_config;
-  payload_config.source.rootfs.path = FLAGS_old_image;
-  payload_config.source.kernel.path = FLAGS_old_kernel;
 
-  payload_config.target.rootfs.path = FLAGS_new_image;
-  payload_config.target.kernel.path = FLAGS_new_kernel;
+  payload_config.is_delta = !FLAGS_old_image.empty() ||
+                            !FLAGS_old_kernel.empty();
+
+  if (payload_config.is_delta) {
+    payload_config.source.partitions.emplace_back(kLegacyPartitionNameRoot);
+    payload_config.source.partitions.back().path = FLAGS_old_image;
+    payload_config.source.partitions.emplace_back(kLegacyPartitionNameKernel);
+    payload_config.source.partitions.back().path = FLAGS_old_kernel;
+  }
+
+  payload_config.target.partitions.emplace_back(kLegacyPartitionNameRoot);
+  payload_config.target.partitions.back().path = FLAGS_new_image;
+  payload_config.target.partitions.emplace_back(kLegacyPartitionNameKernel);
+  payload_config.target.partitions.back().path = FLAGS_new_kernel;
 
   // Use the default soft_chunk_size defined in the config.
   payload_config.hard_chunk_size = FLAGS_chunk_size;
@@ -374,8 +388,6 @@
     CHECK(payload_config.target.LoadImageSize());
   }
 
-  payload_config.is_delta = !FLAGS_old_image.empty();
-
   CHECK(!FLAGS_out_file.empty());
 
   // Ignore failures. These are optional arguments.
@@ -400,10 +412,10 @@
 
   if (payload_config.is_delta) {
     // Avoid opening the filesystem interface for full payloads.
-    CHECK(payload_config.target.rootfs.OpenFilesystem());
-    CHECK(payload_config.target.kernel.OpenFilesystem());
-    CHECK(payload_config.source.rootfs.OpenFilesystem());
-    CHECK(payload_config.source.kernel.OpenFilesystem());
+    for (PartitionConfig& part : payload_config.target.partitions)
+      CHECK(part.OpenFilesystem());
+    for (PartitionConfig& part : payload_config.source.partitions)
+      CHECK(part.OpenFilesystem());
   }
 
   payload_config.major_version = FLAGS_major_version;
@@ -413,14 +425,15 @@
     // Autodetect minor_version by looking at the update_engine.conf in the old
     // image.
     if (payload_config.is_delta) {
-      CHECK(payload_config.source.rootfs.fs_interface);
+      payload_config.minor_version = kInPlaceMinorPayloadVersion;
       chromeos::KeyValueStore store;
       uint32_t minor_version;
-      if (payload_config.source.rootfs.fs_interface->LoadSettings(&store) &&
-          utils::GetMinorVersion(store, &minor_version)) {
-        payload_config.minor_version = minor_version;
-      } else {
-        payload_config.minor_version = kInPlaceMinorPayloadVersion;
+      for (const PartitionConfig& part : payload_config.source.partitions) {
+        if (part.fs_interface && part.fs_interface->LoadSettings(&store) &&
+            utils::GetMinorVersion(store, &minor_version)) {
+          payload_config.minor_version = minor_version;
+          break;
+        }
       }
     } else {
       payload_config.minor_version = kFullPayloadMinorVersion;
diff --git a/payload_generator/payload_generation_config.cc b/payload_generator/payload_generation_config.cc
index 2de696a..66928d3 100644
--- a/payload_generator/payload_generation_config.cc
+++ b/payload_generator/payload_generation_config.cc
@@ -67,19 +67,15 @@
 
 bool ImageConfig::ValidateIsEmpty() const {
   TEST_AND_RETURN_FALSE(ImageInfoIsEmpty());
-
-  TEST_AND_RETURN_FALSE(rootfs.path.empty());
-  TEST_AND_RETURN_FALSE(rootfs.size == 0);
-  TEST_AND_RETURN_FALSE(kernel.path.empty());
-  TEST_AND_RETURN_FALSE(kernel.size == 0);
-  return true;
+  return partitions.empty();
 }
 
 bool ImageConfig::LoadImageSize() {
-  TEST_AND_RETURN_FALSE(!rootfs.path.empty());
-  rootfs.size = utils::FileSize(rootfs.path);
-  if (!kernel.path.empty())
-    kernel.size = utils::FileSize(kernel.path);
+  for (PartitionConfig& part : partitions) {
+    if (part.path.empty())
+      continue;
+    part.size = utils::FileSize(part.path);
+  }
   return true;
 }
 
@@ -94,12 +90,11 @@
 
 bool PayloadGenerationConfig::Validate() const {
   if (is_delta) {
-    TEST_AND_RETURN_FALSE(source.rootfs.ValidateExists());
-    TEST_AND_RETURN_FALSE(source.rootfs.size % block_size == 0);
-
-    if (!source.kernel.path.empty()) {
-      TEST_AND_RETURN_FALSE(source.kernel.ValidateExists());
-      TEST_AND_RETURN_FALSE(source.kernel.size % block_size == 0);
+    for (const PartitionConfig& part : source.partitions) {
+      if (!part.path.empty()) {
+        TEST_AND_RETURN_FALSE(part.ValidateExists());
+        TEST_AND_RETURN_FALSE(part.size % block_size == 0);
+      }
     }
 
     // Check for the supported minor_version values.
@@ -116,17 +111,19 @@
   }
 
   // In all cases, the target image must exists.
-  TEST_AND_RETURN_FALSE(target.rootfs.ValidateExists());
-  TEST_AND_RETURN_FALSE(target.kernel.ValidateExists());
-  TEST_AND_RETURN_FALSE(target.rootfs.size % block_size == 0);
-  TEST_AND_RETURN_FALSE(target.kernel.size % block_size == 0);
+  for (const PartitionConfig& part : target.partitions) {
+    TEST_AND_RETURN_FALSE(part.ValidateExists());
+    TEST_AND_RETURN_FALSE(part.size % block_size == 0);
+    if (minor_version == kInPlaceMinorPayloadVersion &&
+        part.name == kLegacyPartitionNameRoot)
+      TEST_AND_RETURN_FALSE(rootfs_partition_size >= part.size);
+  }
 
   TEST_AND_RETURN_FALSE(hard_chunk_size == -1 ||
                         hard_chunk_size % block_size == 0);
   TEST_AND_RETURN_FALSE(soft_chunk_size % block_size == 0);
 
   TEST_AND_RETURN_FALSE(rootfs_partition_size % block_size == 0);
-  TEST_AND_RETURN_FALSE(rootfs_partition_size >= target.rootfs.size);
 
   return true;
 }
diff --git a/payload_generator/payload_generation_config.h b/payload_generator/payload_generation_config.h
index b7bd484..8ee8887 100644
--- a/payload_generator/payload_generation_config.h
+++ b/payload_generator/payload_generation_config.h
@@ -80,8 +80,7 @@
   ImageInfo image_info;
 
   // The updated partitions.
-  PartitionConfig rootfs = PartitionConfig{kLegacyPartitionNameRoot};
-  PartitionConfig kernel = PartitionConfig{kLegacyPartitionNameKernel};
+  std::vector<PartitionConfig> partitions;
 };
 
 // The PayloadGenerationConfig struct encapsulates all the configuration to