Amend the apex info for ota package

We have already logged the compressed apexes in the target-files.
Because we want to support the apex metrics during OTA update, also
include the uncompressed apexes in the META/apex_info.pb.

For incremental OTA packages, include the source apex version for
each apex package as well.

Bug: 190244686
Test: unit test
Change-Id: I5cf2647c56c4feb5517f9a81aa1e9abc52515bf1
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index c583d01..28a2f5a 100644
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -687,8 +687,10 @@
               os.path.join(OPTIONS.input_tmp, "IMAGES",
                            "{}.img".format(partition_name))))
 
+
 def AddApexInfo(output_zip):
-  apex_infos = GetApexInfoFromTargetFiles(OPTIONS.input_tmp, 'system')
+  apex_infos = GetApexInfoFromTargetFiles(OPTIONS.input_tmp, 'system',
+                                          compressed_only=False)
   apex_metadata_proto = ota_metadata_pb2.ApexMetadata()
   apex_metadata_proto.apex_info.extend(apex_infos)
   apex_info_bytes = apex_metadata_proto.SerializeToString()
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 229f7e9..bac5e3f 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -1166,14 +1166,12 @@
     else:
       logger.warning("Cannot find care map file in target_file package")
 
-  # Copy apex_info.pb over to generated OTA package.
-  try:
-    apex_info_entry = target_zip.getinfo("META/apex_info.pb")
-    with target_zip.open(apex_info_entry, "r") as zfp:
-      common.ZipWriteStr(output_zip, "apex_info.pb", zfp.read(),
-                         compress_type=zipfile.ZIP_STORED)
-  except KeyError:
-    logger.warning("target_file doesn't contain apex_info.pb %s", target_file)
+  # Add the source apex version for incremental ota updates, and write the
+  # result apex info to the ota package.
+  ota_apex_info = ota_utils.ConstructOtaApexInfo(target_zip, source_file)
+  if ota_apex_info is not None:
+    common.ZipWriteStr(output_zip, "apex_info.pb", ota_apex_info,
+                       compress_type=zipfile.ZIP_STORED)
 
   common.ZipClose(target_zip)
 
diff --git a/tools/releasetools/ota_metadata.proto b/tools/releasetools/ota_metadata.proto
index ed9d0c3..689ce80 100644
--- a/tools/releasetools/ota_metadata.proto
+++ b/tools/releasetools/ota_metadata.proto
@@ -72,6 +72,8 @@
   int64 version = 2;
   bool is_compressed = 3;
   int64 decompressed_size = 4;
+  // Used in OTA
+  int64 source_version = 5;
 }
 
 // Just a container to hold repeated apex_info, so that we can easily serialize
diff --git a/tools/releasetools/ota_utils.py b/tools/releasetools/ota_utils.py
index 104f02f..28c246b 100644
--- a/tools/releasetools/ota_utils.py
+++ b/tools/releasetools/ota_utils.py
@@ -569,3 +569,45 @@
 
   SignFile(temp_zip_name, output_zip_name, OPTIONS.package_key, pw,
            whole_file=True)
+
+
+def ConstructOtaApexInfo(target_zip, source_file=None):
+  """If applicable, add the source version to the apex info."""
+
+  def _ReadApexInfo(input_zip):
+    if "META/apex_info.pb" not in input_zip.namelist():
+      logger.warning("target_file doesn't contain apex_info.pb %s", input_zip)
+      return None
+
+    with input_zip.open("META/apex_info.pb", "r") as zfp:
+      return zfp.read()
+
+  target_apex_string = _ReadApexInfo(target_zip)
+  # Return early if the target apex info doesn't exist or is empty.
+  if not target_apex_string:
+    return target_apex_string
+
+  # If the source apex info isn't available, just return the target info
+  if not source_file:
+    return target_apex_string
+
+  with zipfile.ZipFile(source_file, "r", allowZip64=True) as source_zip:
+    source_apex_string = _ReadApexInfo(source_zip)
+  if not source_apex_string:
+    return target_apex_string
+
+  source_apex_proto = ota_metadata_pb2.ApexMetadata()
+  source_apex_proto.ParseFromString(source_apex_string)
+  source_apex_versions = {apex.package_name: apex.version for apex in
+                          source_apex_proto.apex_info}
+
+  # If the apex package is available in the source build, initialize the source
+  # apex version.
+  target_apex_proto = ota_metadata_pb2.ApexMetadata()
+  target_apex_proto.ParseFromString(target_apex_string)
+  for target_apex in target_apex_proto.apex_info:
+    name = target_apex.package_name
+    if name in source_apex_versions:
+      target_apex.source_version = source_apex_versions[name]
+
+  return target_apex_proto.SerializeToString()
diff --git a/tools/releasetools/test_ota_from_target_files.py b/tools/releasetools/test_ota_from_target_files.py
index 661712a..51def30 100644
--- a/tools/releasetools/test_ota_from_target_files.py
+++ b/tools/releasetools/test_ota_from_target_files.py
@@ -24,7 +24,7 @@
 import test_utils
 from ota_utils import (
     BuildLegacyOtaMetadata, CalculateRuntimeDevicesAndFingerprints,
-    FinalizeMetadata, GetPackageMetadata, PropertyFiles)
+    ConstructOtaApexInfo, FinalizeMetadata, GetPackageMetadata, PropertyFiles)
 from ota_from_target_files import (
     _LoadOemDicts, AbOtaPropertyFiles,
     GetTargetFilesZipForCustomImagesUpdates,
@@ -295,6 +295,35 @@
     uncompressed_apex_size = os.path.getsize(original_apex_filepath)
     self.assertEqual(apex_infos[0].decompressed_size, uncompressed_apex_size)
 
+  @staticmethod
+  def construct_tf_with_apex_info(infos):
+    apex_metadata_proto = ota_metadata_pb2.ApexMetadata()
+    apex_metadata_proto.apex_info.extend(infos)
+
+    output = common.MakeTempFile(suffix='.zip')
+    with zipfile.ZipFile(output, 'w') as zfp:
+      common.ZipWriteStr(zfp, "META/apex_info.pb",
+                         apex_metadata_proto.SerializeToString())
+    return output
+
+  def test_ConstructOtaApexInfo_incremental_package(self):
+    infos = [ota_metadata_pb2.ApexInfo(package_name='com.android.apex.1',
+                                       version=1000, is_compressed=False),
+             ota_metadata_pb2.ApexInfo(package_name='com.android.apex.2',
+                                       version=2000, is_compressed=True)]
+    target_file = self.construct_tf_with_apex_info(infos)
+
+    with zipfile.ZipFile(target_file) as target_zip:
+      info_bytes = ConstructOtaApexInfo(target_zip, source_file=target_file)
+    apex_metadata_proto = ota_metadata_pb2.ApexMetadata()
+    apex_metadata_proto.ParseFromString(info_bytes)
+
+    info_list = apex_metadata_proto.apex_info
+    self.assertEqual(2, len(info_list))
+    self.assertEqual('com.android.apex.1', info_list[0].package_name)
+    self.assertEqual(1000, info_list[0].version)
+    self.assertEqual(1000, info_list[0].source_version)
+
   def test_GetPackageMetadata_retrofitDynamicPartitions(self):
     target_info = common.BuildInfo(self.TEST_TARGET_INFO_DICT, None)
     common.OPTIONS.retrofit_dynamic_partitions = True