releasetools: Support verifying AVB signed images with chained partitions.

For example, verify a target_files.zip that has system AVB-signed as a
chained partition.

  $ build/make/tools/releasetools/validate_target_files.py \
      signed-target_files-4904652.zip \
      --verity_key verifiedboot_pub.pem \
      --avb_system_key_path system_pub.pem

Note that verifiedboot_pub.pem should be the key (either public or
private) to verify vbmeta.img, and 'system_pub.pem' should be the key
(either public or private) for the chained partition of system.

testdata/testkey.key is the private key converted from
testdata/testkey.pk8 for testing purpose (`openssl pkcs8 -in
testdata/testkey.pk8 -inform DER -out testdata/testkey.key -nocrypt`).

Bug: 63706333
Test: python -m unittest test_common
Test: python -m unittest test_add_img_to_target_files
Test: `m dist` on aosp_walleye-userdebug; Run validate_target_files.py
      on the generated target_files.zip.
Test: Set up walleye with chained system partition; `m dist`; Run
      validate_target_files.py on the generated target_files.zip.
Change-Id: I38517ab39baf8a5bc1a6062fab2fe229b68e897d

tools/releasetools/add_img_to_target_files.py

@@ -373,22 +373,9 @@ def AppendVBMetaArgsForPartition(cmd, 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"]
-    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, pubkey_path)])
+    chained_partition_arg = common.GetAvbChainedPartitionArg(
+        partition, OPTIONS.info_dict)
+    cmd.extend(["--chain_partition", chained_partition_arg])
   else:
     cmd.extend(["--include_descriptors_from_image", image])
 

tools/releasetools/common.py

@@ -371,6 +371,40 @@ def AppendAVBSigningArgs(cmd, partition):
     cmd.extend(["--salt", avb_salt])
 
 
+def GetAvbChainedPartitionArg(partition, info_dict, key=None):
+  """Constructs and returns the arg to build or verify a chained partition.
+
+  Args:
+    partition: The partition name.
+    info_dict: The info dict to look up the key info and rollback index
+        location.
+    key: The key to be used for building or verifying the partition. Defaults to
+        the key listed in info_dict.
+
+  Returns:
+    A string of form "partition:rollback_index_location:key" that can be used to
+    build or verify vbmeta image.
+
+  Raises:
+    AssertionError: When it fails to extract the public key with avbtool.
+  """
+  if key is None:
+    key = info_dict["avb_" + partition + "_key_path"]
+  avbtool = os.getenv('AVBTOOL') or info_dict["avb_avbtool"]
+  pubkey_path = MakeTempFile(prefix="avb-", suffix=".pubkey")
+  proc = Run(
+      [avbtool, "extract_public_key", "--key", key, "--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 = info_dict[
+      "avb_" + partition + "_rollback_index_location"]
+  return "{}:{}:{}".format(partition, rollback_index_location, pubkey_path)
+
+
 def _BuildBootableImage(sourcedir, fs_config_file, info_dict=None,
                         has_ramdisk=False, two_step_image=False):
   """Build a bootable image from the specified sourcedir.

tools/releasetools/test_common.py

@@ -524,6 +524,9 @@ class CommonApkUtilsTest(unittest.TestCase):
 
 class CommonUtilsTest(unittest.TestCase):
 
+  def setUp(self):
+    self.testdata_dir = test_utils.get_testdata_dir()
+
   def tearDown(self):
     common.Cleanup()
 
@@ -730,6 +733,56 @@ class CommonUtilsTest(unittest.TestCase):
           AssertionError, common.GetSparseImage, 'system', tempdir, input_zip,
           False)
 
+  def test_GetAvbChainedPartitionArg(self):
+    pubkey = os.path.join(self.testdata_dir, 'testkey.pubkey.pem')
+    info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_system_key_path': pubkey,
+        'avb_system_rollback_index_location': 2,
+    }
+    args = common.GetAvbChainedPartitionArg('system', info_dict).split(':')
+    self.assertEqual(3, len(args))
+    self.assertEqual('system', args[0])
+    self.assertEqual('2', args[1])
+    self.assertTrue(os.path.exists(args[2]))
+
+  def test_GetAvbChainedPartitionArg_withPrivateKey(self):
+    key = os.path.join(self.testdata_dir, 'testkey.key')
+    info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_product_key_path': key,
+        'avb_product_rollback_index_location': 2,
+    }
+    args = common.GetAvbChainedPartitionArg('product', info_dict).split(':')
+    self.assertEqual(3, len(args))
+    self.assertEqual('product', args[0])
+    self.assertEqual('2', args[1])
+    self.assertTrue(os.path.exists(args[2]))
+
+  def test_GetAvbChainedPartitionArg_withSpecifiedKey(self):
+    info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_system_key_path': 'does-not-exist',
+        'avb_system_rollback_index_location': 2,
+    }
+    pubkey = os.path.join(self.testdata_dir, 'testkey.pubkey.pem')
+    args = common.GetAvbChainedPartitionArg(
+        'system', info_dict, pubkey).split(':')
+    self.assertEqual(3, len(args))
+    self.assertEqual('system', args[0])
+    self.assertEqual('2', args[1])
+    self.assertTrue(os.path.exists(args[2]))
+
+  def test_GetAvbChainedPartitionArg_invalidKey(self):
+    pubkey = os.path.join(self.testdata_dir, 'testkey_with_passwd.x509.pem')
+    info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_system_key_path': pubkey,
+        'avb_system_rollback_index_location': 2,
+    }
+    self.assertRaises(
+        AssertionError, common.GetAvbChainedPartitionArg, 'system', info_dict)
+
 
 class InstallRecoveryScriptFormatTest(unittest.TestCase):
   """Checks the format of install-recovery.sh.

tools/releasetools/testdata/testkey.key

@@ -0,0 +1,28 @@
+-----BEGIN PRIVATE KEY-----
+MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC+O/I7YvBaCZA3
+KrvP7ErTh6DS3cAvjLY2GkAA7NWcXIICsVwWMtMZAMO+Rtk/o3XY7r53Amg8ue2e
+b0D5Wc8gUVEeDQdRZJscz5CTwmC/b/YBWQBSPknWv23hf7ZYjR5HMk/XlmOfrylA
+oaZzJyKvLSBu+Zi9cvZlSZObnoOYR8JQJEhjYgYn8JwycV4i1VTQvEqxXkyW3kE7
+RW/8JXgRqI4vDIKehm5SFi2jt0eU7/ju/8f3OGQkLng4DV2QPfwQ+A7kad+EYVI1
+dBGYkNNesWB3o75A7jJQ1fyVg/XQzOKZSki1lrTm3rw0AOrBiXdPudbO+4L2vgip
+kPI9/bVNAgMBAAECggEABGjBSY0Wgw+7rvunlL8mUNbQ7HJFVRTO2FwtZZgXr2MZ
+hFR2DPGqoOa6ortjp6zzO071TS7aGaY5krWDbQQe3+Hinm6w37sUOUu6TyJvOaCv
+tAJLFpzo+zg+pL5gDJdgv0e0QAv1TSszKpNUl1Ct5h+Go+vXFXUHrvtQl4fKBwqA
+efxcd3R4z3p/3Cl2ZYIRz9I7UXUZZYwJE7bDNDz3aFZ1jUoELGmhe1O5w0hJY1q6
+PxuOM9bL60yDn0vu0eiCjaPlHeHyGe9pQ1aQLEuwQz9zpWC01dWPVkLmny7HDygC
+VBsdg8MNlzJQ1WV2en11BH72IqZ59U8pD0xEB7a4BQKBgQDxenfXyYZw4AKcaJlP
+ncJmsx/wcgEvWNxiI4etArXES4VIyP2OlSw+q9JbOOpaSk8TJP5PNfUkgTbC4B2y
+gh/AobJP5b7Wn5LrsHc3GY6CzF1i8T4xXQRxnaKWE86SOmZQlEmyCnpyCmfFVuaR
+E8p8CPW/gQLhpxSlQdGZ0bYLiwKBgQDJrJRDdyaI/Isusog/OwKHVGBU6CmRa5tM
+gx+GIlxheqhuDqnBkr0h1kL20Zi80JeG7vKWr+dwfqkEarfdTe+juwlIuQ3MEuXL
+AbsKNuaU1naOqOLm9rjZgRtR7oNLVH5AbkKMaJz1zM6YiMl54FEDX7ZVY8b6q1Kz
+YXT3sGi9hwKBgBsNa0ujagpPLjuzhCllNRgoTRW0z+kr/VSJQnPhb9eT1lS3H6DP
+mWtT+Hb7w1VmKcGtTUg2dUYnq6jdTrZm2YPNGZrV1DFbIyyAUnq7xDlnB7dD64HA
+N/U6gbJqeaPsIvY4BqGJhvorrEBxYdcy7mZC4rUXkOkSvL9exkqDMe/NAoGARaHU
+v0aQg5PO6pyx9kMFqHw1lptiXtdsk4pihAmxI+cZ6IYfjrp/mwNDs7zCo87RwsEV
++Xlay7iv2tqOCVczerDFj9p1LRUJSoKadfhmvNUfsjoVvfFJ+a9eI3fa1VOjE9P+
+HkSwjR3d50Sza+VLk4Kkje8ZcMtejpkDrdG3GFkCgYBXHqciwlFn5nMPFRe8v426
+6YBiUtzCQCZxDtMeeZYCJslFfjrqPXNUcU/flxWwaikjFsLJEtl7aT3Hpdi5I7T8
+yCYkUWqAAh7twEYTOeG6v/tEa/PmsBjZXPD2zkCp76EQmv3gbvsH4F/nA55gT/GR
+2in6XS/4rHBvjn5gF6MFyg==
+-----END PRIVATE KEY-----

tools/releasetools/validate_target_files.py

@@ -315,9 +315,19 @@ def ValidateVerifiedBootImages(input_tmp, info_dict, options):
     key = options['verity_key']
     if key is None:
       key = info_dict['avb_vbmeta_key_path']
+
     # avbtool verifies all the images that have descriptors listed in vbmeta.
     image = os.path.join(input_tmp, 'IMAGES', 'vbmeta.img')
     cmd = ['avbtool', 'verify_image', '--image', image, '--key', key]
+
+    # Append the args for chained partitions if any.
+    for partition in common.AVB_PARTITIONS:
+      key_name = 'avb_' + partition + '_key_path'
+      if info_dict.get(key_name) is not None:
+        chained_partition_arg = common.GetAvbChainedPartitionArg(
+            partition, info_dict, options[key_name])
+        cmd.extend(["--expected_chain_partition", chained_partition_arg])
+
     proc = common.Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
     stdoutdata, _ = proc.communicate()
     assert proc.returncode == 0, \
@@ -339,8 +349,13 @@ def main():
   parser.add_argument(
       '--verity_key',
       help='the verity public key to verify the bootable images (Verified '
-           'Boot 1.0), or the vbmeta image (Verified Boot 2.0), where '
+           'Boot 1.0), or the vbmeta image (Verified Boot 2.0, aka AVB), where '
            'applicable')
+  for partition in common.AVB_PARTITIONS:
+    parser.add_argument(
+        '--avb_' + partition + '_key_path',
+        help='the public or private key in PEM format to verify AVB chained '
+             'partition of {}'.format(partition))
   parser.add_argument(
       '--verity_key_mincrypt',
       help='the verity public key in mincrypt format to verify the system '