#!/bin/sh
# Copyright 2006-2016 Canonical Ltd.
# Copyright 2022-2023 Simon McVittie

# wrapper.sh [OPTIONS] -- COMMAND [ARG...]
# See README.md for details

set -eu

debug () {
    :
}

log () {
    echo "$0: $*" >&2
}

errpath=
outpath=
script_pid_file=

cleanup () {
    if [ -n "$script_pid_file" ]; then
        rm -f "$script_pid_file"
    fi
}

trap cleanup EXIT INT QUIT PIPE

getopt_temp="debug"
getopt_temp="$getopt_temp,artifacts:"
getopt_temp="$getopt_temp,chdir:"
getopt_temp="$getopt_temp,env:"
getopt_temp="$getopt_temp,make-executable:"
getopt_temp="$getopt_temp,script-pid-file:"
getopt_temp="$getopt_temp,source-profile"
getopt_temp="$getopt_temp,stderr:"
getopt_temp="$getopt_temp,stdout:"
getopt_temp="$getopt_temp,tmp:"
getopt_temp="$getopt_temp,unset-env:"
getopt_temp="$(getopt -o '' --long "$getopt_temp" -n "$0" -- "$@")"
eval "set -- $getopt_temp"
unset getopt_temp

while [ "$#" -gt 0 ]; do
    case "$1" in
        (--artifacts)
            debug "creating AUTOPKGTEST_ARTIFACTS: $2"
            # shellcheck disable=SC2174
            mkdir -p -m 1777 -- "$2"
            export AUTOPKGTEST_ARTIFACTS="$2"
            export ADT_ARTIFACTS="$2"
            shift 2
            ;;

        (--chdir)
            debug "changing to directory: $2"
            cd "$2"
            shift 2
            ;;

        (--debug)
            debug () {
                log "$*"
            }
            shift
            ;;

        (--env)
            case "$2" in
                (?*=*)
                    debug "setting environment: $2"
                    export "${2?}"
                    ;;
                (*)
                    log "--env requires VAR=VALUE as argument" >&2
                    exit 255
                    ;;
            esac
            shift 2
            ;;

        (--make-executable)
            debug "marking as executable: $2"
            chmod +x -- "$2"
            shift 2
            ;;

        (--script-pid-file)
            debug "will create pid file: $2"
            script_pid_file="$2"
            shift 2
            ;;

        (--source-profile)
            debug "pretending to be a login shell"
            USER="$(id -nu)"
            export USER
            # /etc/profile, ~/.profile might fail to be sourced, and might
            # contain references to undefined environment variables
            set +eu
            if [ -e /etc/profile ]; then
                # shellcheck source=/dev/null
                . /etc/profile
            fi
            if [ -e ~/.profile ]; then
                # shellcheck source=/dev/null
                . ~/.profile
            fi
            set -eu
            shift
            ;;

        (--stderr)
            # Make sure not to write this in the short form, some of our
            # unit tests assert that debug output doesn't contain it
            debug "will write standard error to $2"
            errpath="$2"
            shift 2
            ;;

        (--stdout)
            debug "will write stdout to $2"
            outpath="$2"
            shift 2
            ;;

        (--tmp)
            debug "creating AUTOPKGTEST_TMP: $2"
            # shellcheck disable=SC2174
            mkdir -p -m 755 -- "$2"
            export AUTOPKGTEST_TMP="$2"
            export ADTTMP="$2"
            shift 2
            ;;

        (--unset-env)
            debug "unsetting environment: $2"
            unset "$2"
            shift 2
            ;;

        (--)
            shift
            break
            ;;

        (-*)
            log "Unknown option: $1"
            exit 255
            ;;

        (*)
            break
            ;;
    esac
done

if [ "$#" -lt 1 ]; then
    log "A command is required" >&2
    exit 255
fi

debug "command to run: $*"

tmp="$(mktemp -d)"
mkfifo "$tmp/err"
mkfifo "$tmp/out"

if [ -n "$outpath" ]; then
    debug "copying $tmp/out to stdout and file: $outpath"
    touch "$outpath"
    tee -a -- "$outpath" < "$tmp/out" &
else
    debug "copying $tmp/out to stdout only"
    cat < "$tmp/out" &
fi
out_pid="$!"

if [ -n "$errpath" ]; then
    debug "copying $tmp/err to standard error and file: $outpath"
    touch "$errpath"
    tee -a -- "$errpath" < "$tmp/err" >&2 &
else
    debug "copying $tmp/err to standard error only"
    cat < "$tmp/err" >&2 &
fi
err_pid="$!"

if [ -n "$script_pid_file" ]; then
    debug "writing script pid $$ to $script_pid_file"
    rm -f "$script_pid_file"
    set -C
    echo "$$" > "$script_pid_file"
    set +C
fi

exit_status=0
# We have to use exec in a subshell instead of running the test in the
# obvious way, to avoid this shell printing a message like "Terminated"
# or "Killed" to $tmp/err if it gets killed by a signal, which autopkgtest
# would interpret as failure when not using allow-stderr.
( set +x; exec >> "$tmp/out" 2>> "$tmp/err"; exec "$@" ) || exit_status="$?"

# The naive implementation here would be to iterate through /proc/[0-9]*/fd/*
# calling readlink on each one. However, that starts a readlink executable
# per fd (hooray shell scripting), which in practice is surprisingly slow,
# particularly on weak hardware (see 4115f7f5 "adt-virt-ssh: Speed up
# adt-ssh-wrapper"). So instead we use find(1) to amortize the process
# startup cost.

# If our temp directory contains \[]*? (unlikely), escape them
debug "checking for leaked background processes..."
tmp_escaped="$(
    echo "$tmp" | sed -E -e 's/\\/\\\\/g' -e 's/\[/\\[/g' -e 's/]/\\]/g' -e 's/([*?])/\\\1/g'
)"

kill="$(
    cd /proc
    find [0-9]*/fd \
        -lname "$tmp_escaped"/out -o -lname "$tmp_escaped"/err \
        2>/dev/null |
    sed -e 's#/fd/.*##' |
    sort -u |
    grep -v -F -e "$out_pid" -e "$err_pid" |
    tr '\n' ' '
)"

if [ -n "$kill" ]; then
    log "Killing leaked background processes: $kill"
    # Intentionally word-splitting
    # shellcheck disable=SC2086
    ps $kill >&2 || :
    # shellcheck disable=SC2086
    kill -CONT $kill >&2 || :
    # shellcheck disable=SC2086
    kill -9 $kill >&2 || :
fi

debug "waiting for tee/cat subprocesses..."
wait "$out_pid" "$err_pid" || :

debug "cleaning up..."
rm -fr "$tmp"
cleanup
debug "Exit status: $exit_status"
exit "$exit_status"
