update_engine_client_android: Add --allocate option

Test: update_engine_client --allocate --header=... --metadata=...
Test: run it twice with the same hash in header; the second time
      does nothing.
Test: run it twice with different hash; the second time re-allocates
Test: apply OTA after allocation.

Bug: 138808058

Change-Id: I3b125bd7e3e340f2f46dd3c58f5b9128ba70e045
diff --git a/update_engine_client_android.cc b/update_engine_client_android.cc
index 6863799..7766f79 100644
--- a/update_engine_client_android.cc
+++ b/update_engine_client_android.cc
@@ -72,6 +72,8 @@
   // Called whenever the UpdateEngine daemon dies.
   void UpdateEngineServiceDied();
 
+  static std::vector<android::String16> ParseHeaders(const std::string& arg);
+
   // Copy of argc and argv passed to main().
   int argc_;
   char** argv_;
@@ -123,15 +125,16 @@
   DEFINE_string(headers,
                 "",
                 "A list of key-value pairs, one element of the list per line. "
-                "Used when --update is passed.");
+                "Used when --update or --allocate is passed.");
 
   DEFINE_bool(verify,
               false,
               "Given payload metadata, verify if the payload is applicable.");
+  DEFINE_bool(allocate, false, "Given payload metadata, allocate space.");
   DEFINE_string(metadata,
                 "/data/ota_package/metadata",
                 "The path to the update payload metadata. "
-                "Used when --verify is passed.");
+                "Used when --verify or --allocate is passed.");
 
   DEFINE_bool(suspend, false, "Suspend an ongoing update and exit.");
   DEFINE_bool(resume, false, "Resume a suspended update.");
@@ -200,6 +203,25 @@
     return ExitWhenIdle(status);
   }
 
+  if (FLAGS_allocate) {
+    auto headers = ParseHeaders(FLAGS_headers);
+    int64_t ret = 0;
+    Status status = service_->allocateSpaceForPayload(
+        android::String16{FLAGS_metadata.data(), FLAGS_metadata.size()},
+        headers,
+        &ret);
+    if (status.isOk()) {
+      if (ret == 0) {
+        LOG(INFO) << "Successfully allocated space for payload.";
+      } else {
+        LOG(INFO) << "Insufficient space; required " << ret << " bytes.";
+      }
+    } else {
+      LOG(INFO) << "Allocation failed.";
+    }
+    return ExitWhenIdle(status);
+  }
+
   if (FLAGS_follow) {
     // Register a callback object with the service.
     callback_ = new UECallback(this);
@@ -212,12 +234,7 @@
   }
 
   if (FLAGS_update) {
-    std::vector<std::string> headers = base::SplitString(
-        FLAGS_headers, "\n", base::KEEP_WHITESPACE, base::SPLIT_WANT_NONEMPTY);
-    std::vector<android::String16> and_headers;
-    for (const auto& header : headers) {
-      and_headers.push_back(android::String16{header.data(), header.size()});
-    }
+    auto and_headers = ParseHeaders(FLAGS_headers);
     Status status = service_->applyPayload(
         android::String16{FLAGS_payload.data(), FLAGS_payload.size()},
         FLAGS_offset,
@@ -261,6 +278,17 @@
   QuitWithExitCode(1);
 }
 
+std::vector<android::String16> UpdateEngineClientAndroid::ParseHeaders(
+    const std::string& arg) {
+  std::vector<std::string> headers = base::SplitString(
+      arg, "\n", base::KEEP_WHITESPACE, base::SPLIT_WANT_NONEMPTY);
+  std::vector<android::String16> and_headers;
+  for (const auto& header : headers) {
+    and_headers.push_back(android::String16{header.data(), header.size()});
+  }
+  return and_headers;
+}
+
 }  // namespace internal
 }  // namespace chromeos_update_engine