#!/usr/local/bin/bash
#
# Mount non-boot-environment datasets.
#
# We can't leave datasets outside the boot environment (for example,
# jails or bhyve VMs) as canmount=on because then every boot
# environment's external datasets would all attempt to mount every
# time. To work around this, we mark those datasets as canmount=noauto
# and run this script to mount datasets under the root of our boot
# environment. This script depends heavily on my zfs dataset structure
# so it needs to be improved to be robust enough for different
# layouts. An example of my layout is:
#
## NAME                                                                                               MOUNTPOINT                  CANMOUNT  TA:BEMOUNT
## zroot                                                                                              none                        off       -
## zroot/global                                                                                       /global                     on        -
## zroot/freebsd                                                                                      none                        on        -
## zroot/freebsd/13.1-RELEASE                                                                         none                        on        -
## zroot/freebsd/13.1-RELEASE/be                                                                      none                        on        -
## zroot/freebsd/13.1-RELEASE/be/main                                                                 /                           noauto    -
## zroot/freebsd/13.1-RELEASE/jails                                                                   none                        on        -
## zroot/freebsd/13.1-RELEASE/jails/foo                                                               /jail/foo                   noauto    on
## zroot/freebsd/13.1-RELEASE/jails/bar                                                               /jail/bar                   noauto    on
## zroot/freebsd/13.1-RELEASE/jails/baz                                                               /jail/baz                   noauto    on
## zroot/freebsd/13.1-RELEASE/vm-bhyve                                                                /vm                         noauto    on
## zroot/linux                                                                                        none                        on        -
## zroot/linux/arch                                                                                   none                        on        -
## zroot/linux/arch/be                                                                                none                        on        -
## zroot/linux/arch/be/main                                                                           /                           noauto    -
set -euo pipefail
IFS=$'\n\t'
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

function main {
    local all_zfs_datasets=$(zfs list -Hp -o 'name,mountpoint,canmount,ta:bemount,mounted')
    local root_dataset=$(find_root_dataset "$all_zfs_datasets")
    local datasets_to_mount=$(find_datasets_to_mount_for_boot_environment "$all_zfs_datasets" "$root_dataset")
    if [ -n "$datasets_to_mount" ]; then
        mount_datasets "$datasets_to_mount"
    fi
}

function reverse_lines {
    sed '1!x;H;1h;$!d;g'
}

function find_dataset {
    local all_zfs_datasets="$1"
    local dataset_name="$2"
    while read dataset; do
        local ds_name=$(awk '{print $1}'<<<"$dataset")
        if [ "$ds_name" = "$dataset_name" ]; then
            echo "$dataset"
            return
        fi
    done<<<"$all_zfs_datasets"
}

function find_root_dataset {
    local all_zfs_datasets="$1"
    while read dataset; do
        local ds_name=$(awk '{print $1}'<<<"$dataset")
        local ds_mountpoint=$(awk '{print $2}'<<<"$dataset")
        # local ds_canmount=$(awk '{print $3}'<<<"$dataset")
        # local ds_bemount=$(awk '{print $4}'<<<"$dataset")
        local ds_mounted=$(awk '{print $5}'<<<"$dataset")
        if [ "$ds_mounted" = "yes" ] && [ "$ds_mountpoint" = "/" ]; then
            echo "$ds_name"
            return
        fi
    done<<<"$all_zfs_datasets"
}

function find_datasets_to_mount_for_boot_environment {
    local all_zfs_datasets="$1"
    local root_dataset="$2"
    # This is a consequence of my layout for zfs datasets. I should
    # make this more robust. Perhaps a zfs property like search up
    # from dataset mounted at / until you find a dataset with property
    # ta:bemountroot="on"?
    local be_root_name="${root_dataset%/*/*}"
    local be_root_dataset=$(find_dataset "$all_zfs_datasets" "$be_root_name")

    while read dataset; do
        local ds_name=$(awk '{print $1}'<<<"$dataset")
        # local ds_mountpoint=$(awk '{print $2}'<<<"$dataset")
        local ds_canmount=$(awk '{print $3}'<<<"$dataset")
        local ds_bemount=$(awk '{print $4}'<<<"$dataset")
        local ds_mounted=$(awk '{print $5}'<<<"$dataset")

        case "$ds_name" in
            "${be_root_name}/"*) ;;
            *)    continue ;;
        esac

        if [ "$ds_bemount" != "on" ]; then
            continue
        fi

        if [ "$ds_mounted" != "no" ]; then
            continue
        fi

        if [ "$ds_canmount" != "noauto" ]; then
            continue
        fi
        echo "$dataset"

    done<<<"$all_zfs_datasets"
}

function mount_datasets {
    local datasets_to_mount=$(reverse_lines<<<"$1")
    while read dataset; do
        local ds_name=$(awk '{print $1}'<<<"$dataset")
        local ds_mountpoint=$(awk '{print $2}'<<<"$dataset")
        local ds_canmount=$(awk '{print $3}'<<<"$dataset")
        local ds_bemount=$(awk '{print $4}'<<<"$dataset")
        local ds_mounted=$(awk '{print $5}'<<<"$dataset")

        mount -v -t zfs "$ds_name" "$ds_mountpoint"

    done<<<"$datasets_to_mount"
}

main