#!/usr/bin/env bash
#
# omv-microvm-run — supervise a single Firecracker microVM.
# Invoked as `ExecStart=/usr/sbin/omv-microvm-run %i` by omv-microvm@.service.
#
# Usage: omv-microvm-run <vm-name>

set -eu

. /usr/share/openmediavault/scripts/helper-functions

VMNAME="${1:?usage: omv-microvm-run <vm-name>}"

_log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ${VMNAME}: $*"
}

# SuccessExitStatus=1 (unit file) covers dtach's own exit(1) on SIGTERM, so
# any other exit-1 before we hand off to dtach must be remapped to exit 2
# or it'd be misread as a clean stop.
STARTED=0
on_exit() {
    local ec=$?
    if [ "$STARTED" -eq 0 ] && [ "$ec" -eq 1 ]; then
        exit 2
    fi
}
trap on_exit EXIT

die() {
    _log "ERROR: $*"
    exit 2
}

VM_XPATH="/config/services/microvm/vms/vm[name='${VMNAME}']"
omv_config_exists "$VM_XPATH" || die "VM '${VMNAME}' not found in configuration."
ENABLE=$(omv_config_get "${VM_XPATH}/enable")
omv_checkyesno "$ENABLE" || die "VM '${VMNAME}' is disabled."

VCPUS=$(omv_config_get "${VM_XPATH}/vcpus")
MEMORY_MIB=$(omv_config_get "${VM_XPATH}/memory_mib")
IMAGE_REF=$(omv_config_get "${VM_XPATH}/imageref")
NETWORKREF=$(omv_config_get "${VM_XPATH}/networkref")
MACADDR=$(omv_config_get "${VM_XPATH}/macaddr")
BOOTARGS=$(omv_config_get "${VM_XPATH}/bootargs")
USE_JAILER=0
omv_checkyesno "$(omv_config_get "${VM_XPATH}/jailer")" && USE_JAILER=1
JAIL_UID=$(omv_config_get "${VM_XPATH}/jail_uid")

[ -n "$IMAGE_REF" ] || die "No image configured."
[ -n "$NETWORKREF" ] || die "No network configured."

NETWORK_XPATH="/config/services/microvm/networks/network[name='${NETWORKREF}']"
omv_config_exists "$NETWORK_XPATH" || die "Network '${NETWORKREF}' not found in configuration."
NET_TYPE=$(omv_config_get "${NETWORK_XPATH}/type")

SF_UUID=$(omv_config_get "/config/services/microvm/sharedfolderref")
omv_isuuid "$SF_UUID" || die "No shared folder configured for VM storage."
SF_PATH=$(omv_get_sharedfolder_path "$SF_UUID")
[ -n "$SF_PATH" ] || die "Could not resolve shared folder path."

VM_DIR="${SF_PATH}/vms/${VMNAME}"
mkdir -p "$VM_DIR"

# Dropped by omv-microvm-snapshot-restore before starting the unit. Skips
# cold-boot provisioning below; the ExecStartPost restore hook loads the
# snapshot's memory/device state instead.
RESTORE_MARKER="${VM_DIR}/.restore_from"
RESTORE_MODE=0
if [ -f "$RESTORE_MARKER" ]; then
    RESTORE_MODE=1
    _log "Restore marker present — skipping boot provisioning, waiting for snapshot load."
fi

# Each VM gets its own writable kernel/rootfs copy on first boot, so it no
# longer depends on the image afterwards.
ROOTFS_PATH="${VM_DIR}/rootfs.ext4"
KERNEL_PATH="${VM_DIR}/vmlinux"

