avbtool: add_hashtree_footer: Add --setup_as_rootfs_from_kernel option.

This option makes it possible to generate system.img with kernel
command-line descriptors for setting up the partition with dm-verity
as the root filesystem. This is different from the existing option
--setup_rootfs_from_kernel which takes a path to system.img and adds
the kernel command-line descriptors to e.g. vbmeta.img.

This option can be used when using chained partitions for system.img.

Also fix README.md to use --include_descriptors_from_image and not
--include_descriptors_from_footer.

Bug: 38304536
Test: New unit tests + all unit tests pass.
Change-Id: I6285877cdb3b63a7c9117c270459d1fbb93e3309
diff --git a/README.md b/README.md
index 31d579b..3e60e37 100644
--- a/README.md
+++ b/README.md
@@ -297,7 +297,7 @@
         --output OUTPUT                                                            \
         [--algorithm ALGORITHM] [--key /path/to/key_used_for_signing_or_pub_key]   \
         [--public_key_metadata /path/to/pkmd.bin] [--rollback_index NUMBER]        \
-        [--include_descriptors_from_footer /path/to/image.bin]                     \
+        [--include_descriptors_from_image /path/to/image.bin]                      \
         [--setup_rootfs_from_kernel /path/to/image.bin]                            \
         [--chain_partition part_name:rollback_index_location:/path/to/key1.bin]    \
         [--signing_helper /path/to/external/signer]                                \
@@ -312,7 +312,7 @@
         [--algorithm ALGORITHM] [--key /path/to/key_used_for_signing_or_pub_key]   \
         [--public_key_metadata /path/to/pkmd.bin] [--rollback_index NUMBER]        \
         [--hash_algorithm HASH_ALG] [--salt HEX]                                   \
-        [--include_descriptors_from_footer /path/to/image.bin]                     \
+        [--include_descriptors_from_image /path/to/image.bin]                      \
         [--setup_rootfs_from_kernel /path/to/image.bin]                            \
         [--output_vbmeta_image OUTPUT_IMAGE] [--do_not_append_vbmeta_image]        \
         [--signing_helper /path/to/external/signer]                                \
@@ -328,8 +328,9 @@
         [--algorithm ALGORITHM] [--key /path/to/key_used_for_signing_or_pub_key]   \
         [--public_key_metadata /path/to/pkmd.bin] [--rollback_index NUMBER]        \
         [--hash_algorithm HASH_ALG] [--salt HEX] [--block_size SIZE]               \
-        [--include_descriptors_from_footer /path/to/image.bin]                     \
+        [--include_descriptors_from_image /path/to/image.bin]                      \
         [--setup_rootfs_from_kernel /path/to/image.bin]                            \
+        [--setup_as_rootfs_from_kernel]                                            \
         [--output_vbmeta_image OUTPUT_IMAGE] [--do_not_append_vbmeta_image]        \
         [--do_not_generate_fec] [--fec_num_roots FEC_NUM_ROOTS]                    \
         [--signing_helper /path/to/external/signer]                                \
diff --git a/avbtool b/avbtool
index 7c18d46..1b459f7 100755
--- a/avbtool
+++ b/avbtool
@@ -2039,11 +2039,11 @@
     data_size += h.auxiliary_data_block_size
     return image.read(data_size)
 
-  def _get_cmdline_descriptors_for_dm_verity(self, image):
+  def _get_cmdline_descriptors_for_hashtree_descriptor(self, ht):
     """Generate kernel cmdline descriptors for dm-verity.
 
     Arguments:
-      image: An ImageHandler (vbmeta or footer) with a hashtree descriptor.
+      ht: A AvbHashtreeDescriptor
 
     Returns:
       A list with two AvbKernelCmdlineDescriptor with dm-verity kernel cmdline
@@ -2055,17 +2055,6 @@
 
     """
 
-    (_, _, descriptors, _) = self._parse_image(image)
-
-    ht = None
-    for desc in descriptors:
-      if isinstance(desc, AvbHashtreeDescriptor):
-        ht = desc
-        break
-
-    if not ht:
-      raise AvbError('No hashtree descriptor in given image')
-
     c = 'dm="1 vroot none ro 1,'
     c += '0'  # start
     c += ' {}'.format((ht.image_size / 512))  # size (# sectors)
