AU: Add support for signing of update payloads after they're generated.

The change adds 2 methods -- one that takes an unsigned payload and a raw
signature size and returns the hash that needs to be signed, and another one
that takes an unsigned payload and a raw signature and generates the signed
payload.

BUG=chromium-os:10872
TEST=unit tests

Change-Id: I65bbe72a1ec67e603e78508c33893695b7de0e6a

Review URL: http://codereview.chromium.org/6265001
diff --git a/delta_diff_generator.cc b/delta_diff_generator.cc
index 4f71c49..0d1e4ce 100644
--- a/delta_diff_generator.cc
+++ b/delta_diff_generator.cc
@@ -1426,26 +1426,11 @@
   // Signatures appear at the end of the blobs. Note the offset in the
   // manifest
   if (!private_key_path.empty()) {
-    LOG(INFO) << "Making room for signature in file";
-    manifest.set_signatures_offset(next_blob_offset);
-    LOG(INFO) << "set? " << manifest.has_signatures_offset();
-    // Add a dummy op at the end to appease older clients
-    DeltaArchiveManifest_InstallOperation* dummy_op =
-        manifest.add_kernel_install_operations();
-    dummy_op->set_type(DeltaArchiveManifest_InstallOperation_Type_REPLACE);
-    dummy_op->set_data_offset(next_blob_offset);
-    manifest.set_signatures_offset(next_blob_offset);
     uint64_t signature_blob_length = 0;
     TEST_AND_RETURN_FALSE(
         PayloadSigner::SignatureBlobLength(private_key_path,
                                            &signature_blob_length));
-    dummy_op->set_data_length(signature_blob_length);
-    manifest.set_signatures_size(signature_blob_length);
-    Extent* dummy_extent = dummy_op->add_dst_extents();
-    // Tell the dummy op to write this data to a big sparse hole
-    dummy_extent->set_start_block(kSparseHole);
-    dummy_extent->set_num_blocks((signature_blob_length + kBlockSize - 1) /
-                                 kBlockSize);
+    AddSignatureOp(next_blob_offset, signature_blob_length, &manifest);
   }
 
   TEST_AND_RETURN_FALSE(InitializePartitionInfos(old_kernel_part,
@@ -1597,6 +1582,27 @@
   return true;
 }
 
+void DeltaDiffGenerator::AddSignatureOp(uint64_t signature_blob_offset,
+                                        uint64_t signature_blob_length,
+                                        DeltaArchiveManifest* manifest) {
+  LOG(INFO) << "Making room for signature in file";
+  manifest->set_signatures_offset(signature_blob_offset);
+  LOG(INFO) << "set? " << manifest->has_signatures_offset();
+  // Add a dummy op at the end to appease older clients
+  DeltaArchiveManifest_InstallOperation* dummy_op =
+      manifest->add_kernel_install_operations();
+  dummy_op->set_type(DeltaArchiveManifest_InstallOperation_Type_REPLACE);
+  dummy_op->set_data_offset(signature_blob_offset);
+  manifest->set_signatures_offset(signature_blob_offset);
+  dummy_op->set_data_length(signature_blob_length);
+  manifest->set_signatures_size(signature_blob_length);
+  Extent* dummy_extent = dummy_op->add_dst_extents();
+  // Tell the dummy op to write this data to a big sparse hole
+  dummy_extent->set_start_block(kSparseHole);
+  dummy_extent->set_num_blocks((signature_blob_length + kBlockSize - 1) /
+                               kBlockSize);
+}
+
 const char* const kBsdiffPath = "bsdiff";
 const char* const kBspatchPath = "bspatch";
 const char* const kDeltaMagic = "CrAU";
diff --git a/delta_diff_generator.h b/delta_diff_generator.h
index d02de48..79683e5 100644
--- a/delta_diff_generator.h
+++ b/delta_diff_generator.h
@@ -232,6 +232,12 @@
       Vertex::Index vertex,
       std::vector<DeltaDiffGenerator::Block>* blocks);
 