if [ "$RESTORE_MODE" -eq 0 ]; then
    if [ ! -f "$ROOTFS_PATH" ] || [ ! -f "$KERNEL_PATH" ]; then
        # Composite key: "name-arch", or "name-arch-version" when a
        # version is set — must match imageRef() in the RPC service.
        IMAGE_XPATH="/config/services/microvm/images/image[(version='' and concat(name,'-',arch)='${IMAGE_REF}') or (version!='' and concat(name,'-',arch,'-',version)='${IMAGE_REF}')]"
        omv_config_exists "$IMAGE_XPATH" || die "Image '${IMAGE_REF}' not found."

        BASE_KERNEL_PATH=$(omv_config_get "${IMAGE_XPATH}/kernel_path")
        BASE_ROOTFS_PATH=$(omv_config_get "${IMAGE_XPATH}/rootfs_path")
    fi

    if [ ! -f "$ROOTFS_PATH" ]; then
        [ -f "$BASE_ROOTFS_PATH" ] || die "Rootfs image missing: ${BASE_ROOTFS_PATH}"
        _log "Copying base rootfs (first boot) ..."
        cp --reflink=auto "$BASE_ROOTFS_PATH" "$ROOTFS_PATH"
    fi

    if [ ! -f "$KERNEL_PATH" ]; then
        [ -f "$BASE_KERNEL_PATH" ] || die "Kernel image missing: ${BASE_KERNEL_PATH}"
        _log "Copying kernel (first boot) ..."
        cp --reflink=auto "$BASE_KERNEL_PATH" "$KERNEL_PATH"
    fi
fi

# /run rather than the shared folder: AF_UNIX paths are capped at ~108
# bytes (SUN_LEN), which shared-folder paths can easily exceed.
RUN_DIR="/run/openmediavault-microvm/${VMNAME}"
mkdir -p "$RUN_DIR"
SOCKET_PATH="${RUN_DIR}/firecracker.socket"
DTACH_SOCK="${RUN_DIR}/console.dtach"
CONSOLE_LOG="${VM_DIR}/console.log"
VMM_LOG="${VM_DIR}/vmm.log"
CONFIG_PATH="${VM_DIR}/vm_config.json"
rm -f "$SOCKET_PATH" "${RUN_DIR}/jailed"
: > "$CONSOLE_LOG"
: > "$VMM_LOG"

# Jailer chroot: <base>/<exec-file-name>/<id>/root, laid out by the
# jailer itself. The id only allows [a-zA-Z0-9-], hence a hash rather than
# the VM name. Must match omv-microvm-jail-teardown. /var/lib rather than
# /run: the jailer mknods /dev/kvm etc. inside it, which a nodev mount
# (like /run) would make unusable.
JAIL_BASE="/var/lib/openmediavault-microvm/jail"
JAIL_ID="mvm-$(echo -n "$VMNAME" | md5sum | cut -c1-16)"
JAIL_ROOT="${JAIL_BASE}/firecracker/${JAIL_ID}/root"
# Where a jailed VM's warm snapshots are written (bind-mounted into the
# chroot as /snapshot-out), before omv-microvm-snapshot-create moves them
# to their real destination. On the VM's own storage so that move is
# usually a rename, not a copy of the whole memory dump.
JAIL_STAGE="${VM_DIR}/.jail-snapshot"

# Leftovers from a run that wasn't cleaned up (host crash), or from before
# the jailer setting was switched off.
omv-microvm-jail-teardown "$VMNAME" || die "Could not clean up the previous jailer chroot."

# Deterministic short name (IFNAMSIZ=15) so omv-microvm-stop-cleanup can
# find and remove the same device.
TAP_NAME="mvm-$(echo -n "$VMNAME" | md5sum | cut -c1-10)"

