diff --git a/Ext4Crypt.cpp b/Ext4Crypt.cpp
index 9c79098..150014c 100644
--- a/Ext4Crypt.cpp
+++ b/Ext4Crypt.cpp
@@ -381,16 +381,14 @@
 
     std::string tmp_path = std::string() + path + "/tmp_mnt";
 
-    // ext4enc:TODO add retry logic
-    rc = umount(tmp_path.c_str());
+    rc = wait_and_unmount(tmp_path.c_str(), true);
     if (rc) {
         SLOGE("umount %s failed with rc %d, msg %s",
               tmp_path.c_str(), rc, strerror(errno));
         return rc;
     }
 
-    // ext4enc:TODO add retry logic
-    rc = umount(path);
+    rc = wait_and_unmount(path, true);
     if (rc) {
         SLOGE("umount %s failed with rc %d, msg %s",
               path, rc, strerror(errno));
diff --git a/cryptfs.c b/cryptfs.c
index 13ccd6c..5598671 100644
--- a/cryptfs.c
+++ b/cryptfs.c
@@ -1446,7 +1446,7 @@
     return encrypt_master_key(passwd, salt, key_buf, master_key, crypt_ftr);
 }
 
-static int wait_and_unmount(char *mountpoint, bool kill)
+int wait_and_unmount(const char *mountpoint, bool kill)
 {
     int i, err, rc;
 #define WAIT_UNMOUNT_COUNT 20
diff --git a/cryptfs.h b/cryptfs.h
index e149061..7f082cc 100644
--- a/cryptfs.h
+++ b/cryptfs.h
@@ -26,6 +26,7 @@
  * partition.
  */
 
+#include <stdbool.h>
 #include <cutils/properties.h>
 
 /* The current cryptfs version */
@@ -225,6 +226,8 @@
 extern "C" {
 #endif
 
+  int wait_and_unmount(const char *mountpoint, bool kill);
+
   typedef int (*kdf_func)(const char *passwd, const unsigned char *salt,
                           unsigned char *ikey, void *params);