diff --git a/lib/common-functions b/lib/common-functions
index 9ac62bd34..6681dc782 100644
--- a/lib/common-functions
+++ b/lib/common-functions
@@ -36,7 +36,7 @@ function mk_build_dir () {
     sudo mount -t tmpfs tmpfs $TMP_IMAGE_DIR
   fi
   sudo chown $(id -u):$(id -g) $TMP_BUILD_DIR $TMP_IMAGE_DIR
-  trap cleanup EXIT
+  trap trap_cleanup EXIT
   echo Building in $TMP_BUILD_DIR
   export TMP_IMAGE_PATH=$TMP_IMAGE_DIR/image.raw
   export TMP_HOOKS_PATH=$TMP_BUILD_DIR/hooks
diff --git a/lib/img-functions b/lib/img-functions
index 56d5edc80..64079bd4c 100644
--- a/lib/img-functions
+++ b/lib/img-functions
@@ -28,6 +28,12 @@ function unmount_image () {
     fi
 }
 
+function trap_cleanup() {
+    exitval=$?
+    cleanup
+    exit $exitval
+}
+
 function cleanup () {
     unmount_image
     cleanup_dirs