AU: Add options to delta_generator to hash and sign generated payloads.

generate_delta_main.cc

Index: generate_delta_main.cc
diff --git a/generate_delta_main.cc b/generate_delta_main.cc
index 6e9e5530b1f0ac650fcc93f4a0a18aaf71e7b9f5..9fc7fd7786d9c0117f485c0cde9148a9137ac363 100644
--- a/generate_delta_main.cc
+++ b/generate_delta_main.cc
@@ -19,6 +19,7 @@
 
 #include "update_engine/delta_diff_generator.h"
 #include "update_engine/delta_performer.h"
+#include "update_engine/payload_signer.h"
 #include "update_engine/prefs.h"
 #include "update_engine/subprocess.h"
 #include "update_engine/terminator.h"
@@ -31,14 +32,18 @@ DEFINE_string(new_dir, "",
               "Directory where the new rootfs is loop mounted read-only");
 DEFINE_string(old_image, "", "Path to the old rootfs");
 DEFINE_string(new_image, "", "Path to the new rootfs");
-DEFINE_string(out_file, "", "Path to output file");
 DEFINE_string(old_kernel, "", "Path to the old kernel partition image");
 DEFINE_string(new_kernel, "", "Path to the new kernel partition image");
+DEFINE_string(in_file, "",
+              "Path to input delta payload file used to hash/sign payloads "
+              "and apply delta over old_image (for debugging)");
+DEFINE_string(out_file, "", "Path to output delta payload file");
+DEFINE_string(out_hash_file, "", "Path to output hash file");
 DEFINE_string(private_key, "", "Path to private key in .pem format");
-DEFINE_string(apply_delta, "",
-              "If set, apply delta over old_image. (For debugging.)");
 DEFINE_string(prefs_dir, "/tmp/update_engine_prefs",
-              "Preferences directory, used with apply_delta.");
+              "Preferences directory, used with apply_delta");
+DEFINE_int32(signature_size, 0, "Raw signature size used for hash calculation");
+DEFINE_string(signature_file, "", "Raw signature file to sign payload with");
 
 // This file contains a simple program that takes an old path, a new path,
 // and an output file as arguments and the path to an output file and
@@ -58,6 +63,79 @@ bool IsDir(const char* path) {
   return S_ISDIR(stbuf.st_mode);
 }
 
