Allow top-level vbmeta struct to be in 'boot' partition.

If there is no 'vbmeta' partition try to load the top-level vbmeta
struct from the end of 'boot' via a footer.

Two use-cases come to mind

 - bring-up when the partition table doesn't yet mention vbmeta; and
 - upgrades where it's not feasible to change the partition table

Bug: None
Test: New unit tests and all unit tests pass.
Change-Id: Id0c6c0f95ce157ffbeb0692d3c9547f49ab58640
diff --git a/avbtool b/avbtool
index 9850634..6ebbd8c 100755
--- a/avbtool
+++ b/avbtool
@@ -1902,7 +1902,7 @@
 
     Arguments:
       output: File to write the image to.
-      chain_partitions: List of partitions to chain.
+      chain_partitions: List of partitions to chain or None.
       algorithm_name: Name of algorithm to use.
       key_path: Path to key to use or None.
       public_key_metadata_path: Path to public key metadata or None.
@@ -1922,26 +1922,10 @@
     """
 
     descriptors = []
-
-    # Insert chained partition descriptors.
-    if chain_partitions:
-      for cp in chain_partitions:
-        cp_tokens = cp.split(':')
-        if len(cp_tokens) != 3:
-          raise AvbError('Malformed chained partition "{}".'.format(cp))
-        desc = AvbChainPartitionDescriptor()
-        desc.partition_name = cp_tokens[0]
-        desc.rollback_index_location = int(cp_tokens[1])
-        if desc.rollback_index_location < 1:
-          raise AvbError('Rollback index location must be 1 or larger.')
-        file_path = cp_tokens[2]
-        desc.public_key = open(file_path, 'rb').read()
-        descriptors.append(desc)
-
     vbmeta_blob = self._generate_vbmeta_blob(
         algorithm_name, key_path, public_key_metadata_path, descriptors,
-        rollback_index, flags, props, props_from_file, kernel_cmdlines,
-        setup_rootfs_from_kernel,
+        chain_partitions, rollback_index, flags, props, props_from_file,
+        kernel_cmdlines, setup_rootfs_from_kernel,
         include_descriptors_from_image, signing_helper, release_string,
         append_to_release_string)
 
@@ -1951,6 +1935,7 @@
 
   def _generate_vbmeta_blob(self, algorithm_name, key_path,
                             public_key_metadata_path, descriptors,
+                            chain_partitions,
                             rollback_index, flags, props, props_from_file,
                             kernel_cmdlines,
                             setup_rootfs_from_kernel,
@@ -1971,6 +1956,7 @@
       key_path: The path to the .pem file used to sign the blob.
       public_key_metadata_path: Path to public key metadata or None.
       descriptors: A list of descriptors to insert or None.
+      chain_partitions: List of partitions to chain or None.
       rollback_index: The rollback index to use.
       flags: Flags to use in the image.
       props: Properties to insert (List of strings of the form 'key:value').
@@ -1998,11 +1984,28 @@
     except KeyError:
       raise AvbError('Unknown algorithm with name {}'.format(algorithm_name))
 
+    if not descriptors:
+      descriptors = []
+
+    # Insert chained partition descriptors, if any
+    if chain_partitions:
+      for cp in chain_partitions:
+        cp_tokens = cp.split(':')
+        if len(cp_tokens) != 3:
+          raise AvbError('Malformed chained partition "{}".'.format(cp))
+        desc = AvbChainPartitionDescriptor()
+        desc.partition_name = cp_tokens[0]
+        desc.rollback_index_location = int(cp_tokens[1])
+        if desc.rollback_index_location < 1:
+          raise AvbError('Rollback index location must be 1 or larger.')
+        file_path = cp_tokens[2]
+        desc.public_key = open(file_path, 'rb').read()
+        descriptors.append(desc)
+
     # Descriptors.
     encoded_descriptors = bytearray()
-    if descriptors:
-      for desc in descriptors:
-        encoded_descriptors.extend(desc.encode())
+    for desc in descriptors:
+      encoded_descriptors.extend(desc.encode())
 
     # Add properties.
     if props:
@@ -2153,8 +2156,9 @@
     write_rsa_key(output, key)
 
   def add_hash_footer(self, image_filename, partition_size, partition_name,
-                      hash_algorithm, salt, algorithm_name, key_path,
-                      public_key_metadata_path, rollback_index, props,
+                      hash_algorithm, salt, chain_partitions, algorithm_name,
+                      key_path,
+                      public_key_metadata_path, rollback_index, flags, props,
                       props_from_file, kernel_cmdlines,
                       setup_rootfs_from_kernel,
                       include_descriptors_from_image, signing_helper,
@@ -2168,10 +2172,12 @@
       partition_name: Name of partition (without A/B suffix).
       hash_algorithm: Hash algorithm to use.
       salt: Salt to use as a hexadecimal string or None to use /dev/urandom.
+      chain_partitions: List of partitions to chain.
       algorithm_name: Name of algorithm to use.
       key_path: Path to key to use or None.
       public_key_metadata_path: Path to public key metadata or None.
       rollback_index: Rollback index.
+      flags: Flags value to use in the image.
       props: Properties to insert (List of strings of the form 'key:value').
       props_from_file: Properties to insert (List of strings 'key:<path>').
       kernel_cmdlines: Kernel cmdlines to insert (list of strings).
@@ -2250,14 +2256,11 @@
       h_desc.salt = salt
       h_desc.digest = digest
 
-      # Flags are only allowed on top-level vbmeta struct.
-      flags = 0
-
       # Generate the VBMeta footer.
       vbmeta_blob = self._generate_vbmeta_blob(
           algorithm_name, key_path, public_key_metadata_path, [h_desc],
-          rollback_index, flags, props, props_from_file, kernel_cmdlines,
-          setup_rootfs_from_kernel,
+          chain_partitions, rollback_index, flags, props, props_from_file,
+          kernel_cmdlines, setup_rootfs_from_kernel,
           include_descriptors_from_image, signing_helper, release_string,
           append_to_release_string)
 
@@ -2310,8 +2313,9 @@
 
   def add_hashtree_footer(self, image_filename, partition_size, partition_name,
                           generate_fec, fec_num_roots, hash_algorithm,
-                          block_size, salt, algorithm_name, key_path,
-                          public_key_metadata_path, rollback_index,
+                          block_size, salt, chain_partitions, algorithm_name,
+                          key_path,
+                          public_key_metadata_path, rollback_index, flags,
                           props, props_from_file, kernel_cmdlines,
                           setup_rootfs_from_kernel,
                           include_descriptors_from_image,
@@ -2332,10 +2336,12 @@
       hash_algorithm: Hash algorithm to use.
       block_size: Block size to use.
       salt: Salt to use as a hexadecimal string or None to use /dev/urandom.
+      chain_partitions: List of partitions to chain.
       algorithm_name: Name of algorithm to use.
       key_path: Path to key to use or None.
       public_key_metadata_path: Path to public key metadata or None.
       rollback_index: Rollback index.
+      flags: Flags value to use in the image.
       props: Properties to insert (List of strings of the form 'key:value').
       props_from_file: Properties to insert (List of strings 'key:<path>').
       kernel_cmdlines: Kernel cmdlines to insert (list of strings).
@@ -2478,15 +2484,12 @@
         ht_desc.fec_offset = fec_offset
         ht_desc.fec_size = len(fec_data)
 
-      # Flags are only allowed on top-level vbmeta struct.
-      flags = 0
-
       # Generate the VBMeta footer and add padding as needed.
       vbmeta_offset = tree_offset + len_hashtree_and_fec
       vbmeta_blob = self._generate_vbmeta_blob(
           algorithm_name, key_path, public_key_metadata_path, [ht_desc],
-          rollback_index, flags, props, props_from_file, kernel_cmdlines,
-          setup_rootfs_from_kernel,
+          chain_partitions, rollback_index, flags, props, props_from_file,
+          kernel_cmdlines, setup_rootfs_from_kernel,
           include_descriptors_from_image, signing_helper, release_string,
           append_to_release_string)
       padding_needed = (round_to_multiple(len(vbmeta_blob), image.block_size) -
@@ -2859,6 +2862,31 @@
                             metavar='IMAGE',
                             action='append',
                             type=argparse.FileType('rb'))
+    # These are only allowed from top-level vbmeta and boot-in-lieu-of-vbmeta.
+    sub_parser.add_argument('--chain_partition',
+                            help='Allow signed integrity-data for partition',
+                            metavar='PART_NAME:ROLLBACK_SLOT:KEY_PATH',
+                            action='append')
+    sub_parser.add_argument('--flags',
+                            help='VBMeta flags',
+                            type=parse_number,
+                            default=0)
+    sub_parser.add_argument('--set_hashtree_disabled_flag',
+                            help='Set the HASHTREE_DISABLED flag',
+                            action='store_true')
+
+  def _fixup_common_args(self, args):
+    """Common fixups needed by subcommands.
+
+    Arguments:
+      args: Arguments to modify.
+
+    Returns:
+      The modified arguments.
+    """
+    if args.set_hashtree_disabled_flag:
+      args.flags |= AVB_VBMETA_IMAGE_FLAGS_HASHTREE_DISABLED
+    return args
 
   def run(self, argv):
     """Command-line processor.