+  // Adds to |manifest| a dummy operation that points to a signature blob
+  // located at the specified offset/length.
+  static void AddSignatureOp(uint64_t signature_blob_offset,
+                             uint64_t signature_blob_length,
+                             DeltaArchiveManifest* manifest);
+
  private:
   // This should never be constructed
   DISALLOW_IMPLICIT_CONSTRUCTORS(DeltaDiffGenerator);
diff --git a/delta_performer.cc b/delta_performer.cc
index 6eb8aea..117e8f5 100644
--- a/delta_performer.cc
+++ b/delta_performer.cc
@@ -194,6 +194,43 @@
 
 }  // namespace {}
 
+DeltaPerformer::MetadataParseResult DeltaPerformer::ParsePayloadMetadata(
+    const std::vector<char>& payload,
+    DeltaArchiveManifest* manifest,
+    uint64_t* metadata_size) {
+  if (payload.size() < strlen(kDeltaMagic) +
+      kDeltaVersionLength + kDeltaProtobufLengthLength) {
+    // Don't have enough bytes to know the protobuf length.
+    return kMetadataParseInsufficientData;
+  }
+  if (memcmp(payload.data(), kDeltaMagic, strlen(kDeltaMagic)) != 0) {
+    LOG(ERROR) << "Bad payload format -- invalid delta magic.";
+    return kMetadataParseError;
+  }
+  uint64_t protobuf_length;
+  COMPILE_ASSERT(sizeof(protobuf_length) == kDeltaProtobufLengthLength,
+                 protobuf_length_size_mismatch);
+  memcpy(&protobuf_length,
+         &payload[strlen(kDeltaMagic) + kDeltaVersionLength],
+         kDeltaProtobufLengthLength);
+  protobuf_length = be64toh(protobuf_length);  // switch big endian to host
+  if (payload.size() < strlen(kDeltaMagic) + kDeltaVersionLength +
+      kDeltaProtobufLengthLength + protobuf_length) {
+    return kMetadataParseInsufficientData;
+  }
+  // We have the full proto buffer in |payload|. Parse it.
+  const int offset = strlen(kDeltaMagic) + kDeltaVersionLength +
+      kDeltaProtobufLengthLength;
+  if (!manifest->ParseFromArray(&payload[offset], protobuf_length)) {
+    LOG(ERROR) << "Unable to parse manifest in update file.";
+    return kMetadataParseError;
+  }
+  *metadata_size = strlen(kDeltaMagic) + kDeltaVersionLength +
+      kDeltaProtobufLengthLength + protobuf_length;
+  return kMetadataParseSuccess;
+}
+
+
 // Wrapper around write. Returns bytes written on success or
 // -errno on error.
 // This function performs as many actions as it can, given the amount of