+void CalculatePayloadHashForSigning() {
+  LOG(INFO) << "Calculating payload hash for signing.";
+  LOG_IF(FATAL, FLAGS_in_file.empty())
+      << "Must pass --in_file to calculate hash for signing.";
+  LOG_IF(FATAL, FLAGS_out_hash_file.empty())
+      << "Must pass --out_hash_file to calculate hash for signing.";
+  LOG_IF(FATAL, FLAGS_signature_size <= 0)
+      << "Must pass --signature_size to calculate hash for signing.";
+  vector<char> hash;
+  CHECK(PayloadSigner::HashPayloadForSigning(
+      FLAGS_in_file, FLAGS_signature_size, &hash));
+  CHECK(utils::WriteFile(
+      FLAGS_out_hash_file.c_str(), hash.data(), hash.size()));
+  LOG(INFO) << "Done calculating payload hash for signing.";
+}
+
+void SignPayload() {
+  LOG(INFO) << "Signing payload.";
+  LOG_IF(FATAL, FLAGS_in_file.empty())
+      << "Must pass --in_file to sign payload.";
+  LOG_IF(FATAL, FLAGS_out_file.empty())
+      << "Must pass --out_file to sign payload.";
+  LOG_IF(FATAL, FLAGS_signature_file.empty())
+      << "Must pass --signature_file to sign payload.";
+  vector<char> signature;
+  CHECK(utils::ReadFile(FLAGS_signature_file, &signature));
+  CHECK(PayloadSigner::AddSignatureToPayload(
+      FLAGS_in_file, signature, FLAGS_out_file));
+  LOG(INFO) << "Done signing payload.";
+}
+
+void ApplyDelta() {
+  LOG(INFO) << "Applying delta.";
+  LOG_IF(FATAL, FLAGS_old_image.empty())
+      << "Must pass --old_image to apply delta.";
+  Prefs prefs;
+  LOG(INFO) << "Setting up preferences under: " << FLAGS_prefs_dir;
+  LOG_IF(ERROR, !prefs.Init(FilePath(FLAGS_prefs_dir)))
+      << "Failed to initialize preferences.";
+  // Get original checksums
+  LOG(INFO) << "Calculating original checksums";
+  PartitionInfo kern_info, root_info;
+  CHECK(DeltaDiffGenerator::InitializePartitionInfo(true,  // is_kernel
+                                                    FLAGS_old_kernel,
+                                                    &kern_info));
+  CHECK(DeltaDiffGenerator::InitializePartitionInfo(false,  // is_kernel
+                                                    FLAGS_old_image,
+                                                    &root_info));
+  vector<char> kern_hash(kern_info.hash().begin(),
+                         kern_info.hash().end());
+  vector<char> root_hash(root_info.hash().begin(),
+                         root_info.hash().end());
+  DeltaPerformer performer(&prefs);
+  performer.set_current_kernel_hash(kern_hash);
+  performer.set_current_rootfs_hash(root_hash);
+  CHECK_EQ(performer.Open(FLAGS_old_image.c_str(), 0, 0), 0);
+  CHECK(performer.OpenKernel(FLAGS_old_kernel.c_str()));
+  vector<char> buf(1024 * 1024);
+  int fd = open(FLAGS_in_file.c_str(), O_RDONLY, 0);
+  CHECK_GE(fd, 0);
+  ScopedFdCloser fd_closer(&fd);
+  for (off_t offset = 0;; offset += buf.size()) {
+    ssize_t bytes_read;
+    CHECK(utils::PReadAll(fd, &buf[0], buf.size(), offset, &bytes_read));
+    if (bytes_read == 0)
+      break;
+    CHECK_EQ(performer.Write(&buf[0], bytes_read), bytes_read);
+  }
+  CHECK_EQ(performer.Close(), 0);
+  DeltaPerformer::ResetUpdateProgress(&prefs, false);
+  LOG(INFO) << "Done applying delta.";
+}
+
 int Main(int argc, char** argv) {
   g_thread_init(NULL);
   google::ParseCommandLineFlags(&argc, &argv, true);
@@ -68,46 +146,16 @@ int Main(int argc, char** argv) {
                        logging::LOG_ONLY_TO_SYSTEM_DEBUG_LOG,
                        logging::DONT_LOCK_LOG_FILE,
                        logging::APPEND_TO_OLD_LOG_FILE);
-  if (!FLAGS_apply_delta.empty()) {
-    if (FLAGS_old_image.empty()) {
-      LOG(FATAL) << "Must pass --old_image with --apply_delta.";
-    }
-    Prefs prefs;
-    LOG(INFO) << "Setting up preferences under: " << FLAGS_prefs_dir;
-    LOG_IF(ERROR, !prefs.Init(FilePath(FLAGS_prefs_dir)))
-        << "Failed to initialize preferences.";
-    // Get original checksums
-    LOG(INFO) << "Calculating original checksums";
-    PartitionInfo kern_info, root_info;
-    CHECK(DeltaDiffGenerator::InitializePartitionInfo(true,  // is_kernel
-                                                      FLAGS_old_kernel,
-                                                      &kern_info));
-    CHECK(DeltaDiffGenerator::InitializePartitionInfo(false,  // is_kernel
-                                                      FLAGS_old_image,
-                                                      &root_info));
-    vector<char> kern_hash(kern_info.hash().begin(),
-                           kern_info.hash().end());
-    vector<char> root_hash(root_info.hash().begin(),
-                           root_info.hash().end());
-    DeltaPerformer performer(&prefs);
-    performer.set_current_kernel_hash(kern_hash);
-    performer.set_current_rootfs_hash(root_hash);
-    CHECK_EQ(performer.Open(FLAGS_old_image.c_str(), 0, 0), 0);
-    CHECK(performer.OpenKernel(FLAGS_old_kernel.c_str()));
-    vector<char> buf(1024 * 1024);
-    int fd = open(FLAGS_apply_delta.c_str(), O_RDONLY, 0);
-    CHECK_GE(fd, 0);
-    ScopedFdCloser fd_closer(&fd);
-    for (off_t offset = 0;; offset += buf.size()) {
-      ssize_t bytes_read;
-      CHECK(utils::PReadAll(fd, &buf[0], buf.size(), offset, &bytes_read));
-      if (bytes_read == 0)
-        break;
-      CHECK_EQ(performer.Write(&buf[0], bytes_read), bytes_read);
-    }
-    CHECK_EQ(performer.Close(), 0);
-    DeltaPerformer::ResetUpdateProgress(&prefs, false);
-    LOG(INFO) << "done applying delta.";
+  if (FLAGS_signature_size > 0 || !FLAGS_out_hash_file.empty()) {
+    CalculatePayloadHashForSigning();
+    return 0;
+  }
+  if (!FLAGS_signature_file.empty()) {
+    SignPayload();
+    return 0;
+  }
+  if (!FLAGS_in_file.empty()) {
+    ApplyDelta();
     return 0;
   }
   CHECK(!FLAGS_new_image.empty());