add basic ZFS support

docs/configuration.md

@@ -194,6 +194,7 @@ These bools simply import `ugrd.kmod.no<category>` modules during `build_pre`.
 
 Additional modules include:
 
+* `ugrd.fs.bcachefs` - Adds the bcachefs module and binary for mounting.
 * `ugrd.fs.btrfs` - Helps with multi-device BTRFS mounts, subvolume selection.
 * `ugrd.fs.fakeudev` - Makes 'fake' udev entries for DM devices.
 * `ugrd.fs.cpio` - Packs the build dir into a CPIO archive with PyCPIO.
@@ -202,6 +203,7 @@ Additional modules include:
 * `ugrd.fs.mdraid` - For MDRAID mounts.
 * `ugrd.fs.resume` - Handles resume from hibernation.
 * `ugrd.fs.test_image` - Creates a test rootfs for automated testing.
+* `ugrd.fs.zfs` - Adds basic ZFS support.
 
 #### ugrd.fs.mounts
 

hooks/installkernel/52-ugrd.install

@@ -32,7 +32,12 @@ main() {
 
 	[[ ${EUID} -eq 0 ]] || die "Please run this script as root"
 
-	ugrd --no-rotate --kver "${ver}" "${initrd}" || die "Failed to generate initramfs"
+	ugrd --no-rotate --kver "${ver}" "${initrd}"
+	case $? in
+		0) einfo "Generated initramfs for kernel: ${ver}";;
+		77) ewarn "Missing ZFS kernel module for kernel: ${ver}" && exit 77;;
+		*) die "Failed to generate initramfs for kernel ${ver}";;
+	esac
 }
 
 main

hooks/kernel-install/52-ugrd.install

@@ -12,4 +12,9 @@ KERNEL_VERSION="${2:?}"
 # only run when the COMMAND is add, and fewer than 5 arguments are passed
 [ "${COMMAND}" = "add" ] && [ "${#}" -lt 5 ] || exit 0
 
-ugrd "$([ "${KERNEL_INSTALL_VERBOSE}" = 1 ] && echo --debug)" --no-rotate --kver "${KERNEL_VERSION}" "${KERNEL_INSTALL_STAGING_AREA}/initrd" || exit 1
+ugrd "$([ "${KERNEL_INSTALL_VERBOSE}" = 1 ] && echo --debug)" --no-rotate --kver "${KERNEL_VERSION}" "${KERNEL_INSTALL_STAGING_AREA}/initrd"
+case $? in
+    0) ;;
+    77) echo "Missing ZFS kernel module for kernel: ${KERNEL_VERSION}"; exit 77 ;;
+    *) exit 1 ;;
+esac

readme.md

@@ -103,6 +103,7 @@ The following root filesystems have been tested:
 The following filesystems have limited support:
 
 * BCACHEFS
+* ZFS
 
 Additionally, the following filesystems have been tested for non-root mounts:
 

src/ugrd/fs/mounts.py

@@ -33,6 +33,8 @@ def _resolve_dev(self, device_path) -> str:
 
     Takes the device path, such as /dev/root, and resolves it to a device indexed in blkid.
     If the device is an overlayfs, resolves the lowerdir device.
