diff --git a/init/builtins.cpp b/init/builtins.cpp
index ed24026..7c66de5 100644
--- a/init/builtins.cpp
+++ b/init/builtins.cpp
@@ -356,51 +356,64 @@
 // mkdir <path> [mode] [owner] [group]
 static Result<void> do_mkdir(const BuiltinArguments& args) {
     mode_t mode = 0755;
-    if (args.size() >= 3) {
-        mode = std::strtoul(args[2].c_str(), 0, 8);
-    }
+    Result<uid_t> uid = -1;
+    Result<gid_t> gid = -1;
 
-    if (!make_dir(args[1], mode)) {
-        /* chmod in case the directory already exists */
-        if (errno == EEXIST) {
-            if (fchmodat(AT_FDCWD, args[1].c_str(), mode, AT_SYMLINK_NOFOLLOW) == -1) {
-                return ErrnoError() << "fchmodat() failed";
-            }
-        } else {
-            return ErrnoErrorIgnoreEnoent() << "mkdir() failed";
-        }
-    }
-
-    if (args.size() >= 4) {
-        auto uid = DecodeUid(args[3]);
-        if (!uid) {
-            return Error() << "Unable to decode UID for '" << args[3] << "': " << uid.error();
-        }
-        Result<gid_t> gid = -1;
-
-        if (args.size() == 5) {
+    switch (args.size()) {
+        case 5:
             gid = DecodeUid(args[4]);
             if (!gid) {
                 return Error() << "Unable to decode GID for '" << args[4] << "': " << gid.error();
             }
-        }
-
-        if (lchown(args[1].c_str(), *uid, *gid) == -1) {
-            return ErrnoError() << "lchown failed";
-        }
-
-        /* chown may have cleared S_ISUID and S_ISGID, chmod again */
-        if (mode & (S_ISUID | S_ISGID)) {
-            if (fchmodat(AT_FDCWD, args[1].c_str(), mode, AT_SYMLINK_NOFOLLOW) == -1) {
-                return ErrnoError() << "fchmodat failed";
+            FALLTHROUGH_INTENDED;
+        case 4:
+            uid = DecodeUid(args[3]);
+            if (!uid) {
+                return Error() << "Unable to decode UID for '" << args[3] << "': " << uid.error();
             }
+            FALLTHROUGH_INTENDED;
+        case 3:
+            mode = std::strtoul(args[2].c_str(), 0, 8);
+            FALLTHROUGH_INTENDED;
+        case 2:
+            break;
+        default:
+            return Error() << "Unexpected argument count: " << args.size();
+    }
+    std::string target = args[1];
+    struct stat mstat;
+    if (lstat(target.c_str(), &mstat) != 0) {
+        if (errno != ENOENT) {
+            return ErrnoError() << "lstat() failed on " << target;
+        }
+        if (!make_dir(target, mode)) {
+            return ErrnoErrorIgnoreEnoent() << "mkdir() failed on " << target;
+        }
+        if (lstat(target.c_str(), &mstat) != 0) {
+            return ErrnoError() << "lstat() failed on new " << target;
         }
     }
-
+    if (!S_ISDIR(mstat.st_mode)) {
+        return Error() << "Not a directory on " << target;
+    }
+    bool needs_chmod = (mstat.st_mode & ~S_IFMT) != mode;
+    if ((*uid != static_cast<uid_t>(-1) && *uid != mstat.st_uid) ||
+        (*gid != static_cast<gid_t>(-1) && *gid != mstat.st_gid)) {
+        if (lchown(target.c_str(), *uid, *gid) == -1) {
+            return ErrnoError() << "lchown failed on " << target;
+        }
+        // chown may have cleared S_ISUID and S_ISGID, chmod again
+        needs_chmod = true;
+    }
+    if (needs_chmod) {
+        if (fchmodat(AT_FDCWD, target.c_str(), mode, AT_SYMLINK_NOFOLLOW) == -1) {
+            return ErrnoError() << "fchmodat() failed on " << target;
+        }
+    }
     if (fscrypt_is_native()) {
-        if (fscrypt_set_directory_policy(args[1].c_str())) {
+        if (fscrypt_set_directory_policy(target)) {
             return reboot_into_recovery(
-                {"--prompt_and_wipe_data", "--reason=set_policy_failed:"s + args[1]});
+                    {"--prompt_and_wipe_data", "--reason=set_policy_failed:"s + target});
         }
     }
     return {};
diff --git a/init/fscrypt_init_extensions.cpp b/init/fscrypt_init_extensions.cpp
index 956228a..0f5a864 100644
--- a/init/fscrypt_init_extensions.cpp
+++ b/init/fscrypt_init_extensions.cpp
@@ -16,18 +16,9 @@
 
 #include "fscrypt_init_extensions.h"
 
-#include <android-base/file.h>
-#include <android-base/logging.h>
-#include <android-base/stringprintf.h>
-#include <android-base/strings.h>
-#include <cutils/properties.h>
-#include <cutils/sockets.h>
 #include <dirent.h>
 #include <errno.h>
 #include <fts.h>
-#include <fscrypt/fscrypt.h>
-#include <keyutils.h>
-#include <logwrap/logwrap.h>
 #include <sys/mount.h>
 #include <sys/stat.h>
 #include <unistd.h>
@@ -35,15 +26,22 @@
 #include <string>
 #include <vector>
 
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <android-base/stringprintf.h>
+#include <android-base/strings.h>
+#include <cutils/properties.h>
+#include <cutils/sockets.h>
+#include <fscrypt/fscrypt.h>
+#include <keyutils.h>
+#include <logwrap/logwrap.h>
 
 #define TAG "fscrypt"
 
-static int set_system_de_policy_on(char const* dir);
+static int set_system_de_policy_on(const std::string& dir);
 
-int fscrypt_install_keyring()
-{
-    key_serial_t device_keyring = add_key("keyring", "fscrypt", 0, 0,
-                                          KEY_SPEC_SESSION_KEYRING);
+int fscrypt_install_keyring() {
+    key_serial_t device_keyring = add_key("keyring", "fscrypt", 0, 0, KEY_SPEC_SESSION_KEYRING);
 
     if (device_keyring == -1) {
         PLOG(ERROR) << "Failed to create keyring";
@@ -56,8 +54,8 @@
 }
 
 // TODO(b/139378601): use a single central implementation of this.
-static void delete_dir_contents(const char* dir) {
-    char* const paths[2] = {const_cast<char*>(dir), nullptr};
+static void delete_dir_contents(const std::string& dir) {
+    char* const paths[2] = {const_cast<char*>(dir.c_str()), nullptr};
     FTS* fts = fts_open(paths, FTS_PHYSICAL | FTS_NOCHDIR | FTS_XDEV, nullptr);
     FTSENT* cur;
     while ((cur = fts_read(fts)) != nullptr) {
@@ -65,7 +63,7 @@
             PLOG(ERROR) << "fts_read";
             break;
         }
-        if (strcmp(dir, cur->fts_path) == 0) {
+        if (dir == cur->fts_path) {
             continue;
         }
         switch (cur->fts_info) {
@@ -96,16 +94,15 @@
     }
 }
 
-int fscrypt_set_directory_policy(const char* dir)
-{
+int fscrypt_set_directory_policy(const std::string& dir) {
     const std::string prefix = "/data/";
 
-    if (!dir || strncmp(dir, prefix.c_str(), prefix.size())) {
+    if (!android::base::StartsWith(dir, prefix)) {
         return 0;
     }
 
     // Special-case /data/media/obb per b/64566063
-    if (strcmp(dir, "/data/media/obb") == 0) {
+    if (dir == "/data/media/obb") {
         // Try to set policy on this directory, but if it is non-empty this may fail.
         set_system_de_policy_on(dir);
         return 0;
@@ -115,7 +112,7 @@
     // To make this less restrictive, consider using a policy file.
     // However this is overkill for as long as the policy is simply
     // to apply a global policy to all /data folders created via makedir
-    if (strchr(dir + prefix.size(), '/')) {
+    if (dir.find_first_of('/', prefix.size()) != std::string::npos) {
         return 0;
     }
 
@@ -132,7 +129,7 @@
         "apex", "preloads", "app-staging",
         "gsi",
     };
-    for (const auto& d: directories_to_exclude) {
+    for (const auto& d : directories_to_exclude) {
         if ((prefix + d) == dir) {
             LOG(INFO) << "Not setting policy on " << dir;
             return 0;
@@ -144,7 +141,7 @@
     }
     // Empty these directories if policy setting fails.
     std::vector<std::string> wipe_on_failure = {
-        "rollback", "rollback-observer",  // b/139193659
+            "rollback", "rollback-observer",  // b/139193659
     };
     for (const auto& d : wipe_on_failure) {
         if ((prefix + d) == dir) {
@@ -157,7 +154,7 @@
     return err;
 }
 
-static int set_system_de_policy_on(char const* dir) {
+static int set_system_de_policy_on(const std::string& dir) {
     std::string ref_filename = std::string("/data") + fscrypt_key_ref;
     std::string policy;
     if (!android::base::ReadFileToString(ref_filename, &policy)) {
@@ -179,14 +176,13 @@
     }
 
     LOG(INFO) << "Setting policy on " << dir;
-    int result = fscrypt_policy_ensure(dir, policy.c_str(), policy.length(),
-                                       modes[0].c_str(),
-                                       modes.size() >= 2 ?
-                                            modes[1].c_str() : "aes-256-cts");
+    int result =
+            fscrypt_policy_ensure(dir.c_str(), policy.c_str(), policy.length(), modes[0].c_str(),
+                                  modes.size() >= 2 ? modes[1].c_str() : "aes-256-cts");
     if (result) {
-        LOG(ERROR) << android::base::StringPrintf(
-            "Setting %02x%02x%02x%02x policy on %s failed!",
-            policy[0], policy[1], policy[2], policy[3], dir);
+        LOG(ERROR) << android::base::StringPrintf("Setting %02x%02x%02x%02x policy on %s failed!",
+                                                  policy[0], policy[1], policy[2], policy[3],
+                                                  dir.c_str());
         return -1;
     }
 
diff --git a/init/fscrypt_init_extensions.h b/init/fscrypt_init_extensions.h
index 2b6c46e..2163ef6 100644
--- a/init/fscrypt_init_extensions.h
+++ b/init/fscrypt_init_extensions.h
@@ -14,20 +14,9 @@
  * limitations under the License.
  */
 
-#ifndef _FSCRYPT_INIT_EXTENSIONS_H_
-#define _FSCRYPT_INIT_EXTENSIONS_H_
+#pragma once
 
-#include <sys/cdefs.h>
-#include <stdbool.h>
-#include <cutils/multiuser.h>
+#include <string>
 
-__BEGIN_DECLS
-
-// These functions assume they are being called from init
-// They will not operate properly outside of init
 int fscrypt_install_keyring();
-int fscrypt_set_directory_policy(const char* path);
-
-__END_DECLS
-
-#endif // _FSCRYPT_INIT_EXTENSIONS_H_
+int fscrypt_set_directory_policy(const std::string& dir);