@@ -203,37 +240,17 @@
   buffer_.insert(buffer_.end(), c_bytes, c_bytes + count);
 
   if (!manifest_valid_) {
-    if (buffer_.size() < strlen(kDeltaMagic) +
-        kDeltaVersionLength + kDeltaProtobufLengthLength) {
-      // Don't have enough bytes to know the protobuf length.
-      return count;
-    }
-    if (memcmp(buffer_.data(), kDeltaMagic, strlen(kDeltaMagic)) != 0) {
-      LOG(ERROR) << "Bad payload format -- invalid delta magic.";
+    MetadataParseResult result = ParsePayloadMetadata(buffer_,
+                                                      &manifest_,
+                                                      &manifest_metadata_size_);
+    if (result == kMetadataParseError) {
       return -EINVAL;
     }
-    uint64_t protobuf_length;
-    COMPILE_ASSERT(sizeof(protobuf_length) == kDeltaProtobufLengthLength,
-                   protobuf_length_size_mismatch);
-    memcpy(&protobuf_length,
-           &buffer_[strlen(kDeltaMagic) + kDeltaVersionLength],
-           kDeltaProtobufLengthLength);
-    protobuf_length = be64toh(protobuf_length);  // switch big endian to host
-    if (buffer_.size() < strlen(kDeltaMagic) + kDeltaVersionLength +
-        kDeltaProtobufLengthLength + protobuf_length) {
+    if (result == kMetadataParseInsufficientData) {
       return count;
     }
-    // We have the full proto buffer in buffer_. Parse it.
-    const int offset = strlen(kDeltaMagic) + kDeltaVersionLength +
-        kDeltaProtobufLengthLength;
-    if (!manifest_.ParseFromArray(&buffer_[offset], protobuf_length)) {
-      LOG(ERROR) << "Unable to parse manifest in update file.";
-      return -EINVAL;
-    }
     // Remove protobuf and header info from buffer_, so buffer_ contains
     // just data blobs
-    manifest_metadata_size_ = strlen(kDeltaMagic) + kDeltaVersionLength +
-        kDeltaProtobufLengthLength + protobuf_length;
     DiscardBufferHeadBytes(manifest_metadata_size_);
     LOG_IF(WARNING, !prefs_->SetInt64(kPrefsManifestMetadataSize,
                                       manifest_metadata_size_))
diff --git a/delta_performer.h b/delta_performer.h
index 8a9914d..c693333 100644
--- a/delta_performer.h
+++ b/delta_performer.h
@@ -25,6 +25,12 @@
 
 class DeltaPerformer : public FileWriter {
  public:
+  enum MetadataParseResult {
+    kMetadataParseSuccess,
+    kMetadataParseError,
+    kMetadataParseInsufficientData,
+  };
+
   DeltaPerformer(PrefsInterface* prefs)
       : prefs_(prefs),
         fd_(-1),
@@ -100,6 +106,17 @@
   // success, false otherwise.
   static bool ResetUpdateProgress(PrefsInterface* prefs, bool quick);
 
+  // Attempts to parse the update metadata starting from the beginning of
+  // |payload| into |manifest|. On success, sets |metadata_size| to the total
+  // metadata bytes (including the delta magic and metadata size fields), and
+  // returns kMetadataParseSuccess. Returns kMetadataParseInsufficientData if
+  // more data is needed to parse the complete metadata. Returns
+  // kMetadataParseError if the metadata can't be parsed given the payload.
+  static MetadataParseResult ParsePayloadMetadata(
+      const std::vector<char>& payload,
+      DeltaArchiveManifest* manifest,
+      uint64_t* metadata_size);
+
   void set_current_kernel_hash(const std::vector<char>& hash) {
     current_kernel_hash_ = hash;
   }
diff --git a/delta_performer_unittest.cc b/delta_performer_unittest.cc
index 51299fb..0d71a6e 100755
--- a/delta_performer_unittest.cc
+++ b/delta_performer_unittest.cc
@@ -92,8 +92,11 @@
   TEST_AND_RETURN_FALSE_ERRNO(return_code == 0);
   return true;
 }
+}  // namespace {}
 
-void DoSmallImageTest(bool full_kernel, bool full_rootfs, bool noop) {
+namespace {
+void DoSmallImageTest(bool full_kernel, bool full_rootfs, bool noop,
+                      bool post_sign) {
   string a_img, b_img;
   EXPECT_TRUE(utils::MakeTempFile("/tmp/a_img.XXXXXX", &a_img, NULL));
   ScopedPathUnlinker a_img_unlinker(a_img);
@@ -208,7 +211,33 @@
             full_kernel ? "" : old_kernel,
             new_kernel,
             delta_path,
-            kUnittestPrivateKeyPath));
+            post_sign ? "" : kUnittestPrivateKeyPath));
+  }
+
+  if (post_sign) {
+    int signature_size;
+    {
+      const vector<char> data(1, 'x');
+      vector<char> hash;
+      ASSERT_TRUE(OmahaHashCalculator::RawHashOfData(data, &hash));
+      vector<char> signature;
+      ASSERT_TRUE(PayloadSigner::SignHash(hash,
+                                          kUnittestPrivateKeyPath,
+                                          &signature));
+      signature_size = signature.size();
+    }
+
+    vector<char> hash;
+    ASSERT_TRUE(PayloadSigner::HashPayloadForSigning(delta_path,
+                                                     signature_size,
+                                                     &hash));
+    vector<char> signature;
+    ASSERT_TRUE(PayloadSigner::SignHash(hash,
+                                        kUnittestPrivateKeyPath,
+                                        &signature));
+    ASSERT_TRUE(PayloadSigner::AddSignatureToPayload(delta_path,
+                                                     signature,
+                                                     delta_path));
   }
 
   // Read delta into memory.