+
+    If the device is a ZFS device, returns the device path.
     """
     if str(device_path) in self["_blkid_info"]:
         self.logger.debug("Device already resolved to blkid indexed device: %s" % device_path)
@@ -43,8 +45,13 @@ def _resolve_dev(self, device_path) -> str:
     device_path = _resolve_overlay_lower_device(self, mountpoint)
     mountpoint = _resolve_device_mountpoint(self, device_path)  # May have changed if it was an overlayfs
 
+    if self["_mounts"][mountpoint]["fstype"] == "zfs":
+        self.logger.info("Resolved ZFS device: %s" % colorize(device_path, "cyan"))
+        return device_path
+
     mount_dev = self["_mounts"][mountpoint]["device"]
     major, minor = _get_device_id(mount_dev.split(":")[0] if ":" in mount_dev else mount_dev)
+
     for device in self["_blkid_info"]:
         check_major, check_minor = _get_device_id(device)
         if (major, minor) == (check_major, check_minor):
@@ -104,6 +111,31 @@ def _resolve_overlay_lower_device(self, mountpoint) -> dict:
     return self["_mounts"][mountpoint]["device"]
 
 
+def _get_mount_dev_fs_type(self, device: str, raise_exception=True) -> str:
+    """Taking the device of an active mount, returns the filesystem type."""
+    for info in self["_mounts"].values():
+        if info["device"] == device:
+            return info["fstype"]
+    if not device.startswith("/dev/"):
+        # Try again with /dev/ prepended if it wasn't already
+        return _get_mount_dev_fs_type(self, f"/dev/{device}", raise_exception)
+
+    if raise_exception:
+        raise ValueError("No mount found for device: %s" % device)
+    else:
+        self.logger.debug("No mount found for device: %s" % device)
+
+
+def _get_mount_source_type(self, mount: dict, with_val=False) -> str:
+    """Gets the source from the mount config."""
+    for source_type in SOURCE_TYPES:
+        if source_type in mount:
+            if with_val:
+                return source_type, mount[source_type]
+            return source_type
+    raise ValueError("No source type found in mount: %s" % mount)
+
+
 def _merge_mounts(self, mount_name: str, mount_config, mount_class) -> None:
     """Merges the passed mount config with the existing mount."""
     if mount_name not in self[mount_class]:
@@ -181,6 +213,11 @@ def _process_mount(self, mount_name: str, mount_config, mount_class="mounts") ->
             if "ugrd.fs.bcachefs" not in self["modules"]:
                 self.logger.info("Auto-enabling module: %s", colorize("bcachefs", "cyan"))
                 self["modules"] = "ugrd.fs.bcachefs"
+        elif mount_type == "zfs":
+            if "ugrd.fs.zfs" not in self["modules"]:
+                self.logger.info("Auto-enabling module: zfs")
+                self["modules"] = "ugrd.fs.zfs"
+                mount_config["options"].add("zfsutil")
         elif mount_type not in ["proc", "sysfs", "devtmpfs", "squashfs", "tmpfs", "devpts"]:
             self.logger.warning("Unknown mount type: %s" % colorize(mount_type, "red", bold=True))
 
@@ -200,16 +237,6 @@ def _process_late_mounts_multi(self, mount_name: str, mount_config) -> None:
     _process_mount(self, mount_name, mount_config, "late_mounts")
 
 
-def _get_mount_source_type(self, mount: dict, with_val=False) -> str:
-    """Gets the source from the mount config."""
-    for source_type in SOURCE_TYPES:
-        if source_type in mount:
-            if with_val:
-                return source_type, mount[source_type]
-            return source_type
-    raise ValueError("No source type found in mount: %s" % mount)
-
-
 def _get_mount_str(self, mount: dict, pad=False, pad_size=44) -> str:
     """returns the mount source string based on the config,
     the output string should work with fstab and mount commands.
@@ -345,6 +372,44 @@ def get_blkid_info(self, device=None) -> dict:
     return self["_blkid_info"][device] if device else self["_blkid_info"]
 
 
+def get_zpool_info(self, poolname=None) -> Union[dict, None]:
+    """Enumerates ZFS pools and devices, adds them to the zpools dict."""
+    if poolname:  # If a pool name is passed, try to get the pool info
+        if "/" in poolname:
+            # If a dataset is passed, get the pool name only
+            poolname = poolname.split("/")[0]
+        if poolname in self["_zpool_info"]:
+            return self["_zpool_info"][poolname]
+
+    # Always try to get zpool info, but only raise an error if a poolname is passed or the ZFS module is enabled
+    try:
+        pool_info = self._run(["zpool", "list", "-vPH", "-o", "name"]).stdout.decode().strip().split("\n")
+    except FileNotFoundError:
+        if "ugrd.fs.zfs" not in self["modules"]:
+            return self.logger.debug("ZFS pool detection failed, but ZFS module not enabled, skipping.")
+        if poolname:
+            raise AutodetectError("Failed to get zpool list for pool: %s" % colorize(poolname, "red"))
+
+    capture_pool = False
+    for line in pool_info:
+        if not capture_pool:
+            poolname = line  # Get the pool name using the first line
+            self["_zpool_info"][poolname] = {"devices": set()}
+            capture_pool = True
+            continue
+        else:  # Otherwise, add devices listed in the pool
+            if line[0] != "\t":
+                capture_pool = False
+                continue  # Keep going
+            # The device name has a tab before it, and may have a space/tab after it
+            device_name = line[1:].split("\t")[0].strip()
+            self.logger.debug("[%s] Found ZFS device: %s" % (colorize(poolname, "blue"), colorize(device_name, "cyan")))
+            self["_zpool_info"][poolname]["devices"].add(device_name)
+
+    if poolname:  # If a poolname was passed, try return the pool info, raise an error if not found
+        return self["_zpool_info"][poolname]
+
+
 @contains("hostonly", "Skipping init mount autodetection, hostonly mode is disabled.", log_level=30)
 @contains("autodetect_init_mount", "Init mount autodetection disabled, skipping.", log_level=30)
 @contains("init_target", "init_target must be set", raise_exception=True)
