am fc267120: am b7e5ab6f: Merge "Allow system images larger than 2GiB."

* commit 'fc2671201caeac485440c2aa146f35a99002bca2':
  Allow system images larger than 2GiB.

5 files changed, 176 insertions, 34 deletions

diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index e98e4b6..8bbe452 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -30,9 +30,6 @@ if sys.hexversion < 0x02070000:
 
 import errno
 import os
-import re
-import shutil
-import subprocess
 import tempfile
 import zipfile
 
@@ -70,10 +67,8 @@ def AddSystem(output_zip, prefix="IMAGES/", recovery_img=None, boot_img=None):
   block_list = common.MakeTempFile(prefix="system-blocklist-", suffix=".map")
   imgname = BuildSystem(OPTIONS.input_tmp, OPTIONS.info_dict,
                         block_list=block_list)
-  with open(imgname, "rb") as f:
-    common.ZipWriteStr(output_zip, prefix + "system.img", f.read())
-  with open(block_list, "rb") as f:
-    common.ZipWriteStr(output_zip, prefix + "system.map", f.read())
+  common.ZipWrite(output_zip, imgname, prefix + "system.img")
+  common.ZipWrite(output_zip, block_list, prefix + "system.map")
 
 
 def BuildSystem(input_dir, info_dict, block_list=None):
@@ -94,10 +89,8 @@ def AddVendor(output_zip, prefix="IMAGES/"):
   block_list = common.MakeTempFile(prefix="vendor-blocklist-", suffix=".map")
   imgname = BuildVendor(OPTIONS.input_tmp, OPTIONS.info_dict,
                      block_list=block_list)
-  with open(imgname, "rb") as f:
-    common.ZipWriteStr(output_zip, prefix + "vendor.img", f.read())
-  with open(block_list, "rb") as f:
-    common.ZipWriteStr(output_zip, prefix + "vendor.map", f.read())
+  common.ZipWrite(output_zip, imgname, prefix + "vendor.img")
+  common.ZipWrite(output_zip, block_list, prefix + "vendor.map")
 
 
 def BuildVendor(input_dir, info_dict, block_list=None):
@@ -296,7 +289,6 @@ def AddImagesToTargetFiles(filename):
   output_zip.close()
 
 def main(argv):
-
   def option_handler(o, a):
     if o in ("-a", "--add_missing"):
       OPTIONS.add_missing = True
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 39c9b3d..6903dc6 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -781,6 +781,46 @@ class PasswordManager(object):
     return result
 
 
+def ZipWrite(zip_file, filename, arcname=None, perms=0o644,
+             compress_type=None):
+  import datetime
+
+  # http://b/18015246
+  # Python 2.7's zipfile implementation wrongly thinks that zip64 is required
+  # for files larger than 2GiB. We can work around this by adjusting their
+  # limit. Note that `zipfile.writestr()` will not work for strings larger than
+  # 2GiB. The Python interpreter sometimes rejects strings that large (though
+  # it isn't clear to me exactly what circumstances cause this).
+  # `zipfile.write()` must be used directly to work around this.
+  #
+  # This mess can be avoided if we port to python3.
+  saved_zip64_limit = zipfile.ZIP64_LIMIT
+  zipfile.ZIP64_LIMIT = (1 << 32) - 1
+
+  if compress_type is None:
+    compress_type = zip_file.compression
+  if arcname is None:
+    arcname = filename
+
+  saved_stat = os.stat(filename)
+
+  try:
+    # `zipfile.write()` doesn't allow us to pass ZipInfo, so just modify the
+    # file to be zipped and reset it when we're done.
+    os.chmod(filename, perms)
+
+    # Use a fixed timestamp so the output is repeatable.
+    epoch = datetime.datetime.fromtimestamp(0)
+    timestamp = (datetime.datetime(2009, 1, 1) - epoch).total_seconds()
+    os.utime(filename, (timestamp, timestamp))
+
+    zip_file.write(filename, arcname=arcname, compress_type=compress_type)
+  finally:
+    os.chmod(filename, saved_stat.st_mode)
+    os.utime(filename, (saved_stat.st_atime, saved_stat.st_mtime))
+    zipfile.ZIP64_LIMIT = saved_zip64_limit
+
+
 def ZipWriteStr(zip, filename, data, perms=0644, compression=None):
   # use a fixed timestamp so the output is repeatable.
   zinfo = zipfile.ZipInfo(filename=filename,
@@ -1092,19 +1132,21 @@ class BlockDifference:
                           'endif;') % (partition,))
 
   def _WriteUpdate(self, script, output_zip):
-    partition = self.partition
-    with open(self.path + ".transfer.list", "rb") as f:
-      ZipWriteStr(output_zip, partition + ".transfer.list", f.read())
-    with open(self.path + ".new.dat", "rb") as f:
-      ZipWriteStr(output_zip, partition + ".new.dat", f.read())
-    with open(self.path + ".patch.dat", "rb") as f:
-      ZipWriteStr(output_zip, partition + ".patch.dat", f.read(),
-                         compression=zipfile.ZIP_STORED)
-
-    call = (('block_image_update("%s", '
-             'package_extract_file("%s.transfer.list"), '
-             '"%s.new.dat", "%s.patch.dat");\n') %
-            (self.device, partition, partition, partition))
+    ZipWrite(output_zip,
+             '{}.transfer.list'.format(self.path),
+             '{}.transfer.list'.format(self.partition))
+    ZipWrite(output_zip,
+             '{}.new.dat'.format(self.path),
+             '{}.new.dat'.format(self.partition))
+    ZipWrite(output_zip,
+             '{}.patch.dat'.format(self.path),
+             '{}.patch.dat'.format(self.partition),
+             compress_type=zipfile.ZIP_STORED)
+
+    call = ('block_image_update("{device}", '
+            'package_extract_file("{partition}.transfer.list"), '
+            '"{partition}.new.dat", "{partition}.patch.dat");\n'.format(
+                device=self.device, partition=self.partition))
     script.AppendExtra(script._WordWrap(call))
 
   def _HashBlocks(self, source, ranges):
