paygen: Generate the payload hash and metadata hash together.

Due to recent changes in delta_generator, payload hash and metadata hash
are calculated together, calling it twice are just doing repeated work.

BUG=None
TEST=./paygen_payload_lib_unittest

Change-Id: I62c6a3621e2da882a5acd8209c17686152810806
Reviewed-on: https://chromium-review.googlesource.com/310226
Commit-Ready: Alex Deymo <deymo@chromium.org>
Tested-by: Sen Jiang <senj@chromium.org>
Reviewed-by: Matthew Sartori <msartori@chromium.org>
Reviewed-by: Alex Deymo <deymo@chromium.org>

2 files changed, 25 insertions, 71 deletions

diff --git a/lib/paygen/paygen_payload_lib.py b/lib/paygen/paygen_payload_lib.py
index fc97be498..b35ab73b7 100644
--- a/lib/paygen/paygen_payload_lib.py
+++ b/lib/paygen/paygen_payload_lib.py
@@ -306,27 +306,29 @@ class _PaygenPayload(object):
     delta_log = self._RunGeneratorCmd(cmd)
     self._StoreDeltaLog(delta_log)
 
-  def _GenPayloadHash(self):
-    """Generate a hash of payload and metadata.
+  def _GenerateHashes(self):
+    """Generate a payload hash and a metadata hash.
 
     Works from an unsigned update payload.
 
     Returns:
-      payload_hash as a string.
+      payload_hash as a string, metadata_hash as a string.
     """
-    logging.info('Calculating payload hashes on %s.', self.payload_file)
+    logging.info('Calculating hashes on %s.', self.payload_file)
 
     # How big will the signatures be.
     signature_sizes = [str(size) for size in self.PAYLOAD_SIGNATURE_SIZES_BYTES]
 
-    with tempfile.NamedTemporaryFile('rb') as payload_hash_file:
-      cmd = ['delta_generator',
-             '-in_file=' + self.payload_file,
-             '-out_hash_file=' + payload_hash_file.name,
-             '-signature_size=' + ':'.join(signature_sizes)]
+    with tempfile.NamedTemporaryFile('rb') as payload_hash_file, \
+         tempfile.NamedTemporaryFile('rb') as metadata_hash_file:
+      cmd = ['brillo_update_payload', 'hash',
+             '--unsigned_payload', self.payload_file,
+             '--payload_hash_file', payload_hash_file.name,
+             '--metadata_hash_file', metadata_hash_file.name,
+             '--signature_size', ':'.join(signature_sizes)]
 
       self._RunGeneratorCmd(cmd)
-      return payload_hash_file.read()
+      return payload_hash_file.read(), metadata_hash_file.read()
 
   def _MetadataSize(self, payload_file):
     """Discover the metadata size.
@@ -346,30 +348,6 @@ class _PaygenPayload(object):
       payload.Init()
       return payload.data_offset
 
-  def _GenMetadataHash(self):
-    """Generate a hash of payload and metadata.
-
-    Works from an unsigned update payload.
-
-    Returns:
-      metadata_hash as a string.
-    """
-    logging.info('Calculating payload hashes on %s.', self.payload_file)
-
-    # How big will the signatures be.
-    signature_sizes = [str(size) for size in self.PAYLOAD_SIGNATURE_SIZES_BYTES]
-
-    # The out_metadata_hash_file flag requires out_hash_file flag to be set.
-    with tempfile.NamedTemporaryFile('rb') as metadata_hash_file:
-      cmd = ['delta_generator',
-             '-in_file=' + self.payload_file,
-             '-out_hash_file=/dev/null',
-             '-out_metadata_hash_file=' + metadata_hash_file.name,
-             '-signature_size=' + ':'.join(signature_sizes)]
-
-      self._RunGeneratorCmd(cmd)
-      return metadata_hash_file.read()
-
   def _GenerateSignerResultsError(self, format_str, *args):
     """Helper for reporting errors with signer results."""
     msg = format_str % args