@@ -351,19 +380,23 @@
 }
 
 TEST(DeltaPerformerTest, RunAsRootSmallImageTest) {
-  DoSmallImageTest(false, false, false);
+  DoSmallImageTest(false, false, false, false);
 }
 
 TEST(DeltaPerformerTest, RunAsRootFullKernelSmallImageTest) {
-  DoSmallImageTest(true, false, false);
+  DoSmallImageTest(true, false, false, false);
 }
 
 TEST(DeltaPerformerTest, RunAsRootFullSmallImageTest) {
-  DoSmallImageTest(true, true, false);
+  DoSmallImageTest(true, true, false, false);
 }
 
 TEST(DeltaPerformerTest, RunAsRootNoopSmallImageTest) {
-  DoSmallImageTest(false, false, true);
+  DoSmallImageTest(false, false, true, false);
+}
+
+TEST(DeltaPerformerTest, RunAsRootSmallImagePostSignTest) {
+  DoSmallImageTest(false, false, false, true);
 }
 
 TEST(DeltaPerformerTest, BadDeltaMagicTest) {
diff --git a/payload_signer.cc b/payload_signer.cc
index 40bfcfc..cb113da 100644
--- a/payload_signer.cc
+++ b/payload_signer.cc
@@ -8,6 +8,8 @@
 #include <base/string_util.h>
 #include <openssl/pem.h>
 
+#include "update_engine/delta_diff_generator.h"
+#include "update_engine/delta_performer.h"
 #include "update_engine/omaha_hash_calculator.h"
 #include "update_engine/subprocess.h"
 #include "update_engine/update_metadata.pb.h"
@@ -20,9 +22,76 @@
 
 const uint32_t kSignatureMessageVersion = 1;
 
-bool PayloadSigner::SignPayload(const string& unsigned_payload_path,
-                                const string& private_key_path,
-                                vector<char>* out_signature_blob) {
+namespace {
+// Given a raw |signature|, packs it into a protobuf and serializes it into a
+// binary blob. Returns true on success, false otherwise.
+bool ConvertSignatureToProtobufBlob(const vector<char> signature,
+                                    vector<char>* out_signature_blob) {
+  // Pack it into a protobuf
+  Signatures out_message;
+  Signatures_Signature* sig_message = out_message.add_signatures();
+  sig_message->set_version(kSignatureMessageVersion);
+  sig_message->set_data(signature.data(), signature.size());
+
+  // Serialize protobuf
+  string serialized;
+  TEST_AND_RETURN_FALSE(out_message.AppendToString(&serialized));
+  out_signature_blob->insert(out_signature_blob->end(),
+                             serialized.begin(),
+                             serialized.end());
+  LOG(INFO) << "Signature blob size: " << out_signature_blob->size();
+  return true;
+}
+
+// Given an unsigned payload under |payload_path| and the |signature_blob_size|
+// generates an updated payload that includes a dummy signature op in its
+// manifest. Returns true on success, false otherwise.
+bool AddSignatureOpToPayload(const std::string& payload_path,
+                             int signature_blob_size,
+                             vector<char>* out_payload) {
+  const int kProtobufOffset = 20;
+  const int kProtobufSizeOffset = 12;
+
+  vector<char> payload;
+  // Loads the payload and parses the manifest.
+  TEST_AND_RETURN_FALSE(utils::ReadFile(payload_path, &payload));
+  LOG(INFO) << "Original payload size: " << payload.size();
+  uint64_t metadata_size;
+  DeltaArchiveManifest manifest;
+  TEST_AND_RETURN_FALSE(DeltaPerformer::ParsePayloadMetadata(
+      payload, &manifest, &metadata_size) ==
+                        DeltaPerformer::kMetadataParseSuccess);
+  LOG(INFO) << "Metadata size: " << metadata_size;
+  TEST_AND_RETURN_FALSE(!manifest.has_signatures_offset() &&
+                        !manifest.has_signatures_size());
+
+  // Updates the manifest to include the signature operation.
+  DeltaDiffGenerator::AddSignatureOp(payload.size() - metadata_size,
+                                     signature_blob_size,
+                                     &manifest);
+
+  // Updates the payload to include the new manifest.
+  string serialized_manifest;
+  TEST_AND_RETURN_FALSE(manifest.AppendToString(&serialized_manifest));
+  LOG(INFO) << "Updated protobuf size: " << serialized_manifest.size();
+  payload.erase(payload.begin() + kProtobufOffset,
+                payload.begin() + metadata_size);
+  payload.insert(payload.begin() + kProtobufOffset,
+                 serialized_manifest.begin(),
+                 serialized_manifest.end());
+
+  // Updates the protobuf size.
+  uint64_t size_be = htobe64(serialized_manifest.size());
+  memcpy(&payload[kProtobufSizeOffset], &size_be, sizeof(size_be));
+  LOG(INFO) << "Updated payload size: " << payload.size();
+  out_payload->swap(payload);
+  return true;
+}
+}  // namespace {}
+
+bool PayloadSigner::SignHash(const vector<char>& hash,
+                             const string& private_key_path,
+                             vector<char>* out_signature) {
   string sig_path;
   TEST_AND_RETURN_FALSE(
       utils::MakeTempFile("/tmp/signature.XXXXXX", &sig_path, NULL));
@@ -32,18 +101,9 @@
   TEST_AND_RETURN_FALSE(
       utils::MakeTempFile("/tmp/hash.XXXXXX", &hash_path, NULL));
   ScopedPathUnlinker hash_path_unlinker(hash_path);
-
-  vector<char> hash_data;
-  {
-    vector<char> payload;
-    // TODO(adlr): Read file in chunks. Not urgent as this runs on the server.
-    TEST_AND_RETURN_FALSE(utils::ReadFile(unsigned_payload_path, &payload));
-    TEST_AND_RETURN_FALSE(OmahaHashCalculator::RawHashOfData(payload,
-                                                             &hash_data));
-  }
   TEST_AND_RETURN_FALSE(utils::WriteFile(hash_path.c_str(),
-                                         &hash_data[0],
-                                         hash_data.size()));
+                                         hash.data(),
+                                         hash.size()));
 
   // This runs on the server, so it's okay to cop out and call openssl
   // executable rather than properly use the library
@@ -61,19 +121,22 @@
 
   vector<char> signature;
   TEST_AND_RETURN_FALSE(utils::ReadFile(sig_path, &signature));
+  out_signature->swap(signature);
+  return true;
+}
 
