diff --git a/generate-ubi b/generate-ubi
index 511edcb..5c2b1b8 100755
--- a/generate-ubi
+++ b/generate-ubi
@@ -66,6 +66,9 @@
 # Scratch directory for untarring and config file
 scratch_dir=`mktemp -d`
 
+# Make sure scratch directory always gets cleaned up
+trap "{ rm -r ${scratch_dir}; }" EXIT
+
 squashfs_file=${scratch_dir}/${squashfs_file_name}
 manifest_file=${scratch_dir}/${manifest_file_name}
 # Untar tarball
@@ -74,14 +77,12 @@
 # All valid PNOR SquashFS Tarballs have a file named "pnor.xz.squashfs"
 if [ ! -f "${squashfs_file}" ]; then
   echo "No \"${squashfs_file_name}\" file in the tarball!"
-  rm -r "${scratch_dir}"
   exit 1
 fi
 
 # Need the manifest file for calculating the version id
 if [ ! -f "${manifest_file}" ]; then
   echo "No \"${manifest_file_name}\" file in the tarball!"
-  rm -r "${scratch_dir}"
   exit 1
 fi
 
@@ -124,7 +125,7 @@
   dd if=/dev/zero bs=1k count=$image_size_kb | tr '\000' '\377' > $image_dst
 }
 
-# Used to temporary hold the UBI volume
+# Used to temporary hold the UBI image
 tmpfile=$(mktemp ${scratch_dir}/ubinized.XXXXXX)
 
 # Configuration file used to create UBI image
@@ -145,4 +146,3 @@
 dd bs=1k conv=notrunc seek=0 if=${tmpfile} of=${outfile}
 
 echo "PNOR UBI image at ${outfile}"
-rm -r "${scratch_dir}"