@@ -621,9 +686,13 @@ def autodetect_root(self) -> None:
     if self["autodetect_root_dm"]:
         if self["mounts"]["root"]["type"] == "btrfs":
             from ugrd.fs.btrfs import _get_btrfs_mount_devices
+
             # Btrfs volumes may be backed by multiple dm devices
             for device in _get_btrfs_mount_devices(self, "/", root_dev):
                 _autodetect_dm(self, "/", device)
+        elif self["mounts"]["root"]["type"] == "zfs":
+            for device in get_zpool_info(self, root_dev)["devices"]:
+                _autodetect_dm(self, "/", device)
         else:
             _autodetect_dm(self, "/")
 
@@ -636,8 +705,17 @@ def _autodetect_mount(self, mountpoint) -> str:
         self.logger.error("Host mounts:\n%s" % pretty_print(self["_mounts"]))
         raise AutodetectError("auto_mount mountpoint not found in host mounts: %s" % mountpoint)
 
-    mount_device = _resolve_dev(self, self["_mounts"][mountpoint]["device"])
-    mount_info = self["_blkid_info"][mount_device]
+    mountpoint_device = self["_mounts"][mountpoint]["device"]
+    # get the fs type from the device as it appears in /proc/mounts
+    fs_type = _get_mount_dev_fs_type(self, mountpoint_device, raise_exception=False)
+    # resolve the device down to the "real" device path, one that has blkid info
+    mount_device = _resolve_dev(self, mountpoint_device)
+    # blkid may need to be re-run if the mount device is not in the blkid info
+    # zfs devices are not in blkid, so we don't need to check for them
+    if fs_type == "zfs":
+        mount_info = {"type": "zfs", "path": mount_device}
+    else:
+        mount_info = get_blkid_info(self, mount_device)  # Raises an exception if the device is not found
 
     if ":" in mount_device:  # Handle bcachefs
         for alt_devices in mount_device.split(":"):
@@ -646,17 +724,27 @@ def _autodetect_mount(self, mountpoint) -> str:
     else:
         autodetect_mount_kmods(self, mount_device)
 
+    # force the name "root" for the root mount, remove the leading slash for other mounts
     mount_name = "root" if mountpoint == "/" else mountpoint.removeprefix("/")
+
+    # Don't overwrite existing mounts if a source type is already set
     if mount_name in self["mounts"] and any(s_type in self["mounts"][mount_name] for s_type in SOURCE_TYPES):
         return self.logger.warning(
             "[%s] Skipping autodetection, mount config already set:\n%s"
             % (colorize(mountpoint, "yellow"), pretty_print(self["mounts"][mount_name]))
         )
 
     mount_config = {mount_name: {"type": "auto", "options": ["ro"]}}  # Default to auto and ro
-    if mount_type := mount_info.get("type"):
-        self.logger.info("Autodetected mount type: %s" % colorize(mount_type, "cyan"))
-        mount_config[mount_name]["type"] = mount_type.lower()
+    fs_type = mount_info.get("type", fs_type) or "auto"
+    if fs_type == "auto":
+        self.logger.warning("Failed to autodetect mount type for mountpoint:" % (colorize(mountpoint, "yellow")))
+    else:
+        self.logger.info("[%s] Autodetected mount type from device: %s" % (mount_device, colorize(fs_type, "cyan")))
+    mount_config[mount_name]["type"] = fs_type.lower()
+
+    # for zfs mounts, set the path to the pool name
+    if fs_type == "zfs":
+        mount_config[mount_name]["path"] = mount_device
 
     for source_type in SOURCE_TYPES:
         if source := mount_info.get(source_type):
@@ -667,7 +755,8 @@ def _autodetect_mount(self, mountpoint) -> str:
             mount_config[mount_name][source_type] = source
             break
     else:
-        raise AutodetectError("[%s] Failed to autodetect mount source." % mountpoint)
+        if fs_type != "zfs":  # For ZFS, the source is the pool name
+            raise AutodetectError("[%s] Failed to autodetect mount source." % mountpoint)
 
     self["mounts"] = mount_config
     return mount_device
@@ -781,6 +870,10 @@ def _validate_host_mount(self, mount, destination_path=None) -> bool:
             break  # Skip host option validation if this is set
         if option == "ro":  # Allow the ro option to be set in the config
             continue
