#!/usr/local/bin/bash
#
set -euo pipefail
IFS=$'\n\t'
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"

function main {
    if [ "$1" = "start" ]; then
        shift 1
        start_jail "${@}"
    elif [ "$1" = "stop" ]; then
        shift 1
        stop_jail "${@}"
    else
        >&2 echo "Unrecognized command"
        exit 1
    fi
}

function start_jail {
    host_interface_name="$1"
    bridge_name="bridge_${host_interface_name}"
    jail_interface_name="$2"
    ip_range="$3"

    assert_bridge "$host_interface_name" "$bridge_name" "$ip_range"

    bridge_link_name=$(detect_available_link "${bridge_name}")
    ngctl -d -f - <<EOF
mkpeer ${bridge_name}: eiface $bridge_link_name ether
name ${bridge_name}:$bridge_link_name $jail_interface_name
EOF
    ifconfig $(ngctl msg "${jail_interface_name}:" getifname | grep Args | cut -d '"' -f 2) name "${jail_interface_name}" up
}

function stop_jail {
    host_interface_name="$1"
    bridge_name="bridge_${host_interface_name}"
    jail_interface_name="$2"

    if ng_exists "${jail_interface_name}:"; then
        wait_for_interface_to_exist "${jail_interface_name}" 120
        ngctl shutdown "${jail_interface_name}:"
    fi

    if ng_exists "${bridge_name}:"; then
        num_remaining_hooks=$(ng_bridge_get_num_hooks "${bridge_name}:")
        if [ $num_remaining_hooks -eq 1 ]; then
            ngctl shutdown "${bridge_name}:"
            ngctl shutdown "${host_interface_name}:"
        fi
    fi
}

function assert_bridge {
    host_interface_name="$1"
    bridge_name="$2"
    ip_range="$3"

    if ! ng_exists "${bridge_name}:"; then
        ngctl -d -f - <<EOF
mkpeer . eiface hook ether
name .:hook $host_interface_name
EOF
        ngctl -d -f - <<EOF
mkpeer ${host_interface_name}: bridge ether link0
name ${host_interface_name}:ether $bridge_name
EOF
        ifconfig $(ngctl msg "${host_interface_name}:" getifname | grep Args | cut -d '"' -f 2) name "${host_interface_name}" "$ip_range" up
    fi
}

function ng_exists {
    ngctl status "${1}" >/dev/null 2>&1
}

function ng_bridge_get_num_hooks {
    ngctl show "${1}" | grep -oE 'Num hooks: [0-9]+' | sed 's/Num hooks: //g'
}

function detect_available_link {
    bridge_name="$1"
    linknum=1
    while true; do
        link_name="link${linknum}"
        if ! ng_exists "${bridge_name}:${link_name}"; then
            echo "$link_name"
            return
        fi
        (>&2 echo "$link_name failed on $bridge_name")
        linknum=$((linknum + 1))
        if [ "$linknum" -gt 90 ]; then
            (>&2 echo "No available links on bridge $bridge_name")
            exit 1
        fi
    done
}

function wait_for_interface_to_exist {
    # Wait for a vnet interface to exist again as a jail is shutting
    # down. If you delete the netgraph node before the interface
    # device exists, then the interface device will persist in a
    # broken state.
    ifname="$1"
    max_wait_seconds="$2"
    start=$(date +%s)
    while true; do
        now=$(date +%s)
        if [ $((now - start)) -gt $max_wait_seconds ]; then
            (>&2 echo "Waited for at least $max_wait_seconds seconds but the interface $ifname did not appear.")
            return 1;
        fi
        if ifconfig "$ifname" >/dev/null 2>&1; then
            return 0;
        fi
        sleep 2
    done
}

main "${@}"