path: root/nixos/modules/system/etc/build-composefs-dump.py

#!/usr/bin/env python3

"""Build a composefs dump from a Json config

See the man page of composefs-dump for details about the format:
https://github.com/containers/composefs/blob/main/man/composefs-dump.md

Ensure to check the file with the check script when you make changes to it:

./check-build-composefs-dump.sh ./build-composefs_dump.py
"""

import glob
import json
import os
import sys
from enum import Enum
from pathlib import Path
from typing import Any

Attrs = dict[str, Any]

# mkcomposefs hard-limits inline content to LCFS_INLINE_CONTENT_MAX (5000 bytes).
# We stay a bit below that. Files larger than this are served from the basedir
# data-only lower layer via an overlay redirect; files at or below it are
# embedded directly into the erofs metadata image, avoiding the redirect
# indirection at read time and keeping the basedir empty in the common case.
INLINE_CONTENT_MAX = 4096


class FileType(Enum):
    """The filetype as defined by the `st_mode` stat field in octal

    You can check the st_mode stat field of a path in Python with
    `oct(os.stat("/path/").st_mode)`
    """

    directory = "4"
    file = "10"
    symlink = "12"


class ComposefsPath:
    path: str
    size: int
    filetype: FileType
    mode: str
    uid: str
    gid: str
    payload: str
    rdev: str = "0"
    nlink: int = 1
    mtime: str = "1.0"
    content: str = "-"
    digest: str = "-"

    def __init__(
        self,
        attrs: Attrs,
        size: int,
        filetype: FileType,
        mode: str,
        payload: str,
        path: str | None = None,
        content: str = "-",
    ):
        if path is None:
            path = attrs["target"]
        self.path = path
        self.size = size
        self.filetype = filetype
        self.content = content

        match len(mode):
            case 3 | 4:
                # We need to pad the mode, because we will later use the concatentation
                # filetype|mode, which assumes that the mode has a length of 4.
                self.mode = f"{mode:0>4}"
            case _:
                raise ValueError(f"mode should be 3 or 4 octal digits, got: {mode}")

        self.uid = attrs["uid"]
        self.gid = attrs["gid"]
        self.payload = payload

    def write_line(self) -> str:
        line_list = [
            str(self.path),
            str(self.size),
            f"{self.filetype.value}{self.mode}",
            str(self.nlink),
            str(self.uid),
            str(self.gid),
            str(self.rdev),
            str(self.mtime),
            str(self.payload),
            str(self.content),
            str(self.digest),
        ]
        return " ".join(line_list)


def eprint(*args: Any, **kwargs: Any) -> None:
    print(*args, **kwargs, file=sys.stderr)


# Bytes that may appear unescaped in a composefs-dump field. Everything else
# is encoded as \xHH. See composefs-dump(5).
_DUMP_SHORT_ESCAPES: dict[int, str] = {
    ord("\\"): r"\\",
    ord("\n"): r"\n",
    ord("\r"): r"\r",
    ord("\t"): r"\t",
}


def escape_dump_field(data: bytes) -> str:
    """Escape raw bytes for use as a composefs-dump field.

    The dump format separates fields by a single space and lines by a single
    newline, uses '\\' as the escape character and reserves '-' for unset
    optional fields, so all of these (plus non-printable bytes and '=') must
    be escaped.
    """
    if data == b"":
        # An empty CONTENT field would be indistinguishable from two spaces
        # between PAYLOAD and DIGEST; callers must emit '-' for size-0 files
        # instead of inlining them.
        raise ValueError("cannot escape empty content; emit '-' instead")
    if data == b"-":
        # A bare '-' means "unset"; escape it so it round-trips as content.
        return r"\x2d"
    out: list[str] = []
    for b in data:
        if b in _DUMP_SHORT_ESCAPES:
            out.append(_DUMP_SHORT_ESCAPES[b])
        elif b in (ord(" "), ord("=")) or not (0x20 <= b <= 0x7E):
            out.append(f"\\x{b:02x}")
        else:
            out.append(chr(b))
    return "".join(out)


def normalize_path(path: str) -> str:
    return str("/" + os.path.normpath(path).lstrip("/"))


