Allow update_engine to communicate with apexd for size calculation

In this CL, we created a ApexHandlerAndroid that can communicate with
apexd via binders to get better estimate for how much space the
update_engine should reserve for capex decompression.

The size check is placed in update_attempter_android, which is also used
in binary for sideloading OTA. Sideloading binary runs during recovery
when binder calls, apexd and data parition are all not available. As
such, it doesn't make sense to reserve space when sideloading.

Bug: 172911822
Test: atest ApexHandlerAndroidTest
Test: atest UpdateAttempterAndroidTest
Test: manually served ota using `update_engine_client --allocate` and
      observed full.tmp was written in /data/apex/ota_reserved
Change-Id: Iccf3d8c2db24e8d8f3406d0aaa65cbf707c9ae51
diff --git a/Android.bp b/Android.bp
index 4ec17f2..1658533 100644
--- a/Android.bp
+++ b/Android.bp
@@ -332,6 +332,7 @@
         "PlatformProperties",
     ],
     shared_libs: [
+        "apex_aidl_interface-cpp",
         "libandroid_net",
         "libbase",
         "libbinder",
@@ -367,6 +368,7 @@
     srcs: [
         ":libupdate_engine_aidl",
         "common/system_state.cc",
+        "aosp/apex_handler_android.cc",
         "aosp/binder_service_android.cc",
         "aosp/binder_service_stable_android.cc",
         "aosp/daemon_android.cc",
@@ -740,6 +742,7 @@
     test_suites: ["device-tests"],
 
     srcs: [
+        "aosp/apex_handler_android_unittest.cc",
         "aosp/dynamic_partition_control_android_unittest.cc",
         "aosp/update_attempter_android_unittest.cc",
         "certificate_checker_unittest.cc",
diff --git a/aosp/apex_handler_android.cc b/aosp/apex_handler_android.cc
new file mode 100644
index 0000000..cdbc983
--- /dev/null
+++ b/aosp/apex_handler_android.cc
@@ -0,0 +1,98 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include <utility>
+
+#include <base/files/file_util.h>
+
+#include "update_engine/aosp/apex_handler_android.h"
+#include "update_engine/common/utils.h"
+
+namespace chromeos_update_engine {
+
+// Don't change this path... apexd relies on it.
+constexpr const char* kApexReserveSpaceDir = "/data/apex/ota_reserved";
+
+uint64_t ApexHandlerAndroid::CalculateSize(
+    const std::vector<ApexInfo>& apex_infos) const {
+  return CalculateSize(apex_infos, GetApexService());
+}
+
+uint64_t ApexHandlerAndroid::CalculateSize(
+    const std::vector<ApexInfo>& apex_infos,
+    android::sp<android::apex::IApexService> apex_service) const {
+  // The safest option is to allocate space for every compressed APEX
+  uint64_t size_required_default = 0;
+
+  // We might not need to decompress every APEX. Communicate with apexd to get
+  // accurate requirement.
+  int64_t size_from_apexd;
+  android::apex::CompressedApexInfoList compressed_apex_info_list;
+
+  for (const auto& apex_info : apex_infos) {
+    if (!apex_info.is_compressed()) {
+      continue;
+    }
+
+    size_required_default += apex_info.decompressed_size();
+
+    android::apex::CompressedApexInfo compressed_apex_info;
+    compressed_apex_info.moduleName = apex_info.package_name();
+    compressed_apex_info.versionCode = apex_info.version();
+    compressed_apex_info.decompressedSize = apex_info.decompressed_size();
+    compressed_apex_info_list.apexInfos.emplace_back(
+        std::move(compressed_apex_info));
+  }
+  if (size_required_default == 0 || apex_service == nullptr) {
+    return size_required_default;
+  }
+
+  auto result = apex_service->calculateSizeForCompressedApex(
+      compressed_apex_info_list, &size_from_apexd);
+  if (!result.isOk()) {
+    return size_required_default;
+  }
+  return size_from_apexd;
+}
+
+bool ApexHandlerAndroid::AllocateSpace(const uint64_t size_required) const {
+  return AllocateSpace(size_required, kApexReserveSpaceDir);
+}
+
+bool ApexHandlerAndroid::AllocateSpace(const uint64_t size_required,
+                                       const std::string& dir_path) const {
+  if (size_required == 0) {
+    return true;
+  }
+  base::FilePath path{dir_path};
+  // The filename is not important, it just needs to be under
+  // kApexReserveSpaceDir. We call it "full.tmp" because the current space
+  // estimation is simply adding up all decompressed sizes.
+  path = path.Append("full.tmp");
+  return utils::ReserveStorageSpace(path.value().c_str(), size_required);
+}
+
+android::sp<android::apex::IApexService> ApexHandlerAndroid::GetApexService()
+    const {
+  auto binder = android::defaultServiceManager()->waitForService(
+      android::String16("apexservice"));
+  if (binder == nullptr) {
+    return nullptr;
+  }
+  return android::interface_cast<android::apex::IApexService>(binder);
+}
+
+}  // namespace chromeos_update_engine
diff --git a/aosp/apex_handler_android.h b/aosp/apex_handler_android.h
new file mode 100644
index 0000000..aac1cd9
--- /dev/null
+++ b/aosp/apex_handler_android.h
@@ -0,0 +1,48 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#ifndef SYSTEM_UPDATE_ENGINE_AOSP_APEX_HANDLER_ANDROID_H_
+#define SYSTEM_UPDATE_ENGINE_AOSP_APEX_HANDLER_ANDROID_H_
+
+#include <string>
+#include <vector>
+
+#include <android/apex/IApexService.h>
+#include <binder/IServiceManager.h>
+
+#include "update_engine/aosp/apex_handler_interface.h"
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+class ApexHandlerAndroid : virtual public ApexHandlerInterface {
+ public:
+  uint64_t CalculateSize(const std::vector<ApexInfo>& apex_infos) const;
+  bool AllocateSpace(const uint64_t size_required) const;
+
+ private:
+  friend class ApexHandlerAndroidTest;
+  android::sp<android::apex::IApexService> GetApexService() const;
+  uint64_t CalculateSize(
+      const std::vector<ApexInfo>& apex_infos,
+      android::sp<android::apex::IApexService> apex_service) const;
+  bool AllocateSpace(const uint64_t size_required,
+                     const std::string& dir_path) const;
+};
+
+}  // namespace chromeos_update_engine
+
+#endif  // SYSTEM_UPDATE_ENGINE_AOSP_APEX_HANDLER_ANDROID_H_
diff --git a/aosp/apex_handler_android_unittest.cc b/aosp/apex_handler_android_unittest.cc
new file mode 100644
index 0000000..3a99f79
--- /dev/null
+++ b/aosp/apex_handler_android_unittest.cc
@@ -0,0 +1,109 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include <utility>
+#include <filesystem>
+
+#include "update_engine/aosp/apex_handler_android.h"
+
+#include <android-base/file.h>
+#include <android-base/strings.h>
+#include <gtest/gtest.h>
+
+using android::base::EndsWith;
+
+namespace chromeos_update_engine {
+
+namespace fs = std::filesystem;
+
+class ApexHandlerAndroidTest : public ::testing::Test {
+ protected:
+  ApexHandlerAndroidTest() = default;
+
+  android::sp<android::apex::IApexService> GetApexService() const {
+    return apex_handler_.GetApexService();
+  }
+
+  uint64_t CalculateSize(
+      const std::vector<ApexInfo>& apex_infos,
+      android::sp<android::apex::IApexService> apex_service) const {
+    return apex_handler_.CalculateSize(apex_infos, apex_service);
+  }
+
+  bool AllocateSpace(const uint64_t size_required,
+                     const std::string& dir_path) const {
+    return apex_handler_.AllocateSpace(size_required, dir_path);
+  }
+
+  ApexInfo CreateApexInfo(const std::string& package_name,
+                          int version,
+                          bool is_compressed,
+                          int decompressed_size) {
+    ApexInfo result;
+    result.set_package_name(package_name);
+    result.set_version(version);
+    result.set_is_compressed(is_compressed);
+    result.set_decompressed_size(decompressed_size);
+    return std::move(result);
+  }
+
+  ApexHandlerAndroid apex_handler_;
+};
+
+// TODO(b/172911822): Once apexd has more optimized response for CalculateSize,
+//  improve this test
+TEST_F(ApexHandlerAndroidTest, CalculateSize) {
+  std::vector<ApexInfo> apex_infos;
+  ApexInfo compressed_apex_1 = CreateApexInfo("sample1", 1, true, 10);
+  ApexInfo compressed_apex_2 = CreateApexInfo("sample2", 2, true, 20);
+  apex_infos.push_back(compressed_apex_1);
+  apex_infos.push_back(compressed_apex_2);
+  auto apex_service = GetApexService();
+  EXPECT_TRUE(apex_service != nullptr) << "Apexservice not found";
+  int required_size = CalculateSize(apex_infos, apex_service);
+  EXPECT_EQ(required_size, 30);
+}
+
+TEST_F(ApexHandlerAndroidTest, AllocateSpace) {
+  // Allocating 0 space should be a no op
+  TemporaryDir td;
+  EXPECT_TRUE(AllocateSpace(0, td.path));
+  EXPECT_TRUE(fs::is_empty(td.path));
+
+  // Allocating non-zero space should create a file with tmp suffix
+  EXPECT_TRUE(AllocateSpace(2 << 20, td.path));
+  EXPECT_FALSE(fs::is_empty(td.path));
+  int num_of_file = 0;
+  for (const auto& entry : fs::directory_iterator(td.path)) {
+    num_of_file++;
+    EXPECT_TRUE(EndsWith(entry.path().string(), ".tmp"));
+    EXPECT_EQ(fs::file_size(entry.path()), 2u << 20);
+  }
+  EXPECT_EQ(num_of_file, 1);
+
+  // AllocateSpace should be safe to call twice
+  EXPECT_TRUE(AllocateSpace(100, td.path));
+  EXPECT_FALSE(fs::is_empty(td.path));
+  num_of_file = 0;
+  for (const auto& entry : fs::directory_iterator(td.path)) {
+    num_of_file++;
+    EXPECT_TRUE(EndsWith(entry.path().string(), ".tmp"));
+    EXPECT_EQ(fs::file_size(entry.path()), 100u);
+  }
+  EXPECT_EQ(num_of_file, 1);
+}
+
+}  // namespace chromeos_update_engine
diff --git a/aosp/apex_handler_interface.h b/aosp/apex_handler_interface.h
new file mode 100644
index 0000000..c3399b6
--- /dev/null
+++ b/aosp/apex_handler_interface.h
@@ -0,0 +1,36 @@
+//
+// Copyright (C) 2021 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#ifndef SYSTEM_UPDATE_ENGINE_AOSP_APEX_HANDLER_INTERFACE_H_
+#define SYSTEM_UPDATE_ENGINE_AOSP_APEX_HANDLER_INTERFACE_H_
+
+#include <vector>
+
+#include "update_engine/update_metadata.pb.h"
+
+namespace chromeos_update_engine {
+
+class ApexHandlerInterface {
+ public:
+  virtual ~ApexHandlerInterface() = default;
+  virtual uint64_t CalculateSize(
+      const std::vector<ApexInfo>& apex_infos) const = 0;
+  virtual bool AllocateSpace(const uint64_t size_required) const = 0;
+};
+
+}  // namespace chromeos_update_engine
+
+#endif  // SYSTEM_UPDATE_ENGINE_AOSP_APEX_HANDLER_INTERFACE_H_
diff --git a/aosp/daemon_state_android.cc b/aosp/daemon_state_android.cc
index 9bdd175..fc89d73 100644
--- a/aosp/daemon_state_android.cc
+++ b/aosp/daemon_state_android.cc
@@ -18,6 +18,7 @@
 
 #include <base/logging.h>
 
+#include "update_engine/aosp/apex_handler_android.h"
 #include "update_engine/aosp/update_attempter_android.h"
 #include "update_engine/common/boot_control.h"
 #include "update_engine/common/boot_control_stub.h"
@@ -64,8 +65,12 @@
   certificate_checker_->Init();
 
   // Initialize the UpdateAttempter before the UpdateManager.
-  update_attempter_.reset(new UpdateAttempterAndroid(
-      this, prefs_.get(), boot_control_.get(), hardware_.get()));
+  update_attempter_.reset(
+      new UpdateAttempterAndroid(this,
+                                 prefs_.get(),
+                                 boot_control_.get(),
+                                 hardware_.get(),
+                                 std::make_unique<ApexHandlerAndroid>()));
 
   return true;
 }
diff --git a/aosp/sideload_main.cc b/aosp/sideload_main.cc
index 3cbc0c7..bf015c9 100644
--- a/aosp/sideload_main.cc
+++ b/aosp/sideload_main.cc
@@ -154,8 +154,11 @@
     return false;
   }
 
