#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-2.0-only

shopt -s extglob

parseopts() {
    local opt='' optarg='' i='' shortopts="$1"
    local -a longopts=() unused_argv=()

    shift
    while [[ -n "$1" && "$1" != '--' ]]; do
        longopts+=("$1")
        shift
    done
    shift

    longoptmatch() {
        local o longmatch=()
        for o in "${longopts[@]}"; do
            if [[ "${o%:}" == "$1" ]]; then
                longmatch=("$o")
                break
            fi
            [[ "${o%:}" == "$1"* ]] && longmatch+=("$o")
        done

        case "${#longmatch[*]}" in
            1)
                # success, override with opt and return arg req (0 == none, 1 == required)
                opt="${longmatch%:}"
                if [[ "${longmatch[*]}" == *: ]]; then
                    return 1
                else
                    return 0
                fi
                ;;
            0)
                # fail, no match found
                return 255
                ;;
            *)
                # fail, ambiguous match
                printf "%s: option '%s' is ambiguous; possibilities:%s\n" "${0##*/}" \
                    "--$1" "$(printf " '%s'" "${longmatch[@]%:}")"
                return 254
                ;;
        esac
    }

    while (( $# )); do
        case "$1" in
            --) # explicit end of options
                shift
                break
                ;;
            -[!-]*) # short option
                for (( i = 1; i < ${#1}; i++ )); do
                    opt=${1:i:1}

                    # option doesn't exist
                    if [[ $shortopts != *$opt* ]]; then
                        printf "%s: invalid option -- '%s'\n" "${0##*/}" "$opt"
                        OPTRET=(--)
                        return 1
                    fi

                    OPTRET+=("-$opt")
                    # option requires optarg
                    if [[ "$shortopts" == *"${opt}:"* ]]; then
                        # if we're not at the end of the option chunk, the rest is the optarg
                        if (( i < ${#1} - 1 )); then
                            OPTRET+=("${1:i+1}")
                            break
                        # if we're at the end, grab the the next positional, if it exists
                        elif (( i == ${#1} - 1 )) && [[ -n "$2" ]]; then
                            OPTRET+=("$2")
                            shift
                            break
                        # parse failure
                        else
                            printf "%s: option '%s' requires an argument\n" "${0##*/}" "-$opt"
                            OPTRET=(--)
                            return 1
                        fi
                    fi
                done
                ;;
            --?*=* | --?*) # long option
                IFS='=' read -r opt optarg <<<"${1#--}"
                longoptmatch "$opt"
                case $? in
                    0)
                        if [[ -n "$optarg" ]]; then
                            printf "%s: option '--%s' doesn't allow an argument\n" "${0##*/}" "$opt"
                            OPTRET=(--)
                            return 1
                        else
                            OPTRET+=("--$opt")
                        fi
                        ;;
                    1)
                        # --longopt=optarg
                        if [[ -n "$optarg" ]]; then
                            OPTRET+=("--$opt" "$optarg")
                        # --longopt optarg
                        elif [[ -n "$2" ]]; then
                            OPTRET+=("--$opt" "$2")
                            shift
                        else
                            printf "%s: option '--%s' requires an argument\n" "${0##*/}" "$opt"
                            OPTRET=(--)
                            return 1
                        fi
                        ;;
                    254)
                        # ambiguous option -- error was reported for us by longoptmatch()
                        OPTRET=(--)
                        return 1
                        ;;
                    255)
                        # parse failure
                        printf "%s: unrecognized option '%s'\n" "${0##*/}" "--$opt"
                        OPTRET=(--)
                        return 1
                        ;;
                esac
                ;;
            *) # non-option arg encountered, add it as a parameter
                unused_argv+=("$1")
                ;;
        esac
        shift
    done

    # add end-of-opt terminator and any leftover positional parameters
    OPTRET+=('--' "${unused_argv[@]}" "$@")
    unset longoptmatch

    return 0
}

kver_x86() {
    local kver
    local -i offset
    # On x86 (since kernel 1.3.73, 1996), regardless of whether it's
    # an Image, a zImage, or a bzImage: The file header is the same,
    # and contains the kernel_version string.
    #
    # scrape the version out of the kernel image. locate the offset
    # to the version string by reading 2 bytes out of image at at
    # address 0x20E. this leads us to a string of, at most, 128 bytes.
    # read the first word from this string as the kernel version.
    #
    # https://www.kernel.org/doc/html/v6.7/arch/x86/boot.html
    # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/arch/x86/boot/header.S?h=v6.7
    offset="$(od -An -j0x20E -dN2 "$1")" || return

    read -r kver _ < \
        <(dd if="$1" bs=1 count=127 skip=$((offset + 0x200)) 2>/dev/null)

    printf '%s' "$kver"
}

detect_compression() {
    # Detect standard compressed files. Not Linux-kernel specific.

    local file="$1" offset="${2:-0}" bytes

    # The first 8 bytes are enough to detect most formats.
    bytes="$(od -An -t x1 -j "$offset" -N 8 "$file" | tr -d ' ')"
    case "$bytes" in
        'fd377a585a00'*)
            printf 'xz'
            return
            ;;
        '894c5a4f'*)
            printf 'lzop'
            return
            ;;
        '1f8b'*)
            printf 'gzip'
            return
            ;;
        '04224d18'*)
            error 'Newer lz4 stream format detected! This may not boot!'
            printf 'lz4'
            return
            ;;
        '02214c18'*)
            printf 'lz4 -l'
            return
            ;;
        '28b52ffd'*)
            printf 'zstd'
            return
            ;;
        '425a68'*) # 'BZh' in ASCII
            printf 'bzip2'
            return
            ;;
        '5d0000'*)
            # lzma detection sucks and there's really no good way to
            # do it without reading large portions of the stream. this
            # check is good enough for GNU tar, apparently, so it's good
            # enough for me.
            printf 'lzma'
            return
            ;;
        ????????'7a696d67'*) # 4 discarded bytes, then 'zimg' in ASCII
            # Linux kernel "zboot" self-decompressing EFI images
            # (since kernel 6.1)
            # (as of kernel 6.7: only on arm64, loongarch, and riscv)

            # Read 32 bytes (before 0x38) from address 0x18, which is a
            # null-terminated string representing the compressed type.
            #
            # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/drivers/firmware/efi/libstub/zboot-header.S?h=v6.7
            read -rd '' bytes < <(od -An -j0x18 -t a -N32 "$1" | sed 's/ nul//g' | tr -dc '[:alnum:]')
            printf 'zboot %s' "$bytes"
            return
            ;;
    esac

    # Try some formats that require us to sniff bytes at other locations.

    bytes="$(od -An -t x1 -j "$((offset+0x24))" -N4 "$file" | tr -d ' ')"
    if [[ "$bytes" == '18286f01' || "$bytes" == '016f2818' ]]; then # endian-sensitive
        # Linux kernel ARM32 self-decompressing "zImage"
        #
        # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/arch/arm/boot/compressed/vmlinux.lds.S?h=v6.7#n121
        # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/arch/arm/boot/compressed/head.S?h=v6.7#n215
        printf 'ARM zImage'
        return
    fi

    # out of ideas, assuming uncompressed
}

decompress_cat() {
    local file="$1" offset="${2:-0}" size="${3:-}"
    local comp_type reader

    comp_type="$(detect_compression "$file" "$offset")"
    case "$comp_type" in
        '') reader='cat' ;;
        'xz') reader='xzcat' ;;
        'lzop') reader='lzop -d' ;;
        'gzip') reader='zcat' ;;
        'lz4') reader='lz4cat' ;;
        'lz4 -l') reader='lz4cat -l' ;;
        'zstd') reader='zstdcat' ;;
        'bzip2') reader='bzcat' ;;
        'lzma') reader='xzcat' ;;
        'zboot '*)
            _zboot_cat "${comp_type#'zboot '}" "$file" "$offset"
            return
            ;;
        'ARM zImage')
            _arm_zimage_cat "$file" "$offset"
            return
            ;;
        *)
            error 'Unknown compression type: %s' "$comp_type"
            return 1
            ;;
    esac
    if (( offset == 0 )) && [[ -z "$size" ]]; then
        $reader - <"$file"
    elif [[ -z "$size" ]]; then
        tail --bytes=+"$((offset+1))" "$file" | $reader -
    else
        tail --bytes=+"$((offset+1))" "$file" | head --bytes="$size" | $reader -
    fi
}

_zboot_cat() {
    # Linux kernel "zboot" self-decompressing EFI images
    # (since kernel 6.1)
    # (as of kernel 6.7: only on arm64, loongarch, and riscv)

    # See zboot-header.S for offsets. See Makefile.zboot for
    # compression types and size adjustments.
    #
    # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/drivers/firmware/efi/libstub/Makefile.zboot?h=v6.7
    # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/drivers/firmware/efi/libstub/zboot-header.S?h=v6.7

    local comp_type="$1" file="$2" offset="${3:-0}"
    local reader start size size_len

    # Read a pair of u32s from 0x08 (right after the "zimg" magic
    # identifier) for the starting-offset and size of the compressed
    # data.
    start="$(od -An -j "$((offset+0x08))" -t u4 -N4 "$file" | tr -dc '[:alnum:]')"
    size="$(od -An -j "$((offset+0x0c))" -t u4 -N4 "$file" | tr -dc '[:alnum:]')"

    [[ "$start" =~ ^[0-9]+$ ]] || return 1
    [[ "$size" =~ ^[0-9]+$ ]] || return 1

    size_len=4 # corresponds to Makefile.zboot:zboot-size-len-y or zboot-header.S:ZBOOT_SIZE_LEN
    case "$comp_type" in
        'gzip') reader='zcat'
                size_len=0
                ;;
        'lz4') reader='lz4cat' ;;
        'lzma') reader='xzcat' ;;
        'lzo') reader='lzop -d' ;;
        'xzkern') reader='xzcat' ;;
        'zstd22') reader='zstdcat' ;;
        *)
            error 'Unknown zboot compression type: %s' "${comp_type#'zboot '}"
            return 1
            ;;
    esac

    tail --bytes=+"$((offset+start+1))" <"$file" | head --bytes="$((size+size_len))" | $reader -
}