@@ -2111,6 +2100,35 @@
 
     return [desc, desc_no_ht]
 
+  def _get_cmdline_descriptors_for_dm_verity(self, image):
+    """Generate kernel cmdline descriptors for dm-verity.
+
+    Arguments:
+      image: An ImageHandler (vbmeta or footer) with a hashtree descriptor.
+
+    Returns:
+      A list with two AvbKernelCmdlineDescriptor with dm-verity kernel cmdline
+      instructions. There is one for when hashtree is not disabled and one for
+      when it is.
+
+    Raises:
+      AvbError: If  |image| doesn't have a hashtree descriptor.
+
+    """
+
+    (_, _, descriptors, _) = self._parse_image(image)
+
+    ht = None
+    for desc in descriptors:
+      if isinstance(desc, AvbHashtreeDescriptor):
+        ht = desc
+        break
+
+    if not ht:
+      raise AvbError('No hashtree descriptor in given image')
+
+    return self._get_cmdline_descriptors_for_hashtree_descriptor(ht)
+
   def make_vbmeta_image(self, output, chain_partitions, algorithm_name,
                         key_path, public_key_metadata_path, rollback_index,
                         flags, props, props_from_file, kernel_cmdlines,
@@ -2142,10 +2160,11 @@
     """
 
     descriptors = []
+    ht_desc_to_setup = None
     vbmeta_blob = self._generate_vbmeta_blob(
         algorithm_name, key_path, public_key_metadata_path, descriptors,
         chain_partitions, rollback_index, flags, props, props_from_file,
-        kernel_cmdlines, setup_rootfs_from_kernel,
+        kernel_cmdlines, setup_rootfs_from_kernel, ht_desc_to_setup,
         include_descriptors_from_image, signing_helper, release_string,
         append_to_release_string)
 
@@ -2159,6 +2178,7 @@
                             rollback_index, flags, props, props_from_file,
                             kernel_cmdlines,
                             setup_rootfs_from_kernel,
+                            ht_desc_to_setup,
                             include_descriptors_from_image, signing_helper,
                             release_string, append_to_release_string):
     """Generates a VBMeta blob.
@@ -2184,6 +2204,8 @@
       kernel_cmdlines: Kernel cmdlines to insert (list of strings).
       setup_rootfs_from_kernel: None or file to generate
         dm-verity kernel cmdline from.
+      ht_desc_to_setup: If not None, an AvbHashtreeDescriptor to
+        generate dm-verity kernel cmdline descriptors from.
       include_descriptors_from_image: List of file objects for which
         to insert descriptors from.
       signing_helper: Program which signs a hash and return signature.
@@ -2258,7 +2280,7 @@
         desc.value = open(file_path, 'rb').read()
         encoded_descriptors.extend(desc.encode())
 
-    # Add AvbKernelCmdline descriptor for dm-verity, if requested.
+    # Add AvbKernelCmdline descriptor for dm-verity from an image, if requested.
     if setup_rootfs_from_kernel:
       image_handler = ImageHandler(
           setup_rootfs_from_kernel.name)
@@ -2266,6 +2288,13 @@
       encoded_descriptors.extend(cmdline_desc[0].encode())
       encoded_descriptors.extend(cmdline_desc[1].encode())
 
+    # Add AvbKernelCmdline descriptor for dm-verity from desc, if requested.
+    if ht_desc_to_setup:
+      cmdline_desc = self._get_cmdline_descriptors_for_hashtree_descriptor(
+          ht_desc_to_setup)
+      encoded_descriptors.extend(cmdline_desc[0].encode())
+      encoded_descriptors.extend(cmdline_desc[1].encode())
+
     # Add kernel command-lines.
     if kernel_cmdlines:
       for i in kernel_cmdlines:
@@ -2558,10 +2587,11 @@
       h_desc.digest = digest
 
       # Generate the VBMeta footer.
+      ht_desc_to_setup = None
       vbmeta_blob = self._generate_vbmeta_blob(
           algorithm_name, key_path, public_key_metadata_path, [h_desc],
           chain_partitions, rollback_index, flags, props, props_from_file,
-          kernel_cmdlines, setup_rootfs_from_kernel,
+          kernel_cmdlines, setup_rootfs_from_kernel, ht_desc_to_setup,
           include_descriptors_from_image, signing_helper, release_string,
           append_to_release_string)
 
@@ -2619,6 +2649,7 @@
                           public_key_metadata_path, rollback_index, flags,
                           props, props_from_file, kernel_cmdlines,
                           setup_rootfs_from_kernel,
+                          setup_as_rootfs_from_kernel,
                           include_descriptors_from_image,
                           calc_max_image_size, signing_helper,
                           release_string, append_to_release_string,
@@ -2648,6 +2679,8 @@
       kernel_cmdlines: Kernel cmdlines to insert (list of strings).
       setup_rootfs_from_kernel: None or file to generate
         dm-verity kernel cmdline from.
+      setup_as_rootfs_from_kernel: If True, generate dm-verity kernel
+        cmdline to set up rootfs.
       include_descriptors_from_image: List of file objects for which
         to insert descriptors from.
       calc_max_image_size: Don't store the hashtree or footer - instead
@@ -2785,12 +2818,16 @@
         ht_desc.fec_offset = fec_offset
         ht_desc.fec_size = len(fec_data)
 
+      ht_desc_to_setup = None
+      if setup_as_rootfs_from_kernel:
+        ht_desc_to_setup = ht_desc
+
       # 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],
           chain_partitions, rollback_index, flags, props, props_from_file,
-          kernel_cmdlines, setup_rootfs_from_kernel,
+          kernel_cmdlines, setup_rootfs_from_kernel, ht_desc_to_setup,
           include_descriptors_from_image, signing_helper, release_string,
           append_to_release_string)
       padding_needed = (round_to_multiple(len(vbmeta_blob), image.block_size) -
@@ -3308,6 +3345,12 @@
                             help=('Do not append vbmeta struct or footer '
                                   'to the image'),
                             action='store_true')