-  UpdateAttempterAndroid update_attempter(
-      &sideload_daemon_state, &prefs, boot_control.get(), hardware.get());
+  UpdateAttempterAndroid update_attempter(&sideload_daemon_state,
+                                          &prefs,
+                                          boot_control.get(),
+                                          hardware.get(),
+                                          nullptr);
   update_attempter.Init();
 
   TEST_AND_RETURN_FALSE(update_attempter.ApplyPayload(
diff --git a/aosp/update_attempter_android.cc b/aosp/update_attempter_android.cc
index 5523a58..c685855 100644
--- a/aosp/update_attempter_android.cc
+++ b/aosp/update_attempter_android.cc
@@ -71,9 +71,6 @@
 
 namespace chromeos_update_engine {
 
-// Don't change this path... apexd relies on it.
-constexpr const char* kApexReserveSpaceDir = "/data/apex/ota_reserved";
-
 namespace {
 
 // Minimum threshold to broadcast an status update in progress and time.
@@ -136,11 +133,13 @@
     DaemonStateInterface* daemon_state,
     PrefsInterface* prefs,
     BootControlInterface* boot_control,
-    HardwareInterface* hardware)
+    HardwareInterface* hardware,
+    std::unique_ptr<ApexHandlerInterface> apex_handler)
     : daemon_state_(daemon_state),
       prefs_(prefs),
       boot_control_(boot_control),
       hardware_(hardware),
+      apex_handler_android_(std::move(apex_handler)),
       processor_(new ActionProcessor()),
       clock_(new Clock()) {
   metrics_reporter_ = metrics::CreateMetricsReporter(
@@ -967,28 +966,6 @@
   return GetCurrentSlot() == 0 ? 1 : 0;
 }
 
-static uint64_t allocateSpaceForApex(const DeltaArchiveManifest& manifest) {
-  // TODO(b/178696931) call apexd's binder once there is one.
-  uint64_t size_required = 0;
-  for (const auto& apex_info : manifest.apex_info()) {
-    if (apex_info.is_compressed()) {
-      size_required += apex_info.decompressed_size();
-    }
-  }
-  if (size_required == 0) {
-    return 0;
-  }
-  base::FilePath path{kApexReserveSpaceDir};
-  // The filename is not important, it just needs to be under
-  // kApexReserveSpaceDir. We call it "full.tmp" because the current space
-  // estimation is simply adding up all decompressed sizes.
-  path = path.Append("full.tmp");
-  if (!utils::ReserveStorageSpace(path.value().c_str(), size_required)) {
-    return size_required;
-  }
-  return 0;
-}
-
 uint64_t UpdateAttempterAndroid::AllocateSpaceForPayload(
     const std::string& metadata_filename,
     const vector<string>& key_value_pair_headers,
@@ -1002,6 +979,13 @@
     return 0;
   }
 
+  std::vector<ApexInfo> apex_infos(manifest.apex_info().begin(),
+                                   manifest.apex_info().end());
+  uint64_t apex_size_required = 0;
+  if (apex_handler_android_ != nullptr) {
+    apex_size_required = apex_handler_android_->CalculateSize(apex_infos);
+  }
+
   string payload_id = GetPayloadId(headers);
   uint64_t required_size = 0;
   if (!DeltaPerformer::PreparePartitionsForUpdate(prefs_,
@@ -1015,14 +999,17 @@
       return 0;
     } else {
       LOG(ERROR) << "Insufficient space for payload: " << required_size
+                 << " bytes, apex decompression: " << apex_size_required
                  << " bytes";
-      return required_size;
+      return required_size + apex_size_required;
     }
   }
-  const auto apex_size = allocateSpaceForApex(manifest);
-  if (apex_size != 0) {
-    LOG(ERROR) << "Insufficient space for apex decompression: " << apex_size;
-    return required_size + apex_size;
+
+  if (apex_size_required > 0 && apex_handler_android_ != nullptr &&
+      !apex_handler_android_->AllocateSpace(apex_size_required)) {
+    LOG(ERROR) << "Insufficient space for apex decompression: "
+               << apex_size_required << " bytes";
+    return apex_size_required;
   }
 
   LOG(INFO) << "Successfully allocated space for payload.";
diff --git a/aosp/update_attempter_android.h b/aosp/update_attempter_android.h
index 499f8f6..70938bc 100644
--- a/aosp/update_attempter_android.h
+++ b/aosp/update_attempter_android.h
@@ -26,6 +26,7 @@
 #include <android-base/unique_fd.h>
 #include <base/time/time.h>
 
+#include "update_engine/aosp/apex_handler_interface.h"
 #include "update_engine/aosp/service_delegate_android_interface.h"
 #include "update_engine/client_library/include/update_engine/update_status.h"
 #include "update_engine/common/action_processor.h"
@@ -57,7 +58,8 @@
   UpdateAttempterAndroid(DaemonStateInterface* daemon_state,
                          PrefsInterface* prefs,
                          BootControlInterface* boot_control_,
-                         HardwareInterface* hardware_);
+                         HardwareInterface* hardware_,
+                         std::unique_ptr<ApexHandlerInterface> apex_handler);
   ~UpdateAttempterAndroid() override;
 
   // Further initialization to be done post construction.
@@ -205,6 +207,8 @@
   BootControlInterface* boot_control_;
   HardwareInterface* hardware_;
 
+  std::unique_ptr<ApexHandlerInterface> apex_handler_android_;
+
   // Last status notification timestamp used for throttling. Use monotonic
   // TimeTicks to ensure that notifications are sent even if the system clock is
   // set back in the middle of an update.
diff --git a/aosp/update_attempter_android_unittest.cc b/aosp/update_attempter_android_unittest.cc
index 173e943..f799df3 100644
--- a/aosp/update_attempter_android_unittest.cc
+++ b/aosp/update_attempter_android_unittest.cc
@@ -69,7 +69,7 @@
   FakeHardware hardware_;
 
   UpdateAttempterAndroid update_attempter_android_{
-      &daemon_state_, &prefs_, &boot_control_, &hardware_};
+      &daemon_state_, &prefs_, &boot_control_, &hardware_, nullptr};
 
   FakeClock* clock_;
   testing::NiceMock<MockMetricsReporter>* metrics_reporter_;