Fix a bug in computing streaming property of payload.bin

When computing the data offset of an entry in zip file, we used length
of extra field from central directory. That is correct most of the time
but wrong if the extra field in central directory has different length
than the one in local file directory. Since python's zipfile doesn't
provide an API to access local file header, we need to parse local file
header ourselves and extract length of extra field.

An incorrect offset will cause magic mismatch error from update_engine,
as update_engine expects to find uncompressed payload at the recorded
offset.

Test: th, partner verification
Bug: 191443484

Change-Id: Id670cd79b0bd65adffaaa5224ae4f8065d66b358

4 files changed, 95 insertions, 15 deletions

diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index b8c812d80d..f3b58f87af 100644
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -1030,8 +1030,5 @@ if __name__ == '__main__':
   try:
     common.CloseInheritedPipes()
     main(sys.argv[1:])
-  except common.ExternalError:
-    logger.exception("\n   ERROR:\n")
-    sys.exit(1)
   finally:
     common.Cleanup()
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index bcc10bd3f1..1f3022bc69 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -241,7 +241,7 @@ import care_map_pb2
 import common
 import ota_utils
 from ota_utils import (UNZIP_PATTERN, FinalizeMetadata, GetPackageMetadata,
-                       PropertyFiles, SECURITY_PATCH_LEVEL_PROP_NAME)
+                       PropertyFiles, SECURITY_PATCH_LEVEL_PROP_NAME, GetZipEntryOffset)
 import target_files_diff
 from check_target_files_vintf import CheckVintfIfTrebleEnabled
 from non_ab_ota import GenerateNonAbOtaPackage
@@ -603,20 +603,20 @@ class AbOtaPropertyFiles(StreamingPropertyFiles):
     payload, till the end of 'medatada_signature_message'.
     """
     payload_info = input_zip.getinfo('payload.bin')
-    payload_offset = payload_info.header_offset
-    payload_offset += zipfile.sizeFileHeader
-    payload_offset += len(payload_info.extra) + len(payload_info.filename)
-    payload_size = payload_info.file_size
+    (payload_offset, payload_size) = GetZipEntryOffset(input_zip, payload_info)
 
-    with input_zip.open('payload.bin') as payload_fp:
-      header_bin = payload_fp.read(24)
+    # Read the underlying raw zipfile at specified offset
+    payload_fp = input_zip.fp
+    payload_fp.seek(payload_offset)
+    header_bin = payload_fp.read(24)
 
     # network byte order (big-endian)
     header = struct.unpack("!IQQL", header_bin)
 
     # 'CrAU'
     magic = header[0]
-    assert magic == 0x43724155, "Invalid magic: {:x}".format(magic)
+    assert magic == 0x43724155, "Invalid magic: {:x}, computed offset {}" \
+        .format(magic, payload_offset)
 
     manifest_size = header[2]
     metadata_signature_size = header[3]
diff --git a/tools/releasetools/ota_utils.py b/tools/releasetools/ota_utils.py
index 28c246b662..573700930c 100644
--- a/tools/releasetools/ota_utils.py
+++ b/tools/releasetools/ota_utils.py
@@ -16,6 +16,7 @@ import copy
 import itertools
 import logging
 import os
+import struct
 import zipfile
 
 import ota_metadata_pb2
@@ -399,6 +400,35 @@ def CalculateRuntimeDevicesAndFingerprints(default_build_info,
   return device_names, fingerprints
 
 
+def GetZipEntryOffset(zfp, entry_info):
+  """Get offset to a beginning of a particular zip entry
+  Args:
+    fp: zipfile.ZipFile
+    entry_info: zipfile.ZipInfo
+
+  Returns:
+    (offset, size) tuple
+  """
+  # Don't use len(entry_info.extra). Because that returns size of extra
+  # fields in central directory. We need to look at local file directory,
+  # as these two might have different sizes.
+
+  # We cannot work with zipfile.ZipFile instances, we need a |fp| for the underlying file.
+  zfp = zfp.fp
+  zfp.seek(entry_info.header_offset)
+  data = zfp.read(zipfile.sizeFileHeader)
+  fheader = struct.unpack(zipfile.structFileHeader, data)
+  # Last two fields of local file header are filename length and
+  # extra length
+  filename_len = fheader[-2]
+  extra_len = fheader[-1]
+  offset = entry_info.header_offset
+  offset += zipfile.sizeFileHeader
+  offset += filename_len + extra_len
+  size = entry_info.file_size
+  return (offset, size)
+
+
 class PropertyFiles(object):
   """A class that computes the property-files string for an OTA package.
 