if [ "$NET_TYPE" = "nat" ]; then
    SUBNET=$(omv_config_get "${NETWORK_XPATH}/subnet")
    BRIDGE=$(omv-microvm-ensure-nat-network "$NETWORKREF" "$SUBNET")

    # Firecracker guests get no DHCP server on a private NAT network, so
    # each VM's address is handed to the kernel directly via the "ip="
    # boot parameter (standard Linux IP autoconfiguration — no in-guest
    # DHCP client required). The preferred offset is a deterministic hash
    # of the VM name, the same approach used for TAP_NAME above; another
    # VM's reservation on that address bumps it to the next free one.
    NET_ADDR="${SUBNET%/*}"
    PREFIX="${SUBNET#*/}"
    _ip_to_int() { local a b c d; IFS=. read -r a b c d <<< "$1"; echo $(( (a << 24) + (b << 16) + (c << 8) + d )); }
    _int_to_ip() { local n=$1; echo "$(( (n >> 24) & 255 )).$(( (n >> 16) & 255 )).$(( (n >> 8) & 255 )).$(( n & 255 ))"; }
    NET_INT=$(_ip_to_int "$NET_ADDR")
    HOST_BITS=$((32 - PREFIX))
    [ "$HOST_BITS" -ge 2 ] || die "Subnet '${SUBNET}' is too small (need at least a /30)."
    MASK_INT=$(( (0xFFFFFFFF << HOST_BITS) & 0xFFFFFFFF ))
    NETMASK=$(_int_to_ip "$MASK_INT")
    GATEWAY_IP=$(_int_to_ip "$((NET_INT + 1))")
    USABLE_HOSTS=$(( (1 << HOST_BITS) - 3 ))
    HASH=$(echo -n "$VMNAME" | cksum | cut -d' ' -f1)

    # Each VM's address is reserved in a file that outlives the VM's run,
    # so two VMs never share one and a VM keeps its address across boots.
    # Removed with the VM (doCommand delete) or the network
    # (omv-microvm-teardown-nat-network).
    NAT_IPS_BASE="/var/lib/openmediavault-microvm/nat-ips"
    NAT_IPS_DIR="${NAT_IPS_BASE}/${NETWORKREF}"
    OWN_IP_FILE="${NAT_IPS_DIR}/${VMNAME}"
    mkdir -p "$NAT_IPS_DIR"
    # Serialises allocation between VMs starting at the same time.
    exec 9> "${NAT_IPS_BASE}/.lock"
    flock 9
    # Its reservation on a network it was on before.
    for f in "$NAT_IPS_BASE"/*/"$VMNAME"; do
        [ "$f" = "$OWN_IP_FILE" ] || rm -f "$f"
    done
    declare -A TAKEN_IPS=()
    for f in "$NAT_IPS_DIR"/*; do
        [ -f "$f" ] && [ "$f" != "$OWN_IP_FILE" ] || continue
        TAKEN_IPS["$(cat "$f")"]=1
    done
    # Usable offsets are 2 .. USABLE_HOSTS+1 (.1 is the gateway).
    _ip_offset_free() {
        [ "$1" -ge 2 ] && [ "$1" -le $((USABLE_HOSTS + 1)) ] \
            && [ -z "${TAKEN_IPS[$(_int_to_ip "$((NET_INT + $1))")]:-}" ]
    }
    GUEST_IP=""
    if [ -f "$OWN_IP_FILE" ]; then
        PREV_IP=$(cat "$OWN_IP_FILE")
        if [[ "$PREV_IP" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]] \
            && _ip_offset_free "$(( $(_ip_to_int "$PREV_IP") - NET_INT ))"; then
            GUEST_IP="$PREV_IP"
        fi
    fi
    if [ -z "$GUEST_IP" ]; then
        for i in $(seq 0 $((USABLE_HOSTS - 1))); do
            OFFSET=$(( ((HASH + i) % USABLE_HOSTS) + 2 ))
            if _ip_offset_free "$OFFSET"; then
                GUEST_IP=$(_int_to_ip "$((NET_INT + OFFSET))")
                [ "$i" -eq 0 ] || _log "WARNING: preferred address taken by another VM; using ${GUEST_IP}."
                break
            fi
        done
    fi
    [ -n "$GUEST_IP" ] || die "No free address left in subnet '${SUBNET}'."
    printf '%s\n' "$GUEST_IP" > "$OWN_IP_FILE"
    flock -u 9
    exec 9>&-

    # This VM's reservation, if it was on a DHCP network last time.
    rm -f /run/openmediavault-microvm-dhcp/*/hosts/"${VMNAME}"

    if omv_checkyesno "$(omv_config_get "${NETWORK_XPATH}/dhcp")"; then
        # A reservation needs a MAC that stays the same across boots;
        # without one set, the guest would pick a random one each time.
        # Locally administered (02:), derived from the name like TAP_NAME.
        if [ -z "$MACADDR" ]; then
            MACADDR="02:fc:$(echo -n "$VMNAME" | md5sum | sed -E 's/^(..)(..)(..)(..).*/\1:\2:\3:\4/')"
        fi
        # Must match omv-microvm-dhcp. dnsmasq notices new files itself.
        DHCP_HOSTS_DIR="/run/openmediavault-microvm-dhcp/${NETWORKREF}/hosts"
        mkdir -p "$DHCP_HOSTS_DIR"
        printf '%s,%s,infinite\n' "$MACADDR" "$GUEST_IP" > "${DHCP_HOSTS_DIR}/${VMNAME}"
        # --no-block: this runs inside the VM unit's own start job. The
        # guest's address doesn't wait on it anyway (see "ip=" below).
        systemctl start --no-block "omv-microvm-dhcp@${NETWORKREF}.service" \
            || _log "WARNING: could not start the DHCP server for network '${NETWORKREF}'."
        # dns0: the gateway, where dnsmasq also answers DNS.
        BOOTARGS="${BOOTARGS} ip=${GUEST_IP}::${GATEWAY_IP}:${NETMASK}::eth0:off:${GATEWAY_IP}"
        _log "Private NAT network '${NETWORKREF}' (DHCP): guest ${GUEST_IP} (${MACADDR}), gateway/DNS ${GATEWAY_IP}, bridge ${BRIDGE}."
    else
        BOOTARGS="${BOOTARGS} ip=${GUEST_IP}::${GATEWAY_IP}:${NETMASK}::eth0:off"
        _log "Private NAT network '${NETWORKREF}': guest ${GUEST_IP}, gateway ${GATEWAY_IP}, bridge ${BRIDGE}."
    fi
