#!/bin/sh

command -v getarg >/dev/null || . /lib/dracut-lib.sh || . /usr/lib/dracut/modules.d/99base/dracut-lib.sh
command -v getargbool >/dev/null || {
    # Compatibility with older Dracut versions.
    # With apologies to the Dracut developers.
    getargbool() {
        _default="$1"; shift
        ! _b=$(getarg "$@") && [ -z "$_b" ] && _b="$_default"
        if [ -n "$_b" ]; then
            [ "$_b" = "0" ] && return 1
            [ "$_b" = "no" ] && return 1
            [ "$_b" = "off" ] && return 1
        fi
        return 0
    }
}

OLDIFS="${IFS}"
NEWLINE="
"
TAB="	"

ZPOOL_IMPORT_OPTS=""
if getargbool 0 zfs_force -y zfs.force -y zfsforce ; then
    warn "ZFS: Will force-import pools if necessary."
    ZPOOL_IMPORT_OPTS="${ZPOOL_IMPORT_OPTS} -f"
fi

# find_bootfs
#   returns the first dataset with the bootfs attribute.
find_bootfs() {
    IFS="${NEWLINE}"
    for dataset in $(zpool list -H -o bootfs); do
        case "${dataset}" in
            "" | "-")
                continue
                ;;
            "no pools available")
                IFS="${OLDIFS}"
                return 1
                ;;
            *)
                IFS="${OLDIFS}"
                echo "${dataset}"
                return 0
                ;;
        esac
    done

    IFS="${OLDIFS}"
    return 1
}

# import_pool POOL
#   imports the given zfs pool if it isn't imported already.
import_pool() {
    pool="${1}"

    if ! zpool list -H "${pool}" > /dev/null 2>&1; then
        info "ZFS: Importing pool ${pool}..."
        # shellcheck disable=SC2086
        if ! zpool import -N ${ZPOOL_IMPORT_OPTS} "${pool}" ; then
            warn "ZFS: Unable to import pool ${pool}"
            return 1
        fi
    fi

    return 0
}

_mount_dataset_cb() {
    # shellcheck disable=SC2154
    mount -o zfsutil -t zfs "${1}" "${NEWROOT}${2}"
}

# mount_dataset DATASET
#   mounts the given zfs dataset.
mount_dataset() {
    dataset="${1}"
    mountpoint="$(zfs get -H -o value mountpoint "${dataset}")"
    ret=0

    # We need zfsutil for non-legacy mounts and not for legacy mounts.
    if [ "${mountpoint}" = "legacy" ] ; then
        mount -t zfs "${dataset}" "${NEWROOT}" || ret=$?
    else
        mount -o zfsutil -t zfs "${dataset}" "${NEWROOT}" || ret=$?

        if [ "$ret" = "0" ]; then
            for_relevant_root_children "${dataset}" _mount_dataset_cb || ret=$?
        fi
    fi

    return "${ret}"
}

# for_relevant_root_children DATASET EXEC
#   Runs "EXEC dataset mountpoint" for all children of DATASET that are needed for system bringup
#   Used by zfs-generator.sh and friends, too!
for_relevant_root_children() {
    dataset="${1}"
    exec="${2}"

    zfs list -t filesystem -Ho name,mountpoint,canmount -r "${dataset}" |
        (
            _ret=0
            while IFS="${TAB}" read -r dataset mountpoint canmount; do
                [ "$canmount" != "on" ] && continue

                case "$mountpoint" in
                    /etc|/bin|/lib|/lib??|/libx32|/usr)
                        # If these aren't mounted we may not be able to get to the real init at all, or pollute the dataset holding the rootfs
                        "${exec}" "${dataset}" "${mountpoint}" || _ret=$?
                        ;;
                    *)
                        # Up to the real init to remount everything else it might need
                        ;;
                esac
            done
            exit "${_ret}"
        )
}