-  // Pack it into a protobuf
-  Signatures out_message;
-  Signatures_Signature* sig_message = out_message.add_signatures();
-  sig_message->set_version(kSignatureMessageVersion);
-  sig_message->set_data(signature.data(), signature.size());
+bool PayloadSigner::SignPayload(const string& unsigned_payload_path,
+                                const string& private_key_path,
+                                vector<char>* out_signature_blob) {
+  vector<char> hash_data;
+  TEST_AND_RETURN_FALSE(OmahaHashCalculator::RawHashOfFile(
+      unsigned_payload_path, -1, &hash_data) ==
+                        utils::FileSize(unsigned_payload_path));
 
-  // Serialize protobuf
-  string serialized;
-  TEST_AND_RETURN_FALSE(out_message.AppendToString(&serialized));
-  out_signature_blob->insert(out_signature_blob->end(),
-                             serialized.begin(),
-                             serialized.end());
+  vector<char> signature;
+  TEST_AND_RETURN_FALSE(SignHash(hash_data, private_key_path, &signature));
+  TEST_AND_RETURN_FALSE(ConvertSignatureToProtobufBlob(signature,
+                                                       out_signature_blob));
   return true;
 }
 
@@ -154,4 +217,48 @@
   return true;
 }
 
+bool PayloadSigner::HashPayloadForSigning(const std::string& payload_path,
+                                          int signature_size,
+                                          vector<char>* out_hash_data) {
+  // TODO(petkov): Reduce memory usage -- the payload is manipulated in memory.
+
+  // Loads the payload and adds the signature op to it.
+  vector<char> signature(signature_size, 0);
+  vector<char> signature_blob;
+  TEST_AND_RETURN_FALSE(ConvertSignatureToProtobufBlob(signature,
+                                                       &signature_blob));
+  vector<char> payload;
+  TEST_AND_RETURN_FALSE(AddSignatureOpToPayload(payload_path,
+                                                signature_blob.size(),
+                                                &payload));
+  // Calculates the hash on the updated payload. Note that the payload includes
+  // the signature op but doesn't include the signature blob at the end.
+  TEST_AND_RETURN_FALSE(OmahaHashCalculator::RawHashOfData(payload,
+                                                           out_hash_data));
+  return true;
+}
+
+bool PayloadSigner::AddSignatureToPayload(const string& payload_path,
+                                          const vector<char>& signature,
+                                          const string& signed_payload_path) {
+  // TODO(petkov): Reduce memory usage -- the payload is manipulated in memory.
+
+  // Loads the payload and adds the signature op to it.
+  vector<char> signature_blob;
+  TEST_AND_RETURN_FALSE(ConvertSignatureToProtobufBlob(signature,
+                                                       &signature_blob));
+  vector<char> payload;
+  TEST_AND_RETURN_FALSE(AddSignatureOpToPayload(payload_path,
+                                                signature_blob.size(),
+                                                &payload));
+  // Appends the signature blob to the end of the payload and writes the new
+  // payload.
+  payload.insert(payload.end(), signature_blob.begin(), signature_blob.end());
+  LOG(INFO) << "Signed payload size: " << payload.size();
+  TEST_AND_RETURN_FALSE(utils::WriteFile(signed_payload_path.c_str(),
+                                         payload.data(),
+                                         payload.size()));
+  return true;
+}
+
 }  // namespace chromeos_update_engine
