Merge lp:~sil2100/ubuntu-system-image/server_uncompressed-initrd into lp:ubuntu-system-image/server

[Łukasz Zemczak]

Add support for uncompressed/xz-compressed initrd images in our device tarballs. Add additional header stripping to specific devices while repacking the keyring as the initrd have custom 512 bytes at the beginning. Handle the old s-i 2 keyring paths in the recovery.

Add support for uncompressed/xz-compressed initrd images in our device tarballs. Add additional header stripping to specific devices while repacking the keyring as the initrd have custom 512 bytes at the beginning. Handle the old s-i 2 keyring paths in the recovery.

This merge basically has 3 fixes, all related more or less to eachother - all required to fix repacking of keyrings for devices such as krillin or arale.

[Łukasz Zemczak (sil2100)]

Now it's ready for review basically.

[Barry Warsaw (barry)]

Just a few thoughts.

[Łukasz Zemczak (sil2100)]

Hey Barry! Thanks for the code-check. I didn't use subprocess.DEVNULL for two reasons: first, I didn't know about it, and second - the existing code was using open(os.devnull). Would have to change it in the whole testcase then - which we can do of course.

[Barry Warsaw (barry)]

On Feb 12, 2016, at 03:10 PM, Łukasz Zemczak wrote:

>Hey Barry! Thanks for the code-check. I didn't use subprocess.DEVNULL for two
>reasons: first, I didn't know about it, and second - the existing code was
>using open(os.devnull). Would have to change it in the whole testcase then -
>which we can do of course.

Reducing tech-debt would make for a nice sprint task! :)

[Łukasz Zemczak (sil2100)]

Anyway, I'll leave it as is right now since from what I see there's no subprocess.DEVNULL in Python 2.7, while system-image is currently written to be both 2.7 and 3 compatible. os.devnull seems to be the safe choice here.

That being said: the branch in overall got tested by Bhushan Shah (bshah) who reproduced the original issue and said it helped. If you could +1 the code in overall I would gladly merge it in.

[Barry Warsaw (barry)]

A few comments inlined, but other than that, +1

review: Approve

[lp:~sil2100/ubuntu-system-image/server_uncompressed-initrd]

290. By Łukasz Zemczak on 2016-02-17

Changes as per Barry's comments - make the buffer size variable a global constant and use assertIsNone() when it makes sense.

=== modified file 'lib/systemimage/generators.py'
--- lib/systemimage/generators.py	2015-12-14 19:56:16 +0000
+++ lib/systemimage/generators.py	2016-02-17 09:23:34 +0000
@@ -1217,7 +1217,8 @@
 
             if "keyring" in options:
                 if not tools.repack_recovery_keyring(conf, path,
-                                                     options['keyring']):
+                                                     options['keyring'],
+                                                     device_name):
                     if os.path.exists(path):
                         os.remove(path)
                     return None

=== modified file 'lib/systemimage/tests/test_tools.py'
--- lib/systemimage/tests/test_tools.py	2015-12-14 19:56:16 +0000
+++ lib/systemimage/tests/test_tools.py	2016-02-17 09:23:34 +0000
@@ -254,6 +254,31 @@
         os.environ['PATH'] = bin_dir
         self.assertFalse(tools.find_on_path("program"))
 
+    def test_manipulate_recovery_header(self):
+        """Check if stripping and reattaching recovery headers works."""
+        source_path = os.path.join(self.temp_directory, "source")
+        stripped_path = os.path.join(self.temp_directory, "stripped")
+        reattached_path = os.path.join(self.temp_directory, "reattached")
+        
+        header = bytearray(512)
+        contents = b"RECOVERY"
+        for i in range(0, 64):
+            header[i] = i
+        with open(source_path, "wb+") as f:
+            f.write(header)
+            f.write(contents)
+
+        stripped = tools.strip_recovery_header(source_path, stripped_path)
+        self.assertEqual(header, bytes(stripped))
+        with open(stripped_path, "rb") as f:
+            self.assertEqual(f.read(), contents)
+
+        tools.reattach_recovery_header(stripped_path, reattached_path,
+                                       stripped)
+        with open(reattached_path, "rb") as f, open(source_path, "rb") as fs:
+            self.assertEqual(f.read(), fs.read())
+
+
     @unittest.skipUnless(HAS_TEST_KEYS, MISSING_KEYS_WARNING)
     def test_repack_recovery_keyring(self):
         # Generate the keyring tarballs
@@ -321,6 +346,32 @@
                                                    self.temp_directory,
                                       "archive-master")
 
+        tools.reattach_recovery_header(os.path.join(self.temp_directory,
+                                                    "initrd.gz"),
+                                       os.path.join(self.temp_directory,
+                                                    "initrd.header"),
+                                       bytearray(512))
+
+        with open(os.devnull, "w") as devnull:
+            subprocess.call(["abootimg", "--create",
+                             "%s/partitions/recovery.img" %
+                             self.temp_directory,
+                             "-k", "%s/kernel" % self.temp_directory,
+                             "-r", "%s/initrd.header" % self.temp_directory,
+                             "-f", "%s/bootimg.cfg" % self.temp_directory],
+                            stderr=devnull, stdout=devnull)
+
+            subprocess.call(["tar", "Jcf",
+                             "%s/recovery-spec.tar.xz" % self.temp_directory,
+                             "-C", self.temp_directory,
+                             "partitions/"], stderr=devnull, stdout=devnull)
+
+        # Try repacking in case of a recovery with a special header
+        tools.repack_recovery_keyring(self.config, "%s/recovery-spec.tar.xz" %
+                                                   self.temp_directory,
+                                      "archive-master", "krillin")
+
+
     def test_system_image_30_symlinks(self):
         # To support system-image 3.0, generate symlinks for config.d
         version_tarball = "%s/version.tar" % self.temp_directory
@@ -476,3 +527,24 @@
 
         six.assertCountEqual(
             self, [base_image1, base_image2], delta_base)
+
+    def test_guess_file_compression(self):
+        """Check if we can correctly guess compression algorithms."""
+        test_string = "test-string"
+
+        # Simple compress/uncompress
+        test_file = os.path.join(self.temp_directory, "test.txt")
+        with open(test_file, "w+") as fd:
+            fd.write(test_string)
+
+        self.assertIsNone(tools.guess_file_compression(test_file))
+
+        xz_file = os.path.join(self.temp_directory, "test.xz")
+        tools.xz_compress(test_file, xz_file)
+        self.assertEqual(
+            tools.guess_file_compression(xz_file), "xz")
+
+        gzip_file = os.path.join(self.temp_directory, "test.gz")
+        tools.gzip_compress(test_file, gzip_file)
+        self.assertEqual(
+            tools.guess_file_compression(gzip_file), "gzip")

=== modified file 'lib/systemimage/tools.py'
--- lib/systemimage/tools.py	2015-12-14 19:56:16 +0000
+++ lib/systemimage/tools.py	2016-02-17 09:23:34 +0000
@@ -32,6 +32,8 @@
 from systemimage import gpg
 
 
+READ_SIZE = 1024*1024
+
 logger = logging.getLogger(__name__)
 
 
@@ -183,6 +185,27 @@
     gpg.sign_file(config, "image-signing", path.replace(".tar.xz", ".json"))
 
 
+def guess_file_compression(path):
+    """
+        Try to guess through the magic signature in the first bytes of the
+        file.
+    """
+
+    compressions = {
+        b"\x1f\x8b\x08": "gzip",
+        b"\xfd\x37\x7a\x58\x5a\x00": "xz"
+    }
+    length = max(len(x) for x in compressions)
+
+    with open(path, 'rb') as f:
+        start = f.read(length)
+    for magic, compression in compressions.items():
+        if start.startswith(magic):
+            return compression
+
+    return None
+
+
 def gzip_compress(path, destination=None, level=9):
     """
         Compress a file (path) using gzip.
@@ -307,9 +330,47 @@
                        mirror.ssh_key, mirror.ssh_command)
 
 
-def repack_recovery_keyring(conf, path, keyring_name):
+def strip_recovery_header(source_path, dest_path):
+    """
+        Strip the first 512 bytes of the custom header from the source_path
+        file (initrd). Writes the stripped file to dest_path and returns the
+        header contents.
+    """
+
+    with open(source_path, "rb") as source:
+        header_contents = source.read(512)
+        with open(dest_path, "wb") as dest:
+            data = source.read(READ_SIZE)
+            while data:
+                dest.write(data)
+                data = source.read(READ_SIZE)
+    return header_contents
+
+
+def reattach_recovery_header(source_path, dest_path, header_contents):
+    """
+        Reattach the stripped header (in header_contents) in front of the
+        source_path file contents. This writes the end file to dest_path.
+    """
+
+    with open(dest_path, "wb") as dest:
+        dest.write(header_contents)
+        with open(source_path, "rb") as source:
+            data = source.read(READ_SIZE)
+            while data:
+                dest.write(data)
+                data = source.read(READ_SIZE)
+
+
+def repack_recovery_keyring(conf, path, keyring_name, device_name=None):
     tempdir = tempfile.mkdtemp()
 
+    # In case of certain devices, special care of the recovery is needed
+    additional_header = False
+    if device_name in ("krillin", "vegetahd", "arale"):
+        logging.debug("Expecting additional header in recovery image.")
+        additional_header = True
+
     xz_uncompress(path, os.path.join(tempdir, "input.tar"))
 
     input_tarball = tarfile.open(os.path.join(tempdir, "input.tar"), "r:")
@@ -335,10 +396,25 @@
     state_path = os.path.join(tempdir, "fakeroot_state")
 
     with chdir(os.path.join(tempdir, "initrd")):
-        gzip_uncompress(os.path.join(tempdir, "img", "initrd.img"),
-                        os.path.join(tempdir, "img", "initrd"))
-
-        with open(os.path.join(tempdir, "img", "initrd"), "rb") as fd:
+        initrdimg_path = os.path.join(tempdir, "img", "initrd.img")
+        initrd_path = os.path.join(tempdir, "img", "initrd")
+
+        if additional_header:
+            # Remove the 512 header bytes before unpacking
+            tmp_path = os.path.join(tempdir, "img", "initrd.img.tmp")
+            header_contents = strip_recovery_header(initrdimg_path, tmp_path)
+            os.rename(tmp_path, initrdimg_path)
+
+        # The initrd can be either compressed or uncompressed
+        compression = guess_file_compression(initrdimg_path)
+        if compression == "gzip":
+            gzip_uncompress(initrdimg_path, initrd_path)
+        elif compression == "xz":
+            xz_uncompress(initrdimg_path, initrd_path)
+        else:
+            shutil.copyfile(initrdimg_path, initrd_path)
+
+        with open(initrd_path, "rb") as fd:
             with open(os.path.devnull, "w") as devnull:
                 subprocess.call(["fakeroot", "-s", state_path, "cpio", "-i"],
                                 stdin=fd, stdout=devnull, stderr=devnull)
@@ -346,13 +422,18 @@
     # Swap the files
     keyring_path = os.path.join(conf.gpg_keyring_path, keyring_name)
 
+    # Handle two different keyring paths in the recovery
+    dest_keyring_path = os.path.join(tempdir, "initrd", "usr", "share",
+                                     "system-image", keyring_name)
+    if not os.path.exists("%s.tar.xz" % dest_keyring_path):
+        dest_keyring_path = os.path.join(tempdir, "initrd", "etc",
+                                         "system-image", keyring_name)
+
     shutil.copy("%s.tar.xz" % keyring_path,
-                os.path.join(tempdir, "initrd", "usr", "share", "system-image",
-                             "%s.tar.xz" % keyring_name))
+                "%s.tar.xz" % dest_keyring_path)
 
     shutil.copy("%s.tar.xz.asc" % keyring_path,
-                os.path.join(tempdir, "initrd", "usr", "share", "system-image",
-                             "%s.tar.xz.asc" % keyring_name))
+                "%s.tar.xz.asc" % dest_keyring_path)
 
     # Re-generate the initrd
     with chdir(os.path.join(tempdir, "initrd")):
@@ -367,8 +448,19 @@
 
     os.rename(os.path.join(tempdir, "img", "initrd.img"),
               os.path.join(tempdir, "img", "initrd.img.bak"))
-    gzip_compress(os.path.join(tempdir, "img", "initrd"),
-                  os.path.join(tempdir, "img", "initrd.img"))
+
+    if compression == "gzip":
+        gzip_compress(initrd_path, initrdimg_path)
+    elif compression == "xz":
+        xz_compress(initrd_path, initrdimg_path)
+    else:
+        shutil.copyfile(initrd_path, initrdimg_path)
+
+    if additional_header:
+        # Append the previously removed header
+        tmp_path = os.path.join(tempdir, "img", "initrd.img.tmp")
+        reattach_recovery_header(initrdimg_path, tmp_path, header_contents)
+        os.rename(tmp_path, initrdimg_path)
 
     # Rewrite bootimg.cfg
     content = ""