@@ -517,10 +547,7 @@ class PropertyFiles(object):
     def ComputeEntryOffsetSize(name):
       """Computes the zip entry offset and size."""
       info = zip_file.getinfo(name)
-      offset = info.header_offset
-      offset += zipfile.sizeFileHeader
-      offset += len(info.extra) + len(info.filename)
-      size = info.file_size
+      (offset, size) = GetZipEntryOffset(zip_file, info)
       return '%s:%d:%d' % (os.path.basename(name), offset, size)
 
     tokens = []
diff --git a/tools/releasetools/test_ota_utils.py b/tools/releasetools/test_ota_utils.py
new file mode 100644
index 0000000000..9a82e6f9f1
--- /dev/null
+++ b/tools/releasetools/test_ota_utils.py
@@ -0,0 +1,56 @@
+# Copyright (C) 2021 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import unittest
+import io
+import ota_utils
+import zipfile
+
+
+class TestZipEntryOffset(unittest.TestCase):
+  def test_extra_length_differ(self):
+      # This is a magic zip file such that:
+      # 1. It has 1 entry, `file.txt'`
+      # 2. The central directory entry for the entry contains an extra field of#
+      # length 24, while the local file header for the entry contains an extra#
+      # field of length 28.
+      # It is key that the entry contains extra field of different length.
+      # The sole purpose of this test case is make sure our offset computing
+      # logic works in this scenario.
+
+      # This is created by:
+      # touch file.txt
+      # zip -0 test.zip file.txt
+      # Above command may or may not work on all platforms.
+      # Some zip implementation will keep the extra field size consistent.
+      # Some don't
+    magic_zip = b'PK\x03\x04\n\x00\x00\x00\x00\x00nY\xfcR\x00\x00\x00\x00\x00\x00\x00' +\
+        b'\x00\x00\x00\x00\x00\x08\x00\x1c\x00file.txtUT\t\x00\x03' +\
+        b'\xa0s\x01a\xa0s\x01aux\x0b\x00\x01\x04\x88\xc4\t\x00\x04S_\x01\x00' +\
+        b'PK\x01\x02\x1e\x03\n\x00\x00\x00\x00\x00nY\xfcR\x00\x00\x00\x00' +\
+        b'\x00\x00\x00\x00\x00\x00\x00\x00\x08\x00\x18\x00\x00\x00\x00\x00' +\
+        b'\x00\x00\x00\x00\x80\x81\x00\x00\x00\x00file.txt' +\
+        b'UT\x05\x00\x03\xa0s\x01aux\x0b\x00\x01\x04\x88\xc4\t\x00\x04' +\
+        b'S_\x01\x00PK\x05\x06\x00\x00\x00\x00\x01\x00\x01\x00N\x00\x00' +\
+        b'\x00B\x00\x00\x00\x00\x00'
+    # Just making sure we concatenated the bytes correctly
+    self.assertEqual(len(magic_zip), 166)
+    fp = io.BytesIO(magic_zip)
+    with zipfile.ZipFile(fp, 'r') as zfp:
+      self.assertGreater(len(zfp.infolist()), 0)
+      zinfo = zfp.getinfo("file.txt")
+      (offset, size) = ota_utils.GetZipEntryOffset(zfp, zinfo)
+      self.assertEqual(size, zinfo.file_size)
+      self.assertEqual(offset, zipfile.sizeFileHeader+len(zinfo.filename) + 28)