@@ -2891,17 +2919,6 @@
                             type=argparse.FileType('wb'),
                             required=True)
     self._add_common_args(sub_parser)
-    sub_parser.add_argument('--chain_partition',
-                            help='Allow signed integrity-data for partition',
-                            metavar='PART_NAME:ROLLBACK_SLOT:KEY_PATH',
-                            action='append')
-    sub_parser.add_argument('--flags',
-                            help='VBMeta flags',
-                            type=parse_number,
-                            default=0)
-    sub_parser.add_argument('--set_hashtree_disabled_flag',
-                            help='Set the HASHTREE_DISABLED flag',
-                            action='store_true')
     sub_parser.set_defaults(func=self.make_vbmeta_image)
 
     sub_parser = subparsers.add_parser('add_hash_footer',
@@ -3104,8 +3121,7 @@
 
   def make_vbmeta_image(self, args):
     """Implements the 'make_vbmeta_image' sub-command."""
-    if args.set_hashtree_disabled_flag:
-      args.flags |= AVB_VBMETA_IMAGE_FLAGS_HASHTREE_DISABLED
+    args = self._fixup_common_args(args)
     self.avb.make_vbmeta_image(args.output, args.chain_partition,
                                args.algorithm, args.key,
                                args.public_key_metadata, args.rollback_index,
@@ -3119,11 +3135,13 @@
 
   def add_hash_footer(self, args):
     """Implements the 'add_hash_footer' sub-command."""
+    args = self._fixup_common_args(args)
     self.avb.add_hash_footer(args.image.name, args.partition_size,
                              args.partition_name, args.hash_algorithm,
-                             args.salt, args.algorithm, args.key,
+                             args.salt, args.chain_partition, args.algorithm,
+                             args.key,
                              args.public_key_metadata, args.rollback_index,
-                             args.prop, args.prop_from_file,
+                             args.flags, args.prop, args.prop_from_file,
                              args.kernel_cmdline,
                              args.setup_rootfs_from_kernel,
                              args.include_descriptors_from_image,
@@ -3135,14 +3153,15 @@
 
   def add_hashtree_footer(self, args):
     """Implements the 'add_hashtree_footer' sub-command."""
+    args = self._fixup_common_args(args)
     self.avb.add_hashtree_footer(args.image.name if args.image else None,
                                  args.partition_size,
                                  args.partition_name,
                                  args.generate_fec, args.fec_num_roots,
                                  args.hash_algorithm, args.block_size,
-                                 args.salt, args.algorithm, args.key,
-                                 args.public_key_metadata,
-                                 args.rollback_index, args.prop,
+                                 args.salt, args.chain_partition, args.algorithm,
+                                 args.key, args.public_key_metadata,
+                                 args.rollback_index, args.flags, args.prop,
                                  args.prop_from_file,
                                  args.kernel_cmdline,
                                  args.setup_rootfs_from_kernel,