releasetools: Clean up AppendVBMetaArgsForPartition().

Drop the parameter of public_key_dir. It's not necessary to stage all
the extracted public keys into the same temporary dir.

Redirect the stderr output to STDOUT while calling avbtool, so that it
can dump the error message accordingly.

Also remove the check for having empty img_path arg. It doesn't look
like a clean logic to allow the function to be called with invalid /
empty arg. The only caller (i.e. AddVBMeta) already asserts img_path
being valid as well.

Test: python -m unittest test_add_img_to_target_files
Test: `m dist` with aosp_walleye-userdebug
Change-Id: Id58c5ae780ac8a22661ffea629144d4836839175

tools/releasetools/add_img_to_target_files.py

@@ -358,29 +358,39 @@ def AddUserdata(output_zip):
   img.Write()
 
 
-def AppendVBMetaArgsForPartition(cmd, partition, img_path, public_key_dir):
-  if not img_path:
-    return
+def AppendVBMetaArgsForPartition(cmd, partition, image):
+  """Appends the VBMeta arguments for partition.
 
+  It sets up the VBMeta argument by including the partition descriptor from the
+  given 'image', or by configuring the partition as a chained partition.
+
+  Args:
+    cmd: A list of command args that will be used to generate the vbmeta image.
+        The argument for the partition will be appended to the list.
+    partition: The name of the partition (e.g. "system").
+    image: The path to the partition image.
+  """
   # Check if chain partition is used.
   key_path = OPTIONS.info_dict.get("avb_" + partition + "_key_path")
   if key_path:
     # extract public key in AVB format to be included in vbmeta.img
     avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
-    public_key_path = os.path.join(public_key_dir, "%s.avbpubkey" % partition)
-    p = common.Run([avbtool, "extract_public_key", "--key", key_path,
-                    "--output", public_key_path],
-                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    p.communicate()
-    assert p.returncode == 0, \
-        "avbtool extract_public_key fail for partition: %r" % partition
+    pubkey_path = common.MakeTempFile(prefix="avb-", suffix=".pubkey")
+    proc = common.Run(
+        [avbtool, "extract_public_key", "--key", key_path, "--output",
+         pubkey_path],
+        stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    stdoutdata, _ = proc.communicate()
+    assert proc.returncode == 0, \
+        "Failed to extract pubkey for {}:\n{}".format(
+            partition, stdoutdata)
 
     rollback_index_location = OPTIONS.info_dict[
         "avb_" + partition + "_rollback_index_location"]
     cmd.extend(["--chain_partition", "%s:%s:%s" % (
-        partition, rollback_index_location, public_key_path)])
+        partition, rollback_index_location, pubkey_path)])
   else:
-    cmd.extend(["--include_descriptors_from_image", img_path])
+    cmd.extend(["--include_descriptors_from_image", image])
 
 
 def AddVBMeta(output_zip, partitions):
@@ -389,8 +399,8 @@ def AddVBMeta(output_zip, partitions):
   Args:
     output_zip: The output zip file, which needs to be already open.
     partitions: A dict that's keyed by partition names with image paths as
-        values. Only valid partition names are accepted, which include 'boot',
-        'recovery', 'system', 'vendor', 'dtbo'.
+        values. Only valid partition names are accepted, as listed in
+        common.AVB_PARTITIONS.
   """
   img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "vbmeta.img")
   if os.path.exists(img.input_name):
@@ -401,13 +411,12 @@ def AddVBMeta(output_zip, partitions):
   cmd = [avbtool, "make_vbmeta_image", "--output", img.name]
   common.AppendAVBSigningArgs(cmd, "vbmeta")
 
-  public_key_dir = common.MakeTempDir(prefix="avbpubkey-")
   for partition, path in partitions.items():
-    assert partition in common.AVB_PARTITIONS, 'Unknown partition: %s' % (
-        partition,)
-    assert os.path.exists(path), 'Failed to find %s for partition %s' % (
-        path, partition)
-    AppendVBMetaArgsForPartition(cmd, partition, path, public_key_dir)
+    assert partition in common.AVB_PARTITIONS, \
+        'Unknown partition: {}'.format(partition)
+    assert os.path.exists(path), \
+        'Failed to find {} for {}'.format(path, partition)
+    AppendVBMetaArgsForPartition(cmd, partition, path)
 
   args = OPTIONS.info_dict.get("avb_vbmeta_args")
   if args and args.strip():

tools/releasetools/test_add_img_to_target_files.py

@@ -22,7 +22,8 @@ import zipfile
 import common
 import test_utils
 from add_img_to_target_files import (
-    AddCareMapTxtForAbOta, AddPackRadioImages, CheckAbOtaImages, GetCareMap)
+    AddCareMapTxtForAbOta, AddPackRadioImages, AppendVBMetaArgsForPartition,
+    CheckAbOtaImages, GetCareMap)
 from rangelib import RangeSet
 
 
@@ -264,6 +265,31 @@ class AddImagesToTargetFilesTest(unittest.TestCase):
     # The existing entry should be scheduled to be replaced.
     self.assertIn('META/care_map.txt', OPTIONS.replace_updated_files_list)
 
+  def test_AppendVBMetaArgsForPartition(self):
+    OPTIONS.info_dict = {}
+    cmd = []
+    AppendVBMetaArgsForPartition(cmd, 'system', '/path/to/system.img')
+    self.assertEqual(
+        ['--include_descriptors_from_image', '/path/to/system.img'], cmd)
+
+  def test_AppendVBMetaArgsForPartition_vendorAsChainedPartition(self):
+    testdata_dir = test_utils.get_testdata_dir()
+    pubkey = os.path.join(testdata_dir, 'testkey.pubkey.pem')
+    OPTIONS.info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_vendor_key_path': pubkey,
+        'avb_vendor_rollback_index_location': 5,
+    }
+    cmd = []
+    AppendVBMetaArgsForPartition(cmd, 'vendor', '/path/to/vendor.img')
+    self.assertEqual(2, len(cmd))
+    self.assertEqual('--chain_partition', cmd[0])
+    chained_partition_args = cmd[1].split(':')
+    self.assertEqual(3, len(chained_partition_args))
+    self.assertEqual('vendor', chained_partition_args[0])
+    self.assertEqual('5', chained_partition_args[1])
+    self.assertTrue(os.path.exists(chained_partition_args[2]))
+
   def test_GetCareMap(self):
     sparse_image = test_utils.construct_sparse_image([
         (0xCAC1, 6),