#!/bin/bash
#
# Proof test for patches/v1.11.0/0004-vault-atomic-update-user.patch.
#
#   sudo bash proof.sh /path/to/versitygw [/path/to/bao]
#
# Starts a throwaway OpenBao dev server (in memory, 127.0.0.1:${PROOF_BAO_PORT:-18200}) and two
# gateways from the given binary on one KV path, as several instances of one namespace run:
#
#   A  127.0.0.1:${PROOF_PORT:-7391}   reaches OpenBao through vaultproxy.py, which holds every request
#                                  ${PROOF_DELAY:-3} s, as from a distant OpenBao, and can refuse KV writes
#   B  next port                   reaches OpenBao directly
#
# Both use AppRole with the policy leil-cluster gives its gateways and the default 120 s IAM cache
# (PROOF_CACHE_TTL changes it, for a quick try only). Then it rotates an account through A's admin API
# while B signs requests with the old secret past its cache, updates the account from A and B at once,
# and has A's KV write refused. Everything lives under ${PROOF_DIR:-/var/tmp/s3gw-vault-update} and is
# removed at the end; nothing installed is touched. Prints one line per check and exits non-zero if
# any check fails. KEEP=1 leaves the directory (logs and samples) behind. About four minutes.
#
# The account, made by root through A's admin API, on leil-cluster's reserved test band:
#
#   vu-rot   uid 74992  gid 74993  user

set -uo pipefail

VGW=${1:?usage: proof.sh /path/to/versitygw [/path/to/bao]}
BAO=${2:-$(command -v bao)}
HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
W=${PROOF_DIR:-/var/tmp/s3gw-vault-update}
PA=${PROOF_PORT:-7391}
PB=$((PA + 1))
BAO_PORT=${PROOF_BAO_PORT:-18200}
PROXY_PORT=$((BAO_PORT + 2))   # the dev server's cluster listener takes BAO_PORT + 1
TTL=${PROOF_CACHE_TTL:-120}
D=${PROOF_DELAY:-3}
BAO_URL=http://127.0.0.1:$BAO_PORT
KV=s3gw-iam
ACCOUNTS=proof/0x0/accounts
ENTRY=$KV/data/$ACCOUNTS/vu-rot
export S3_REGION=us-east-1

[ -x "$VGW" ] || { echo "FATAL: $VGW is not an executable" >&2; exit 2; }
[ -n "$BAO" ] && [ -x "$BAO" ] || { echo "FATAL: no bao binary (give its path as the second argument)" >&2; exit 2; }
for c in python3 curl; do
    command -v $c >/dev/null || { echo "FATAL: $c is required" >&2; exit 2; }
done
listening() { (exec 3<>"/dev/tcp/127.0.0.1/$1") 2>/dev/null; }
for p in $PA $PB $BAO_PORT $((BAO_PORT + 1)) $PROXY_PORT; do
    if listening $p; then
        echo "FATAL: something already listens on 127.0.0.1:$p (set PROOF_PORT / PROOF_BAO_PORT)" >&2
        exit 2
    fi
done

rm -rf "$W" && mkdir -p "$W/tmp" "$W/proxy" "$W/home" "$W/a-data" "$W/b-data" || exit 2
chmod 0700 "$W"
secret() { python3 -c 'import secrets; print(secrets.token_hex(16))'; }
export BAO_TOKEN ROOT_KEY=vuroot ROOT_SECRET S1 S2 S3 S4
BAO_TOKEN=$(secret) ROOT_SECRET=$(secret) S1=$(secret) S2=$(secret) S3=$(secret) S4=$(secret)

PIDS=()
cleanup() {
    kill "${PIDS[@]}" 2>/dev/null
    wait "${PIDS[@]}" 2>/dev/null
    [ "${KEEP:-0}" = 1 ] && echo "kept $W" || rm -rf "$W"
}
trap cleanup EXIT

now() { date +%s.%N; }
calc() { awk "BEGIN { printf \"%.3f\", $* }"; }
sleep_until() {
    local s
    s=$(calc "$1 - $(now)")
    awk "BEGIN { exit !($s > 0) }" && sleep "$s"
}
wait_port() {   # wait_port <port> <pid> <log>
    for _ in $(seq 1 100); do
        listening "$1" && return 0
        kill -0 "$2" 2>/dev/null || { cat "$3"; echo "FATAL: $3: the process exited" >&2; exit 2; }
        sleep 0.2
    done
    echo "FATAL: nothing listens on $1" >&2
    exit 2
}

printf 'X-Vault-Token: %s\n' "$BAO_TOKEN" >"$W/tmp/bao.hdr"
bao_api() {   # bao_api <method> <path> [json]
    curl -sS -f -X "$1" -H "@$W/tmp/bao.hdr" ${3:+-d "$3"} "$BAO_URL/v1/$2"
}
jget() { python3 -c 'import json, sys
d = json.load(sys.stdin)
for k in sys.argv[1].split("."): d = d[k]
print(d)' "$1"; }

PASS=0 FAIL=0
result() {   # result <id> <ok?> <got> <what> [detail]
    if [ "$2" = 1 ]; then
        PASS=$((PASS + 1)); printf 'PASS  %-3s %s  %s\n' "$1" "$3" "$4"
    else
        FAIL=$((FAIL + 1)); printf 'FAIL  %-3s %s  %s  [%s]\n' "$1" "$3" "$4" "${5:-}"
    fi
}
# check <id> <want status> <what> <A|B> <key> <secret var> <s3req.py args...>
check() {
    local id=$1 want=$2 what=$3 gw=$4 key=$5 sec=$6 got port
    shift 6
    [ "$gw" = A ] && port=$PA || port=$PB
    got=$(S3_ENDPOINT=http://127.0.0.1:$port S3_KEY=$key S3_SECRET=${!sec} \
        python3 "$HERE/../owner-scope/s3req.py" "$@" 2>"$W/tmp/err")
    [ "$got" = "$want" ]; result "$id" $((! $?)) "${got:-none}" "$what" "want $want: $(cat "$W/tmp/err")"
}
# check_cmd <id> <want: ok|refused> <what> <command...>
check_cmd() {
    local id=$1 want=$2 what=$3 got=ok
    shift 3
    "$@" >"$W/tmp/err" 2>&1 || got=refused
    [ "$got" = "$want" ]; result "$id" $((! $?)) "$got" "$what" "want $want: $(tr '\n' ' ' <"$W/tmp/err")"
}
# verdict <id> <what> <command...> -- the command's one line is the evidence, PASS or FAIL
verdict() {
    local id=$1 what=$2 out ok=1
    shift 2
    out=$("$@" 2>&1) || ok=0
    if [ $ok = 1 ]; then
        PASS=$((PASS + 1)); printf 'PASS  %-3s %s: %s\n' "$id" "$what" "$out"
    else
        FAIL=$((FAIL + 1)); printf 'FAIL  %-3s %s: %s\n' "$id" "$what" "$out"
    fi
}
watch() { python3 "$HERE/watch.py" "$@"; }
admin() {   # admin <A|B> <admin args...>
    local port
    [ "$1" = A ] && port=$PA || port=$PB
    shift
    ADMIN_ACCESS_KEY_ID=$ROOT_KEY ADMIN_SECRET_ACCESS_KEY=$ROOT_SECRET \
        "$VGW" admin --region "$S3_REGION" --endpoint-url "http://127.0.0.1:$port" "$@"
}
listed() { admin B list-users | awk '$1 == "vu-rot" { print $2, $3, $4, $5 }' | grep -qx 'user 74992 74993 41'; }

echo "--- setup"
HOME=$W/home BAO_DEV_ROOT_TOKEN_ID=$BAO_TOKEN VAULT_DEV_ROOT_TOKEN_ID=$BAO_TOKEN \
    "$BAO" server -dev -dev-listen-address="127.0.0.1:$BAO_PORT" >"$W/bao.log" 2>&1 &
PIDS+=($!)
wait_port "$BAO_PORT" "${PIDS[-1]}" "$W/bao.log"
for _ in $(seq 1 50); do curl -sf "$BAO_URL/v1/sys/health" >/dev/null && break; sleep 0.2; done
echo "openbao $("$BAO" version | head -1), dev server on 127.0.0.1:$BAO_PORT"
echo "gateway $("$VGW" --version 2>&1 | head -1), sha256 $(sha256sum "$VGW" | cut -c1-16)"

POLICY="path \"$KV/data/$ACCOUNTS/*\" { capabilities = [\"create\", \"read\", \"update\"] }
path \"$KV/metadata/$ACCOUNTS/*\" { capabilities = [\"read\", \"delete\", \"list\"] }
path \"$KV/metadata/$ACCOUNTS\" { capabilities = [\"list\"] }"
setup_bao() {
    bao_api POST "sys/mounts/$KV" '{"type": "kv", "options": {"version": "2"}}' &&
        bao_api POST sys/auth/approle '{"type": "approle"}' &&
        bao_api PUT sys/policies/acl/s3gw-proof "$(printf '%s' "$POLICY" | python3 -c 'import json, sys; print(json.dumps({"policy": sys.stdin.read()}))')" &&
        bao_api POST auth/approle/role/s3gw-proof '{"token_policies": ["s3gw-proof"], "token_ttl": "1h"}' &&
        ROLE_ID=$(bao_api GET auth/approle/role/s3gw-proof/role-id | jget data.role_id) &&
        SECRET_ID=$(bao_api POST auth/approle/role/s3gw-proof/secret-id | jget data.secret_id)
}
setup_bao >"$W/tmp/err" 2>&1
[ -n "${SECRET_ID:-}" ]; result S1 $((! $?)) ok "KV v2 mount $KV, AppRole s3gw-proof with the gateways' policy" "$(tr '\n' ' ' <"$W/tmp/err")"
[ -n "${SECRET_ID:-}" ] || exit 1

python3 "$HERE/vaultproxy.py" "$PROXY_PORT" "$BAO_PORT" "$W/proxy" >"$W/proxy.log" 2>&1 &
PIDS+=($!)
wait_port "$PROXY_PORT" "${PIDS[-1]}" "$W/proxy.log"
gateway() {   # gateway <name> <port> <vault url>
    ROOT_ACCESS_KEY_ID=$ROOT_KEY ROOT_SECRET_ACCESS_KEY=$ROOT_SECRET VGW_IAM_VAULT_ROLE_SECRET=$SECRET_ID \
        "$VGW" --region "$S3_REGION" --port "127.0.0.1:$2" \
        --iam-vault-endpoint-url "$3" --iam-vault-mount-path "$KV" --iam-vault-secret-storage-path "$ACCOUNTS" \
        --iam-vault-auth-method approle --iam-vault-role-id "$ROLE_ID" --iam-cache-ttl "$TTL" \
        posix "$W/$1-data" >"$W/$1.log" 2>&1 &
    PIDS+=($!)
    wait_port "$2" "${PIDS[-1]}" "$W/$1.log"
}
gateway a "$PA" "http://127.0.0.1:$PROXY_PORT"
gateway b "$PB" "$BAO_URL"
echo "gateway A on 127.0.0.1:$PA (OpenBao through the proxy), B on 127.0.0.1:$PB, IAM cache ${TTL} s"

check_cmd S2 ok "create-user vu-rot (uid 74992, gid 74993) through A" \
    admin A create-user --access vu-rot --secret "$S1" --role user --user-id 74992 --group-id 74993
read -r V0 CREATED0 < <(watch meta "$BAO_URL" "$KV/metadata/$ACCOUNTS/vu-rot")
t0=$(now)
check S3 200 "vu-rot signs on B with its secret (B caches it now)"   B vu-rot S1 GET /
t1=$(now)
EXPIRY=$(calc "($t0 + $t1) / 2 + $TTL")
check S4 200 "vu-rot signs on A with its secret"                     A vu-rot S1 GET /

echo "--- rotation: update-user through A while B signs with the old secret past its cache"
END=$(calc "$EXPIRY + 2 * $D + 3")
python3 "$HERE/watch.py" kv "$BAO_URL" "$ENTRY" "$END" "$W/kv-rotation.log" &
KVW=$!
S3_KEY=vu-rot S3_SECRET=$S1 python3 "$HERE/watch.py" s3 "http://127.0.0.1:$PB" "$END" "$W/b-rotation.log" &
S3W=$!
PIDS+=($KVW $S3W)
echo "$D" >"$W/proxy/delay"
# Upstream deletes the entry one delay into the update and makes it again a delay later: put
# that gap across B's cache expiry, so that B has to look the account up while it is open.
START=$(calc "$EXPIRY - 2.5 * $D")
awk "BEGIN { exit !($START < $(now)) }" && echo "      (the cache time is too short to line the rotation up with B's expiry)"
sleep_until "$START"
t=$(now)
check_cmd R1 ok "update-user vu-rot --secret <new> through A" admin A update-user --access vu-rot --secret "$S2"
echo "      took $(calc "$(now) - $t") s with OpenBao $D s away from A; B's cache expires $(calc "$EXPIRY - $(now)") s from now"
check R2 200 "A takes the new secret at once (the update drops A's own cache entry)" A vu-rot S2 GET /
wait $KVW $S3W
verdict R3 "the KV entry never went missing" watch never-missing "$W/kv-rotation.log"
verdict R4 "B never answered InvalidAccessKeyId to the old secret" watch never-invalid "$W/b-rotation.log"
verdict R5 "B took the old secret until its cache expired, and not after" watch cache-expiry "$W/b-rotation.log" "$EXPIRY"
check R6 200 "B takes the new secret after its cache time" B vu-rot S2 GET /
verdict R7 "UserID and GroupID survive" watch expect "$BAO_URL" "$ENTRY" secret=S2 role=user userID=74992 groupID=74993 projectID=0
read -r V1 CREATED1 < <(watch meta "$BAO_URL" "$KV/metadata/$ACCOUNTS/vu-rot")
[ "$V1" = $((V0 + 1)) ] && [ "$CREATED1" = "$CREATED0" ]
result R8 $((! $?)) "v$V0 -> v$V1" "one write: the entry's version went up by one and it kept its creation time" \
    "created $CREATED0, now $CREATED1: the entry was deleted and made again"

echo "--- two updates at once: --project-id 41 through A, and B rotating the secret in between"
(admin A update-user --access vu-rot --project-id 41 >"$W/tmp/a-update" 2>&1; echo $? >"$W/tmp/a-update.rc") &
A_UPDATE=$!
sleep "$(calc "1.5 * $D")"
check_cmd C1 ok "update-user vu-rot --secret <newer> through B, between A's read and A's write" \
    admin B update-user --access vu-rot --secret "$S3"
wait $A_UPDATE
check_cmd C2 ok "A's update-user --project-id 41 finished" test "$(cat "$W/tmp/a-update.rc")" = 0
verdict C3 "neither update was lost" watch expect "$BAO_URL" "$ENTRY" secret=S3 projectID=41 userID=74992 groupID=74993
check_cmd C4 ok "admin list-users on B shows vu-rot as user, 74992, 74993, project 41" listed

echo "--- a KV write that fails leaves the account as it was"
CUR=S3
watch expect "$BAO_URL" "$ENTRY" secret=S3 >/dev/null || CUR=S2
read -r V2 _ < <(watch meta "$BAO_URL" "$KV/metadata/$ACCOUNTS/vu-rot")
echo 0 >"$W/proxy/delay"
printf 'POST /v1/%s/data/' "$KV" >"$W/proxy/refuse"
check_cmd F1 refused "update-user vu-rot --secret <another> through A, its KV write refused on the way" \
    admin A update-user --access vu-rot --secret "$S4"
rm -f "$W/proxy/refuse"
verdict F2 "the entry is unchanged" watch expect "$BAO_URL" "$ENTRY" version="$V2" secret=$CUR userID=74992 groupID=74993
check F3 200 "A still signs vu-rot with its current secret (A caches nothing for it)" A vu-rot $CUR GET /
check F4 200 "so does B"                                                           B vu-rot $CUR GET /

echo
echo "$PASS passed, $FAIL failed"
[ "$FAIL" = 0 ]