From b578c2031ac12809445f7a348ce110bfb99762d0 Mon Sep 17 00:00:00 2001
From: Storm Dragon <stormdragon2976@gmail.com>
Date: Wed, 9 Dec 2020 10:39:57 -0500
Subject: [PATCH] Added cleanup code for if the script exits unexpectedly.

---
 build/build-stormux.sh | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/build/build-stormux.sh b/build/build-stormux.sh
index b4aac96..b74249b 100755
--- a/build/build-stormux.sh
+++ b/build/build-stormux.sh
@@ -18,8 +18,21 @@
 # 02110-1301, USA.
 #
 
-set -e # Don't want to destroy stuff if this goes majorly wrong.
+# keep track of mounted status for exit function
+mounted=1
 
+set -e # Don't want to destroy stuff if this goes majorly wrong.
+trap cleanup EXIT # make sure the script cleans up after itself before closing.
+
+
+cleanup() {
+    if [[ $mounted -eq 0 ]]; then
+        umount -R /mnt
+        partx -d ${loopdev}
+        losetup --detach ${loopdev}
+    fi
+    exit 0
+}
 
 help() {
     echo -e "Usage:\n"
@@ -116,6 +129,8 @@ mkfs.ext4 -F ${loopdev}p2
 mount ${loopdev}p2 /mnt
 mkdir /mnt/boot
 mount ${loopdev}p1 /mnt/boot
+# Things are mounted now, so set mounted to 0 (bash true)
+mounted=0
 wget "${url[$imageVersion]}" -O- | bsdtar -xpf - -C /mnt
 arch-chroot /mnt << EOF
 # set up pacman
@@ -148,6 +163,5 @@ systemctl enable cronie.service fenrirscreenreader.service
 # Update fstab for Raspberry Pi 4.
 [[ $imageVersion -eq 4 ]] && sed -i 's/mmcblk0/mmcblk1/g' /etc/fstab
 EOF
-umount -R /mnt
-partx -d ${loopdev}
-losetup --detach ${loopdev}
+
+exit 0