+        if option == "zfsutil":
+            if self["_mounts"][destination_path]["fstype"] == "zfs":
+                continue
+            raise ValueError("Cannot set 'zfsutil' option for non-zfs mount: %s" % destination_path)
         if option not in host_mount_options:
             raise ValidationError(
                 "Host mount options mismatch. Expected: %s, Found: %s" % (mount["options"], host_mount_options)
@@ -864,13 +957,36 @@ def export_mount_info(self) -> None:
     self["exports"]["MOUNTS_ROOT_TARGET"] = self["mounts"]["root"]["destination"]
 
 
+def autodetect_zfs_device_kmods(self, poolname) -> list[str]:
+    """Gets kmods for all devices in a ZFS pool and adds them to _kmod_auto."""
+    for device in get_zpool_info(self, poolname)["devices"]:
+        if device_kmods := resolve_blkdev_kmod(self, device):
+            self.logger.info(
+                "[%s:%s] Auto-enabling kernel modules for ZFS device: %s"
+                % (
+                    colorize(poolname, "blue"),
+                    colorize(device, "blue", bold=True),
+                    colorize(", ".join(device_kmods), "cyan"),
+                )
+            )
+            self["_kmod_auto"] = device_kmods
+
+
 def autodetect_mount_kmods(self, device) -> None:
     """Autodetects the kernel modules for a block device."""
+    if fs_type := _get_mount_dev_fs_type(self, device, raise_exception=False):
+        # This will fail for most non-zfs devices
+        if fs_type == "zfs":
+            return autodetect_zfs_device_kmods(self, device)
+
     if "/" not in str(device):
         device = f"/dev/{device}"
 
     if device_kmods := resolve_blkdev_kmod(self, device):
-        self.logger.info("Auto-enabling kernel modules for device: %s" % colorize(", ".join(device_kmods), "cyan"))
+        self.logger.info(
+            "[%s] Auto-enabling kernel modules for device: %s"
+            % (colorize(device, "blue"), colorize(", ".join(device_kmods), "cyan"))
+        )
         self["_kmod_auto"] = device_kmods
 
 

src/ugrd/fs/mounts.toml

@@ -29,7 +29,7 @@ late_fstab = "/etc/fstab.late"
 		     "_process_late_mounts_multi"]
 
 [imports.build_enum]
-"ugrd.fs.mounts" = [ "get_mounts_info", "get_virtual_block_info", "get_blkid_info",
+"ugrd.fs.mounts" = [ "get_mounts_info", "get_virtual_block_info", "get_blkid_info", "get_zpool_info",
 		     "autodetect_root", "autodetect_mounts", "autodetect_init_mount" ]
 
 [imports.build_tasks]
@@ -80,6 +80,7 @@ no_fsck = "bool"  # Whether or not to skip fsck on the root device when applicab
 _mounts = "dict"  # The mounts information
 _vblk_info = "dict"  # Virtual block device information
 _blkid_info = "dict"  # The blkid information
+_zpool_info = "dict"  # The zpool information
 
 # Define the base of the root mount
 [mounts.root]

src/ugrd/fs/zfs.py

@@ -0,0 +1,10 @@
+__version__ = "0.2.2"
+
+
+def zpool_import(self) -> str:
+    """ Returns bash lines to import all ZFS pools """
+    return """
+    edebug 'Importing all ZFS pools'
+    export ZPOOL_IMPORT_UDEV_TIMEOUT_MS=0  # Disable udev timeout
+    einfo "$(zpool import -aN)"
+    """

src/ugrd/fs/zfs.toml

@@ -0,0 +1,9 @@
+binaries = ["zfs", "zpool"]
+kmod_init = ["zfs"]
+
+
+[imports.init_main]
+"ugrd.fs.zfs" = [ "zpool_import" ]
+
+[import_order.after]
+zpool_import = ["mount_fstab", "crypt_init", "init_lvm"]

src/ugrd/kmod/kmod.py

@@ -353,6 +353,12 @@ def process_ignored_module(self, module: str) -> None:
             if key == "kmod_init":
                 if module in self["_kmod_modinfo"] and self["_kmod_modinfo"][module]["filename"] == "(builtin)":
                     self.logger.debug("Removing built-in module from kmod_init: %s" % module)
+                elif module == "zfs":
+                    self.logger.critical("ZFS module is required but missing.")
+                    self.logger.critical("Please build/install the required kmods before running this script.")
+                    self.logger.critical("Detected kernel version: %s" % self["kernel_version"])
+                    # https://github.com/projg2/installkernel-gentoo/commit/1c70dda8cd2700e5306d2ed74886b66ad7ccfb42
+                    exit(77)
                 else:
                     raise ValueError("Required module cannot be imported and is not builtin: %s" % module)
             else: