From 9fd886274174b15bc2313d8267df8da9e727ef96 Mon Sep 17 00:00:00 2001
From: David Anderson <dvander@google.com>
Date: Fri, 5 Mar 2021 14:10:55 -0800
Subject: [PATCH] init: only mlock() system pages when performing snapuserd
 transitions.

Bug: 181032115
Test: manual test w/ VABC OTA
Change-Id: Ib4d2856b9b5eaf8688534f9d84edeb64d4b3244d
---
 init/Android.bp               |  2 ++
 init/Android.mk               |  1 +
 init/snapuserd_transition.cpp | 33 ++++++++++++++++++++++++++++++---
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/init/Android.bp b/init/Android.bp
index 3ff17672c..1381c1d25 100644
--- a/init/Android.bp
+++ b/init/Android.bp
@@ -143,6 +143,7 @@ cc_defaults {
         "libcgrouprc_format",
         "liblmkd_utils",
         "libmodprobe",
+        "libprocinfo",
         "libprotobuf-cpp-lite",
         "libpropertyinfoserializer",
         "libpropertyinfoparser",
@@ -308,6 +309,7 @@ cc_binary {
         "libsnapshot_cow",
         "libsnapshot_init",
         "update_metadata-protos",
+        "libprocinfo",
     ],
 
     static_executable: true,
diff --git a/init/Android.mk b/init/Android.mk
index 65ee385de..3c7d95acf 100644
--- a/init/Android.mk
+++ b/init/Android.mk
@@ -130,6 +130,7 @@ LOCAL_STATIC_LIBRARIES := \
     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 19b5c57a0..40467b7d3 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 @@ SnapuserdSelinuxHelper::SnapuserdSelinuxHelper(std::unique_ptr<SnapshotManager>&
     });
 }
 
+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 @@ void SnapuserdSelinuxHelper::StartTransition() {
 
     // 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");