@@ -540,9 +518,7 @@ class _PaygenPayload(object):
       List of payload signatures, List of metadata signatures.
     """
     # Create hashes to sign.
-    # TODO(senj): Calculate the two hashes in one shot.
-    payload_hash = self._GenPayloadHash()
-    metadata_hash = self._GenMetadataHash()
+    payload_hash, metadata_hash = self._GenerateHashes()
 
     # Sign them.
     # pylint: disable=unpacking-non-sequence
diff --git a/lib/paygen/paygen_payload_lib_unittest.py b/lib/paygen/paygen_payload_lib_unittest.py
index 68daf22b2..8c6d18882 100644
--- a/lib/paygen/paygen_payload_lib_unittest.py
+++ b/lib/paygen/paygen_payload_lib_unittest.py
@@ -125,7 +125,7 @@ class PaygenPayloadLibBasicTest(PaygenPayloadLibTest):
 
     if not au_generator_uri_override:
       au_generator_uri_override = gspaths.ChromeosReleases.GeneratorUri(
-          payload.tgt_image.channel, payload.tgt_image.board, '6351.0.0')
+          payload.tgt_image.channel, payload.tgt_image.board, '7587.0.0')
 
     return paygen_payload_lib._PaygenPayload(
         payload=payload,
@@ -403,8 +403,8 @@ class PaygenPayloadLibBasicTest(PaygenPayloadLibTest):
     self.mox.ReplayAll()
     gen._GenerateUnsignedPayload()
 
-  def testGenPayloadHashes(self):
-    """Test _GenPayloadHash via mox."""
+  def testGenerateHashes(self):
+    """Test _GenerateHashes via mox."""
     gen = self._GetStdGenerator()
 
     # Stub out the required functions.
@@ -412,35 +412,16 @@ class PaygenPayloadLibBasicTest(PaygenPayloadLibTest):
                              '_RunGeneratorCmd')
 
     # Record the expected function calls.
-    cmd = ['delta_generator',
-           '-in_file=' + gen.payload_file,
-           mox.IsA(str),
-           '-signature_size=256']
+    cmd = ['brillo_update_payload', 'hash',
+           '--unsigned_payload', gen.payload_file,
+           '--payload_hash_file', mox.IsA(str),
+           '--metadata_hash_file', mox.IsA(str),
+           '--signature_size', '256']
     gen._RunGeneratorCmd(cmd)
 
     # Run the test.
     self.mox.ReplayAll()
-    self.assertEqual(gen._GenPayloadHash(), '')
-
-  def testGenMetadataHashes(self):
-    """Test _GenPayloadHash via mox."""
-    gen = self._GetStdGenerator()
-
-    # Stub out the required functions.
-    self.mox.StubOutWithMock(paygen_payload_lib._PaygenPayload,
-                             '_RunGeneratorCmd')
-
-    # Record the expected function calls.
-    cmd = ['delta_generator',
-           '-in_file=' + gen.payload_file,
-           '-out_hash_file=/dev/null',
-           mox.IsA(str),
-           '-signature_size=256']
-    gen._RunGeneratorCmd(cmd)
-
-    # Run the test.
-    self.mox.ReplayAll()
-    self.assertEqual(gen._GenMetadataHash(), '')
+    self.assertEqual(gen._GenerateHashes(), ('', ''))
 
   def testSignHashes(self):
     """Test _SignHashes via mox."""
@@ -563,9 +544,7 @@ class PaygenPayloadLibBasicTest(PaygenPayloadLibTest):
 
     # Set up stubs.
     self.mox.StubOutWithMock(paygen_payload_lib._PaygenPayload,
-                             '_GenPayloadHash')
-    self.mox.StubOutWithMock(paygen_payload_lib._PaygenPayload,
-                             '_GenMetadataHash')
+                             '_GenerateHashes')
     self.mox.StubOutWithMock(paygen_payload_lib._PaygenPayload,
                              '_SignHashes')
     self.mox.StubOutWithMock(paygen_payload_lib._PaygenPayload,
@@ -574,8 +553,7 @@ class PaygenPayloadLibBasicTest(PaygenPayloadLibTest):
                              '_StoreMetadataSignatures')
 
     # Record expected calls.
-    gen._GenPayloadHash().AndReturn(payload_hash)
-    gen._GenMetadataHash().AndReturn(metadata_hash)
+    gen._GenerateHashes().AndReturn((payload_hash, metadata_hash))
     gen._SignHashes([payload_hash, metadata_hash]).AndReturn(
         (payload_sigs, metadata_sigs))
     gen._InsertPayloadSignatures(payload_sigs)
@@ -760,7 +738,7 @@ class PaygenPayloadLibEndToEndTest(PaygenPayloadLibTest):
         cache=self.cache,
         work_dir=self.tempdir,
         au_generator_uri=gspaths.ChromeosReleases.GeneratorUri(
-            payload.tgt_image.channel, payload.tgt_image.board, '6351.0.0'),
+            payload.tgt_image.channel, payload.tgt_image.board, '7587.0.0'),
         sign=sign)
 
     self.assertTrue(os.path.exists(output_uri))