Add an option to set VABC compression algo

When testing on cuttlefish, space consumption isn't a concern, so
compression can be disabled to save time. We also want to test the
non-compressed code path, as we expect some OEMs to use userspace COW
w/o compression.

If no --vabc_compression_param option specified to
ota_from_target_files, we will use the compression algo specified in
target_files's dynamic_partition_info.txt.

Bug: 219966986
Test: th
Change-Id: Ie18064b15938550ad6fb4efc86fc100c12321dc0

tools/releasetools/ota_from_target_files.py

@@ -241,6 +241,9 @@ A/B OTA specific options
   --spl_downgrade
       Force generate an SPL downgrade OTA. Only needed if target build has an
       older SPL.
+
+  --vabc_compression_param
+      Compression algorithm to be used for VABC. Available options: gz, brotli, none
 """
 
 from __future__ import print_function
@@ -312,6 +315,7 @@ OPTIONS.force_minor_version = None
 OPTIONS.compressor_types = None
 OPTIONS.enable_zucchini = True
 OPTIONS.enable_lz4diff = False
+OPTIONS.vabc_compression_param = None
 
 POSTINSTALL_CONFIG = 'META/postinstall_config.txt'
 DYNAMIC_PARTITION_INFO = 'META/dynamic_partitions_info.txt'
@@ -651,6 +655,24 @@ class AbOtaPropertyFiles(StreamingPropertyFiles):
     return (payload_offset, metadata_total)
 
 
+def ModifyVABCCompressionParam(content, algo):
+  """ Update update VABC Compression Param in dynamic_partitions_info.txt
+  Args:
+    content: The string content of dynamic_partitions_info.txt
+    algo: The compression algorithm should be used for VABC. See
+          https://cs.android.com/android/platform/superproject/+/master:system/core/fs_mgr/libsnapshot/cow_writer.cpp;l=127;bpv=1;bpt=1?q=CowWriter::ParseOptions&sq=
+  Returns:
+    Updated content of dynamic_partitions_info.txt , with custom compression algo
+  """
+  output_list = []
+  for line in content.splitlines():
+    if line.startswith("virtual_ab_compression_method="):
+      continue
+    output_list.append(line)
+  output_list.append("virtual_ab_compression_method="+algo)
+  return "\n".join(output_list)
+
+
 def UpdatesInfoForSpecialUpdates(content, partitions_filter,
                                  delete_keys=None):
   """ Updates info file for secondary payload generation, partial update, etc.
@@ -805,6 +827,27 @@ def ParseInfoDict(target_file_path):
     return common.LoadInfoDict(zfp)
 
 
+def GetTargetFilesZipForCustomVABCCompression(input_file, vabc_compression_param):
+  """Returns a target-files.zip with a custom VABC compression param.
+  Args:
+    input_file: The input target-files.zip path
+    vabc_compression_param: Custom Virtual AB Compression algorithm
+
+  Returns:
+    The path to modified target-files.zip
+  """
+  target_file = common.MakeTempFile(prefix="targetfiles-", suffix=".zip")
+  shutil.copyfile(input_file, target_file)
+  common.ZipDelete(target_file, DYNAMIC_PARTITION_INFO)
+  with zipfile.ZipFile(input_file, 'r', allowZip64=True) as zfp:
+    dynamic_partition_info = zfp.read(DYNAMIC_PARTITION_INFO).decode()
+    dynamic_partition_info = ModifyVABCCompressionParam(
+        dynamic_partition_info, vabc_compression_param)
+    with zipfile.ZipFile(target_file, "a", allowZip64=True) as output_zip:
+      output_zip.writestr(DYNAMIC_PARTITION_INFO, dynamic_partition_info)
+  return target_file
+
+
 def GetTargetFilesZipForPartialUpdates(input_file, ab_partitions):
   """Returns a target-files.zip for partial ota update package generation.
 
@@ -879,6 +922,9 @@ def GetTargetFilesZipForPartialUpdates(input_file, ab_partitions):
       content = input_zip.read(info_file).decode()
       modified_info = UpdatesInfoForSpecialUpdates(
           content, lambda p: p in ab_partitions)
+      if OPTIONS.vabc_compression_param and info_file == DYNAMIC_PARTITION_INFO:
+        modified_info = ModifyVABCCompressionParam(
+            modified_info, OPTIONS.vabc_compression_param)
       common.ZipWriteStr(partial_target_zip, info_file, modified_info)
 
     # TODO(xunchang) handle META/postinstall_config.txt'
@@ -1114,7 +1160,7 @@ def GenerateAbOtaPackage(target_file, output_file, source_file=None):
   # Source build's update_engine must be able to write XOR ops, and target
   # build's snapuserd must be able to interpret XOR ops.
   if not target_info.is_vabc_xor or OPTIONS.disable_vabc or \
-      (source_info is not None and not source_info.is_vabc_xor):
+          (source_info is not None and not source_info.is_vabc_xor):
     logger.info("VABC XOR Not supported, disabling")
     OPTIONS.enable_vabc_xor = False
   additional_args = []
@@ -1132,6 +1178,9 @@ def GenerateAbOtaPackage(target_file, output_file, source_file=None):
     target_file = GetTargetFilesZipForPartialUpdates(target_file,
                                                      OPTIONS.partial)
     additional_args += ["--is_partial_update", "true"]
+  elif OPTIONS.vabc_compression_param:
+    target_file = GetTargetFilesZipForCustomVABCCompression(
+        target_file, OPTIONS.vabc_compression_param)
   elif OPTIONS.skip_postinstall:
     target_file = GetTargetFilesZipWithoutPostinstallConfig(target_file)
   # Target_file may have been modified, reparse ab_partitions
@@ -1166,7 +1215,7 @@ def GenerateAbOtaPackage(target_file, output_file, source_file=None):
                       str(OPTIONS.enable_zucchini).lower()]
 
   if not ota_utils.IsLz4diffCompatible(source_file, target_file):
-    logger.warn(
+    logger.warning(
         "Source build doesn't support lz4diff, or source/target don't have compatible lz4diff versions. Disabling lz4diff.")
     OPTIONS.enable_lz4diff = False
 
@@ -1377,6 +1426,8 @@ def main(argv):
     elif o == "--enable_lz4diff":
       assert a.lower() in ["true", "false"]
       OPTIONS.enable_lz4diff = a.lower() != "false"
+    elif o == "--vabc_compression_param":
+      OPTIONS.vabc_compression_param = a.lower()
     else:
       return False
     return True
@@ -1426,6 +1477,7 @@ def main(argv):
                                  "compressor_types=",
                                  "enable_zucchin=",
                                  "enable_lz4diff=",
+                                 "vabc_compression_param=",
                              ], extra_option_handler=option_handler)
 
   if len(args) != 2: