#!/usr/bin/env bash
# Boot a KASAN Linux guest and run poc.c inside it.
#
# Required env:
#   KERNEL_IMAGE  path to bzImage (KASAN + CONFIG_DM=y)
#   VM_IMAGE      path to guest disk (e.g. bookworm.img)
#   VM_KEY        path to guest SSH private key (e.g. bookworm.id_rsa)
#
# Optional: VM_PORT VM_MEMORY_MB VM_CPUS RUN_TIMEOUT_SEC
#
# Important: do NOT set panic_on_warn — a related bd_*_disk_holder WARNING
# often fires before the KASAN UAF and would abort the guest early.
set -euo pipefail
cd "$(dirname "$0")"

export KERNEL_IMAGE="${KERNEL_IMAGE:?set KERNEL_IMAGE}"
export VM_IMAGE="${VM_IMAGE:?set VM_IMAGE}"
export VM_KEY="${VM_KEY:?set VM_KEY}"
export VM_PORT="${VM_PORT:-10057}"
export VM_MEMORY_MB="${VM_MEMORY_MB:-2048}"
export VM_CPUS="${VM_CPUS:-4}"
export RUN_TIMEOUT_SEC="${RUN_TIMEOUT_SEC:-240}"

REPRO_SRC="$(pwd)/poc.c"
LOG_DIR="${LOG_DIR:-${OUTPUT_DIR:-$(pwd)}}"
mkdir -p "$LOG_DIR"
SERIAL_LOG="$LOG_DIR/vm_serial.log"
DMESG_GREP="$LOG_DIR/repro_dmesg_grep.txt"
SSH_OUT="$LOG_DIR/repro_ssh_output.txt"
PIDFILE="$LOG_DIR/qemu-repro.pid"

SSH_OPTS=(-i "$VM_KEY" -p "$VM_PORT" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes)
SCP_OPTS=(-i "$VM_KEY" -P "$VM_PORT" -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes)

require_file() {
	if [[ ! -f "$1" ]]; then
		echo "missing required file: $1" >&2
		exit 1
	fi
}

stop_qemu() {
	if [[ -f "$PIDFILE" ]]; then
		local pid
		pid="$(cat "$PIDFILE")"
		if [[ "$pid" =~ ^[0-9]+$ ]]; then
			kill "$pid" 2>/dev/null || true
			wait "$pid" 2>/dev/null || true
		fi
		rm -f "$PIDFILE"
	fi
}

wait_for_ssh() {
	local timeout="${1:-180}"
	local start now
	start="$(date +%s)"
	while true; do
		if ssh "${SSH_OPTS[@]}" -o ConnectTimeout=3 root@127.0.0.1 true 2>/dev/null; then
			return 0
		fi
		now="$(date +%s)"
		if (( now - start >= timeout )); then
			echo "guest SSH not ready on port $VM_PORT within ${timeout}s" >&2
			echo "--- last 40 lines of serial ---" >&2
			tail -n 40 "$SERIAL_LOG" >&2 || true
			return 1
		fi
		sleep 2
	done
}

kvm_args=()
if [[ -r /dev/kvm ]]; then
	kvm_args=(-enable-kvm -cpu host)
fi

require_file "$KERNEL_IMAGE"
require_file "$VM_IMAGE"
require_file "$VM_KEY"
require_file "$REPRO_SRC"
chmod 600 "$VM_KEY"

trap 'stop_qemu' EXIT
stop_qemu
: >"$SERIAL_LOG"

echo "[+] Booting KASAN guest (port=$VM_PORT timeout=${RUN_TIMEOUT_SEC}s cpus=$VM_CPUS)"
qemu-system-x86_64 \
	"${kvm_args[@]}" \
	-m "$VM_MEMORY_MB" \
	-smp "$VM_CPUS" \
	-kernel "$KERNEL_IMAGE" \
	-append "root=/dev/sda rw console=ttyS0 earlyprintk=serial oops=panic panic=-1 nomodeset net.ifnames=0 biosdevname=0" \
	-drive "file=$VM_IMAGE,format=raw" \
	-snapshot \
	-netdev "user,id=net0,host=10.0.2.10,hostfwd=tcp::${VM_PORT}-:22" \
	-device virtio-net-pci,netdev=net0 \
	-nographic \
	-pidfile "$PIDFILE" \
	>"$SERIAL_LOG" 2>&1 &

wait_for_ssh 180
scp "${SCP_OPTS[@]}" "$REPRO_SRC" root@127.0.0.1:/tmp/repro.c

GUEST_TIMEOUT=$((RUN_TIMEOUT_SEC > 60 ? RUN_TIMEOUT_SEC - 30 : RUN_TIMEOUT_SEC))

set +e
timeout "$RUN_TIMEOUT_SEC" ssh "${SSH_OPTS[@]}" -o ServerAliveInterval=10 root@127.0.0.1 \
	"modprobe dm_mod loop 2>/dev/null || true; \
	 if ! gcc -O2 -static -o /tmp/repro /tmp/repro.c -lpthread 2>/tmp/gcc.err; then \
	   echo '[!] static link failed, retrying dynamic:'; cat /tmp/gcc.err; \
	   gcc -O2 -o /tmp/repro /tmp/repro.c -lpthread; \
	 fi; \
	 dmesg -C; \
	 timeout $GUEST_TIMEOUT /tmp/repro; \
	 echo __DMESG_START__; dmesg" >"$SSH_OUT" 2>&1
ssh_rc=$?
set -e

{
	echo "ssh_exit_code=$ssh_rc"
	echo
	sed -n '/__DMESG_START__/,$p' "$SSH_OUT" | tail -n +2
	echo
	echo "=== serial log grep ==="
	grep -E 'KASAN|slab-use-after-free|dm_setup_md_queue|bd_link_disk_holder|bd_unlink_disk_holder|BUG:|WARNING:' "$SERIAL_LOG" || true
} >"$DMESG_GREP"

echo
echo "[+] Logs: $SSH_OUT $DMESG_GREP $SERIAL_LOG"
if grep -q 'slab-use-after-free' "$DMESG_GREP" "$SERIAL_LOG" 2>/dev/null; then
	echo "[+] KASAN UAF reproduced"
	exit 0
else
	echo "[!] KASAN not seen — check serial log or increase RUN_TIMEOUT_SEC"
	grep -E 'WARNING:|BUG:|panic|Oops' "$SERIAL_LOG" | head -40 || true
	exit 1
fi