diff --git a/payload_signer.h b/payload_signer.h
index 59affd3..5ebba19 100644
--- a/payload_signer.h
+++ b/payload_signer.h
@@ -7,11 +7,11 @@
 
 #include <string>
 #include <vector>
-#include "base/basictypes.h"
 
-// This function signs a payload with the OS vendor's private key.
-// It takes an update up to the signature blob and returns the signature
-// blob, which should be appended. See update_metadata.proto for more info.
+#include <base/basictypes.h>
+
+// This class encapsulates methods used for payload signing and signature
+// verification. See update_metadata.proto for more info.
 
 namespace chromeos_update_engine {
 
@@ -19,6 +19,17 @@
 
 class PayloadSigner {
  public:
+  // Given a raw |hash| and a private key in |private_key_path| calculates the
+  // raw signature in |out_signature|. Returns true on success, false otherwise.
+  static bool SignHash(const std::vector<char>& hash,
+                       const std::string& private_key_path,
+                       std::vector<char>* out_signature);
+
+  // Given an unsigned payload in |unsigned_payload_path| and a private key in
+  // |private_key_path|, calculates the signature blob into
+  // |out_signature_blob|. Note that the payload must already have an updated
+  // manifest that includes the dummy signature op. Returns true on success,
+  // false otherwise.
   static bool SignPayload(const std::string& unsigned_payload_path,
                           const std::string& private_key_path,
                           std::vector<char>* out_signature_blob);
@@ -34,6 +45,23 @@
                               const std::string& public_key_path,
                               std::vector<char>* out_hash_data);
 
+
+  // Given an unsigned payload in |payload_path| (with no dummy signature op)
+  // and the raw |signature_size| calculates the raw hash that needs to be
+  // signed in |out_hash_data|. Returns true on success, false otherwise.
+  static bool HashPayloadForSigning(const std::string& payload_path,
+                                    int signature_size,
+                                    std::vector<char>* out_hash_data);
+
+  // Given an unsigned payload in |payload_path| (with no dummy signature op)
+  // and the raw |signature| updates the payload to include the signature thus
+  // turning it into a signed payload. The new payload is stored in
+  // |signed_payload_path|. |payload_path| and |signed_payload_path| can point
+  // to the same file. Returns true on success, false otherwise.
+  static bool AddSignatureToPayload(const std::string& payload_path,
+                                    const std::vector<char>& signature,
+                                    const std::string& signed_payload_path);
+
  private:
   // This should never be constructed
   DISALLOW_IMPLICIT_CONSTRUCTORS(PayloadSigner);