#!/bin/sh

usage() {
    cat << EOF
usage: tmpoverlay [OPTIONS] [SOURCE...] DEST

Create a tmpfs-backed overlayfs at DEST starting with SOURCEs. If no SOURCE is
specified, use DEST as the source. To free the memory, simply umount DEST.

options:
  -c, --no-canonicalize  don't canonicalize paths
  -h, --help             print this help
  -o, --overlayfs OPTS   add overlayfs mount options, e.g. redirect_dir/metacopy
  -n, --no-mtab          don't write to /etc/mtab
  -N, --mount-name NAME  source name for mount (default "overlay")
  -t, --tmpfs OPTS       add tmpfs mount options, e.g. size
  -v, --verbose          verbose mode

examples:
  tmpoverlay / /new_root       # make a thin copy of root
  tmpoverlay /etc              # make read-only /etc writable
  tmpoverlay /a /b /c /merged  # merge /a, /b, /c, and a fresh tmpfs
  tmpoverlay /                 # USE WITH CAUTION, see docs
EOF
}

log() {
    fmt=$1
    shift
    printf 'tmpoverlay: ' >&2
    # shellcheck disable=SC2059
    printf "$fmt" "$@" >&2
    printf '\n' >&2
}

logv() {
    if [ -n "$verbose" ]; then
        log "$@"
    fi
}

cmdv() {
    logv '%s ' "$@"
    "$@"
}

unset own_tmpdir
die() {
    r=$?
    trap - INT TERM QUIT
    unset sig
    case $1 in
        -[A-Z]*) sig=${1#-};;
        '') ;;
        *) log "$@"
    esac
    logv 'an error occurred, cleaning up'
    if [ -d "$tmpdir" ] && [ -n "$own_tmpdir" ]; then
        if mountpoint -q "$tmpdir"; then
            logv 'unmounting tmpdir'
            # shellcheck disable=SC2086
            umount $no_mtab "$tmpdir"
        fi
        logv 'deleting tmpdir'
        rmdir "$tmpdir"
    fi
    if [ -n "$sig" ]; then
        logv 'reraising SIG%s' "$1"
        kill "-$1" "$$" && read -r _
    fi
    logv 'exiting'
    if [ $r = 0 ]; then
        exit 1
    fi
    exit $r
}

canon() {
    if [ -n "$no_canon" ]; then
        printf '%s\n' "$1"
    else
        realpath "$1"
    fi
}

my_getopt() {
    getopt \
        -l no-canonicalize \
        -l help \
        -l overlayfs: \
        -l no-mtab \
        -l mount-name: \
        -l tmpfs: \
        -l verbose \
        -n tmpoverlay \
        -- \
        cho:nN:t:v \
        "$@"
}

args=$(my_getopt "$@") || { usage >&2; exit 1; }
eval set -- "$args"
unset args

unset no_canon extra_ovl_opts no_mtab mount_name tmpfs_opts verbose
while true; do
    case "$1" in
        -c|--no-canonicalize) no_canon=-c; shift;;
        -h|--help) usage; exit 0;;
        -o|--overlayfs)
            [ -n "$2" ] && extra_ovl_opts="$extra_ovl_opts,$2"
            shift 2
            ;;
        -n|--no-mtab) no_mtab=-n; shift;;
        -N|--mount-name) mount_name=$2; shift 2;;
        -t|--tmpfs)
            [ -n "$2" ] && tmpfs_opts="${tmpfs_opts:+$tmpfs_opts,}$2"
            shift 2
            ;;
        -v|--verbose) verbose=-v; shift;;
        --) shift; break;;
        *) die "getopt failure"
    esac
done

