init: only mlock() system pages when performing snapuserd transitions.

Bug: 181032115
Test: manual test w/ VABC OTA
Change-Id: Ib4d2856b9b5eaf8688534f9d84edeb64d4b3244d
diff --git a/init/Android.bp b/init/Android.bp
index 3ff1767..1381c1d 100644
--- a/init/Android.bp
+++ b/init/Android.bp
@@ -143,6 +143,7 @@
         "libcgrouprc_format",
         "liblmkd_utils",
         "libmodprobe",
+        "libprocinfo",
         "libprotobuf-cpp-lite",
         "libpropertyinfoserializer",
         "libpropertyinfoparser",
@@ -308,6 +309,7 @@
         "libsnapshot_cow",
         "libsnapshot_init",
         "update_metadata-protos",
+        "libprocinfo",
     ],
 
     static_executable: true,
diff --git a/init/Android.mk b/init/Android.mk
index 65ee385..3c7d95a 100644
--- a/init/Android.mk
+++ b/init/Android.mk
@@ -130,6 +130,7 @@
     libsnapshot_cow \
     libsnapshot_init \
     update_metadata-protos \
+    libprocinfo \
 
 LOCAL_SANITIZE := signed-integer-overflow
 # First stage init is weird: it may start without stdout/stderr, and no /proc.
diff --git a/init/snapuserd_transition.cpp b/init/snapuserd_transition.cpp
index 19b5c57..40467b7 100644
--- a/init/snapuserd_transition.cpp
+++ b/init/snapuserd_transition.cpp
@@ -24,6 +24,7 @@
 
 #include <filesystem>
 #include <string>
+#include <string_view>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
@@ -34,6 +35,7 @@
 #include <libsnapshot/snapshot.h>
 #include <libsnapshot/snapuserd_client.h>
 #include <private/android_filesystem_config.h>
+#include <procinfo/process_map.h>
 #include <selinux/android.h>
 
 #include "block_dev_initializer.h"
@@ -157,6 +159,33 @@
     });
 }
 
+static void LockAllSystemPages() {
+    bool ok = true;
+    auto callback = [&](const android::procinfo::MapInfo& map) -> void {
+        if (!ok || android::base::StartsWith(map.name, "/dev/") ||
+            !android::base::StartsWith(map.name, "/")) {
+            return;
+        }
+        auto start = reinterpret_cast<const void*>(map.start);
+        auto len = map.end - map.start;
+        if (!len) {
+            return;
+        }
+        if (mlock(start, len) < 0) {
+            LOG(ERROR) << "mlock failed, " << start << " for " << len << " bytes.";
+            ok = false;
+        }
+    };
+
+    if (!android::procinfo::ReadProcessMaps(getpid(), callback) || !ok) {
+        LOG(FATAL) << "Could not process /proc/" << getpid() << "/maps file for init, "
+                   << "falling back to mlockall().";
+        if (mlockall(MCL_CURRENT) < 0) {
+            LOG(FATAL) << "mlockall failed";
+        }
+    }
+}
+
 void SnapuserdSelinuxHelper::StartTransition() {
     LOG(INFO) << "Starting SELinux transition of snapuserd";
 
@@ -170,9 +199,7 @@
 
     // We cannot access /system after the transition, so make sure init is
     // pinned in memory.
-    if (mlockall(MCL_CURRENT) < 0) {
-        LOG(FATAL) << "mlockall failed";
-    }
+    LockAllSystemPages();
 
     argv_.emplace_back("snapuserd");
     argv_.emplace_back("-no_socket");