# initramfs-tools boot driver -- sourced by ``/init`` when the kernel
# cmdline carries ``boot=nbdboot``. Defines ``mountroot`` (the
# contract initramfs-tools' /init calls to populate ``/root`` before
# pivoting + ``exec /sbin/init``).
#
# Sequence inside mountroot():
#
#   1. Parse ``pixie.*`` cmdline params -- nbd endpoint, image
#      export name, overlay size, server URL, MAC.
#   2. modprobe nbd + overlay.
#   3. nbd-client connects to nbdmux; ``/dev/nbd0`` appears.
#   4. partx scans /dev/nbd0; pick the largest partition as the
#      root (good-enough heuristic for nosi disk images; can be
#      overridden via ``pixie.root_part=<devnode>``).
#   5. mount root read-only at /lower.
#   6. mount tmpfs (size = pixie.overlay_size, default 10G) at /upper.
#   7. overlayfs(lower=/lower, upper=/upper/up, work=/upper/work)
#      at ``/root`` (the path initramfs-tools' /init will
#      ``pivot_root`` into).
#   8. Best-effort POST status to pixie.server so the machine's
#      timeline reflects "nbdboot up" before pivot_root.
#
# Failure handling: on any unrecoverable step we drop to a
# ``panic`` (initramfs-tools' single-shell recovery) with the
# reason. The operator sees the message on tty/serial.

# shellcheck disable=SC2034
PREREQ=""
prereqs() { echo "$PREREQ"; }
case "${1:-}" in
    prereqs) prereqs; exit 0 ;;
esac

# Diagnostic trace helper -- prints to kernel log buffer (visible in
# dmesg + on the boot console) so the operator can tell whether
# ``mountroot`` was even entered and, if so, which step failed. Cheap
# to leave in place; the boot output already carries kernel spew.
_nbdboot_trace() { echo "nbdboot: $*" >/dev/kmsg 2>/dev/null || echo "nbdboot: $*"; }

# initramfs-tools' ``panic`` spawns ``sh -i`` on /dev/console and
# RETURNS if that shell exits. On serial without a controlling
# terminal the shell exits immediately, so ``panic`` becomes a no-op
# and mountroot merrily continues through every subsequent broken
# step, papering over the real failure. Use ``_nbdboot_die`` instead
# to trace, best-effort POST, and then hang PID 1 so the operator
# can actually read the message on the boot console.
_nbdboot_die() {
    _nbdboot_trace "FATAL: $*"
    _post_status "nbdboot.die"
    # Replace this shell with a permanent sleep so /init can't
    # continue past a failed step. The kernel keeps PID 1 alive and
    # the console stays readable.
    exec /bin/busybox sleep 2147483647
}

_nbdboot_trace "script sourced (BOOT=nbdboot)"

# initramfs-tools' standard library (provides ``log_begin_msg``,
# ``log_end_msg``, ``panic``, ``configure_networking``, etc.)
# shellcheck disable=SC1091
. /scripts/functions

_get_cmdline() {
    # ``$1`` is the parameter name (without ``=``). Reads /proc/cmdline
    # and prints just the value, or empty if the parameter is absent.
    sed -n "s/.*\\<$1=\\([^ ]*\\).*/\\1/p" /proc/cmdline
}

_post_status() {
    # Best-effort POST. ``$1`` = status string. Silent on failure;
    # the boot continues regardless.
    server="$(_get_cmdline pixie.server)"
    mac="$(_get_cmdline pixie.mac)"
    [ -n "$server" ] && [ -n "$mac" ] || return 0
    # busybox-static's ``wget`` does plain HTTP POST when we use
    # ``--post-data``; the body is irrelevant to pixie (server reads
    # ``status`` from a query param shape historically). For the
    # nbdboot variant we just touch the endpoint -- the URL alone
    # carries the signal.
    /bin/busybox wget -q -O /dev/null \
        --post-data="status=$1" \
        "${server}/pxe/${mac}/status" || true
}

