Allow dtb image to be empty in vendor_boot image

Currently, building vendor_boot image will check size of
dtb image and abort if not specified. We want to allow this condition
and make user to build dtb whether in vendor_boot or vendor_kernel_boot.

Bug: 222429295
Bug: 214409109
Test: atest --host mkbootimg_test
Signed-off-by: Lucas Wei <lucaswei@google.com>
Change-Id: I7a96ddc700d2ca43bd1063fd634ef8af76e94005

3 files changed, 78 insertions, 5 deletions

diff --git a/mkbootimg.py b/mkbootimg.py
index 5c65e2c..ec29581 100755
--- a/mkbootimg.py
+++ b/mkbootimg.py
@@ -141,9 +141,6 @@ def write_header_v3_and_above(args):
 
 
 def write_vendor_boot_header(args):
-    if filesize(args.dtb) == 0:
-        raise ValueError('DTB image must not be empty.')
-
     if args.header_version > 3:
         vendor_ramdisk_size = args.vendor_ramdisk_total_size
         vendor_boot_header_size = VENDOR_BOOT_IMAGE_HEADER_V4_SIZE
diff --git a/tests/mkbootimg_test.py b/tests/mkbootimg_test.py
index 28f47f0..e691e30 100644
--- a/tests/mkbootimg_test.py
+++ b/tests/mkbootimg_test.py
@@ -760,6 +760,80 @@ class MkbootimgTest(unittest.TestCase):
             self.assertEqual(raw_vendor_cmdline,
                              vendor_cmdline.encode() + b'\x00')
 
+    def test_vendor_boot_v4_without_dtb(self):
+        """Tests building vendor_boot version 4 without dtb image."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
+            ramdisk = generate_test_file(
+                os.path.join(temp_out_dir, 'ramdisk'), 0x1000)
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--header_version', '4',
+                '--vendor_boot', vendor_boot_img,
+                '--vendor_ramdisk', ramdisk,
+            ]
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', vendor_boot_img,
+                '--out', os.path.join(temp_out_dir, 'out'),
+            ]
+            expected_output = [
+                'boot magic: VNDRBOOT',
+                'vendor boot image header version: 4',
+                'dtb size: 0',
+            ]
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            output = [line.strip() for line in result.stdout.splitlines()]
+            if not subsequence_of(expected_output, output):
+                msg = '\n'.join([
+                    'Unexpected unpack_bootimg output:',
+                    'Expected:',
+                    ' ' + '\n '.join(expected_output),
+                    '',
+                    'Actual:',
+                    ' ' + '\n '.join(output),
+                ])
+                self.fail(msg)
+
+    def test_unpack_vendor_boot_image_v4_without_dtb(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity when no dtb image."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
+            vendor_boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'vendor_boot.img.reconstructed')
+            ramdisk = generate_test_file(
+                os.path.join(temp_out_dir, 'ramdisk'), 0x121212)
+
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--header_version', '4',
+                '--vendor_boot', vendor_boot_img,
+                '--vendor_ramdisk', ramdisk,
+            ]
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', vendor_boot_img,
+                '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
+            ]
+            subprocess.run(mkbootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--vendor_boot', vendor_boot_img_reconstructed,
+            ]
+            unpack_format_args = shlex.split(result.stdout)
+            mkbootimg_cmds.extend(unpack_format_args)
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
+                'reconstructed vendor_boot image differ from the original')
+
 
 # I don't know how, but we need both the logger configuration and verbosity
 # level > 2 to make atest work. And yes this line needs to be at the very top
diff --git a/unpack_bootimg.py b/unpack_bootimg.py
index 437408f..9a515ee 100755
--- a/unpack_bootimg.py
+++ b/unpack_bootimg.py
@@ -353,7 +353,8 @@ class VendorBootImageInfoFormatter:
         args.extend(['--vendor_cmdline', self.cmdline])
         args.extend(['--board', self.product_name])
 
-        args.extend(['--dtb', os.path.join(self.image_dir, 'dtb')])
+        if self.dtb_size > 0:
+          args.extend(['--dtb', os.path.join(self.image_dir, 'dtb')])
 
         if self.header_version > 3:
             args.extend(['--vendor_bootconfig',
@@ -451,7 +452,8 @@ def unpack_vendor_boot_image(args):
 
     dtb_offset = page_size * (num_boot_header_pages + num_boot_ramdisk_pages
                              ) # header + vendor_ramdisk
-    image_info_list.append((dtb_offset, info.dtb_size, 'dtb'))
+    if info.dtb_size > 0:
+      image_info_list.append((dtb_offset, info.dtb_size, 'dtb'))
 
     create_out_dir(args.out)
     for offset, size, name in image_info_list: