diff --git a/tmp b/tmp
index d08bc88..ff585d1 100755
--- a/tmp
+++ b/tmp
@@ -6,19 +6,25 @@
 # Distributed under terms of the MIT license.
 #
 
-## Define variables
+# Define variables
 name="$(basename "$0" .sh)"
 tmpdir="$(mktemp -d)"
-## Define functions
+# Verify them
+[ -z ${tmpdir+x} ] && exit 2
+# Define functions
 function log() {
 	[ -z ${1+x} ] && return 1
 	printf "${name}: $1\n"
 }
-## Do the do
+function cleanup() {
+	rm -rf "${tmpdir}"
+}
+trap "cleanup" EXIT
+# Do the do
 log "This folder will be removed when this shell exits"
 (
 cd "${tmpdir}"
 $SHELL
 )
-rm -rf "${tmpdir}"
+cleanup