+    # This is different from --setup_rootfs_from_kernel insofar that
+    # it doesn't take an IMAGE, the generated cmdline will be for the
+    # hashtree we're adding.
+    sub_parser.add_argument('--setup_as_rootfs_from_kernel',
+                            action='store_true',
+                            help='Adds kernel cmdline for setting up rootfs')
     self._add_common_args(sub_parser)
     sub_parser.set_defaults(func=self.add_hashtree_footer)
 
@@ -3510,6 +3553,7 @@
                                  args.prop_from_file,
                                  args.kernel_cmdline,
                                  args.setup_rootfs_from_kernel,
+                                 args.setup_as_rootfs_from_kernel,
                                  args.include_descriptors_from_image,
                                  args.calc_max_image_size, args.signing_helper,
                                  args.internal_release_string,
diff --git a/test/avb_slot_verify_unittest.cc b/test/avb_slot_verify_unittest.cc
index 76ce2f5..a51daf2 100644
--- a/test/avb_slot_verify_unittest.cc
+++ b/test/avb_slot_verify_unittest.cc
@@ -45,6 +45,7 @@
   }
 
   void CmdlineWithHashtreeVerification(bool hashtree_verification_on);
+  void CmdlineWithChainedHashtreeVerification(bool hashtree_verification_on);
 
   FakeAvbOps ops_;
 };
@@ -1407,6 +1408,181 @@
   CmdlineWithHashtreeVerification(true);
 }
 
+void AvbSlotVerifyTest::CmdlineWithChainedHashtreeVerification(
+    bool hashtree_verification_on) {
+  const size_t system_size = 1028 * 1024;
+  const size_t system_partition_size = 1536 * 1024;
+
+  // Generate a 1028 KiB file with known content.
+  std::vector<uint8_t> contents;
+  contents.resize(system_size);
+  for (size_t n = 0; n < system_size; n++)
+    contents[n] = uint8_t(n);
+  base::FilePath system_path = testdir_.Append("system_a.img");
+  EXPECT_EQ(system_size,
+            static_cast<const size_t>(
+                base::WriteFile(system_path,
+                                reinterpret_cast<const char*>(contents.data()),
+                                contents.size())));
+
+  // Check that we correctly generate dm-verity kernel cmdline
+  // snippets, if requested.
+  EXPECT_COMMAND(0,
+                 "./avbtool add_hashtree_footer --salt d00df00d --image %s "
+                 "--partition_size %d --partition_name foobar "
+                 "--algorithm SHA256_RSA2048 "
+                 "--key test/data/testkey_rsa2048.pem "
+                 "--internal_release_string \"\" "
+                 "--do_not_generate_fec "
+                 "--setup_as_rootfs_from_kernel",
+                 system_path.value().c_str(),
+                 (int)system_partition_size);
+
+  EXPECT_EQ(
+      "Footer version:           1.0\n"
+      "Image size:               1572864 bytes\n"
+      "Original image size:      1052672 bytes\n"
+      "VBMeta offset:            1069056\n"
+      "VBMeta size:              1664 bytes\n"
+      "--\n"
+      "Minimum libavb version:   1.0\n"
+      "Header Block:             256 bytes\n"
+      "Authentication Block:     320 bytes\n"
+      "Auxiliary Block:          1088 bytes\n"
+      "Algorithm:                SHA256_RSA2048\n"
+      "Rollback Index:           0\n"
+      "Flags:                    0\n"
+      "Release String:           ''\n"
+      "Descriptors:\n"
+      "    Hashtree descriptor:\n"
+      "      Version of dm-verity:  1\n"
+      "      Image Size:            1052672 bytes\n"
+      "      Tree Offset:           1052672\n"
+      "      Tree Size:             16384 bytes\n"
+      "      Data Block Size:       4096 bytes\n"
+      "      Hash Block Size:       4096 bytes\n"
+      "      FEC num roots:         0\n"
+      "      FEC offset:            0\n"
+      "      FEC size:              0 bytes\n"
+      "      Hash Algorithm:        sha1\n"
+      "      Partition Name:        foobar\n"
+      "      Salt:                  d00df00d\n"
+      "      Root Digest:           e811611467dcd6e8dc4324e45f706c2bdd51db67\n"
+      "    Kernel Cmdline descriptor:\n"
+      "      Flags:                 1\n"
+      "      Kernel Cmdline:        'dm=\"1 vroot none ro 1,0 2056 verity 1 "
+      "PARTUUID=$(ANDROID_SYSTEM_PARTUUID) PARTUUID=$(ANDROID_SYSTEM_PARTUUID) "
+      "4096 4096 257 257 sha1 e811611467dcd6e8dc4324e45f706c2bdd51db67 "
+      "d00df00d 2 $(ANDROID_VERITY_MODE) ignore_zero_blocks\" root=/dev/dm-0'\n"
+      "    Kernel Cmdline descriptor:\n"
+      "      Flags:                 2\n"
+      "      Kernel Cmdline:        "
+      "'root=PARTUUID=$(ANDROID_SYSTEM_PARTUUID)'\n",
+      InfoImage(system_path));
+
+  base::FilePath pk_path = testdir_.Append("testkey_rsa2048.avbpubkey");
+  EXPECT_COMMAND(
+      0,
+      "./avbtool extract_public_key --key test/data/testkey_rsa2048.pem"
+      " --output %s",
+      pk_path.value().c_str());
+
+  GenerateVBMetaImage(
+      "vbmeta_a.img",
+      "SHA256_RSA2048",
+      4,
+      base::FilePath("test/data/testkey_rsa2048.pem"),
+      base::StringPrintf("--kernel_cmdline should_be_in_both=1 "
+                         "--algorithm SHA256_RSA2048 "
+                         "--flags %d "
+                         "--chain_partition system:1:%s "
+                         "--internal_release_string \"\"",
+                         hashtree_verification_on
+                             ? 0
+                             : AVB_VBMETA_IMAGE_FLAGS_HASHTREE_DISABLED,
+                         pk_path.value().c_str()));
+
+  EXPECT_EQ(
+      base::StringPrintf("Minimum libavb version:   1.0\n"
+                         "Header Block:             256 bytes\n"
+                         "Authentication Block:     320 bytes\n"
+                         "Auxiliary Block:          1216 bytes\n"
+                         "Algorithm:                SHA256_RSA2048\n"
+                         "Rollback Index:           4\n"
+                         "Flags:                    %d\n"
+                         "Release String:           ''\n"
+                         "Descriptors:\n"
+                         "    Chain Partition descriptor:\n"
+                         "      Partition Name:          system\n"
+                         "      Rollback Index Location: 1\n"
+                         "      Public key (sha1):       "
+                         "cdbb77177f731920bbe0a0f94f84d9038ae0617d\n"
+                         "    Kernel Cmdline descriptor:\n"
+                         "      Flags:                 0\n"
+                         "      Kernel Cmdline:        'should_be_in_both=1'\n",
+                         hashtree_verification_on
+                             ? 0
+                             : AVB_VBMETA_IMAGE_FLAGS_HASHTREE_DISABLED),
+      InfoImage(vbmeta_image_path_));
+
+  ops_.set_expected_public_key(
+      PublicKeyAVB(base::FilePath("test/data/testkey_rsa2048.pem")));
+
+  // Check that avb_slot_verify() picks the cmdline descriptors based
+  // on their flags value... note that these descriptors are in the
+  // 'system' partition.
+  AvbSlotVerifyData* slot_data = NULL;
+  const char* requested_partitions[] = {"boot", NULL};
+  EXPECT_EQ(AVB_SLOT_VERIFY_RESULT_OK,
+            avb_slot_verify(ops_.avb_ops(),
+                            requested_partitions,
+                            "_a",
+                            AVB_SLOT_VERIFY_FLAGS_NONE,
+                            AVB_HASHTREE_ERROR_MODE_RESTART_AND_INVALIDATE,
+                            &slot_data));
+  EXPECT_NE(nullptr, slot_data);
+  if (hashtree_verification_on) {
+    EXPECT_EQ(
+        "dm=\"1 vroot none ro 1,0 2056 verity 1 "
+        "PARTUUID=1234-fake-guid-for:system_a "
+        "PARTUUID=1234-fake-guid-for:system_a 4096 4096 257 257 sha1 "
+        "e811611467dcd6e8dc4324e45f706c2bdd51db67 d00df00d 2 "
+        "restart_on_corruption ignore_zero_blocks\" root=/dev/dm-0 "
+        "should_be_in_both=1 "
+        "androidboot.vbmeta.device=PARTUUID=1234-fake-guid-for:vbmeta_a "
+        "androidboot.vbmeta.avb_version=1.0 "
+        "androidboot.vbmeta.device_state=locked "
+        "androidboot.vbmeta.hash_alg=sha256 androidboot.vbmeta.size=3456 "
+        "androidboot.vbmeta.digest="
+        "5ee1669b112625322657b83ec932c73dad9b0222011b5aa3e8273f4e0ee025dc "
+        "androidboot.vbmeta.invalidate_on_error=yes "
+        "androidboot.veritymode=enforcing",
+        std::string(slot_data->cmdline));
+  } else {
+    // NOTE: androidboot.veritymode is 'disabled', not 'enforcing' and
+    // androidboot.vbmeta.invalidate_on_error isn't set.
+    EXPECT_EQ(
+        "root=PARTUUID=1234-fake-guid-for:system_a should_be_in_both=1 "
+        "androidboot.vbmeta.device=PARTUUID=1234-fake-guid-for:vbmeta_a "
+        "androidboot.vbmeta.avb_version=1.0 "
+        "androidboot.vbmeta.device_state=locked "
+        "androidboot.vbmeta.hash_alg=sha256 androidboot.vbmeta.size=3456 "
+        "androidboot.vbmeta.digest="
+        "ae792c45a9d898b532ff9625b60043a8d9eae7e6106b9cba31837d50ba40f81c "
+        "androidboot.veritymode=disabled",
+        std::string(slot_data->cmdline));
+  }
+  avb_slot_verify_data_free(slot_data);
+}
+
+TEST_F(AvbSlotVerifyTest, CmdlineWithChainedHashtreeVerificationOff) {
+  CmdlineWithChainedHashtreeVerification(false);
+}
+
+TEST_F(AvbSlotVerifyTest, CmdlineWithChainedHashtreeVerificationOn) {
+  CmdlineWithChainedHashtreeVerification(true);
+}
+
 // In the event that there's no vbmeta partition, we treat the vbmeta
 // struct from 'boot' as the top-level partition. Check that this
 // works.