Support incremental dev option OTA during signing

This allows incremental dev option to be used on release-key devices.
Boot images are signed during the signing process, and hence the dev
option OTAs(which are derivative of boot image) need to be re-generated.
Previously we only re-generate full OTAs, now we support incrementals
too.

Test: th
Bug: 339658378
Change-Id: I4755379b49ff8adf351ccaf76fe38f19c1685e9e

tools/releasetools/Android.bp

@@ -555,6 +555,7 @@ python_binary_host {
         "releasetools_common",
         "ota_metadata_proto",
         "ota_utils_lib",
+        "update_payload",
     ],
 }
 

tools/releasetools/sign_target_files_apks.py

@@ -190,10 +190,10 @@ from xml.etree import ElementTree
 
 import add_img_to_target_files
 import ota_from_raw_img
-import ota_utils
 import apex_utils
 import common
 import payload_signer
+import update_payload
 from payload_signer import SignOtaPackage, PAYLOAD_BIN
 
 
@@ -223,6 +223,7 @@ OPTIONS.vendor_otatools = None
 OPTIONS.allow_gsi_debug_sepolicy = False
 OPTIONS.override_apk_keys = None
 OPTIONS.override_apex_keys = None
+OPTIONS.input_tmp = None
 
 
 AVB_FOOTER_ARGS_BY_PARTITION = {
@@ -583,16 +584,13 @@ def IsBuildPropFile(filename):
 
 def RegenerateKernelPartitions(input_tf_zip: zipfile.ZipFile, output_tf_zip: zipfile.ZipFile, misc_info):
   """Re-generate boot and dtbo partitions using new signing configuration"""
+  files_to_unzip = [
+      "PREBUILT_IMAGES/*", "BOOTABLE_IMAGES/*.img", "*/boot_16k.img", "*/dtbo_16k.img"]
   if OPTIONS.input_tmp is None:
-    OPTIONS.input_tmp = common.UnzipTemp(input_tf_zip.filename, [
-                                "*/boot.img", "*/dtbo.img"])
+    OPTIONS.input_tmp = common.UnzipTemp(input_tf_zip.filename, files_to_unzip)
   else:
-    common.UnzipToDir(input_tf_zip, OPTIONS.input_tmp, [
-                                "*/boot.img", "*/dtbo.img"])
+    common.UnzipToDir(input_tf_zip.filename, OPTIONS.input_tmp, files_to_unzip)
   unzip_dir = OPTIONS.input_tmp
-  image_dir = os.path.join(unzip_dir, "IMAGES")
-  shutil.rmtree(image_dir)
-  os.makedirs(image_dir, exist_ok=True)
 
   boot_image = common.GetBootableImage(
       "IMAGES/boot.img", "boot.img", unzip_dir, "BOOT", misc_info)
@@ -605,33 +603,59 @@ def RegenerateKernelPartitions(input_tf_zip: zipfile.ZipFile, output_tf_zip: zip
   return unzip_dir
 
 
-def RegenerateBootOTA(input_tf_zip: zipfile.ZipFile, output_tf_zip: zipfile.ZipFile, misc_info, filename, input_ota):
-  if filename not in ["VENDOR/boot_otas/boot_ota_4k.zip", "SYSTEM/boot_otas/boot_ota_4k.zip"]:
-    # We only need to re-generate 4K boot OTA, for other OTA packages
-    # simply copy as is
-    with input_tf_zip.open(filename, "r") as in_fp:
-      shutil.copyfileobj(in_fp, input_ota)
-      input_ota.flush()
+def RegenerateBootOTA(input_tf_zip: zipfile.ZipFile, filename, input_ota):
+  with input_tf_zip.open(filename, "r") as in_fp:
+    payload = update_payload.Payload(in_fp)
+  is_incremental = any([part.HasField('old_partition_info')
+                        for part in payload.manifest.partitions])
+  is_boot_ota = filename.startswith(
+      "VENDOR/boot_otas/") or filename.startswith("SYSTEM/boot_otas/")
+  if not is_boot_ota:
+    return
+  is_4k_boot_ota = filename in [
+      "VENDOR/boot_otas/boot_ota_4k.zip", "SYSTEM/boot_otas/boot_ota_4k.zip"]
+  # Only 4K boot image is re-generated, so if 16K boot ota isn't incremental,
+  # we do not need to re-generate
+  if not is_4k_boot_ota and not is_incremental:
     return
-  timestamp = misc_info["build.prop"].GetProp(
-      "ro.system.build.date.utc")
-  unzip_dir = RegenerateKernelPartitions(
-      input_tf_zip, output_tf_zip, misc_info)
-  signed_boot_image = os.path.join(unzip_dir, "IMAGES/boot.img")
-  signed_dtbo_image = os.path.join(unzip_dir, "IMAGES/dtbo.img")
 
+  timestamp = str(payload.manifest.max_timestamp)
+  partitions = [part.partition_name for part in payload.manifest.partitions]
+  unzip_dir = OPTIONS.input_tmp
+  signed_boot_image = os.path.join(unzip_dir, "IMAGES", "boot.img")
   if not os.path.exists(signed_boot_image):
     logger.warn("Need to re-generate boot OTA {} but failed to get signed boot image. 16K dev option will be impacted, after rolling back to 4K user would need to sideload/flash their device to continue receiving OTAs.")
     return
-  logger.info(
-      "Re-generating boot OTA {} with timestamp {}".format(filename, timestamp))
+  signed_dtbo_image = os.path.join(unzip_dir, "IMAGES", "dtbo.img")
+  if "dtbo" in partitions and not os.path.exists(signed_dtbo_image):
+    raise ValueError(
+        "Boot OTA {} has dtbo partition, but no dtbo image found in target files.".format(filename))
+  if is_incremental:
+    signed_16k_boot_image = os.path.join(
+        unzip_dir, "IMAGES", "boot_16k.img")
+    signed_16k_dtbo_image = os.path.join(
+        unzip_dir, "IMAGES", "dtbo_16k.img")
+    if is_4k_boot_ota:
+      if os.path.exists(signed_16k_boot_image):
+        signed_boot_image = signed_16k_boot_image + ":" + signed_boot_image
+      if os.path.exists(signed_16k_dtbo_image):
+        signed_dtbo_image = signed_16k_dtbo_image + ":" + signed_dtbo_image
+    else:
+      if os.path.exists(signed_16k_boot_image):
+        signed_boot_image += ":" + signed_16k_boot_image
+      if os.path.exists(signed_16k_dtbo_image):
+        signed_dtbo_image += ":" + signed_16k_dtbo_image
+
+
   args = ["ota_from_raw_img", "--package_key", OPTIONS.package_key,
           "--max_timestamp", timestamp, "--output", input_ota.name]
-  if os.path.exists(signed_dtbo_image):
+  if "dtbo" in partitions:
     args.extend(["--partition_name", "boot,dtbo",
                 signed_boot_image, signed_dtbo_image])
   else:
     args.extend(["--partition_name", "boot", signed_boot_image])
+  logger.info(
+      "Re-generating boot OTA {} using cmd {}".format(filename, args))
   ota_from_raw_img.main(args)
 
 
@@ -657,6 +681,8 @@ def ProcessTargetFiles(input_tf_zip: zipfile.ZipFile, output_tf_zip: zipfile.Zip
   if misc_info.get('avb_enable') == 'true':
     RewriteAvbProps(misc_info)
 
+  RegenerateKernelPartitions(input_tf_zip, output_tf_zip, misc_info)
+
   for info in input_tf_zip.infolist():
     filename = info.filename
     if filename.startswith("IMAGES/"):
@@ -667,10 +693,10 @@ def ProcessTargetFiles(input_tf_zip: zipfile.ZipFile, output_tf_zip: zipfile.Zip
     if filename.startswith("OTA/") and filename.endswith(".img"):
       continue
 
-    data = input_tf_zip.read(filename)
-    out_info = copy.copy(info)
     (is_apk, is_compressed, should_be_skipped) = GetApkFileInfo(
         filename, compressed_extension, OPTIONS.skip_apks_with_path_prefix)
+    data = input_tf_zip.read(filename)
+    out_info = copy.copy(info)
 
     if is_apk and should_be_skipped:
       # Copy skipped APKs verbatim.
@@ -734,8 +760,7 @@ def ProcessTargetFiles(input_tf_zip: zipfile.ZipFile, output_tf_zip: zipfile.Zip
     elif filename.endswith(".zip") and IsEntryOtaPackage(input_tf_zip, filename):
       logger.info("Re-signing OTA package {}".format(filename))
       with tempfile.NamedTemporaryFile() as input_ota, tempfile.NamedTemporaryFile() as output_ota:
-        RegenerateBootOTA(input_tf_zip, output_tf_zip,
-                          misc_info, filename, input_ota)
+        RegenerateBootOTA(input_tf_zip, filename, input_ota)
 
         SignOtaPackage(input_ota.name, output_ota.name)
         common.ZipWrite(output_tf_zip, output_ota.name, filename,