diff --git a/tools/releasetools/img_from_target_files.py b/tools/releasetools/img_from_target_files.py
index 4b88e73..a9d4cbe 100755
--- a/tools/releasetools/img_from_target_files.py
+++ b/tools/releasetools/img_from_target_files.py
@@ -88,11 +88,13 @@ def main(argv):
       # and all we have to do is copy them to the output zip.
       images = os.listdir(images_path)
       if images:
-        for i in images:
-          if bootable_only and i not in ("boot.img", "recovery.img"): continue
-          if not i.endswith(".img"): continue
-          with open(os.path.join(images_path, i), "r") as f:
-            common.ZipWriteStr(output_zip, i, f.read())
+        for image in images:
+          if bootable_only and image not in ("boot.img", "recovery.img"):
+            continue
+          if not image.endswith(".img"):
+            continue
+          common.ZipWrite(
+              output_zip, os.path.join(images_path, image), image)
         done = True
 
     if not done:
diff --git a/tools/releasetools/ota_from_target_files b/tools/releasetools/ota_from_target_files
index 25309a4..6e0fefc 100755
--- a/tools/releasetools/ota_from_target_files
+++ b/tools/releasetools/ota_from_target_files
@@ -646,10 +646,8 @@ endif;
   WriteMetadata(metadata, output_zip)
 
 
-def WritePolicyConfig(file_context, output_zip):
-  f = open(file_context, 'r');
-  basename = os.path.basename(file_context)
-  common.ZipWriteStr(output_zip, basename, f.read())
+def WritePolicyConfig(file_name, output_zip):
+  common.ZipWrite(output_zip, file_name, os.path.basename(file_name))
 
 
 def WriteMetadata(metadata, output_zip):
diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py
new file mode 100644
index 0000000..f163f92
--- /dev/null
+++ b/tools/releasetools/test_common.py
@@ -0,0 +1,108 @@
+#
+# Copyright (C) 2015 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 os
+import tempfile
+import time
+import unittest
+import zipfile
+
+import common
+
+
+def random_string_with_holes(size, block_size, step_size):
+  data = ["\0"] * size
+  for begin in range(0, size, step_size):
+    end = begin + block_size
+    data[begin:end] = os.urandom(block_size)
+  return "".join(data)
+
+
+class CommonZipTest(unittest.TestCase):
+  def _test_ZipWrite(self, contents, extra_zipwrite_args=None):
+    extra_zipwrite_args = dict(extra_zipwrite_args or {})
+
+    test_file = tempfile.NamedTemporaryFile(delete=False)
+    zip_file = tempfile.NamedTemporaryFile(delete=False)
+
+    test_file_name = test_file.name
+    zip_file_name = zip_file.name
+
+    # File names within an archive strip the leading slash.
+    arcname = extra_zipwrite_args.get("arcname", test_file_name)
+    if arcname[0] == "/":
+      arcname = arcname[1:]
+
+    zip_file.close()
+    zip_file = zipfile.ZipFile(zip_file_name, "w")
+
+    try:
+      test_file.write(contents)
+      test_file.close()
+
+      old_stat = os.stat(test_file_name)
+      expected_mode = extra_zipwrite_args.get("perms", 0o644)
+
+      time.sleep(5)  # Make sure the atime/mtime will change measurably.
+
+      common.ZipWrite(zip_file, test_file_name, **extra_zipwrite_args)
+
+      new_stat = os.stat(test_file_name)
+      self.assertEqual(int(old_stat.st_mode), int(new_stat.st_mode))
+      self.assertEqual(int(old_stat.st_mtime), int(new_stat.st_mtime))
+
+      zip_file.close()
+      zip_file = zipfile.ZipFile(zip_file_name, "r")
+      info = zip_file.getinfo(arcname)
+
+      self.assertEqual(info.date_time, (2009, 1, 1, 0, 0, 0))
+      mode = (info.external_attr >> 16) & 0o777
+      self.assertEqual(mode, expected_mode)
+      self.assertEqual(zip_file.read(arcname), contents)
+    finally:
+      os.remove(test_file_name)
+      os.remove(zip_file_name)
+
+  def test_ZipWrite(self):
+    file_contents = os.urandom(1024)
+    self._test_ZipWrite(file_contents)
+
+  def test_ZipWrite_with_opts(self):
+    file_contents = os.urandom(1024)
+    self._test_ZipWrite(file_contents, {
+        "arcname": "foobar",
+        "perms": 0o777,
+        "compress_type": zipfile.ZIP_DEFLATED,
+    })
+
+  def test_ZipWrite_large_file(self):
+    kilobytes = 1024
+    megabytes = 1024 * kilobytes
+    gigabytes = 1024 * megabytes
+
+    size = int(2 * gigabytes + 1)
+    block_size = 4 * kilobytes
+    step_size = 4 * megabytes
+    file_contents = random_string_with_holes(
+        size, block_size, step_size)
+    self._test_ZipWrite(file_contents, {
+        "compress_type": zipfile.ZIP_DEFLATED,
+    })
+
+  def test_ZipWrite_resets_ZIP64_LIMIT(self):
+    default_limit = (1 << 31) - 1
+    self.assertEqual(default_limit, zipfile.ZIP64_LIMIT)
+    self._test_ZipWrite('')
+    self.assertEqual(default_limit, zipfile.ZIP64_LIMIT)