Merge "mkbootimg: Add '--boot_signature' and mark '--gki-*' as deprecated"
diff --git a/mkbootimg.py b/mkbootimg.py
index 3694f15..05aaf38 100755
--- a/mkbootimg.py
+++ b/mkbootimg.py
@@ -105,6 +105,13 @@
     return dtbo_offset
 
 
+def should_add_legacy_gki_boot_signature(args):
+    if (args.boot_signature is None and args.gki_signing_key and
+            args.gki_signing_algorithm):
+        return True
+    return False
+
+
 def write_header_v3_and_above(args):
     if args.header_version > 3:
         boot_header_size = BOOT_IMAGE_HEADER_V4_SIZE
@@ -127,7 +134,12 @@
                            args.cmdline))
     if args.header_version >= 4:
         # The signature used to verify boot image v4.
-        args.output.write(pack('I', BOOT_IMAGE_V4_SIGNATURE_SIZE))
+        boot_signature_size = 0
+        if args.boot_signature:
+            boot_signature_size = filesize(args.boot_signature)
+        elif should_add_legacy_gki_boot_signature(args):
+            boot_signature_size = BOOT_IMAGE_V4_SIGNATURE_SIZE
+        args.output.write(pack('I', boot_signature_size))
     pad_file(args.output, BOOT_IMAGE_HEADER_V3_PAGESIZE)
 
 
@@ -536,14 +548,8 @@
                         help='boot image header version')
     parser.add_argument('-o', '--output', type=FileType('wb'),
                         help='output file name')
-    parser.add_argument('--gki_signing_algorithm',
-                        help='GKI signing algorithm to use')
-    parser.add_argument('--gki_signing_key',
-                        help='path to RSA private key file')
-    parser.add_argument('--gki_signing_signature_args', default='',
-                        help='other hash arguments passed to avbtool')
-    parser.add_argument('--gki_signing_avbtool_path', default='avbtool',
-                        help='path to avbtool for boot signature generation')
+    parser.add_argument('--boot_signature', type=FileType('rb'),
+                        help='path to the GKI certificate file')
     parser.add_argument('--vendor_boot', type=FileType('wb'),
                         help='vendor boot output file name')
     parser.add_argument('--vendor_ramdisk', type=FileType('rb'),
@@ -551,6 +557,19 @@
     parser.add_argument('--vendor_bootconfig', type=FileType('rb'),
                         help='path to the vendor bootconfig file')
 
+    gki_2_0_signing_args = parser.add_argument_group(
+        '[DEPRECATED] GKI 2.0 signing arguments')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_algorithm', help='GKI signing algorithm to use')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_key', help='path to RSA private key file')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_signature_args', default='',
+        help='other hash arguments passed to avbtool')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_avbtool_path', default='avbtool',
+        help='path to avbtool for boot signature generation')
+
     args, extra_args = parser.parse_known_args()
     if args.vendor_boot is not None and args.header_version > 3:
         extra_args = parse_vendor_ramdisk_args(args, extra_args)
@@ -576,15 +595,19 @@
     vbmeta partition) via the Android Verified Boot process, when the
     device boots.
     """
-    args.output.flush()  # Flush the buffer for signature calculation.
 
-    # Appends zeros if the signing key is not specified.
-    if not args.gki_signing_key or not args.gki_signing_algorithm:
-        zeros = b'\x00' * BOOT_IMAGE_V4_SIGNATURE_SIZE
-        args.output.write(zeros)
-        pad_file(args.output, pagesize)
+    if args.boot_signature:
+        write_padded_file(args.output, args.boot_signature, pagesize)
         return
 
+    if not should_add_legacy_gki_boot_signature(args):
+        return
+
+    # Fallback to the legacy certificating method.
+
+    # Flush the buffer for signature calculation.
+    args.output.flush()
+
     # Outputs the signed vbmeta to a separate file, then append to boot.img
     # as the boot signature.
     with tempfile.TemporaryDirectory() as temp_out_dir:
diff --git a/tests/mkbootimg_test.py b/tests/mkbootimg_test.py
index ae5cf6b..1e13d55 100644
--- a/tests/mkbootimg_test.py
+++ b/tests/mkbootimg_test.py
@@ -34,8 +34,6 @@
 VENDOR_BOOT_ARGS_OFFSET = 28
 VENDOR_BOOT_ARGS_SIZE = 2048
 
-BOOT_IMAGE_V4_SIGNATURE_SIZE = 4096
-
 TEST_KERNEL_CMDLINE = (
     'printk.devkmsg=on firmware_class.path=/vendor/etc/ init=/init '
     'kfence.sample_interval=500 loop.max_part=7 bootconfig'
@@ -86,7 +84,7 @@
         # C0103: invalid-name for maxDiff.
         self.maxDiff = None  # pylint: disable=C0103
 
-    def _test_boot_image_v4_signature(self, avbtool_path):
+    def _test_legacy_boot_image_v4_signature(self, avbtool_path):
         """Tests the boot_signature in boot.img v4."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             boot_img = os.path.join(temp_out_dir, 'boot.img')