_arm_zimage_cat() {
    # Linux kernel ARM32 self-decompressing "zImage"

    # The end of the file looks like:
    #
    #    input_data: char []<compressed data>
    #    piggy_size: u32  len(<uncompressed_data>)
    #    padding:    char []<zeros>               ; <=4096 bytes
    #    got:        u32  []<global_offsets>
    #    trailer:    char []<whatever>            ; <512 bytes
    #
    # We can find the location of "piggy_size" from the file header,
    # and then we can ignore some zeros to find the GOT. It is
    # reasonable to assume (but still a heuristic) that the largest
    # value in the GOT that is <0x20000 (128KiB) is the location of
    # input_data.
    #
    # Note that piggy_size is *not* at a 4-byte-aligned location.
    #
    # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/arch/arm/boot/compressed/vmlinux.lds.S?h=v6.7
    #
    # About that 0x20000 heuristic: input_data is put after the
    # extraction code, and so its location is basically determined by
    # the size of the extraction code, which will vary based on a
    # number of factors. So, this heuristic relies on a few things:
    #
    #  1. That ${max_extraction_code_size} is less than ~128KiB
    #  2. That ${min_compressed_kernel_size}+${min_extraction_code_size}
    #     is more than ~128KiB.
    #
    # How safe are those assumptions?
    #
    #  1. Using GCC 12.2.0 with kernels 4.15 - 6.8 in a variety of
    #     configurations, I've observed input_data being placed at
    #     addresses in the range 0x152e - 0x10f63 (~5.3KiB - ~68KiB).
    #     This is well within our 128KiB threshold.
    #  2. Let's just say that at the limit min_extraction_code_size =~
    #     0; I think it is safe to assume that even compressed,
    #     kernels are measured in MiB, not KiB. (Though, when writing
    #     tests for this code, it means we must be careful to use
    #     files that are large enough!)

    local file="$1" offset="${2:-0}"

    # Read magic numbers.
    local hdr_endian_indicator cpu_endian_indicator have_table_indicator hdr_endian cpu_endian
    hdr_endian_indicator="$(od -An -t x1 -j "$((offset+0x24))" -N4 "$file" | tr -d ' ')"
    cpu_endian_indicator="$(od -An -t x1 -j "$((offset+0x30))" -N4 "$file" | tr -d ' ')"
    have_table_indicator="$(od -An -t x1 -j "$((offset+0x34))" -N4 "$file" | tr -d ' ')"
    case "$hdr_endian_indicator" in #
        '18286f01') hdr_endian='little' ;;
        '016f2818') hdr_endian='big' ;;
        *)
            error '_arm_zimage_cat called, but file does not look like an ARM zImage'
            return 1
            ;;
    esac
    case "$cpu_endian_indicator" in
        # Prior to v3.17, hdr_endian and cpu_endian were the same, and
        # there was no cpu_endian_indicator.
        '01020304') cpu_endian='little' ;;
        '04030201') cpu_endian='big' ;;
        *)          cpu_endian="$hdr_endian" ;;
    esac
    if [[ "$have_table_indicator" != '45454545' ]]; then
        error 'ARM zImage is too old (before v4.15, 2017)'
        return 1
    fi

    # Read offsets.
    local zboot_rom_text magic_table_addr piggy_size_addr aligned_after_piggy
    zboot_rom_text="$(od --endian="$hdr_endian" -An -t u4 -j "$((offset+0x28))" -N4 "$file" | tr -d ' ')"
    magic_table_addr="$(od --endian="$hdr_endian" -An -t u4 -j "$((offset+0x38))" -N4 "$file" | tr -d ' ')"
    piggy_size_addr="$(od --endian="$hdr_endian" -An -t u4 -j $((offset+magic_table_addr+8)) -N4 "$file" | tr -d ' ')"
    aligned_after_piggy=$((((piggy_size_addr+4+3)/4)*4))

    # Read the padding and GOT.
    # These values have zboot_rom_text added to them.
    local -i input_data_addr=0 in_got=0 val size
    while read -r val; do
        if (( val == 0 )); then
            if (( in_got )); then
                break
            else
                continue
            fi
        fi
        in_got=1
        val="$((val - zboot_rom_text))"
        if (( val > input_data_addr && val < 0x20000 )); then
            input_data_addr=$val
        fi
    done < <(od --endian="$cpu_endian" -An -v -t u4 -j "$((offset+aligned_after_piggy))" "$file" | xargs printf '%s\n')
    if (( input_data_addr == 0 )); then
        error 'Could not find input_data_addr in ARM zImage'
        return 1
    fi
    size="$((piggy_size_addr-input_data_addr))"

    # gzip contains a 4-byte uncompressed-size suffix, so for
    # CONFIG_KERNEL_GZIP it saves 4 bytes by having the input_data
    # extend into the piggy_size.
    #
    # That is: the Makefile says
    #    compress-$(CONFIG_KERNEL_GZIP) = gzip
    # instead of
    #    compress-$(CONFIG_KERNEL_GZIP) = gzip_with_size
    #
    # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/arch/arm/boot/compressed/Makefile?h=v6.7#n79
    if [[ "$(detect_compression "$file" "$((offset+input_data_addr))")" == 'gzip' ]]; then
        size=$((size+4))
    fi

    decompress_cat "$file" "$((offset+input_data_addr))" "$size"
}

kver_generic() {
    # For unknown architectures, we can try to grep the uncompressed or gzipped
    # image for the boot banner.
    # This should work at least for ARM when run on /boot/Image, or RISC-V on
    # gzipped /boot/vmlinuz-linuz. On other architectures it may be worth trying
    # rather than bailing, and inform the user if none was found.

    local kver=''

    # Loosely grep for `linux_banner`:
    # https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/init/version-timestamp.c?h=v6.7#n28
    read -r _ _ kver _ < <(decompress_cat "$1" | grep -m1 -aoE 'Linux version .(\.[-_[:alnum:]+]+)+')

    printf '%s' "$kver"
}

kver() {
    # this is intentionally very loose. only ensure that we're
    # dealing with some sort of string that starts with something
    # resembling dotted decimal notation. remember that there's no
    # requirement for CONFIG_LOCALVERSION to be set.
    local kver re='^[[:digit:]]+(\.[[:digit:]]+)+'

    local arch
    arch="$(uname -m)"
    if [[ $arch == @(i?86|x86_64) ]]; then
        kver="$(kver_x86 "$1")"
    else
        kver="$(kver_generic "$1")"
    fi

    [[ "$kver" =~ $re ]] || return 1

    printf '%s' "$kver"
}

plain() {
    local mesg="$1"; shift
    # shellcheck disable=SC2059
    printf "    $_color_bold$mesg$_color_none\n" "$@" >&1
}

quiet() {
    # _optquiet is assigned in mkinitcpio
    # shellcheck disable=SC2154
    (( _optquiet )) || plain "$@"
}

msg() {
    local mesg="$1"; shift
    # shellcheck disable=SC2059
    printf "$_color_green==>$_color_none $_color_bold$mesg$_color_none\n" "$@" >&1
}

msg2() {
    local mesg="$1"; shift
    # shellcheck disable=SC2059
    printf "  $_color_blue->$_color_none $_color_bold$mesg$_color_none\n" "$@" >&1
}

warning() {
    local mesg="$1"; shift
    # shellcheck disable=SC2059
    printf "$_color_yellow==> WARNING:$_color_none $_color_bold$mesg$_color_none\n" "$@" >&2
}

error() {
    local mesg="$1"; shift
    # shellcheck disable=SC2059
    printf "$_color_red==> ERROR:$_color_none $_color_bold$mesg$_color_none\n" "$@" >&2
    return 1
}

die() {
    error "$@"
    cleanup 1
}

