From 65029a2801f42b0609360f56ef8a98c5ba316e8d Mon Sep 17 00:00:00 2001
From: Kelvin Zhang <zhangkelvin@google.com>
Date: Tue, 3 Nov 2020 10:07:51 -0500
Subject: [PATCH] Include complete boot image on builds containing gki apex

Test: treehugger
Bug: 171906589
Change-Id: I231dac425ab14602492346006e331e4666173614
---
 tools/releasetools/ota_from_target_files.py | 51 +++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 869b713796..3c98c80307 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -217,10 +217,13 @@ from __future__ import print_function
 
 import logging
 import multiprocessing
+import os
 import os.path
+import re
 import shlex
 import shutil
 import struct
+import subprocess
 import sys
 import zipfile
 
@@ -968,6 +971,50 @@ def GeneratePartitionTimestampFlagsDowngrade(pre_partition_state, post_partition
     ",".join([key + ":" + val for (key, val) in partition_timestamps.items()])
     ]
 
+def IsSparseImage(filepath):
+  with open(filepath, 'rb') as fp:
+    # Magic for android sparse image format
+    # https://source.android.com/devices/bootloader/images
+    return fp.read(4) == b'\x3A\xFF\x26\xED'
+
+def SupportsMainlineGkiUpdates(target_file):
+  """Return True if the build supports MainlineGKIUpdates.
+
+  This function scans the product.img file in IMAGES/ directory for
+  pattern |*/apex/com.android.gki.*.apex|. If there are files
+  matching this pattern, conclude that build supports mainline
+  GKI and return True
+
+  Args:
+    target_file: Path to a target_file.zip, or an extracted directory
+  Return:
+    True if thisb uild supports Mainline GKI Updates.
+  """
+  if target_file is None:
+    return False
+  if os.path.isfile(target_file):
+    target_file = common.UnzipTemp(target_file, ["IMAGES/product.img"])
+  if not os.path.isdir(target_file):
+    assert os.path.isdir(target_file), \
+        "{} must be a path to zip archive or dir containing extracted"\
+        " target_files".format(target_file)
+  image_file = os.path.join(target_file, "IMAGES", "product.img")
+
+  if not os.path.isfile(image_file):
+    return False
+
+  if IsSparseImage(image_file):
+    # Unsparse the image
+    tmp_img = common.MakeTempFile(suffix=".img")
+    subprocess.check_output(["simg2img", image_file, tmp_img])
+    image_file = tmp_img
+
+  cmd = ["debugfs_static", "-R", "ls -p /apex", image_file]
+  output = subprocess.check_output(cmd).decode()
+
+  pattern = re.compile(r"com\.android\.gki\..*\.apex")
+  return pattern.search(output) is not None
+
 def GenerateAbOtaPackage(target_file, output_file, source_file=None):
   """Generates an Android OTA package that has A/B update payload."""
   # Stage the output zip package for package signing.
@@ -1039,6 +1086,10 @@ def GenerateAbOtaPackage(target_file, output_file, source_file=None):
 
   additional_args += ["--max_timestamp", max_timestamp]
 
+  if SupportsMainlineGkiUpdates(source_file):
+    logger.info("Detected build with mainline GKI, include full boot image.")
+    additional_args.extend(["--full_boot", "true"])
+
   payload.Generate(
       target_file,
       source_file,