def leading_directories(path: str) -> list[str]:
    """Return the leading directories of path

    Given the path "alsa/conf.d/50-pipewire.conf", for example, this function
    returns `[ "alsa", "alsa/conf.d" ]`.
    """
    parents = list(Path(path).parents)
    parents.reverse()
    # remove the implicit `.` from the start of a relative path or `/` from an
    # absolute path
    del parents[0]
    return [str(i) for i in parents]


def add_leading_directories(
    target: str, attrs: Attrs, paths: dict[str, ComposefsPath]
) -> None:
    """Add the leading directories of a target path to the composefs paths

    mkcomposefs expects that all leading directories are explicitly listed in
    the dump file. Given the path "alsa/conf.d/50-pipewire.conf", for example,
    this function adds "alsa" and "alsa/conf.d" to the composefs paths.
    """
    path_components = leading_directories(target)
    for component in path_components:
        composefs_path = ComposefsPath(
            attrs,
            path=component,
            size=4096,
            filetype=FileType.directory,
            mode="0755",
            payload="-",
        )
        paths[component] = composefs_path


def main() -> None:
    """Build a composefs dump from a Json config

    This config describes the files that the final composefs image is supposed
    to contain.
    """
    config_file = sys.argv[1]
    if not config_file:
        eprint("No config file was supplied.")
        sys.exit(1)

    with open(config_file, "rb") as f:
        config = json.load(f)

    if not config:
        eprint("Config is empty.")
        sys.exit(1)

    eprint("Building composefs dump...")

    paths: dict[str, ComposefsPath] = {}
    for attrs in config:
        # Normalize the target path to work around issues in how targets are
        # declared in `environment.etc`.
        attrs["target"] = normalize_path(attrs["target"])

        target = attrs["target"]
        source = attrs["source"]
        mode = attrs["mode"]

        if "*" in source:  # Path with globbing
            glob_sources = glob.glob(source)
            for glob_source in glob_sources:
                basename = os.path.basename(glob_source)
                glob_target = f"{target}/{basename}"

                composefs_path = ComposefsPath(
                    attrs,
                    path=glob_target,
                    size=100,
                    filetype=FileType.symlink,
                    mode="0777",
                    payload=glob_source,
                )

                paths[glob_target] = composefs_path
                add_leading_directories(glob_target, attrs, paths)
        else:  # Without globbing
            if mode == "symlink" or mode == "direct-symlink":
                composefs_path = ComposefsPath(
                    attrs,
                    # A high approximation of the size of a symlink
                    size=100,
                    filetype=FileType.symlink,
                    mode="0777",
                    payload=source,
                )
            elif os.path.isdir(source):
                composefs_path = ComposefsPath(
                    attrs,
                    size=4096,
                    filetype=FileType.directory,
                    mode=mode,
                    payload=source,
                )
            else:
                size = os.stat(source).st_size
                if size <= INLINE_CONTENT_MAX:
                    # Inline small files directly into the erofs image so they
                    # do not need to be served from the basedir data layer via
                    # an overlay redirect. Empty files need neither payload nor
                    # content; mkcomposefs treats size=0 as an empty inline
                    # file.
                    if size > 0:
                        with open(source, "rb") as fh:
                            raw = fh.read()
                        content = escape_dump_field(raw)
                        size = len(raw)
                    else:
                        content = "-"
                    composefs_path = ComposefsPath(
                        attrs,
                        size=size,
                        filetype=FileType.file,
                        mode=mode,
                        payload="-",
                        content=content,
                    )
                else:
                    composefs_path = ComposefsPath(
                        attrs,
                        size=size,
                        filetype=FileType.file,
                        mode=mode,
                        # payload needs to be relative path in this case
                        payload=target.lstrip("/"),
                    )
            paths[target] = composefs_path
            add_leading_directories(target, attrs, paths)

    composefs_dump = ["/ 4096 40755 1 0 0 0 0.0 - - -"]  # Root directory
    for key in sorted(paths):
        composefs_path = paths[key]
        eprint(composefs_path.path)
        composefs_dump.append(composefs_path.write_line())

    print("\n".join(composefs_dump))


if __name__ == "__main__":
    main()