Fix unsupported fs.squashfs extraction in sonic-installer (sonic-net#1366)

Abstracted away the access to the path of the rootfs via a contextmanager.
In Arista secureboot, the rootfs is never extracted on the flash. Instead it's mounted directly from within the signed SWI.
The update_sonic_environment function however always assume that the rootfs to be at the same place.

- How I did it

To alleviate this restriction, a new context manager to obtain the rootfs is introduced.
The choice of a context manager rather than a function is entirely based on error management and cleanup.
Mounting a squashfs from a swi file requires the use of losetup which makes the rootfs available under /dev/loopX
Once done or on error, we need to free this resource which becomes free when using a contextmanager.

Author: Staphylo

sonic_installer/bootloader/aboot.py

@@ -9,6 +9,7 @@
 import subprocess
 import sys
 import zipfile
+from contextlib import contextmanager
 
 import click
 
@@ -18,7 +19,9 @@
    HOST_PATH,
    IMAGE_DIR_PREFIX,
    IMAGE_PREFIX,
+   ROOTFS_NAME,
    run_command,
+   run_command_or_raise,
 )
 from .bootloader import Bootloader
 
@@ -34,7 +37,7 @@ def isSecureboot():
     global _secureboot
     if _secureboot is None:
         with open('/proc/cmdline') as f:
-           m  = re.search(r"secure_boot_enable=[y1]", f.read())
+            m = re.search(r"secure_boot_enable=[y1]", f.read())
         _secureboot = bool(m)
     return _secureboot
 
@@ -179,3 +182,25 @@ def base64Decode(cls, text):
     def detect(cls):
         with open('/proc/cmdline') as f:
             return 'Aboot=' in f.read()
+
+    def _get_swi_file_offset(self, swipath, filename):
+        with zipfile.ZipFile(swipath) as swi:
+            with swi.open(filename) as f:
+                return f._fileobj.tell() # pylint: disable=protected-access
+
+    @contextmanager
+    def get_rootfs_path(self, image_path):
+        rootfs_path = os.path.join(image_path, ROOTFS_NAME)
+        if os.path.exists(rootfs_path) and not isSecureboot():
+            yield rootfs_path
+            return
+
+        swipath = os.path.join(image_path, DEFAULT_SWI_IMAGE)
+        offset = self._get_swi_file_offset(swipath, ROOTFS_NAME)
+        loopdev = subprocess.check_output(['losetup', '-f']).rstrip()
+
+        try:
+            run_command_or_raise(['losetup', '-o', str(offset), loopdev, swipath])
+            yield loopdev
+        finally:
+            run_command_or_raise(['losetup', '-d', loopdev])

sonic_installer/bootloader/bootloader.py

@@ -2,12 +2,14 @@
 Abstract Bootloader class
 """
 
+from contextlib import contextmanager
 from os import path
 
 from ..common import (
    HOST_PATH,
    IMAGE_DIR_PREFIX,
    IMAGE_PREFIX,
+   ROOTFS_NAME,
 )
 
 class Bootloader(object):
@@ -68,3 +70,7 @@ def get_image_path(cls, image):
         prefix = path.join(HOST_PATH, IMAGE_DIR_PREFIX)
         return image.replace(IMAGE_PREFIX, prefix)
 
+    @contextmanager
+    def get_rootfs_path(self, image_path):
+        """returns the path to the squashfs"""
+        yield path.join(image_path, ROOTFS_NAME)

sonic_installer/common.py

@@ -13,6 +13,7 @@
 HOST_PATH = '/host'
 IMAGE_PREFIX = 'SONiC-OS-'
 IMAGE_DIR_PREFIX = 'image-'
+ROOTFS_NAME = 'fs.squashfs'
 
 # Run bash command and print output to stdout
 def run_command(command):

sonic_installer/main.py

@@ -11,7 +11,7 @@
 from swsscommon.swsscommon import SonicV2Connector
 
 from .bootloader import get_bootloader
-from .common import run_command, run_command_or_raise
+from .common import run_command, run_command_or_raise, IMAGE_PREFIX
 from .exception import SonicRuntimeException
 
 SYSLOG_IDENTIFIER = "sonic-installer"
@@ -218,8 +218,7 @@ def print_deprecation_warning(deprecated_cmd_or_subcmd, new_cmd_or_subcmd):
                 fg="red", err=True)
     click.secho("Please use '{}' instead".format(new_cmd_or_subcmd), fg="red", err=True)
 
-
-def update_sonic_environment(click, binary_image_version):
+def update_sonic_environment(click, bootloader, binary_image_version):
     """Prepare sonic environment variable using incoming image template file. If incoming image template does not exist
        use current image template file.
     """
@@ -234,38 +233,38 @@ def umount_next_image_fs(mount_point):
     SONIC_ENV_TEMPLATE_FILE = os.path.join("usr", "share", "sonic", "templates", "sonic-environment.j2")
     SONIC_VERSION_YML_FILE = os.path.join("etc", "sonic", "sonic_version.yml")
 
-    sonic_version = re.sub("SONiC-OS-", '', binary_image_version)
-    new_image_dir = os.path.join('/', "host", "image-{0}".format(sonic_version))
-    new_image_squashfs_path = os.path.join(new_image_dir, "fs.squashfs")
+    sonic_version = re.sub(IMAGE_PREFIX, '', binary_image_version)
+    new_image_dir = bootloader.get_image_path(binary_image_version)
     new_image_mount = os.path.join('/', "tmp", "image-{0}-fs".format(sonic_version))
     env_dir = os.path.join(new_image_dir, "sonic-config")
     env_file = os.path.join(env_dir, "sonic-environment")
 
-    try:
-        mount_next_image_fs(new_image_squashfs_path, new_image_mount)
-
-        next_sonic_env_template_file = os.path.join(new_image_mount, SONIC_ENV_TEMPLATE_FILE)
-        next_sonic_version_yml_file = os.path.join(new_image_mount, SONIC_VERSION_YML_FILE)
-
-        sonic_env = run_command_or_raise([
-                "sonic-cfggen",
-                "-d",
-                "-y",
-                next_sonic_version_yml_file,
-                "-t",
-                next_sonic_env_template_file,
-        ])
-        os.mkdir(env_dir, 0o755)
-        with open(env_file, "w+") as ef:
-            print(sonic_env, file=ef)
-        os.chmod(env_file, 0o644)
-    except SonicRuntimeException as ex:
-        echo_and_log("Warning: SONiC environment variables are not supported for this image: {0}".format(str(ex)), LOG_ERR, fg="red")
-        if os.path.exists(env_file):
-            os.remove(env_file)
-            os.rmdir(env_dir)
-    finally:
-        umount_next_image_fs(new_image_mount)
+    with bootloader.get_rootfs_path(new_image_dir) as new_image_squashfs_path:
+        try:
+            mount_next_image_fs(new_image_squashfs_path, new_image_mount)
+
+            next_sonic_env_template_file = os.path.join(new_image_mount, SONIC_ENV_TEMPLATE_FILE)
+            next_sonic_version_yml_file = os.path.join(new_image_mount, SONIC_VERSION_YML_FILE)
+
+            sonic_env = run_command_or_raise([
+                    "sonic-cfggen",
+                    "-d",
+                    "-y",
+                    next_sonic_version_yml_file,
+                    "-t",
+                    next_sonic_env_template_file,
+            ])
+            os.mkdir(env_dir, 0o755)
+            with open(env_file, "w+") as ef:
+                print(sonic_env, file=ef)
+            os.chmod(env_file, 0o644)
+        except SonicRuntimeException as ex:
+            echo_and_log("Warning: SONiC environment variables are not supported for this image: {0}".format(str(ex)), LOG_ERR, fg="red")
+            if os.path.exists(env_file):
+                os.remove(env_file)
+                os.rmdir(env_dir)
+        finally:
+            umount_next_image_fs(new_image_mount)
 
 # Main entrypoint
 @click.group(cls=AliasedGroup)
@@ -332,7 +331,7 @@ def install(url, force, skip_migration=False):
         else:
             run_command('config-setup backup')
 
-        update_sonic_environment(click, binary_image_version)
+        update_sonic_environment(click, bootloader, binary_image_version)
 
     # Finally, sync filesystem
     run_command("sync;sync;sync")