Merge "Clean up temporary files on run_qemu.sh exit"
diff --git a/aarch64/run_qemu.sh b/aarch64/run_qemu.sh
index f4f2edb..d628518 100755
--- a/aarch64/run_qemu.sh
+++ b/aarch64/run_qemu.sh
@@ -27,6 +27,7 @@
 default_var GIC		3
 default_var GDB		0
 default_var VERBOSE	0
+default_var KEEP_TEMP	0
 default_var TIMEOUT	""
 
 KiB=1024
@@ -36,7 +37,7 @@
 function usage() {
 	cat <<EOF
 
-Usage: $0 [-h] [-v]
+Usage: $0 [-h] [-v] [-K]
        [-e QEMU] [-L ROM_DIR] [-k KERNEL] [-r ROOTFS] [-R DRIVE] [-T TEMPLATE]
        [-c CPU] [-s NUM_CPUS] [-m MEM] [-g GIC] [-G]
        [-t TIMEOUT]
@@ -55,9 +56,17 @@
     -g    version of GIC (defaults to ${DEFAULT_GIC})
     -G    enable debugging of emulated system with GDB
     -t    kill QEMU after given number of seconds
+    -K    keep temp files
 EOF
 }
 
+# Note: Due to a bug in older versions of Bash, use '${array[@]+"${array[@]}"}'
+# to expand potentially empty arrays. '${array[@]}' is treated as undefined.
+TMP_FILES=()
+function cleanup() {
+	rm -f ${TMP_FILES[@]+"${TMP_FILES[@]}"}
+}
+
 function file_size() {
 	stat --format=%s "$1"
 }
@@ -75,9 +84,9 @@
 function template_overlay() {
 	local in="$1"
 	local out="$2"
-	local addr="$3"
-	local size="$4"
-	local tmp="$(mktemp)"
+	local tmp="$3"
+	local addr="$4"
+	local size="$5"
 
 	# Convert input DTB back to source.
 	dtc -I dtb -O dts -o "${tmp}" "${in}"
@@ -100,7 +109,7 @@
 APPEND=()
 EXTRA_RO_MOUNTS=()
 
-while getopts ":e:L:k:r:R:T:c:s:m:g:t:vGh" OPT; do
+while getopts ":e:L:k:r:R:T:c:s:m:g:t:vGKh" OPT; do
 	case "${OPT}" in
 	e)	QEMU="${OPTARG}"		;;
 	L)	ROM_DIR="${OPTARG}"		;;
@@ -115,6 +124,7 @@
 	t)	TIMEOUT="${OPTARG}"		;;
 	v)	VERBOSE=1			;;
 	G)	GDB=1				;;
+	K)	KEEP_TEMP=1			;;
 	h)
 		usage
 		exit 0
@@ -138,12 +148,14 @@
 	exit 1
 fi
 
+if [ "${KEEP_TEMP}" -ne 1 ]; then
+	trap cleanup EXIT
+fi
+
 if [ -n "${TIMEOUT}" ]; then
 	CMD+=(timeout -k 1s --foreground "${TIMEOUT}")
 fi
 
-# Note: Due to a bug in older versions of Bash, use '${array[@]+"${array[@]}"}'
-# to expand potentially empty arrays. '${array[@]}' is treated as undefined.
 CMD+=("${QEMU}")
 CMD+=(-M virt)
 CMD+=(-machine virtualization=true -machine virt,gic-version=${GIC})
@@ -173,14 +185,19 @@
 CMD+=(-append "${APPEND[*]}")
 
 if [ -n "${TEMPLATE}" ]; then
-	# Dump the QEMU DTB.
 	QEMU_DTB="$(mktemp)"
+	QEMU_PATCHED_DTB="$(mktemp)"
+	TMP_DTS="$(mktemp)"
+
+	# Mark files for deletion on EXIT.
+	TMP_FILES+=("${QEMU_DTB}" "${QEMU_PATCHED_DTB}" "${TMP_DTS}")
+
+	# Dump the QEMU DTB.
 	"${CMD[@]}" -machine dumpdtb="${QEMU_DTB}" > /dev/null
 
 	# Compile the overlayed DTB with dummy values to determine its size.
-	TMP_DTB="$(mktemp)"
-	template_overlay "${QEMU_DTB}" "${TMP_DTB}" 0x0 0x0
-	TMP_DTB_SIZE=$(file_size "${TMP_DTB}")
+	template_overlay "${QEMU_DTB}" "${QEMU_PATCHED_DTB}" "${TMP_DTS}" 0x0 0x0
+	TMP_DTB_SIZE=$(file_size "${QEMU_PATCHED_DTB}")
 
 	# From QEMU's hw/arm/boot.c:
 	# RAM always starts at 1GiB PA offset. The kernel is placed there.
@@ -191,8 +208,8 @@
 	TEMPLATE_ADDR=$(hex $(align_up_pow2 $TEMPLATE_ADDR 64*KiB))
 	TEMPLATE_SIZE=$(hex $(file_size "${TEMPLATE}"))
 
-	QEMU_PATCHED_DTB="$(mktemp)"
-	template_overlay "${QEMU_DTB}" "${QEMU_PATCHED_DTB}" "${TEMPLATE_ADDR}" "${TEMPLATE_SIZE}"
+	template_overlay "${QEMU_DTB}" "${QEMU_PATCHED_DTB}" "${TMP_DTS}" \
+			 "${TEMPLATE_ADDR}" "${TEMPLATE_SIZE}"
 
 	CMD+=(-dtb "${QEMU_PATCHED_DTB}")
 	CMD+=(-device loader,file="${TEMPLATE}",addr=${TEMPLATE_ADDR},force-raw=true)
@@ -202,4 +219,6 @@
 	set -x
 fi
 
-exec "${CMD[@]}"
+# Invoke QEMU and then propagate its exit code.
+# We do this instead of `exec` to delete TMP_FILES in the EXIT handler.
+set +e; "${CMD[@]}"; exit $?