update_engine: Fix a bug with P2P updates when there are multiple payload

It seems currently, the implementaion opens the file descriptor for
writing the P2P file only once even if there are multiple payloads. This
is a bug for example where there are DLCs being updated in addition to
the platform update. This CL closes the FD everytime a payload is
sucessfully downloaded and opens it every time a new payload is starting
to be downloaded.

Unfortunately this can't be tested easily locally. So we need to land
this patch and then use DUTs in the lab to perform a p2p updates (using
canary builds) when DLCs are available to make sure this CL is correct.

BUG=b:172219304
TEST=cros_workon_make --board reef --test update_engine

Change-Id: If02035937288be84eb66a5c3abda87bce6ecdb09
Reviewed-on: https://chromium-review.googlesource.com/c/aosp/platform/system/update_engine/+/2776259
Tested-by: Amin Hassani <[email protected]>
Reviewed-by: Jae Hoon Kim <[email protected]>
Reviewed-by: Vyshu Khota <[email protected]>
Commit-Queue: Amin Hassani <[email protected]>
(cherry picked from commit 3774de980a5e2eb815178b6a1fa8257dbd767de9)
Reviewed-on: https://chromium-review.googlesource.com/c/aosp/platform/system/update_engine/+/2809342
Auto-Submit: Amin Hassani <[email protected]>
Commit-Queue: Jae Hoon Kim <[email protected]>
diff --git a/cros/download_action_chromeos.cc b/cros/download_action_chromeos.cc
index ea84f2c..0cc9f81 100644
--- a/cros/download_action_chromeos.cc
+++ b/cros/download_action_chromeos.cc
@@ -80,6 +80,7 @@
 
   // Don't use p2p from this point onwards.
   p2p_file_id_.clear();
+  p2p_visible_ = false;
 }
 
 bool DownloadActionChromeos::SetupP2PSharingFd() {
@@ -122,8 +123,7 @@
                                             size_t length,
                                             off_t file_offset) {
   if (p2p_sharing_fd_ == -1) {
-    if (!SetupP2PSharingFd())
-      return;
+    return;
   }
 
   // Check that the file is at least |file_offset| bytes long - if
@@ -255,6 +255,10 @@
   }
 
   if (SystemState::Get() != nullptr) {
+    // Close any previous P2P sharing. It will be a no-op if there were no
+    // previous file sharing.
+    CloseP2PSharingFd(false);
+
     const PayloadStateInterface* payload_state =
         SystemState::Get()->payload_state();
     string file_id = utils::CalculateP2PFileId(payload_->hash, payload_->size);
@@ -262,7 +266,9 @@
       // If we're sharing the update, store the file_id to convey
       // that we should write to the file.
       p2p_file_id_ = file_id;
-      LOG(INFO) << "p2p file id: " << p2p_file_id_;
+      LOG(INFO) << "Sharing p2p file with id: " << p2p_file_id_;
+      if (!SetupP2PSharingFd())
+        return;
     } else {
       // Even if we're not sharing the update, it could be that
       // there's a partial file from a previous attempt with the same
diff --git a/cros/download_action_chromeos_unittest.cc b/cros/download_action_chromeos_unittest.cc
index e56c25c..fb78744 100644
--- a/cros/download_action_chromeos_unittest.cc
+++ b/cros/download_action_chromeos_unittest.cc
@@ -534,9 +534,10 @@
   void SetupDownload(off_t starting_offset) {
     start_at_offset_ = starting_offset;
     // Prepare data 10 kB of data.
-    data_.clear();
-    for (unsigned int i = 0; i < 10 * 1000; i++)
-      data_ += 'a' + (i % 25);
+    data_.resize(10 * 1000);
+    std::generate(data_.begin(), data_.end(), [i = 0]() mutable {
+      return 'a' + (i++ % 26);
+    });
 
     // Setup p2p.
     FakeP2PManagerConfiguration* test_conf = new FakeP2PManagerConfiguration();
@@ -704,6 +705,69 @@
   EXPECT_EQ(0, p2p_manager_->CountSharedFiles());
 }
 
+TEST_F(P2PDownloadActionTest, MultiplePayload) {
+  SetupDownload(0);
+  EXPECT_CALL(*FakeSystemState::Get()->mock_payload_state(),
+              GetUsingP2PForSharing())
+      .WillRepeatedly(Return(true));
+
+  EXPECT_CALL(*FakeSystemState::Get()->mock_payload_state(), NextPayload())
+      .WillOnce(Return(true));
+
+  MockFileWriter mock_file_writer;
+  EXPECT_CALL(mock_file_writer, Close()).WillRepeatedly(Return(0));
+  EXPECT_CALL(mock_file_writer, Write(_, _, _))
+      .WillRepeatedly(
+          DoAll(SetArgPointee<2>(ErrorCode::kSuccess), Return(true)));
+
+  InstallPlan install_plan;
+  install_plan.payloads.push_back(
+      {.size = data_.length() / 4,
+       .hash = {'1', '1', '1', '1', 'h', 'a', 's', 'h'}});
+  install_plan.payloads.push_back(
+      {.size = (data_.length() / 4) * 3,
+       .hash = {'2', '2', '2', '2', 'h', 'a', 's', 'h'}});
+
+  auto feeder_action = std::make_unique<ObjectFeederAction<InstallPlan>>();
+  feeder_action->set_obj(install_plan);
+  auto download_action = std::make_unique<DownloadActionChromeos>(
+      FakeSystemState::Get()->prefs(),
+      FakeSystemState::Get()->boot_control(),
+      FakeSystemState::Get()->hardware(),
+      new MockHttpFetcher(data_.c_str(), data_.length(), nullptr),
+      /*interactive=*/false);
+
+  download_action->SetTestFileWriter(&mock_file_writer);
+  BondActions(feeder_action.get(), download_action.get());
+  processor_.EnqueueAction(std::move(feeder_action));
+  processor_.EnqueueAction(std::move(download_action));
+
+  loop_.PostTask(
+      FROM_HERE,
+      base::Bind(
+          [](ActionProcessor* processor) { processor->StartProcessing(); },
+          base::Unretained(&processor_)));
+  loop_.Run();
+  EXPECT_FALSE(loop_.PendingTasks());
+
+  EXPECT_EQ(2, p2p_manager_->CountSharedFiles());
+  size_t offset = 0;
+  for (auto& payload : install_plan.payloads) {
+    string file_id = utils::CalculateP2PFileId(payload.hash, payload.size);
+    LOG(INFO) << "test" << file_id;
+    EXPECT_EQ(payload.size, p2p_manager_->FileGetSize(file_id));
+    string file_content;
+    EXPECT_TRUE(
+        ReadFileToString(p2p_manager_->FileGetPath(file_id), &file_content));
+    EXPECT_EQ(data_.substr(0, payload.size), file_content);
+    offset += payload.size;
+  }
+
+  // We don't use the |delegate_| in this test. So just sets it's processing
+  // done to true so it doesn't complain on destruction.
+  delegate_.processing_done_called_ = true;
+}
+
 class RestrictedTimeIntervalDownloadActionTest : public ::testing::Test {
  protected:
   void SetUp() override {