mountroot() {
    _nbdboot_trace "mountroot() entered"
    nbd_url="$(_get_cmdline pixie.nbd)"
    image="$(_get_cmdline pixie.image)"
    overlay_size="$(_get_cmdline pixie.overlay_size)"
    root_part_override="$(_get_cmdline pixie.root_part)"
    # Persistent mode gate. ``pixie.persist=1`` on the kernel cmdline
    # tells this pivot the NBD export IS the machine's private
    # writable qcow2 (backed by the shared base blob on pixie's side
    # via qemu-nbd + backing_file). Skip the tmpfs+overlayfs setup
    # and mount /dev/nbd0's root partition read-write straight at
    # /root; writes land in the qcow2 and persist across reboots.
    persist="$(_get_cmdline pixie.persist)"

    # Defaults: 10 GiB tmpfs is enough for typical CI workloads
    # writing logs + scratch state; operator caps with pixie.overlay_size.
    : "${overlay_size:=10G}"

    _nbdboot_trace "cmdline nbd=${nbd_url:-<unset>} image=${image:-<unset>} overlay_size=${overlay_size} root_part=${root_part_override:-<auto>}"

    if [ -z "$nbd_url" ] || [ -z "$image" ]; then
        _nbdboot_die "missing pixie.nbd or pixie.image on kernel cmdline"
    fi

    # nbd_url is ``tcp://<host>:<port>``. Strip the scheme and split.
    nbd_host="${nbd_url#tcp://}"
    nbd_host="${nbd_host%%:*}"
    nbd_port="${nbd_url##*:}"

    _nbdboot_trace "modprobe nbd (nbds_max=1 max_part=16) + overlay"
    log_begin_msg "nbdboot: modprobe nbd + overlay"
    # ``max_part=16`` is essential: without it the nbd module leaves
    # /dev/nbd0 as a single flat device and NEVER exposes the
    # partition device nodes we scan for below. Debian's default is
    # ``max_part=0``.
    #
    # Linux keeps module parameters from the FIRST load and silently
    # ignores params passed to subsequent modprobe calls. If nbd got
    # pre-loaded during initramfs boot (e.g. by udev on a modalias),
    # our ``max_part=16`` would drop on the floor and partition
    # nodes would never appear. Force a clean reload.
    rmmod nbd 2>/dev/null || true
    modprobe nbd nbds_max=1 max_part=16 || _nbdboot_die "modprobe nbd failed"
    modprobe overlay || _nbdboot_die "modprobe overlay failed"
    _nbdboot_trace "nbd sysfs max_part=$(cat /sys/module/nbd/parameters/max_part 2>/dev/null || echo '<unreadable>')"
    log_end_msg

    # iPXE releases its DHCP lease when it chainloads the kernel, so
    # userspace has to redo network setup before nbd-client can talk
    # to nbdmux. configure_networking() (from /scripts/functions) does
    # DHCP via ipconfig(8) on the auto-detected NIC. Fatal on
    # failure -- nbd-client cannot succeed without a route to the
    # server.
    _nbdboot_trace "configure_networking (DHCP via ipconfig)"
    log_begin_msg "nbdboot: configure_networking"
    configure_networking || _nbdboot_die "configure_networking failed (DHCP)"
    log_end_msg
    _nbdboot_trace "network up: $(/bin/busybox ip -o -4 addr show 2>/dev/null | /bin/busybox awk '{print $2, $4}' | /bin/busybox tr '\n' ';')"

    # ``-persist`` keeps nbd-client's supervisor process alive after
    # the initial attach so it reconnects if the TCP session drops.
    # Ubuntu's systemd-networkd resets the network stack a few seconds
    # into userspace boot, which was killing the socket the initramfs
    # opened and leaving /dev/nbd0 dead. jbd2 on the /boot partition
    # then queues writes forever (observed 120s hung-task warnings).
    # With -persist the client reconnects transparently, and
    # nbdkit's cow filter is per-nbdkit-instance not per-connection,
    # so the writable overlay survives the drop.
    _nbdboot_trace "nbd-client -persist ${nbd_host}:${nbd_port} -name ${image}"
    log_begin_msg "nbdboot: nbd-client ${nbd_host}:${nbd_port} -name ${image}"
    if ! nbd-client -persist "$nbd_host" "$nbd_port" -name "$image" /dev/nbd0; then
        _post_status "nbdboot.nbd_connect_failed"
        _nbdboot_die "nbd-client failed to connect to ${nbd_host}:${nbd_port}"
    fi
    log_end_msg

    # ``nbd-client`` returns as soon as the client-side attach ioctl
    # is done, but the kernel commits the capacity change + partition
    # scan asynchronously via a workqueue. Racing partprobe before
    # the capacity event lands leaves /dev/nbd0 sized 0 and produces
    # NO partition nodes at all (been there). Poll for size > 0
    # before scanning; give up after ~10s.
    i=0
    while [ "$i" -lt 100 ]; do
        nbd0_size="$(/bin/busybox blockdev --getsize64 /dev/nbd0 2>/dev/null || echo 0)"
        [ "${nbd0_size:-0}" -gt 0 ] && break
        /bin/busybox sleep 0.1
        i=$((i + 1))
    done
    _nbdboot_trace "nbd0 capacity ready after ${i} tick(s): ${nbd0_size:-0} bytes"

    udevadm settle --timeout=10 || true
    /bin/busybox blockdev --rereadpt /dev/nbd0 2>/dev/null || true
    udevadm settle --timeout=10 || true

    # First try the kernel's own async partition scan. On Linux 6.12
    # with this nbd driver it consistently returns no partition nodes
    # for GPT disks (BLKRRPART returns 0 but the scanner gives up
    # somewhere between reading LBA 0 and validating the GPT header
    # on LBA 1). Poll briefly so we don't over-wait if the answer is
    # "not going to happen".
    i=0
    while [ "$i" -lt 20 ]; do
        [ -b /dev/nbd0p1 ] && break
        /bin/busybox sleep 0.1
        i=$((i + 1))
    done
    if [ ! -b /dev/nbd0p1 ]; then
        # Kernel didn't produce partition nodes. Use partx from
        # userspace: it reads the partition table itself (via
        # libblkid) and installs each partition on the disk via
        # BLKPG_ADD_PARTITION ioctls, bypassing the in-kernel
        # scanner. Same end state (``/dev/nbd0pN`` nodes appear
        # for udev to notice) without a loop-device wrapper.
        _px_out="$(/usr/bin/partx --add --verbose /dev/nbd0 2>&1)"; _px_rc=$?
        _nbdboot_trace "partx --add rc=${_px_rc} out='${_px_out}'"
        udevadm settle --timeout=10 || true
    fi
    _nbdboot_trace "nbd0 nodes after scan: $(/bin/busybox ls -1 /dev/nbd0* 2>/dev/null | /bin/busybox tr '\n' ' ')"

    if [ -n "$root_part_override" ]; then
        root_part="$root_part_override"
    elif [ -b /dev/nbd0p1 ]; then
        # Native partition scan worked -- pick the largest partition.
        root_part="$(
            /bin/busybox find /dev -maxdepth 1 -name 'nbd0p*' -print0 \
            | xargs -0 -I{} /bin/busybox sh -c 'echo "$(blockdev --getsize64 "$1") $1"' _ {} \
            | sort -n | tail -1 | cut -d' ' -f2
        )"
    else
        # No partition table (or scan didn't fire) -- treat the
        # whole device as a raw filesystem.
        root_part=/dev/nbd0
    fi
    _nbdboot_trace "picked root_part=${root_part:-<none>}"
    if [ -z "$root_part" ] || [ ! -e "$root_part" ]; then
        _post_status "nbdboot.no_root_partition"
        _nbdboot_die "could not pick a root partition on /dev/nbd0"
    fi

    mkdir -p /lower /upper /root
    _nbdboot_trace "$(/sbin/blkid -p ${root_part} 2>&1)"
    # Try common in-kernel filesystems explicitly, then fall back to
    # auto-detect. modprobe each candidate first because initramfs
    # kernel modules aren't preloaded and mount(2) with -t auto only
    # tries filesystems whose modules are already registered.
    modprobe ext4 2>/dev/null || true
    modprobe xfs 2>/dev/null || true
    modprobe btrfs 2>/dev/null || true
    _mnt_out=""; _mnt_rc=1
    if [ "$persist" = "1" ]; then
        # Persistent path: NBD serves this machine's private qcow2.
        # Mount it read-write straight at /root; the qcow2's backing_file
        # points at the shared base blob on pixie's side, so unchanged
        # bytes are still shared but writes stay in the qcow2. No
        # tmpfs upper, no overlayfs.
        _mnt_rc=1
        for fstype in ext4 xfs btrfs auto; do
            _nbdboot_trace "persist: mount -t ${fstype} -o rw ${root_part} -> /root"
            _mnt_out="$(mount -t "$fstype" -o rw "$root_part" /root 2>&1)"; _mnt_rc=$?
            [ "$_mnt_rc" -eq 0 ] && break
            _nbdboot_trace "  -> rc=${_mnt_rc} out='${_mnt_out}'"
        done
        if [ "$_mnt_rc" -ne 0 ]; then
            _post_status "nbdboot.mount_root_failed"
            _nbdboot_die "persist: failed to mount ${root_part} rw as ext4/xfs/btrfs/auto"
        fi
        _nbdboot_trace "persist: mounted ${root_part} rw at /root"
        # For the persistent path, fall through to the fstab-rewrite +
        # service-masking prep below. Those steps target /root regardless
        # of whether it's an overlay merge or a plain mount, and they're
        # idempotent -- the second boot re-masks already-masked units +
        # rewrites an already-rewritten fstab without side effects.
    else
        for fstype in ext4 xfs btrfs auto; do
            _nbdboot_trace "mount -t ${fstype} -o ro ${root_part} -> /lower"
            _mnt_out="$(mount -t "$fstype" -o ro "$root_part" /lower 2>&1)"; _mnt_rc=$?
            [ "$_mnt_rc" -eq 0 ] && break
            _nbdboot_trace "  -> rc=${_mnt_rc} out='${_mnt_out}'"
        done
        if [ "$_mnt_rc" -ne 0 ]; then
            _post_status "nbdboot.mount_root_failed"
            _nbdboot_die "failed to mount ${root_part} as ext4/xfs/btrfs/auto"
        fi
        _nbdboot_trace "mounted ${root_part}; ls /lower:  $(/bin/busybox ls /lower 2>&1 | /bin/busybox tr '\n' ' ')"

        _nbdboot_trace "tmpfs(${overlay_size}) -> /upper"
        log_begin_msg "nbdboot: tmpfs(${overlay_size}) at /upper"
        if ! mount -t tmpfs -o "size=${overlay_size}" tmpfs /upper; then
            _post_status "nbdboot.mount_tmpfs_failed"
            _nbdboot_die "failed to mount tmpfs at /upper"
        fi
        mkdir -p /upper/up /upper/work
        log_end_msg

        _nbdboot_trace "overlay(lower=/lower,upper=/upper/up,work=/upper/work) -> /root"
        log_begin_msg "nbdboot: overlayfs at /root"
        if ! mount -t overlay overlay \
            -o "lowerdir=/lower,upperdir=/upper/up,workdir=/upper/work" \
            /root
        then
            _post_status "nbdboot.mount_overlay_failed"
            _nbdboot_die "failed to mount overlayfs at /root"
        fi
        log_end_msg
    fi

    # Rewrite the guest's /etc/fstab in the overlay upper before we
    # hand control back to /init. Cloud disk images (nosi's Ubuntu +
    # Debian) list /boot and /boot/efi as separate ext4/vfat mounts
    # off partitions we can only reach through the loop stack over
    # nbd. Those mounts have NO overlayfs wrapper, so ext4 journal
    # writes to /boot land straight in nbdkit's cow filter through
    # loop + nbd, and something in that chain wedges jbd2 within
    # 120s (observed on GIGABYTE MC12-LE0 booting nosi's
    # ubuntu-2604-headless). Root is fine because mountroot() above
    # already wrapped it in overlay + tmpfs.
    #
    # Strip /boot and /boot/efi from fstab: the bootloader has
    # already consumed /boot before we ever ran, and nothing in
    # userspace needs to write to it. The overlay upper wins over
    # the lower's fstab because /root is the overlay merge and /init
    # pivots into that.
    # Overlay upperdir is /upper/up (see the ``mount -t overlay``
    # invocation above); writing to /upper/etc/fstab lands OUTSIDE
    # the overlay and the pivoted root still sees the lower's fstab.
    # The strip has to go into the actual upperdir.
    # ``prep_root`` is the path we write to when preparing the
    # first-boot state (fstab, service masks, resolv.conf). For the
    # ephemeral overlay path, this is the overlay upperdir (writes
    # here are the ones the pivoted root actually sees). For the
    # persistent path, /root is already the writable mount, so we
    # write directly there. Same downstream code, one variable.
    if [ "$persist" = "1" ]; then
        prep_root=/root
        prep_fstab_src=/root/etc/fstab
    else
        prep_root=/upper/up
        prep_fstab_src=/lower/etc/fstab
    fi

    if [ -e "$prep_fstab_src" ]; then
        mkdir -p "$prep_root/etc"
        /bin/busybox awk '$2 != "/boot" && $2 != "/boot/efi" { print }' \
            "$prep_fstab_src" > "$prep_root/etc/fstab"
        _nbdboot_trace "fstab: rewrote /etc/fstab under ${prep_root} (dropped /boot, /boot/efi) -- $(wc -l < "$prep_root/etc/fstab") lines"
    fi

    # Prevent Ubuntu/Debian's systemd-networkd + cloud-init +
    # NetworkManager from touching the NIC the initramfs already
    # configured. Under nbdboot the root filesystem lives on that
    # NIC (NBD over TCP), so anything that flushes the IP + re-DHCPs
    # -- which networkd + netplan + cloud-init's network stage +
    # NetworkManager all do by default -- tears down the socket the
    # kernel was reading root pages from, then wedges every process
    # on the ensuing D-state page fault because the reconnect needs
    # an IP that needs networkd that needs pages from a dead NBD.
    # Circular deadlock, no hung-task warning (nothing is runnable
    # long enough to trip 120s), silent console freeze after journald
    # flush.
    #
    # Ubuntu 26.04 Server switched from systemd-networkd to
    # NetworkManager by default; a nbdboot into a fresh nosi
    # ubuntu-2604-headless was burning 60 seconds on
    # NetworkManager-wait-online.service before its own "online"
    # check timed out (systemd-analyze critical-chain showed
    # ``NetworkManager-wait-online.service @3.946s +59.988s``). Mask
    # the whole NM stack: the initrd-owned NIC survives userspace
    # untouched.
    #
    # Mask by symlink-to-/dev/null in the overlay upper's systemd
    # unit dir. Fresh cloud install case: no runtime override needed
    # since we don't get a chance to reconfigure networking anyway.
    # Also masks the .socket units so socket-activation can't
    # resurrect them.
    mkdir -p ${prep_root}/etc/systemd/system
    for unit in \
        systemd-networkd.service \
        systemd-networkd.socket \
        systemd-networkd-wait-online.service \
        NetworkManager.service \
        NetworkManager-wait-online.service \
        NetworkManager-dispatcher.service \
        nm-cloud-setup.service \
        nm-cloud-setup.timer \
        cloud-init.service \
        cloud-init-local.service \
        cloud-config.service \
        cloud-final.service \
    ; do
        ln -sf /dev/null "${prep_root}/etc/systemd/system/${unit}"
    done
    _nbdboot_trace "network: masked systemd-networkd + NetworkManager + cloud-init units so the initramfs-owned NIC survives userspace"

    # With networkd + NetworkManager + cloud-init all masked, nothing
    # in userspace populates /etc/resolv.conf. Propagate the
    # DHCP-supplied DNS servers + search domain from initramfs-tools'
    # /run/net-*.conf into the overlay upper so ``apt-get update`` /
    # ``curl`` / ``ping <host>`` work in the pivoted root. First-match wins: the loop breaks on
    # the first ``/run/net-*.conf`` that carries a non-empty
    # ``DNSSERVERS`` line, which matches the single interface the
    # initramfs actually configured.
    for net_conf in /run/net-*.conf; do
        [ -e "$net_conf" ] || continue
        # shellcheck source=/dev/null
        . "$net_conf" 2>/dev/null || continue
        [ -n "${DNSSERVERS:-}" ] || continue
        {
            echo "# Written by nbdboot initramfs from ${net_conf}."
            echo "# nbdboot masks systemd-networkd + NetworkManager + cloud-init so nothing"
            echo "# in userspace would otherwise reseed this file."
            [ -n "${DOMAINSEARCH:-}" ] && echo "search ${DOMAINSEARCH}"
            for _ns in ${DNSSERVERS}; do
                echo "nameserver ${_ns}"
            done
        } > "${prep_root}/etc/resolv.conf"
        _nbdboot_trace "network: wrote /etc/resolv.conf from ${net_conf} (nameservers: ${DNSSERVERS})"
        break
    done

    _nbdboot_trace "mountroot() done -- returning to /init for pivot_root"
    _post_status "nbdboot.up"
    # initramfs-tools' /init pivot_root's into /root and exec's
    # /sbin/init from here.
}