@@ -162,15 +160,16 @@
 
             self.assertEqual(result.stdout, expected_boot_signature_info)
 
-    def test_boot_image_v4_signature_without_avbtool_path(self):
+    def test_legacy_boot_image_v4_signature_without_avbtool_path(self):
         """Boot signature generation without --gki_signing_avbtool_path."""
-        self._test_boot_image_v4_signature(avbtool_path=None)
+        self._test_legacy_boot_image_v4_signature(avbtool_path=None)
 
-    def test_boot_image_v4_signature_with_avbtool_path(self):
+    def test_legacy_boot_image_v4_signature_with_avbtool_path(self):
         """Boot signature generation with --gki_signing_avbtool_path."""
-        self._test_boot_image_v4_signature(avbtool_path=self._avbtool_path)
+        self._test_legacy_boot_image_v4_signature(
+            avbtool_path=self._avbtool_path)
 
-    def test_boot_image_v4_signature_exceed_size(self):
+    def test_legacy_boot_image_v4_signature_exceed_size(self):
         """Tests the boot signature size exceeded in a boot image version 4."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             boot_img = os.path.join(temp_out_dir, 'boot.img')
@@ -205,7 +204,7 @@
                 self.assertIn('ValueError: boot sigature size is > 4096',
                               e.stderr)
 
-    def test_boot_image_v4_signature_zeros(self):
+    def test_boot_image_v4_signature_empty(self):
         """Tests no boot signature in a boot image version 4."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             boot_img = os.path.join(temp_out_dir, 'boot.img')
@@ -235,11 +234,8 @@
             subprocess.run(mkbootimg_cmds, check=True)
             subprocess.run(unpack_bootimg_cmds, check=True)
 
-            boot_signature = os.path.join(
-                temp_out_dir, 'out', 'boot_signature')
-            with open(boot_signature) as f:
-                zeros = '\x00' * BOOT_IMAGE_V4_SIGNATURE_SIZE
-                self.assertEqual(f.read(), zeros)
+            boot_signature = os.path.join(temp_out_dir, 'out', 'boot_signature')
+            self.assertFalse(os.path.exists(boot_signature))
 
     def test_vendor_boot_v4(self):
         """Tests vendor_boot version 4."""
@@ -418,6 +414,48 @@
                 filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
                 'reconstructed vendor_boot image differ from the original')
 
+    def test_unpack_boot_image_v4(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img = os.path.join(temp_out_dir, 'boot.img')
+            boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'boot.img.reconstructed')
+            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
+                                        0x1000)
+            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
+                                         0x1000)
+            boot_signature = generate_test_file(
+                os.path.join(temp_out_dir, 'boot_signature'), 0x800)
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--header_version', '4',
+                '--kernel', kernel,
+                '--ramdisk', ramdisk,
+                '--cmdline', TEST_KERNEL_CMDLINE,
+                '--boot_signature', boot_signature,
+                '--output', boot_img,
+            ]
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', 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',
+                '--out', boot_img_reconstructed,
+            ]
+            mkbootimg_cmds.extend(shlex.split(result.stdout))
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(boot_img, boot_img_reconstructed),
+                'reconstructed boot image differ from the original')
+
     def test_unpack_boot_image_v3(self):
         """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
diff --git a/unpack_bootimg.py b/unpack_bootimg.py
index 2b176e5..ae59429 100755
--- a/unpack_bootimg.py
+++ b/unpack_bootimg.py
@@ -53,6 +53,8 @@
 
 
 def format_os_version(os_version):
+    if os_version == 0:
+        return None
     a = os_version >> 14
     b = os_version >> 7 & ((1<<7) - 1)
     c = os_version & ((1<<7) - 1)
@@ -60,6 +62,8 @@
 
 
 def format_os_patch_level(os_patch_level):
+    if os_patch_level == 0:
+        return None
     y = os_patch_level >> 4
     y += 2000
     m = os_patch_level & ((1<<4) - 1)
@@ -130,12 +134,18 @@
     def format_mkbootimg_argument(self):
         args = []
         args.extend(['--header_version', str(self.header_version)])
-        args.extend(['--os_version', self.os_version])
-        args.extend(['--os_patch_level', self.os_patch_level])
+        if self.os_version:
+            args.extend(['--os_version', self.os_version])
+        if self.os_patch_level:
+            args.extend(['--os_patch_level', self.os_patch_level])
 
         args.extend(['--kernel', os.path.join(self.image_dir, 'kernel')])
         args.extend(['--ramdisk', os.path.join(self.image_dir, 'ramdisk')])
 
+        if self.header_version >= 4 and self.boot_signature_size > 0:
+            args.extend(['--boot_signature',
+                         os.path.join(self.image_dir, 'boot_signature')])
+
         if self.header_version <= 2:
             if self.second_size > 0:
                 args.extend(['--second',