map() {
    local r=0
    for _ in "${@:2}"; do
        # shellcheck disable=SC1105,SC2210,SC2035
        "$1" "$_" || (( $# > 255 ? r=1 : ++r ))
    done
    return "$r"
}

arrayize_config() {
    set -f
    [[ ${MODULES@a} != *a* ]] && IFS=' ' read -r -a MODULES <<<"$MODULES"
    [[ ${BINARIES@a} != *a* ]] && IFS=' ' read -r -a BINARIES <<<"$BINARIES"
    [[ ${FILES@a} != *a* ]] && IFS=' ' read -r -a FILES <<<"$FILES"
    [[ ${HOOKS@a} != *a* ]] && IFS=' ' read -r -a HOOKS <<<"$HOOKS"
    [[ ${COMPRESSION_OPTIONS@a} != *a* ]] && IFS=' ' read -r -a COMPRESSION_OPTIONS <<<"$COMPRESSION_OPTIONS"
    set +f
}

in_array() {
    # Search for an element in an array.
    #   $1: needle
    #   ${@:2}: haystack

    local item='' needle="$1"; shift

    for item in "$@"; do
        [[ "$item" == "$needle" ]] && return 0 # Found
    done
    return 1 # Not Found
}

index_of() {
    # get the array index of an item. sets the global var _idx with
    # index and returns 0 if found, otherwise returns 1.
    local item="$1"; shift

    for (( _idx=1; _idx <= $#; _idx++ )); do
        if [[ "$item" == "${!_idx}" ]]; then
            (( --_idx ))
            return 0
        fi
    done

    # not found
    unset _idx
    return 1
}

funcgrep() {
    awk -v funcmatch="$1" '
        /^[[:space:]]*[[:alnum:]_]+[[:space:]]*\([[:space:]]*\)/ {
            match($1, funcmatch)
            print substr($1, RSTART, RLENGTH)
        }' "$2"
}

list_hookpoints() {
    local funcs script

    # _d_hooks is assigned in mkinitcpio
    # shellcheck disable=SC2154
    script="$(PATH="$_d_hooks" type -P "$1")" || return 0

    mapfile -t funcs < <(funcgrep '^run_[[:alnum:]_]+' "$script")

    echo
    msg "This hook has runtime scripts:"
    in_array run_earlyhook "${funcs[@]}" && msg2 "early hook"
    in_array run_hook "${funcs[@]}" && msg2 "pre-mount hook"
    in_array run_latehook "${funcs[@]}" && msg2 "post-mount hook"
    in_array run_cleanuphook "${funcs[@]}" && msg2 "cleanup hook"
    in_array run_emergencyhook "${funcs[@]}" && msg2 "emergency hook"
}

modprobe() {
    # _optmoduleroot is assigned in mkinitcpio
    # shellcheck disable=SC2154
    command modprobe -d "$_optmoduleroot" -S "$KERNELVERSION" "$@"
}

all_modules() {
    # Add modules to the initcpio, filtered by grep.
    #   $@: filter arguments to grep
    #   -f FILTER: ERE to filter found modules

    local -i count=0
    local mod='' OPTIND='' OPTARG='' modfilter=()

    while getopts ':f:' flag; do
        [[ "$flag" = "f" ]] && modfilter+=("$OPTARG")
    done
    shift $(( OPTIND - 1 ))

    # _d_kmoduledir is assigned in mkinitcpio
    # shellcheck disable=SC2154
    while read -r -d '' mod; do
        (( ++count ))

        for f in "${modfilter[@]}"; do
            [[ "$mod" =~ $f ]] && continue 2
        done

        mod="${mod##*/}"
        mod="${mod%.ko*}"
        printf '%s\n' "${mod//-/_}"
    done < <(LC_ALL=C.UTF-8 find "$_d_kmoduledir" -name '*.ko*' -print0 2>/dev/null | grep -EZz "$@")

    (( count ))
}

add_all_modules() {
    # Add modules to the initcpio.
    #   $@: arguments to all_modules

    local mod
    local -a mods

    mapfile -t mods < <(all_modules "$@")
    map add_module "${mods[@]}"

    return $(( !${#mods[*]} ))
}

add_checked_modules() {
    # Add modules to the initcpio, filtered by the list of autodetected
    # modules.
    #   $@: arguments to all_modules

    local mod
    local -a mods

    # _autodetect_cache is declared in mkinitcpio and assigned in install/autodetect
    # shellcheck disable=SC2154
    if (( ${#_autodetect_cache[*]} )); then
        mapfile -t mods < <(all_modules "$@" | grep -xFf <(printf '%s\n' "${!_autodetect_cache[@]}"))
    else
        mapfile -t mods < <(all_modules "$@")
    fi

    map add_module "${mods[@]}"

    return $(( !${#mods[*]} ))
}

add_firmware() {
    # add a firmware file to the image.
    #   $1: firmware path fragment

    local fw fwpath
    local -a fwfile
    local -i r=1

    for fw; do
        # _d_fwpath is assigned in mkinitcpio
        # shellcheck disable=SC2154
        for fwpath in "${_d_fwpath[@]}"; do
            # shellcheck disable=SC2153
            if ! compgen -G "${BUILDROOT}${fwpath}/${fw}?(.*)" &>/dev/null; then
                # modinfo has firmware entries with globs, read entries into an array
                if read -r fwfile < <(compgen -G "${fwpath}/${fw}?(.*)"); then
                    map add_file "${fwfile[@]}" && r=0
                    break
                fi
            else
                r=0
                break
            fi
        done
    done

    return "$r"
}

add_module() {
    # Add a kernel module to the initcpio image. Dependencies will be
    # discovered and added.
    #   $1: module name

    local target='' module='' softdeps=() deps=() field='' value='' firmware=()
    local ign_errors=0 found=0

    [[ "$KERNELVERSION" == 'none' ]] && return 0

    if [[ "$1" == *\? ]]; then
        ign_errors=1
        set -- "${1%?}"
    fi

    target="${1%.ko*}" target="${target//-/_}"

    # skip expensive stuff if this module has already been added
    (( _addedmodules["$target"] == 1 )) && return

    while IFS=':= ' read -r -d '' field value; do
        case "$field" in
            filename)
                # Only add modules with filenames that look like paths (e.g.
                # it might be reported as "(builtin)"). We'll defer actually
                # checking whether or not the file exists -- any errors can be
                # handled during module install time.
                if [[ "$value" == /* ]]; then
                    found=1
                    module="${value##*/}" module="${module%.ko*}"
                    quiet "adding module: %s (%s)" "$module" "$value"
                    _modpaths["$value"]=1
                    _addedmodules["${module//-/_}"]=1
                fi
                ;;
            depends)
                IFS=',' read -r -a deps <<<"$value"
                map add_module "${deps[@]}"
                ;;
            firmware)
                firmware+=("$value")
                ;;
            softdep)
                read -ra softdeps <<<"$value"
                for module in "${softdeps[@]}"; do
                    [[ $module == *: ]] && continue
                    add_module "$module?"
                done
                ;;
        esac
    done < <(modinfo -b "$_optmoduleroot" -k "$KERNELVERSION" -0 "$target" 2>/dev/null)

    if (( !found )); then
        (( ign_errors || _addedmodules["$target"] )) && return 0
        error "module not found: '%s'" "$target"
        return 1
    fi

    if (( ${#firmware[*]} )); then
        add_firmware "${firmware[@]}" ||
            warning "Possibly missing firmware for module: '%s'" "$target"
    fi

    # handle module quirks
    case "$target" in
        fat)
            add_module "nls_ascii?" # from CONFIG_FAT_DEFAULT_IOCHARSET
            add_module "nls_cp437?" # from CONFIG_FAT_DEFAULT_CODEPAGE
            ;;
        ocfs2)
            add_module "configfs?"
            ;;
        btrfs)
            add_module "libcrc32c?"
            ;;
        f2fs)
            add_module "crypto-crc32?"
            ;;
        ext4)
            add_module "crypto-crc32c?"
            ;;
    esac
}

add_full_dir() {
    # Add a directory and all its contents, recursively, to the initcpio image.
    # No parsing is performed and the contents of the directory is added as is.
    #   $1: path to directory
    #   $2: glob pattern to filter file additions (optional)
    #   $3: path prefix that will be stripped off from the image path (optional)

    local f='' filter="${2:-*}" strip_prefix="$3"

    if [[ -n "$1" && -d "$1" ]]; then
        add_dir "$1"

        for f in "$1"/*; do
            if [[ -L "$f" ]]; then
                # Explicit glob matching
                # shellcheck disable=SC2053
                if [[ "$f" == $filter ]]; then
                    add_symlink "${f#"${strip_prefix}"}" "$(readlink "$f")"
                fi
            elif [[ -d "$f" ]]; then
                add_full_dir "$f" "$filter" "$strip_prefix"
            elif [[ -f "$f" ]]; then
                # Explicit glob matching
                # shellcheck disable=SC2053
                if [[ "$f" == $filter ]]; then
                    add_file "$f" "${f#"${strip_prefix}"}"
                fi
            fi
        done
    fi
}

add_dir() {
    # add a directory (with parents) to $BUILDROOT
    #   $1: pathname on initcpio
    #   $2: mode (optional)

    if [[ -z "$1" || "$1" != /?* ]]; then
        return 1
    fi

    local path="$1" mode="${2:-755}"

    # shellcheck disable=SC2153
    if [[ -d "${BUILDROOT}${1}" ]]; then
        # ignore dir already exists
        return 0
    fi

    quiet "adding dir: %s" "$path"
    command install -dm"${mode}" "${BUILDROOT}${path}"
}

add_dir_early() {
    # add a directory (with parents) to $EARLYROOT
    #   $1: pathname on initcpio
    #   $2: mode (optional)

    # $EARLYROOT is assigned by mkinitcpio
    # shellcheck disable=SC2153
    BUILDROOT="$EARLYROOT" add_dir "$@" || return
}

add_symlink() {
    # Add a symlink to the initcpio image. There is no checking done
    # to ensure that the target of the symlink exists.
    #   $1: pathname of symlink on image
    #   $2: absolute path to target of symlink (optional, can be read from $1)

    local name="$1" target="${2:-$1}" linkobject

    (( $# == 1 || $# == 2 )) || return 1

    # find out the link target
    if [[ "$name" == "$target" ]]; then
        linkobject="$(LC_ALL=C.UTF-8 find "$target" -prune -printf '%l')"
        # use relative path if the target is a file in the same directory as the link
        # anything more would lead to the insanity of parsing each element in its path
        if [[ "$linkobject" != *'/'* && ! -L "${name%/*}/${linkobject}" ]]; then
            target="$linkobject"
        else
            target="$(realpath -eq -- "$target")"
        fi
    elif [[ -L "$target" ]]; then
        target="$(realpath -eq -- "$target")"
    fi
    if [[ -z "$target" ]]; then
        error "invalid symlink: '%s'" "$name"
        return 1
    fi

    add_dir "${name%/*}"

    if [[ -L "${BUILDROOT}${1}" ]]; then
        quiet "overwriting symlink %s -> %s" "$name" "$target"
    else
        quiet "adding symlink: %s -> %s" "$name" "$target"
    fi
    ln -sfn "$target" "${BUILDROOT}${name}"
}

add_file() {
    # Add a plain file to the initcpio image. No parsing is performed and only
    # the singular file is added.
    #   $1: path to file
    #   $2: destination on initcpio (optional, defaults to same as source)
    #   $3: mode

    # determine source and destination
    local src="$1" dest="${2:-$1}" mode="$3" srcrealpath destrealpath

    # Treat '-' as /dev/stdin
    if [[ "$src" == '-' ]]; then
        src='/dev/stdin'
    fi

    if [[ -c "$src" || -p "$src" ]]; then
        if [[ "$src" == "$dest" || "$dest" == *'/' ]]; then
            error "no destination file specified for: '%s'" "$src"
            return 1
        fi
        if [[ -z "$mode" ]]; then
            error "no file mode specified for: '%s'" "$dest"
            return 1
        fi
    elif [[ ! -f "$src" ]]; then
        error "file not found: '%s'" "$src"
        return 1
    fi

    # Handle cases where the parent of the destination is a symlink.
    # What if there is a symlink higher up in the file path, you ask?
    # We simply hope that never happens...
    if [[ -z "$mode" && "$src" == "$dest" && -L "${dest%/*}" ]]; then
        destrealpath="$(realpath -- "${dest%/*}")"
        if [[ ! -e "${BUILDROOT}${dest%/*}" ]]; then
            add_dir "${destrealpath}"
            add_symlink "${dest%/*}"
        fi

        # Rewrite destination to reflect real path of target
        dest="${destrealpath}/${dest##*/}"
    fi

    # check if $src is a symlink
    if [[ -L "$src" && ! -c "$src" && ! -p "$src" ]]; then
        srcrealpath="$(realpath -- "$src")
        if [[ "$srcrealpath" != "$dest" ]]; then
            # add the target file
            add_file "$srcrealpath" "$srcrealpath" "$mode"
            # create the symlink
            add_symlink "$dest" "$src"
            return
        fi
    fi

    # cp does not create directories leading to the destination and install
    # does not create the last directory if the destination is a directory.
    [[ ! -e "${BUILDROOT}${dest%/*}" && ( -z "$mode" || "$dest" == *'/' ) ]] && add_dir "${dest%/*}"

    # if destination ends with a slash, then use the source file name
    if [[ -e "${BUILDROOT}${dest/%\//\/${src##*/}}" ]]; then
        quiet 'overwriting file: %s' "${dest/%\//\/${src##*/}}"
    else
        quiet 'adding file: %s' "${dest/%\//\/${src##*/}}"
    fi
    if [[ -z "$mode" ]]; then
        command cp --remove-destination --preserve=mode,ownership "$src" "${BUILDROOT}${dest}"
    else
        command install -Dm"$mode" "$src" "${BUILDROOT}${dest}"
    fi
}

add_file_early() {
    # Add a plain file to $EARLYROOT. No parsing is performed and only
    # the singular file is added.
    #   $1: path to file
    #   $2: destination on initcpio (optional, defaults to same as source)
    #   $3: mode

    # $EARLYROOT is assigned by mkinitcpio
    # shellcheck disable=SC2153
    BUILDROOT="$EARLYROOT" add_file "$@" || return
}

add_runscript() {
    # Adds a runtime script to the initcpio image. The name is derived from the
    # script which calls it as the basename of the caller.

    local fn script hookname="${BASH_SOURCE[1]##*/}"
    local -a funcs

    if ! script="$(PATH="$_d_hooks" type -P "$hookname")"; then
        error "runtime script for '%s' not found" "$hookname"
        return
    fi

    if [[ -L "$script" ]]; then
        script="$(realpath -- "$script")"
    fi
    add_binary "$script" "/hooks/$hookname" 755

    mapfile -t funcs < <(funcgrep '^run_[[:alnum:]_]+' "$script")

    for fn in "${funcs[@]}"; do
        case $fn in
            run_earlyhook)
                _runhooks['early']+=" $hookname"
                ;;
            run_hook)
                _runhooks['hooks']+=" $hookname"
                ;;
            run_latehook)
                _runhooks['late']+=" $hookname"
                ;;
            run_cleanuphook)
                _runhooks['cleanup']="$hookname ${_runhooks['cleanup']}"
                ;;
            run_emergencyhook)
                _runhooks['emergency']="$hookname ${_runhooks['emergency']}"
                ;;
        esac
    done
}

add_binary() {
    # Add a binary file to the initcpio image. library dependencies will
    # be discovered and added.
    #   $1: path to binary
    #   $2: destination on initcpio (optional, defaults to same as source)
    #   $3: mode (optional)

    local line='' regex='' binary='' dest='' mode='' sodep='' resolved='' shebang='' interpreter='' args=()

    if [[ "${1:0:1}" != '/' ]]; then
        if ! binary="$(type -P "$1")"; then
            error "binary not found: '%s'" "$1"
            return 1
        fi
    else
        binary="$1"
    fi

    dest="${2:-$binary}"
    mode="$3"

    args=("$binary" "$dest")

    if [[ -n "$mode" ]]; then
        args+=("$mode")
    fi

    add_file "${args[@]}" || return 1

    # non-binaries
    if ! lddout="$(ldd "$binary" 2>/dev/null)"; then
        # detect if the file has a shebang
        if IFS='' LC_ALL=C.UTF-8 read -rn2 -d '' shebang <"$binary" && [[ "$shebang" == '#!' ]]; then
            read -r shebang <"$binary"
            interpreter="${shebang##\#\!*([[:space:]])}"
            # strip /usr/bin/env and warn if it is missing
            if [[ "$interpreter" == '/usr/bin/env'* ]]; then
                [[ -e "${BUILDROOT}/usr/bin/env" ]] || warning "Possibly missing '/usr/bin/env' for script: %s" "$binary"
                interpreter="${interpreter##'/usr/bin/env'+([[:space:]])}"
            fi
            # strip parameters
            interpreter="${interpreter%%[[:space:]]*}"
            # check if the interpreter exists in BUILDROOT
            if [[ "$interpreter" != '/'* ]] && PATH="${BUILDROOT}/usr/local/sbin:${BUILDROOT}/usr/local/bin:${BUILDROOT}/usr/bin" type -P "$interpreter" &>/dev/null; then
                :
            elif [[ -e "${BUILDROOT}/${interpreter}" ]]; then
                :
            else
                warning "Possibly missing '%s' for script: %s" "$interpreter" "$binary"
            fi
        fi
        return 0
    fi

    # resolve sodeps
    regex='^(|.+ )(/.+) \(0x[a-fA-F0-9]+\)'
    while read -r line; do
        if [[ "$line" =~ $regex ]]; then
            sodep="${BASH_REMATCH[2]}"
        elif [[ "$line" = *'not found' ]]; then
            error "binary dependency '%s' not found for '%s'" "${line%% *}" "$1"
            (( ++_builderrors ))
            continue
        fi

        if [[ -f "$sodep" && ! -e "${BUILDROOT}${sodep}" ]]; then
            add_file "$sodep" "$sodep"
        fi
    done <<<"$lddout"

    return 0
}

add_udev_rule() {
    # Add an udev rules file to the initcpio image. Dependencies on binaries
    # will be discovered and added.
    #   $1: path to rules file (or name of rules file)

    local rules="$1" rule=() key='' value='' binary=''

    if [[ "${rules:0:1}" != '/' ]]; then
        rules="$(PATH='/usr/lib/udev/rules.d:/lib/udev/rules.d' type -P "$rules")"
    fi
    if [[ -z "$rules" ]]; then
        # complain about not found rules
        return 1
    fi

    add_file "$rules" /usr/lib/udev/rules.d/"${rules##*/}"

    while IFS=, read -ra rule; do
        # skip empty lines, comments
        # rule is an array, but we are only checking if it's an empty string
        # shellcheck disable=SC2128
        [[ -z "$rule" || "$rule" == @(+([[:space:]])|#*) ]] && continue

        for pair in "${rule[@]}"; do
            IFS=' =' read -r key value <<<"$pair"
            case "$key" in
                'RUN{program}' | 'RUN+' | 'IMPORT{program}' | 'ENV{REMOVE_CMD}')
                    # strip quotes
                    binary="${value//[\"\']/}"
                    # just take the first word as the binary name
                    binary="${binary%% *}"
                    [[ "${binary:0:1}" == '$' ]] && continue
                    if [[ "${binary:0:1}" != '/' ]]; then
                        binary="$(PATH='/usr/lib/udev:/lib/udev' type -P "$binary")"
                    fi
                    add_binary "$binary"
                    ;;
            esac
        done
    done <"$rules"
}

parse_config() {
    # parse key global variables set by the config file.

    map add_module "${MODULES[@]}"
    map add_binary "${BINARIES[@]}"
    map add_file "${FILES[@]}"

    tee "$BUILDROOT/buildconfig" <"$1" | {
        # When MODULES is not an array (but instead implicitly converted at
        # startup), sourcing the config causes the string value of MODULES
        # to be assigned as MODULES[0]. Avoid this by explicitly unsetting
        # MODULES before re-sourcing the config.
        unset MODULES

        # shellcheck disable=SC1091
        . /dev/stdin

        # arrayize MODULES if necessary.
        [[ ${MODULES@a} != *a* ]] && read -ra MODULES <<<"${MODULES//-/_}"

        for mod in "${MODULES[@]%\?}"; do
            mod="${mod//-/_}"
            # only add real modules (2 == builtin)
            (( _addedmodules["$mod"] == 1 )) && add+=("$mod")
        done
        (( ${#add[*]} )) && printf 'MODULES="%s"\n' "${add[*]}"

        printf '%s="%s"\n' \
            'EARLYHOOKS' "${_runhooks['early']# }" \
            'HOOKS' "${_runhooks['hooks']# }" \
            'LATEHOOKS' "${_runhooks['late']# }" \
            'CLEANUPHOOKS' "${_runhooks['cleanup']% }" \
            'EMERGENCYHOOKS' "${_runhooks['emergency']% }"
    } >"$BUILDROOT/config"
}

initialize_buildroot() {
    # creates a temporary directory for the buildroot and initialize it with a
    # basic set of necessary directories and symlinks

    local kernver="$1" generatedir="$2" workdir arch buildroot osreleasefile root
    arch="$(uname -m)"

    if ! workdir="$(mktemp -d --tmpdir mkinitcpio.XXXXXX 2>/dev/null)"; then
        error 'Failed to create temporary working directory in %s' "${TMPDIR:-/tmp}"
        return 1
    fi
    earlyroot="$workdir/early"
    buildroot="${generatedir:-$workdir/root}"

    if [[ ! -w "${generatedir:-$workdir}" ]]; then
        error 'Unable to write to build root: %s' "$buildroot"
        return 1
    fi

    # early root structure
    install -dm755 "$earlyroot"
    # this flag file is used by some tools
    echo 1 > "${earlyroot}/early_cpio"

    # base directory structure
    for root in "$buildroot" "$earlyroot"; do
        install -dm755 "$root"/{new_root,proc,sys,dev,run,tmp,var,etc,usr/{local{,/bin,/sbin,/lib},lib,bin}}
        ln -s "usr/lib" "$root/lib"
        ln -s "bin"     "$root/usr/sbin"
        ln -s "usr/bin" "$root/bin"
        ln -s "usr/bin" "$root/sbin"
        ln -s "../run"  "$root/var/run"

        case "$arch" in
            x86_64)
                ln -s "lib"     "$root/usr/lib64"
                ln -s "usr/lib" "$root/lib64"
                ;;
        esac
    done

    # mkinitcpio version stamp
    # shellcheck disable=SC2154
    printf '%s' "$version" >"$buildroot/VERSION"

    # kernel module dir
    [[ "$kernver" != 'none' ]] && install -dm755 "$buildroot/usr/lib/modules/$kernver/kernel"

    # mount tables
    ln -s ../proc/self/mounts "$buildroot/etc/mtab"
    : >"$buildroot/etc/fstab"

    # add os-release and initrd-release for systemd
    if [[ -e /etc/os-release ]]; then
        if [[ -L /etc/os-release ]]; then
            osreleasefile="$(realpath -- /etc/os-release)"
            install -Dm0644 "$osreleasefile" "${buildroot}${osreleasefile}"
            cp -adT /etc/os-release "${buildroot}/etc/os-release"
            cp -adT /etc/os-release "${buildroot}/etc/initrd-release"
        else
            install -Dm0644 /etc/os-release "${buildroot}/etc/os-release"
            ln -sT os-release "${buildroot}/etc/initrd-release"
        fi
    else
        : >"$buildroot/etc/initrd-release"
    fi

    # add a blank ld.so.conf to keep ldconfig happy
    : >"$buildroot/etc/ld.so.conf"

    printf '%s' "$workdir"
}

run_build_hook() {
    local hook="$1" script='' resolved=''
    # shellcheck disable=SC2034
    local MODULES=() BINARIES=() FILES=() SCRIPT=''

    # find script in install dirs
    # _d_install is assigned in mkinitcpio
    # shellcheck disable=SC2154
    if ! script="$(PATH="$_d_install" type -P "$hook")"; then
        error "Hook '$hook' cannot be found"
        return 1
    fi

    # check for deprecation
    if resolved="$(readlink -e "$script")" && [[ "${script##*/}" != "${resolved##*/}" ]]; then
        warning "Hook '%s' is deprecated. Replace it with '%s' in your config" \
            "${script##*/}" "${resolved##*/}"
        script="$resolved"
    fi

    # source
    unset -f build
    # shellcheck disable=SC1090
    if ! . "$script"; then
        error 'Failed to read %s' "$script"
        return 1
    fi

    if ! declare -f build >/dev/null; then
        error "Hook '%s' has no build function" "${script}"
        return 1
    fi

    # run
    if (( _optquiet )); then
        msg2 "Running build hook: [%s]" "${script##*/}"
    else
        msg2 "Running build hook: [%s]" "$script"
    fi
    build

    # if we made it this far, return successfully. Hooks can
    # do their own error catching if it's severe enough, and
    # we already capture errors from the add_* functions.
    return 0
}

try_enable_color() {
    local colors

    if ! colors="$(tput colors 2>/dev/null)"; then
        warning "Failed to enable color. Check your TERM environment variable"
        return
    fi

    if (( colors > 0 )) && tput setaf 0 &>/dev/null; then
        _color_none="$(tput sgr0)"
        _color_bold="$(tput bold)"
        _color_blue="$_color_bold$(tput setaf 4)"
        _color_green="$_color_bold$(tput setaf 2)"
        _color_red="$_color_bold$(tput setaf 1)"
        _color_yellow="$_color_bold$(tput setaf 3)"
    fi
}

install_modules() {
    local m
    local -a xz_comp gz_comp zst_comp

    [[ "$KERNELVERSION" == 'none' ]] && return 0

    if (( $# == 0 )); then
        warning "No modules were added to the image. This is probably not what you want."
        return 0
    fi

    for m in "$@"; do
        add_file "$m"
        # unzip modules prior to recompression
        if [[ "$MODULES_DECOMPRESS" == 'yes' ]]; then
            case "$m" in
                *.xz)
                    xz_comp+=("$BUILDROOT/$m")
                    ;;
                *.gz)
                    gz_comp+=("$BUILDROOT/$m")
                    ;;
                *.zst)
                    zst_comp+=("$BUILDROOT/$m")
                    ;;
            esac
        fi
    done

    (( ${#xz_comp[*]} )) && xz -d "${xz_comp[@]}"
    (( ${#gz_comp[*]} )) && gzip -d "${gz_comp[@]}"
    (( ${#zst_comp[*]} )) && zstd -d --rm -q "${zst_comp[@]}"

    msg "Generating module dependencies"
    map add_file "$_d_kmoduledir"/modules.{builtin,builtin.modinfo,order}
    depmod -b "$BUILDROOT" "$KERNELVERSION"

    # remove all non-binary module.* files (except devname for on-demand module loading)
    rm "${BUILDROOT}${_d_kmoduledir}"/modules.!(*.bin|devname|softdep)
}

find_module_from_symbol() {
    # Find a module based off on the symbol
    #   $1: symbol to find
    #   $2: the directory to look at
    #
    # The directory can either be a:
    #   absolute directory with a leading /
    #   A subdirectory with a = prefix, like =drivers/hid

    local moduledir symbols="$1" directories=("${@:2}")
    for dir in "${directories[@]}"; do
        case "${dir::1}" in
            = )
                moduledir="/lib/modules/$KERNELVERSION/kernel/${dir:1}"
                ;;
            /*)
                moduledir="$dir"
                ;;
        esac


        module_to_stdout() {
            case "$1" in
                *.xz) xz -d "$1" -c ;;
                *.gz) gzip -d "$1" -c ;;
                *.zst) zstd -q -d "$1" -c ;;
                *) cat "$1" ;;
            esac
        }

        while read -r -d '' mod; do
            if module_to_stdout "$mod" | grep -Eq "^($symbols)" &> /dev/null; then
                mod=${mod##*/}
                mod="${mod%.ko*}"
                printf '%s\n' "${mod//-/_}"
            fi
        done < <(LC_ALL=C.UTF-8 find "$moduledir" -name '*.ko*' -print0 2>/dev/null)
    done
}

add_all_modules_from_symbol() {
    local symbol="$1"
    local -a mods paths=("${@:2}")

    mapfile -t mods < <(find_module_from_symbol "$symbol" "${paths[@]}")
    if (( ! ${#mods[@]} )); then
        paths=("${paths[@]#=}") paths=("${paths[@]/#/\'}") paths=("${paths[@]/%/\',}")
        local path_string="${paths[*]}"
        warning "No module containing the symbol '%s' found in: %s" "$symbol" "${path_string%,}"
        return 1
    fi

    # add_checked_modules_from_symbol
    if [[ "${FUNCNAME[1]}" == 'add_checked_modules_from_symbol' ]]; then
        # _autodetect_cache is declared in mkinitcpio and assigned in install/autodetect
        # shellcheck disable=SC2154
        if (( ${#_autodetect_cache[@]} )); then
            mapfile -t mods < <(printf '%s\n' "${mods[@]}" | grep -xFf <(printf '%s\n' "${!_autodetect_cache[@]}"))
            # Do not fail if no autodetected module has the symbol
            if (( ! ${#mods[@]} )); then
                return 0
            fi
        fi
    fi

    map add_module "${mods[@]}"
}


add_checked_modules_from_symbol() {
    if ! add_all_modules_from_symbol "$@"; then
        return 1
    fi
}

# Add TPM2 PCR signature to the UKI
# $1: path to private key
# $2: path to public key
# $3: PCR selection (comma-separated list)
add_tpm2_signature() {
    local privkey="$1" pubkey="$2" pcrs="$3"
    local pcrsig_json
    
    # Generate PCR signature using systemd-measure
    if ! pcrsig_json="$(systemd-measure sign --private-key="$privkey" --public-key="$pubkey" --pcr="$pcrs")"; then
        error "Failed to generate TPM2 PCR signature"
        return 1
    fi
    
    # Store signature for later use
    printf '%s' "$pcrsig_json" > "$BUILDROOT/pcrsig.json"
    add_file "$pubkey" "/pcrpkey"
}

# Add TPM2 PCR policy to LUKS2 header
# $1: device
# $2: path to public key
# $3: PCR selection
# $4: path to signature
add_tpm2_luks_policy() {
    local device="$1" pubkey="$2" pcrs="$3" sig="$4"
    
    if ! systemd-cryptenroll --tpm2-public-key="$pubkey" \
                            --tpm2-public-key-pcrs="$pcrs" \
                            --tpm2-signature="$sig" "$device"; then
        error "Failed to add TPM2 PCR policy to LUKS2 header"
        return 1
    fi
}

if [[ "$0" == *'/functions' && "$1" == 'run_mkinitcpio_func' ]]; then
    shift
    if declare -F "$1" >/dev/null; then
        "$@"
    else
        exit 2
    fi
fi

# vim: set ft=sh ts=4 sw=4 et:
