#!/usr/bin/env bash
# GitLab Runner custom executor: each job in a fresh exe.dev sandbox.
# From https://exe.dev/deploy/gitlab-runner.
#
#   config   pick the job VM's name
#   prepare  cp BASE_VM to that name as a sandbox, wait for it to answer
#   run      run one of the job's scripts on it, over SSH
#   cleanup  delete it
#   sweep    delete job VMs older than MAX_HOURS, in case a cleanup never ran
#
# Settings come from /etc/gitlab-runner/exe.env. A job can ask for a bigger
# VM with the CI/CD variable EXE_CPU (2, 4, 8, ... up to MAX_CPU), and gets
# 2 GiB of memory per vCPU.
set -uo pipefail
# shellcheck source=/dev/null
. /etc/gitlab-runner/exe.env
: "${BASE_VM:?set BASE_VM in /etc/gitlab-runner/exe.env}"
CP_FLAGS=${CP_FLAGS:---cpu=2 --memory=4GB}
JOB_PREFIX=${JOB_PREFIX:-$(hostname)-job-}
SSH_KEY=${SSH_KEY:-/etc/gitlab-runner/exe/id_ed25519}
KNOWN_HOSTS=${KNOWN_HOSTS:-/etc/gitlab-runner/exe/known_hosts}
MAX_HOURS=${MAX_HOURS:-6}
MAX_CPU=${MAX_CPU:-8}

ssh_() {
  ssh -i "$SSH_KEY" -o IdentitiesOnly=yes -o UserKnownHostsFile="$KNOWN_HOSTS" \
      -o StrictHostKeyChecking=yes -o BatchMode=yes -o ConnectTimeout=15 \
      -o ServerAliveInterval=30 "$@"
}
exe() { ssh_ exe.dev "$@"; }

case "${1:-}" in
config)
  # A name exe.dev doesn't have yet; prepare, run and cleanup get it as $EXE_VM.
  vm="${JOB_PREFIX}${CUSTOM_ENV_CI_JOB_ID}-$(openssl rand -hex 2)"
  printf '{"builds_dir":"/home/exedev/builds","cache_dir":"/home/exedev/cache",'
  printf '"builds_dir_is_shared":false,"hostname":"%s",' "$vm"
  printf '"driver":{"name":"exe.dev sandbox","version":"1"},"job_env":{"EXE_VM":"%s"}}\n' "$vm"
  ;;
prepare)
  if [ -n "${CUSTOM_ENV_CI_JOB_IMAGE:-}" ]; then
    echo "Note: image: $CUSTOM_ENV_CI_JOB_IMAGE is ignored; the job runs on the VM itself."
  fi
  flags=$CP_FLAGS
  if [ -n "${CUSTOM_ENV_EXE_CPU:-}" ]; then
    case "$CUSTOM_ENV_EXE_CPU" in
    2 | 4 | 8 | 16 | 32) cpu=$CUSTOM_ENV_EXE_CPU ;;
    *) cpu=0 ;;
    esac
    if [ "$cpu" -eq 0 ] || [ "$cpu" -gt "$MAX_CPU" ]; then
      echo "EXE_CPU must be 2, 4, 8, ... up to MAX_CPU=$MAX_CPU"
      exit "$BUILD_FAILURE_EXIT_CODE"
    fi
    flags="--cpu=$cpu --memory=$((cpu * 2))GB"
  fi
  exe rm "$EXE_VM" >/dev/null 2>&1   # left over from a failed attempt
  t0=$SECONDS
  # shellcheck disable=SC2086
  if ! out=$(exe cp "$BASE_VM" "$EXE_VM" --sandbox $flags 2>&1); then
    echo "cp $BASE_VM $EXE_VM failed: $out"
    exit "$SYSTEM_FAILURE_EXIT_CODE"
  fi
  for _ in $(seq 30); do
    if ssh_ "$EXE_VM.exe.xyz" true 2>/dev/null; then
      echo "Running on $EXE_VM, a copy of $BASE_VM with $flags ($((SECONDS - t0))s)"
      exit 0
    fi
    sleep 1
  done
  echo "$EXE_VM did not answer over SSH"
  exit "$SYSTEM_FAILURE_EXIT_CODE"
  ;;
run)
  # $2 is a script the runner wrote, $3 the stage (get_sources, step_script...).
  # Copy it over and run it with stdin closed, so job commands can't eat it.
  # shellcheck disable=SC2016
  ssh_ "$EXE_VM.exe.xyz" 'f=$(mktemp); cat >"$f"; bash "$f" </dev/null; rc=$?; rm -f "$f"; exit $rc' <"$2"
  rc=$?
  if [ "$rc" -eq 255 ]; then
    exit "$SYSTEM_FAILURE_EXIT_CODE"   # SSH itself failed
  elif [ "$rc" -ne 0 ]; then
    echo "$rc" >"$BUILD_EXIT_CODE_FILE"
    exit "$BUILD_FAILURE_EXIT_CODE"
  fi
  ;;
cleanup)
  exe rm "$EXE_VM" >/dev/null && echo "deleted $EXE_VM"
  ;;
sweep)
  if ! vms=$(exe ls --json 2>&1); then
    echo "ssh exe.dev ls failed: $vms" >&2
    exit 1
  fi
  printf '%s\n' "$vms" | jq -r --arg p "$JOB_PREFIX" --argjson h "$MAX_HOURS" '.vms[]
      | select(.vm_name | startswith($p))
      | select((.created_at | sub("\\.[0-9]+"; "") | fromdateiso8601) < now - $h * 3600)
      | .vm_name' |
    while read -r vm; do exe rm "$vm" >/dev/null && echo "swept $vm"; done
  ;;
*)
  echo "usage: $0 config|prepare|run|cleanup|sweep" >&2
  exit 2
  ;;
esac