[ $# != 0 ] || { log 'no paths specified'; usage >&2; exit 1; }

unset lowerdir
while [ "$#" != 1 ]; do
    d=$(canon "$1")
    if [ -h "$d" ] || ! [ -d "$d" ]; then
        die 'source "%s" is not a directory' "$d"
    fi
    lowerdir=${lowerdir+:}$d
    shift
done

dest=$(canon "$1")
[ -d "$dest" ] || die 'destination "%s" is not a directory' "$dest"
[ "$dest" != / ] || log 'overmounting root, use with caution'

if [ -z "$lowerdir" ]; then
    lowerdir=$dest
fi

trap 'die -INT' INT
trap 'die -TERM' TERM
trap 'die -QUIT' QUIT

if [ -z "$tmpdir" ]; then
    logv 'creating tmpdir'
    own_tmpdir=1
    tmpdir=$(umask 077; mktemp -dt tmpoverlay.XXXXXXXXXX) || exit 1
    logv 'created tmpdir: %s' "$tmpdir"
fi
mount -t tmpfs ${tmpfs_opts:+-o "$tmpfs_opts"} ${verbose:+-v} tmpfs "$tmpdir" || die

upperdir="$tmpdir/upper"
workdir="$tmpdir/work"
tmpmnt="$tmpdir/tmpmnt"
logv 'creating dirs'
mkdir "$upperdir" "$workdir" "$tmpmnt" || die

ovl_opts="lowerdir=$lowerdir,upperdir=$upperdir,workdir=$workdir"
logv 'testing overlay options'
mount -n -t overlay -o "$ovl_opts" overlay "$tmpmnt" || die "overlayfs is not supported"
umount "$tmpmnt"
if [ -n "$extra_ovl_opts" ]; then
    ovl_opts="$ovl_opts,$extra_ovl_opts"
    mount -n -t overlay -o "$ovl_opts" overlay "$tmpmnt" || die "invalid extra overlayfs options"
    umount "$tmpmnt"
fi
chk_ovl_opt() {
    if [ -n "$2" ]; then
        case ",$ovl_opts," in
            *",$1=$2,"*) return 0
        esac
        [ "$2" = on ] && val=Y || val=N
    else
        val=Y
    fi
    f=/sys/module/overlay/parameters/$1
    [ -e "$f" ] && [ "$(cat "$f")" = "$val" ]
}
try_ovl_opt() {
    # returns 0 iff option is/gets enabled
    case ",$ovl_opts," in
        *",$1=off,"*) return 1;;
        *",$1[=,]"*) return 0
    esac
    if chk_ovl_opt "$1"; then
        logv 'skipping %s'
        return
    fi
    logv 'trying %s' "$1${2+=$2}"
    new_ovl_opts="$ovl_opts,$1${2+=$2}"
    mount -n -t overlay -o "$new_ovl_opts" "overlay" "$tmpmnt" 2>/dev/null || return
    umount "$tmpmnt" || die
    # clear out workdir/work/incompat/volatile and upperdir index
    rm -r "$workdir" "$upperdir" || die
    mkdir "$upperdir" "$workdir" || die
    ovl_opts="$new_ovl_opts"
}
try_ovl_opt index on
# redirect_dir and metacopy are unsafe with untrusted non-bottom layers
# nfs_export conflicts with metacopy
[ "${lowerdir#*:}" = "$lowerdir" ] && \
    ! chk_ovl_opt userxattr on && \
    try_ovl_opt redirect_dir on && \
    { chk_ovl_opt nfs_export on || try_ovl_opt metacopy on; }
try_ovl_opt volatile

logv 'copying lowerdir owner/perms to upperdir'
lastlowerdir=${lowerdir##*:}
# stat -c isn't posix -.-
ls=$(ls -dn "$lastlowerdir/.") || die
tmp=${ls#* * }
owner=${tmp%% *}
tmp=${tmp#* }
group=${tmp%% *}
chown "$owner:$group" "$upperdir" || die
mode=${ls%% *}
[ "${#mode}" = 10 ] || die "bad ls permission format"
mode=${mode#?}
umode=${mode%??????}
ugmode=${mode%???}
gmode=${ugmode#???}
omode=${mode#??????}
chmod "u=$umode,g=$gmode,o=$omode" "$upperdir" || die
# -m - covers ACLs (system.posix_acl_access) and file caps
# (security.capability). theoretically someone might have get/setcap and/or
# get/setfacl but not get/setxattr, but this is unlikely since libcap/acl
# require attr.
logv 'copying root xattrs'
if attrs=$(cd "$lastlowerdir" && getfattr -d -m - . 2>/dev/null); then
    if [ -n "$attrs" ]; then
        printf '%s\n' "$attrs" | (cd "$upperdir"; setfattr --restore=-) || die
    fi
else
    log 'getfattr failed, skipping xattrs'
fi

logv 'mounting overlay'
mount_opts="$no_canon $no_mtab $verbose"
# shellcheck disable=SC2086
cmdv mount $mount_opts -t overlay -o "$ovl_opts" "${mount_name-overlay}" "$dest" || die

logv 'cleaning up'
umount "$tmpdir" || exit
rmdir "$tmpdir" || exit

logv 'done'