# 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)"

    # 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
    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

    # 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.
    if [ -e /lower/etc/fstab ]; then
        mkdir -p /upper/up/etc
        /bin/busybox awk '$2 != "/boot" && $2 != "/boot/efi" { print }' \
            /lower/etc/fstab > /upper/up/etc/fstab
        _nbdboot_trace "fstab: rewrote /etc/fstab in overlay upper (dropped /boot, /boot/efi) -- $(wc -l < /upper/up/etc/fstab) lines"
    fi

    # Prevent Ubuntu/Debian's systemd-networkd + cloud-init 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 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.
    #
    # 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 unit so socket-activation can't
    # resurrect networkd.
    mkdir -p /upper/up/etc/systemd/system
    for unit in \
        systemd-networkd.service \
        systemd-networkd.socket \
        systemd-networkd-wait-online.service \
        cloud-init.service \
        cloud-init-local.service \
        cloud-config.service \
        cloud-final.service \
    ; do
        ln -sf /dev/null "/upper/up/etc/systemd/system/${unit}"
    done
    _nbdboot_trace "network: masked systemd-networkd + cloud-init units so the initramfs-owned NIC survives userspace"

    # With networkd + cloud-init 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 + 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
        } > /upper/up/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.
}