# ask_for_password
#
# Wraps around plymouth ask-for-password and adds fallback to tty password ask
# if plymouth is not present.
#
# --cmd command
#   Command to execute. Required.
# --prompt prompt
#   Password prompt. Note that function already adds ':' at the end.
#   Recommended.
# --tries n
#   How many times repeat command on its failure.  Default is 3.
# --ply-[cmd|prompt|tries]
#   Command/prompt/tries specific for plymouth password ask only.
# --tty-[cmd|prompt|tries]
#   Command/prompt/tries specific for tty password ask only.
# --tty-echo-off
#   Turn off input echo before tty command is executed and turn on after.
#   It's useful when password is read from stdin.
ask_for_password() {
    ply_tries=3
    tty_tries=3
    while [ "$#" -gt 0 ]; do
        case "$1" in
            --cmd) ply_cmd="$2"; tty_cmd="$2"; shift;;
            --ply-cmd) ply_cmd="$2"; shift;;
            --tty-cmd) tty_cmd="$2"; shift;;
            --prompt) ply_prompt="$2"; tty_prompt="$2"; shift;;
            --ply-prompt) ply_prompt="$2"; shift;;
            --tty-prompt) tty_prompt="$2"; shift;;
            --tries) ply_tries="$2"; tty_tries="$2"; shift;;
            --ply-tries) ply_tries="$2"; shift;;
            --tty-tries) tty_tries="$2"; shift;;
            --tty-echo-off) tty_echo_off=yes;;
            *) echo "ask_for_password(): wrong opt '$1'" >&2;;
        esac
        shift
    done

    {
        flock -s 9

        # Prompt for password with plymouth, if installed and running.
        if plymouth --ping 2>/dev/null; then
            plymouth ask-for-password \
                --prompt "$ply_prompt" --number-of-tries="$ply_tries" | \
                eval "$ply_cmd"
            ret=$?
        else
            if [ "$tty_echo_off" = yes ]; then
                stty_orig="$(stty -g)"
                stty -echo
            fi

            i=1
            while [ "$i" -le "$tty_tries" ]; do
                [ -n "$tty_prompt" ] && \
                    printf "%s [%i/%i]:" "$tty_prompt" "$i" "$tty_tries" >&2
                eval "$tty_cmd" && ret=0 && break
                ret=$?
                i=$((i+1))
                [ -n "$tty_prompt" ] && printf '\n' >&2
            done
            unset i
            [ "$tty_echo_off" = yes ] && stty "$stty_orig"
        fi
    } 9>/.console_lock

    [ "$ret" -ne 0 ] && echo "Wrong password" >&2
    return "$ret"
}

# Parse root=, rootfstype=, return them decoded and normalised to zfs:AUTO for auto, plain dset for explicit
#
# True if ZFS-on-root, false if we shouldn't
#
# Supported values:
#   root=
#   root=zfs
#   root=zfs:
#   root=zfs:AUTO
#
#   root=ZFS=data/set
#   root=zfs:data/set
#   root=zfs:ZFS=data/set (as a side-effect; allowed but undocumented)
#
#   rootfstype=zfs AND root=data/set <=> root=data/set
#   rootfstype=zfs AND root=         <=> root=zfs:AUTO
#
# '+'es in explicit dataset decoded to ' 's.
decode_root_args() {
    if [ -n "$rootfstype" ]; then
        [ "$rootfstype" = zfs ]
        return
    fi

    root=$(getarg root=)
    rootfstype=$(getarg rootfstype=)

    # shellcheck disable=SC2249
    case "$root" in
        ""|zfs|zfs:|zfs:AUTO)
            root=zfs:AUTO
            rootfstype=zfs
            return 0
            ;;

        ZFS=*|zfs:*)
            root="${root#zfs:}"
            root="${root#ZFS=}"
            root=$(echo "$root" | tr '+' ' ')
            rootfstype=zfs
            return 0
            ;;
    esac

    if [ "$rootfstype" = "zfs" ]; then
        case "$root" in
            "") root=zfs:AUTO ;;
            *)  root=$(echo "$root" | tr '+' ' ') ;;
        esac
        return 0
    fi

    return 1
}