else
    rm -f /run/openmediavault-microvm-dhcp/*/hosts/"${VMNAME}" \
        /var/lib/openmediavault-microvm/nat-ips/*/"${VMNAME}"
    BRIDGE=$(omv_config_get "${NETWORK_XPATH}/bridge")
    [ -n "$BRIDGE" ] || die "Network '${NETWORKREF}' has no bridge device configured."
fi

if [ "$USE_JAILER" -eq 1 ]; then
    # Owned by the jail user so unprivileged Firecracker may attach to it;
    # an existing device's owner can't be changed, so always recreate it.
    if ip link show "$TAP_NAME" >/dev/null 2>&1; then
        ip link del "$TAP_NAME"
    fi
    ip tuntap add dev "$TAP_NAME" mode tap user "$JAIL_UID" group "$JAIL_UID"
elif ! ip link show "$TAP_NAME" >/dev/null 2>&1; then
    ip tuntap add dev "$TAP_NAME" mode tap
fi
ip link set "$TAP_NAME" master "$BRIDGE"
ip link set "$TAP_NAME" up

NETIF_JSON="{\"iface_id\":\"eth0\",\"host_dev_name\":\"${TAP_NAME}\"}"
if [ -n "$MACADDR" ]; then
    NETIF_JSON="{\"iface_id\":\"eth0\",\"host_dev_name\":\"${TAP_NAME}\",\"guest_mac\":\"${MACADDR}\"}"
fi

DISK_ROWS=$(omv-microvm-list-disks "$VMNAME") || die "Could not list data disks."

# The paths Firecracker itself is given — inside the chroot when jailed.
FC_KERNEL_PATH="$KERNEL_PATH"
FC_ROOTFS_PATH="$ROOTFS_PATH"
FC_CONFIG_PATH="$CONFIG_PATH"
FC_SOCKET_PATH="$SOCKET_PATH"
FC_VMM_LOG="$VMM_LOG"
fc_disk_path() {
    if [ "$USE_JAILER" -eq 1 ]; then
        echo "/disks/${1}.img"
    else
        echo "$2"
    fi
}

if [ "$USE_JAILER" -eq 1 ]; then
    [ -x /usr/local/bin/jailer ] || die "Jailer binary missing: /usr/local/bin/jailer"
    [ "${JAIL_UID:-0}" -gt 0 ] 2>/dev/null || die "No jail user id allocated; save the VM's settings again."

    # Bind mounts rather than hard links, which can't cross filesystems
    # (the chroot is on the root fs, disks are on shared folders).
    # ExecStopPost's omv-microvm-jail-teardown undoes all of them, also
    # when this script fails partway.
    _jail_bind() {
        # _jail_bind <host-path> <path-in-chroot> <readonly 0|1>
        local target="${JAIL_ROOT}${2}"
        chown "${JAIL_UID}:${JAIL_UID}" "$1" || die "Could not hand ${1} to the jail user."
        mkdir -p "$(dirname "$target")"
        touch "$target"
        mount --bind "$1" "$target" || die "Could not bind-mount ${1} into the jail."
        if [ "$3" -eq 1 ]; then
            mount -o remount,bind,ro "$target" || die "Could not make ${2} read-only in the jail."
        fi
    }

    mkdir -p "${JAIL_ROOT}/run"
    chown "${JAIL_UID}:${JAIL_UID}" "${JAIL_ROOT}/run"

    _jail_bind "$KERNEL_PATH" /vmlinux 1
    _jail_bind "$ROOTFS_PATH" /rootfs.ext4 0
    _jail_bind "$VMM_LOG" /vmm.log 0
    while IFS=$'\t' read -r disk_name disk_path disk_readonly _ _; do
        [ -n "$disk_name" ] || continue
        [ -f "$disk_path" ] || die "Data disk '${disk_name}' missing: ${disk_path}"
        _jail_bind "$disk_path" "$(fc_disk_path "$disk_name" "$disk_path")" "$disk_readonly"
    done <<< "$DISK_ROWS"

    rm -rf "$JAIL_STAGE"
    mkdir -m 0700 "$JAIL_STAGE"
    chown "${JAIL_UID}:${JAIL_UID}" "$JAIL_STAGE"
    mkdir -p "${JAIL_ROOT}/snapshot-out"
    mount --bind "$JAIL_STAGE" "${JAIL_ROOT}/snapshot-out" || die "Could not bind-mount the snapshot staging directory into the jail."

    if [ "$RESTORE_MODE" -eq 1 ]; then
        # The restore hook loads it from /snapshot-in. Handed back to root
        # by the hook once loaded.
        RESTORE_SRC=$(cat "$RESTORE_MARKER")
        chown "${JAIL_UID}:${JAIL_UID}" "${RESTORE_SRC}/memfile" "${RESTORE_SRC}/vmstate" \
            || die "Could not hand the snapshot to the jail user."
        mkdir -p "${JAIL_ROOT}/snapshot-in"
        mount --bind "$RESTORE_SRC" "${JAIL_ROOT}/snapshot-in" || die "Could not bind-mount the snapshot into the jail."
        mount -o remount,bind,ro "${JAIL_ROOT}/snapshot-in"
    fi

    FC_KERNEL_PATH="/vmlinux"
    FC_ROOTFS_PATH="/rootfs.ext4"
    FC_CONFIG_PATH="/vm_config.json"
    FC_SOCKET_PATH="/run/firecracker.socket"
    FC_VMM_LOG="/vmm.log"

    # Everything else (snapshot, restore hook, cleanup) keeps using the
    # usual socket path; the marker tells them the VM is jailed.
    ln -sfn "${JAIL_ROOT}${FC_SOCKET_PATH}" "$SOCKET_PATH"
    printf '%s' "$JAIL_ROOT" > "${RUN_DIR}/jailed"
    _log "Jailed as uid ${JAIL_UID} in ${JAIL_ROOT}."
fi

if [ "$RESTORE_MODE" -eq 0 ]; then
    # Data disks follow the rootfs in name order, so the guest sees them as
    # /dev/vdb, /dev/vdc, ... — ext4 ones are also labelled with the disk
    # name, which is the stabler thing to mount by.
    DATA_DRIVES_JSON="[]"
    while IFS=$'\t' read -r disk_name disk_path disk_readonly _ _; do
        [ -n "$disk_name" ] || continue
        [ -f "$disk_path" ] || die "Data disk '${disk_name}' missing: ${disk_path}"
        DATA_DRIVES_JSON=$(jq -c \
            --arg id "$disk_name" \
            --arg path "$(fc_disk_path "$disk_name" "$disk_path")" \
            --argjson ro "$([ "$disk_readonly" -eq 1 ] && echo true || echo false)" \
            '. + [{"drive_id": $id, "path_on_host": $path, "is_root_device": false, "is_read_only": $ro}]' \
            <<< "$DATA_DRIVES_JSON")
        _log "Attaching data disk '${disk_name}' ($([ "$disk_readonly" -eq 1 ] && echo ro || echo rw)): ${disk_path}"
    done <<< "$DISK_ROWS"

    jq -n \
        --arg kernel "$FC_KERNEL_PATH" \
        --arg bootargs "$BOOTARGS" \
        --arg rootfs "$FC_ROOTFS_PATH" \
        --argjson vcpus "$VCPUS" \
        --argjson mem "$MEMORY_MIB" \
        --argjson netif "$NETIF_JSON" \
        --argjson datadrives "$DATA_DRIVES_JSON" \
        '{
            "boot-source": {
                "kernel_image_path": $kernel,
                "boot_args": $bootargs
            },
            "drives": ([{
                "drive_id": "rootfs",
                "path_on_host": $rootfs,
                "is_root_device": true,
                "is_read_only": false
            }] + $datadrives),
            "network-interfaces": [$netif],
            "machine-config": {
                "vcpu_count": $vcpus,
                "mem_size_mib": $mem
            }
        }' > "$CONFIG_PATH"

    if [ "$USE_JAILER" -eq 1 ]; then
        install -o "$JAIL_UID" -g "$JAIL_UID" -m 0600 "$CONFIG_PATH" "${JAIL_ROOT}${FC_CONFIG_PATH}"
    fi
fi

FC_ARGS="--api-sock '${FC_SOCKET_PATH}' --log-path '${FC_VMM_LOG}' --level Info"
if [ "$RESTORE_MODE" -eq 0 ]; then
    FC_ARGS="${FC_ARGS} --config-file '${FC_CONFIG_PATH}'"
fi
if [ "$USE_JAILER" -eq 1 ]; then
    # No --daemonize/--new-pid-ns: the jailer must exec Firecracker in
    # place, so it stays attached to the console pty below and in this
    # unit's cgroup.
    FC_CMD="/usr/local/bin/jailer --id '${JAIL_ID}' --exec-file /usr/local/bin/firecracker --uid ${JAIL_UID} --gid ${JAIL_UID} --chroot-base-dir '${JAIL_BASE}' -- ${FC_ARGS}"
else
    FC_CMD="/usr/local/bin/firecracker ${FC_ARGS}"
fi

_log "Starting Firecracker (vcpus=${VCPUS} mem=${MEMORY_MIB}MiB bridge=${BRIDGE} tap=${TAP_NAME} restore=${RESTORE_MODE} jailer=${USE_JAILER}) ..."

# dtach lets omv-microvm-console attach to the live console later; script
# records the session to console.log regardless of whether anyone's
# attached. KillMode=control-group handles shutdown for the whole chain.
STARTED=1
exec dtach -N "$DTACH_SOCK" script -qfc "$FC_CMD" "$CONSOLE_LOG"
