# Installation script for Root-on-ZFS on a disk for Raspberry Pi 4B
#
# For educational purposes only.
# Author is not liable for any damages caused by incorrect handling of this tool.
#
# Author: Libor Martinek

set -eu

# Absolute path to this script.
THIS="$(realpath $0)"

# --- BEGIN ENV VARS ---
#
# Name of the disk that will have the FreeBSD on ZFS installed.
RPIZFS_DISK="${RPIZFS_DISK:-da0}"
# FreeBSD image to install.
RPIZFS_IMG="${RPIZFS_IMG:-FreeBSD-15.1-RELEASE-arm64-aarch64-RPI.img}"
# Where to mount the ZFS pool on the newly installed FreeBSD.
RPIZFS_POOLMOUNTPOINT="${RPIZFS_POOLMOUNTPOINT:-/zroot}"
# Name of the ZFS pool on the newly installed FreeBSD.
RPIZFS_POOLNAME="${RPIZFS_POOLNAME:-zroot}"
# Where to download the decompressed FreeBSD image from. May be left as dummy value when
# the 'RPIZFS_IMG' image is already present in the 'RPIZFS_TMPDIR'.
RPIZFS_REMOTE_URL="${RPIZFS_REMOTE_URL:-"http://192.0.2.1:8080"}"
# Where to download to or look for the FreeBSD image. Please note that '/tmp' will work only on
# Raspberries that have more than ~5GB RAM, which is the size of the downloaded image! That is
# because by default, '/tmp' directory is mounted as a temporary filesystem (tmpfs), so
# contents are held entirely in RAM.
RPIZFS_TMPDIR="${RPIZFS_TMPDIR:-/tmp}"
#
# --- END ENV VARS ---

# Private or derived variables:
rpizfs_altroot="/rpizfs-altroot";                readonly rpizfs_altroot
rpizfs_diskpath="/dev/${RPIZFS_DISK}";           readonly rpizfs_diskpath
rpizfs_imgpath="${RPIZFS_TMPDIR}/${RPIZFS_IMG}"; readonly rpizfs_imgpath
rpizfs_mntdir="/mnt";                            readonly rpizfs_mntdir
rpizfs_partition="${RPIZFS_DISK}s2";             readonly rpizfs_partition

help() {
    # Print help, public function documentation and environment variables used
    # to configure the script.
    echo "$(basename ${THIS}) - Install Root-on-ZFS on a disk for Raspberry Pi 4B"
    echo
    echo "Following functions can be executed:"
    echo
    print_function_docs
    echo "---"
    echo
    echo "Following environment variables can be used to configure the script behaviour:"
    echo
    print_config_vars_docs
}

function_docs_awk=$(cat <<'EOF'
BEGIN {
    in_function_body = 0;
}

# Public function header encountered, so print its name and start accepting comments.
/^[a-zA-Z][a-zA-Z0-9_]+\(\) \{/ {
    in_function_body = 1;
    sub(/\(\) \{/, "", $0);
    print $0;
}

# Non-comment line in function body encountered, so reset state.
/^[ ]+[a-zA-Z0-9_:]+/ {
    if (in_function_body == 1) {
        in_function_body = 0;
        printf "\n";
    }
}

# Comment line in function body encountered, so strip the leading '#' and print it.
/^[ ]*\# / {
    if (in_function_body == 1) {
        sub(/^[ ]*\# /, "    ", $0);
        print $0;
    }
}

# End of function body encountered, so reset state.
/^\}$/ {
    if (in_function_body == 1) {
        in_function_body = 0;
        printf "\n";
    }
}
EOF
)
readonly function_docs_awk

print_function_docs() {
    # Print all public functions and their documentation.
    awk "${function_docs_awk}" "${THIS}"
}

config_vars_awk=$(cat <<'EOF'
BEGIN {
    comment_acc = "";
    in_env_section = 0;
}

# Comment encountered. If we are in the env section, append the line to
# the 'comment_acc' without the leading '#'.
/^# [a-zA-Z0-9\-_]/ {
    if (in_env_section == 1) {
        sub(/^# /, "    ", $0);
        # Yeah, following is string concatenation in awk...
        comment_acc = comment_acc $0 "\n";
    }
}

# Special beginning marker encountered, accumulation can start.
$0 == "# --- BEGIN ENV VARS ---" {
    in_env_section = 1;
}

# Special ending marker encountered, no point in continuing further.
$0 == "# --- END ENV VARS ---" {
    exit 0;
}

# Variable assignment encountered. Print its name and all preceding comment
# lines accumulated. Reset the accumulator.
/^RPIZFS_[a-zA-Z0-9-_]+\=/ {
    sub(/\=.*$/, "", $0);
    print $0;
    print comment_acc;
    comment_acc = "";
}
EOF
)
readonly config_vars_awk

print_config_vars_docs() {
    # Print all 'RPIZFS_' variables and their documentation that may be configured
    # by the user.
    awk "${config_vars_awk}" "${THIS}"
}

print_config_values() {
    # Print the values of all 'RPIZFS_' variables that be configured by the user.
    # They are in shell format - "name=value" without spaces - so can be dumped
    # to a file and then sourced from it again.
    for varname in $(grep -E -o -e '^RPIZFS_\w+' "${THIS}"); do
        echo -n "$varname="
        eval echo "\$${varname}"
    done
}

check_system() {
    # Check necessary preconditions of the installation, like
    # existence of the disk, permissions, etc.
    echo -n "Checking existence of '${rpizfs_imgpath}' ... "
    if test -f "${rpizfs_imgpath}"; then
        echo "OK"
    else
        echo "will need to be downloaded."
    fi

    echo -n "Checking existence of '${RPIZFS_DISK}' disk ... "
    if test -e "${rpizfs_diskpath}"; then
        echo "OK"
    else
        echo "does not exist! Exiting."
        exit 1
    fi

    echo -n "Checking ability to write to '${rpizfs_mntdir}' directory ... "
    if test -w "${rpizfs_mntdir}"; then
        echo "OK"
    else
        echo "not writable! Exiting."
        exit 1
    fi
}

_fetch_image() {
    # Fetch the installation image. No-op if the image is already present.
    if test -f "${rpizfs_imgpath}"; then
        echo "Image already downloaded. Skipping"
    else
        local remote_image_uri="${RPIZFS_REMOTE_URL}/${RPIZFS_IMG}"
        fetch "${remote_image_uri}" -o "${rpizfs_imgpath}"
    fi
}

_dd_to_disk() {
    # Burn the FreeBSD install image to the disk.
    dd \
        if="${rpizfs_imgpath}" \
        of="${rpizfs_diskpath}" \
        bs=4M \
        iflag=direct \
        oflag=direct \
        status=progress
}

_resize_partitions() {
    # Resize and modified partitions created by the '_dd_to_disk' function.
    gpart resize -i 2 "${RPIZFS_DISK}"
    gpart modify -i 1 "${rpizfs_partition}"
    gpart resize \
        -i 1 \
        -a 1M \
        "${rpizfs_partition}"
}

_create_rootzfs() {
    # Create a ZFS pool and the filesystem hierarchy within it.
    :
    # https://openzfs.github.io/openzfs-docs/man/master/8/zpool-create.8.html
    zpool create \
        -f \
        -O compress=lz4 \
        -O atime=off \
        -o ashift=12 \
        -o altroot="${rpizfs_altroot}" \
        -m none \
        "${RPIZFS_POOLNAME}" \
        "${rpizfs_partition}"

    # https://openzfs.github.io/openzfs-docs/man/master/8/zfs-create.8.html
    zfs create -o mountpoint=none "${RPIZFS_POOLNAME}/ROOT"
    zfs create -o mountpoint=/ "${RPIZFS_POOLNAME}/ROOT/default"
    zfs create -o mountpoint=/tmp -o exec=on -o setuid=off "${RPIZFS_POOLNAME}/tmp"
    zfs create -o mountpoint=/usr -o canmount=off "${RPIZFS_POOLNAME}/usr"
    zfs create -o setuid=off "${RPIZFS_POOLNAME}/usr/ports"
    zfs create "${RPIZFS_POOLNAME}/usr/src"
    zfs create -o mountpoint=/var -o canmount=off "${RPIZFS_POOLNAME}/var"
    zfs create -o exec=off -o setuid=off "${RPIZFS_POOLNAME}/var/audit"
    zfs create -o exec=off -o setuid=off "${RPIZFS_POOLNAME}/var/crash"
    zfs create -o exec=off -o setuid=off "${RPIZFS_POOLNAME}/var/log"
    zfs create -o atime=on "${RPIZFS_POOLNAME}/var/mail"
    zfs create -o setuid=off "${RPIZFS_POOLNAME}/var/tmp"

    zfs set mountpoint="${RPIZFS_POOLMOUNTPOINT}" "${RPIZFS_POOLNAME}"
    zpool set bootfs="${RPIZFS_POOLNAME}/ROOT/default" "${RPIZFS_POOLNAME}"
    zfs set canmount=noauto "${RPIZFS_POOLNAME}/ROOT/default"
}

_copy_image_to_rootzfs() {
    # Create a virtual disk from the install image and copy it to the filesystem
    # created by the '_create_rootzfs' function.
    local md_disk_name="md123";                        readonly md_disk_name
    local md_partition_disk="/dev/${md_disk_name}s2a"; readonly md_partition_disk
    local mnt_dir="${rpizfs_mntdir}/rpi-rootzfs";      readonly mnt_dir

    # Create a "virtual disk" from the install image.
    # https://man.freebsd.org/cgi/man.cgi?mdconfig
    mdconfig -a -t vnode -f "${rpizfs_imgpath}" -u "${md_disk_name}"

    mkdir -p "${mnt_dir}"
    mount "${md_partition_disk}" "${mnt_dir}"

    # The critical step - the mounted install image is archived in its entirety,
    # and then unarchived in the ZFS file system hierarchy created by the
    # '_create_rootzfs' function.
    ( cd "${mnt_dir}" && tar cf - . ) | ( cd "${rpizfs_altroot}" && tar xf - )

    # Clean-up: Unmount the virtual disk, destroy it and remove the mount
    #   directory.
    umount "${mnt_dir}"
    mdconfig -d -u "${md_disk_name}"
    rm -rf "${mnt_dir}"
}

_adjust_boot_conf() {
    # Adjust configuration so the newly installed Raspberry Pi boots correctly
    # on ZFS.
    :
    # https://man.freebsd.org/cgi/man.cgi?loader.conf
    echo 'zfs_load="YES"' >> "${rpizfs_altroot}/boot/loader.conf"
    # Following allows time synchronization running as a non-root user.
    # https://man.freebsd.org/cgi/man.cgi?query=mac_ntpd&sektion=4&format=html
    echo 'mac_ntpd_load="YES"' >> "${rpizfs_altroot}/boot/loader.conf"
    # I enable Wireguard kernel module by default.
    # https://man.freebsd.org/cgi/man.cgi?query=if_wg&sektion=4&format=html
    echo 'if_wg_load="YES"' >> "${rpizfs_altroot}/boot/loader.conf"
    # Comment out mounting UFS file system.
    sed -i '' '/ufs/ s/^/#/' "${rpizfs_altroot}/etc/fstab"
    # Comment out mounting tmpfs on '/tmp' directory. There is a separate ZFS
    # dataset for it.
    sed -i '' '/tmpfs/ s/^/#/' "${rpizfs_altroot}/etc/fstab"
}

_create_rc_conf() {
    # Create a basic '/etc/rc.conf' for the newly installed Raspberry Pi.
    :
    # https://man.freebsd.org/cgi/man.cgi?rc.conf
    echo 'hostname="rpi-rootzfs"' > "${rpizfs_altroot}/etc/rc.conf"
    echo 'sshd_enable="YES"' >> "${rpizfs_altroot}/etc/rc.conf"
    echo 'sendmail_enable="NONE"' >> "${rpizfs_altroot}/etc/rc.conf"
    echo 'ifconfig_DEFAULT="DHCP"' >> "${rpizfs_altroot}/etc/rc.conf"
    # https://docs.freebsd.org/en/books/handbook/config/#configtuning-syslog
    # That's the default, but have it here anyway.
    echo 'syslogd_flags="-s"' >> "${rpizfs_altroot}/etc/rc.conf"

    # Set dumpdev to "AUTO" to enable crash dumps, "NO" to disable
    echo 'dumpdev="NO"' >> "${rpizfs_altroot}/etc/rc.conf"
    echo 'powerd_enable="YES"' >> "${rpizfs_altroot}/etc/rc.conf"
    echo 'zfs_enable="YES"' >> "${rpizfs_altroot}/etc/rc.conf"

    # If the 'mac_ntpd_load' is loaded, this enables running the 'ntpd' daemon
    # as a non-root user and synchronizes time on startup.
    echo 'ntpd_enable="YES"' >> "${rpizfs_altroot}/etc/rc.conf"
    echo 'ntpd_sync_on_start="YES"' >> "${rpizfs_altroot}/etc/rc.conf"
    echo 'ntpd_user="ntpd"' >> "${rpizfs_altroot}/etc/rc.conf"
}

_export_zpool() {
    # "Unmount" the ZFS pool created.
    :
    # https://openzfs.github.io/openzfs-docs/man/master/8/zpool-export.8.html
    zpool export "${RPIZFS_POOLNAME}"
}


install() {
    # Check system and then execute the whole installation procedure.
    echo "Installation procedure initialized."
    echo "Variables used:"
    print_config_values

    echo "Checking system:"
    check_system

    # Fetch the uncompressed FreeBSD installation image ...
    _fetch_image
    # ... burn it to the external disk ...
    _dd_to_disk
    # ... tweak the disk partitions ...
    _resize_partitions
    # ... create a filesystem hierarchy in a new ZFS pool on the external disk ...
    _create_rootzfs
    # ... copy the FreeBSD installation image to the filesystem in the pool ...
    _copy_image_to_rootzfs
    # ... adjust '/boot/loader.conf' and `/etc/fstab`, so RPi boots correctly from ZFS ...
    _adjust_boot_conf
    # ... create '/etc/rc.conf' on RPi ...
    _create_rc_conf
    # ... and finally, unmount the ZFS pool.
    _export_zpool

    echo
    echo "Installation successful."
    echo "You can shutdown the RPi ('shutdown now'), remove the microSD card and start it again."
}

# "$0" Dispatch pattern
# https://www.oilshell.org/blog/2021/07/blog-backlog-1.html#shell-programming-patterns
"$@"