Make uprobestats probe setUidTempAllowlistStateLSP

Allow uprobestats to run on user (public) builds when the target class
is on the allowlist. Add a BPF program for instrumeting
setUidTempAllowlistStateLSP. Add code to output the stats to statsd.

Flag: NONE the whole uprobestats program is already flagged.
Change-Id: Ief5f438e55f8bab4b413160eec243fa87077d315
Bug: 296108553
Test: hello_uprobestats -t -n test_setUidTempAllowlistStateLSP
diff --git a/src/Android.bp b/src/Android.bp
index 44d97a5..69ae4ec 100644
--- a/src/Android.bp
+++ b/src/Android.bp
@@ -22,6 +22,7 @@
         "Bpf.cpp",
         "ConfigResolver.cpp",
         "Process.cpp",
+        "Guardrail.cpp",
         "config.proto",
     ],
     header_libs: [
@@ -66,6 +67,7 @@
     required: [
         "BitmapAllocation.o",
         "GenericInstrumentation.o",
+        "ProcessManagement.o",
     ],
     proto: {
         type: "lite",
@@ -114,3 +116,25 @@
         canonical_path_from_root: false,
     },
 }
+
+cc_test {
+    name: "libuprobestats_test",
+    srcs: [
+        "config.proto",
+        "Guardrail-test.cpp",
+    ],
+    shared_libs: [
+        "libbase",
+    ],
+    static_libs: [
+        "libbase",
+        "libgtest",
+        "liblog",
+        "libprotoutil",
+        "libuprobestats",
+    ],
+    proto: {
+        type: "lite",
+        static: true,
+    },
+}
diff --git a/src/Bpf.cpp b/src/Bpf.cpp
index 6821cfd..9ff8920 100644
--- a/src/Bpf.cpp
+++ b/src/Bpf.cpp
@@ -95,6 +95,8 @@
                                              int timeoutMs);
 template std::vector<CallTimestamp> pollRingBuf(const char *mapPath,
                                                 int timeoutMs);
+template std::vector<SetUidTempAllowlistStateRecord>
+pollRingBuf(const char *mapPath, int timeoutMs);
 
 std::vector<int32_t> consumeRingBuf(const char *mapPath) {
   auto result = android::bpf::BpfRingbuf<uint64_t>::Create(mapPath);
diff --git a/src/Bpf.h b/src/Bpf.h
index 8ac2042..16a64d9 100644
--- a/src/Bpf.h
+++ b/src/Bpf.h
@@ -36,6 +36,11 @@
   unsigned long timestampNs;
 };
 
+struct SetUidTempAllowlistStateRecord {
+  __u64 uid;
+  bool onAllowlist;
+};
+
 template <typename T>
 std::vector<T> pollRingBuf(const char *mapPath, int timeoutMs);
 
diff --git a/src/ConfigResolver.cpp b/src/ConfigResolver.cpp
index 611e816..68dcd1f 100644
--- a/src/ConfigResolver.cpp
+++ b/src/ConfigResolver.cpp
@@ -143,4 +143,4 @@
 
 } // namespace config_resolver
 } // namespace uprobestats
-} // namespace android
\ No newline at end of file
+} // namespace android
diff --git a/src/Guardrail-test.cpp b/src/Guardrail-test.cpp
new file mode 100644
index 0000000..968c534
--- /dev/null
+++ b/src/Guardrail-test.cpp
@@ -0,0 +1,66 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <gtest/gtest.h>
+
+#include "Guardrail.h"
+
+namespace android {
+namespace uprobestats {
+
+class GuardrailTest : public ::testing::Test {};
+
+TEST_F(GuardrailTest, EverythingAllowedOnUserDebugAndEng) {
+  ::uprobestats::protos::UprobestatsConfig config;
+  config.add_tasks()->add_probe_configs()->set_method_signature(
+      "void com.android.server.am.SomeClass.doWork()");
+  EXPECT_TRUE(guardrail::isAllowed(config, "userdebug"));
+  EXPECT_TRUE(guardrail::isAllowed(config, "eng"));
+}
+
+TEST_F(GuardrailTest, OomAdjusterAllowed) {
+  ::uprobestats::protos::UprobestatsConfig config;
+  config.add_tasks()->add_probe_configs()->set_method_signature(
+      "void com.android.server.am.OomAdjuster.setUidTempAllowlistStateLSP(int, "
+      "boolean)");
+  config.add_tasks()->add_probe_configs()->set_method_signature(
+      "void "
+      "com.android.server.am.OomAdjuster$$ExternalSyntheticLambda0.accept(java."
+      "lang.Object)");
+  EXPECT_TRUE(guardrail::isAllowed(config, "user"));
+  EXPECT_TRUE(guardrail::isAllowed(config, "userdebug"));
+  EXPECT_TRUE(guardrail::isAllowed(config, "eng"));
+}
+
+TEST_F(GuardrailTest, DisallowOomAdjusterWithSuffix) {
+  ::uprobestats::protos::UprobestatsConfig config;
+  config.add_tasks()->add_probe_configs()->set_method_signature(
+      "void com.android.server.am.OomAdjusterWithSomeSuffix.doWork()");
+  EXPECT_FALSE(guardrail::isAllowed(config, "user"));
+}
+
+TEST_F(GuardrailTest, DisallowedMethodInSecondTask) {
+  ::uprobestats::protos::UprobestatsConfig config;
+  config.add_tasks()->add_probe_configs()->set_method_signature(
+      "void com.android.server.am.OomAdjuster.setUidTempAllowlistStateLSP(int, "
+      "boolean)");
+  config.add_tasks()->add_probe_configs()->set_method_signature(
+      "void com.android.server.am.disallowedClass.doWork()");
+  EXPECT_FALSE(guardrail::isAllowed(config, "user"));
+}
+
+} // namespace uprobestats
+} // namespace android
diff --git a/src/Guardrail.cpp b/src/Guardrail.cpp
new file mode 100644
index 0000000..a55411e
--- /dev/null
+++ b/src/Guardrail.cpp
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <android-base/strings.h>
+#include <config.pb.h>
+#include <string>
+
+namespace android {
+namespace uprobestats {
+namespace guardrail {
+
+using std::string;
+
+namespace {
+
+constexpr std::array kAllowedMethodPrefixes = {
+    "com.android.server.am.CachedAppOptimizer",
+    "com.android.server.am.OomAdjuster",
+    "com.android.server.am.OomAdjusterModernImpl",
+};
+
+} // namespace
+
+bool isAllowed(const ::uprobestats::protos::UprobestatsConfig &config,
+               const string &buildType) {
+  if (buildType != "user") {
+    return true;
+  }
+  for (const auto &task : config.tasks()) {
+    for (const auto &probeConfig : task.probe_configs()) {
+      const string &methodSignature = probeConfig.method_signature();
+      std::vector<string> components =
+          android::base::Split(methodSignature, " ");
+      if (components.size() < 2) {
+        return false;
+      }
+      const string &fullMethodName = components[1];
+      bool allowed = false;
+      for (const std::string allowedPrefix : kAllowedMethodPrefixes) {
+        if (android::base::StartsWith(fullMethodName, allowedPrefix + ".") ||
+            android::base::StartsWith(fullMethodName, allowedPrefix + "$")) {
+          allowed = true;
+          break;
+        }
+      }
+      if (!allowed) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+} // namespace guardrail
+} // namespace uprobestats
+} // namespace android
diff --git a/src/Guardrail.h b/src/Guardrail.h
new file mode 100644
index 0000000..30c5591
--- /dev/null
+++ b/src/Guardrail.h
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <config.pb.h>
+#include <string>
+
+namespace android {
+namespace uprobestats {
+namespace guardrail {
+
+bool isAllowed(const ::uprobestats::protos::UprobestatsConfig &config,
+               const std::string &buildType);
+
+} // namespace guardrail
+} // namespace uprobestats
+} // namespace android
diff --git a/src/UprobeStats.cpp b/src/UprobeStats.cpp
index c45f1b5..ad660cc 100644
--- a/src/UprobeStats.cpp
+++ b/src/UprobeStats.cpp
@@ -30,6 +30,7 @@
 
 #include "Bpf.h"
 #include "ConfigResolver.h"
+#include "Guardrail.h"
 #include <stats_event.h>
 
 using namespace android::uprobestats;
@@ -38,6 +39,8 @@
     std::string("GenericInstrumentation_call_detail");
 const std::string kGenericBpfMapTimestamp =
     std::string("GenericInstrumentation_call_timestamp");
+const std::string kProcessManagementMap =
+    std::string("ProcessManagement_output_buf");
 const int kJavaArgumentRegisterOffset = 2;
 const bool kDebug = true;
 
@@ -48,10 +51,6 @@
     }                                                                          \
   } while (0)
 
-bool isUserBuild() {
-  return android::base::GetProperty("ro.build.type", "unknown") == "user";
-}
-
 bool isUprobestatsEnabled() {
   return android::uprobestats::flags::enable_uprobestats();
 }
@@ -135,6 +134,27 @@
         AStatsEvent_release(event);
         LOG_IF_DEBUG("successfully wrote atom id: " << atom_id);
       }
+    } else if (mapPath.find(kProcessManagementMap) != std::string::npos) {
+      LOG_IF_DEBUG("Polling for SetUidTempAllowlistStateRecord result");
+      auto result = bpf::pollRingBuf<bpf::SetUidTempAllowlistStateRecord>(
+          mapPath.c_str(), timeoutMs);
+      for (auto value : result) {
+        LOG_IF_DEBUG("SetUidTempAllowlistStateRecord result... uid: "
+                     << value.uid << " onAllowlist: " << value.onAllowlist
+                     << " mapPath: " << mapPath);
+        if (!args.taskConfig.has_statsd_logging_config()) {
+          LOG_IF_DEBUG("no statsd logging config");
+          continue;
+        }
+        auto statsd_logging_config = args.taskConfig.statsd_logging_config();
+        int atom_id = statsd_logging_config.atom_id();
+        AStatsEvent *event = AStatsEvent_obtain();
+        AStatsEvent_setAtomId(event, atom_id);
+        AStatsEvent_writeInt32(event, value.uid);
+        AStatsEvent_writeBool(event, value.onAllowlist);
+        AStatsEvent_write(event);
+        AStatsEvent_release(event);
+      }
     } else {
       LOG_IF_DEBUG("Polling for i64 result");
       auto result = bpf::pollRingBuf<uint64_t>(mapPath.c_str(), timeoutMs);
@@ -149,12 +169,6 @@
 }
 
 int main(int argc, char **argv) {
-  if (isUserBuild()) {
-    // TODO(296108553): See if we could avoid shipping this binary on user
-    // builds.
-    LOG(ERROR) << "uprobestats disabled on user build. Exiting.";
-    return 1;
-  }
   if (!isUprobestatsEnabled()) {
     LOG(ERROR) << "uprobestats disabled by flag. Exiting.";
     return 1;
@@ -170,6 +184,11 @@
     LOG(ERROR) << "Failed to parse uprobestats config: " << argv[1];
     return 1;
   }
+  if (!guardrail::isAllowed(config.value(), android::base::GetProperty(
+                                                "ro.build.type", "unknown"))) {
+    LOG(ERROR) << "uprobestats probing config disallowed on this device.";
+    return 1;
+  }
   auto resolvedTask = config_resolver::resolveSingleTask(config.value());
   if (!resolvedTask.has_value()) {
     LOG(ERROR) << "Failed to parse task";
diff --git a/src/bpf_progs/Android.bp b/src/bpf_progs/Android.bp
index c474469..1b7bf63 100644
--- a/src/bpf_progs/Android.bp
+++ b/src/bpf_progs/Android.bp
@@ -13,3 +13,14 @@
     srcs: ["GenericInstrumentation.c"],
     sub_dir: "uprobestats",
 }
+
+bpf {
+    name: "ProcessManagement.o",
+    srcs: ["ProcessManagement.c"],
+    btf: true,
+    cflags: [
+        "-Wall",
+        "-Werror",
+    ],
+    sub_dir: "uprobestats",
+}
diff --git a/src/bpf_progs/ProcessManagement.c b/src/bpf_progs/ProcessManagement.c
new file mode 100644
index 0000000..b8623d6
--- /dev/null
+++ b/src/bpf_progs/ProcessManagement.c
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <bpf_helpers.h>
+#include <linux/bpf.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+// TODO: import this struct from generic header, access registers via generic
+// function
+struct pt_regs {
+  unsigned long regs[16];
+  unsigned long pc;
+  unsigned long pr;
+  unsigned long sr;
+  unsigned long gbr;
+  unsigned long mach;
+  unsigned long macl;
+  long tra;
+};
+
+struct SetUidTempAllowlistStateRecord {
+  __u64 uid;
+  bool onAllowlist;
+};
+
+DEFINE_BPF_RINGBUF_EXT(output_buf, struct SetUidTempAllowlistStateRecord, 4096,
+                       AID_UPROBESTATS, AID_UPROBESTATS, 0600, "", "", PRIVATE,
+                       BPFLOADER_MIN_VER, BPFLOADER_MAX_VER, LOAD_ON_ENG,
+                       LOAD_ON_USER, LOAD_ON_USERDEBUG);
+
+DEFINE_BPF_PROG("uprobe/set_uid_temp_allowlist_state", AID_UPROBESTATS,
+                AID_UPROBESTATS, BPF_KPROBE2)
+(struct pt_regs *ctx) {
+  struct SetUidTempAllowlistStateRecord *output = bpf_output_buf_reserve();
+  if (output == NULL)
+    return 1;
+  output->uid = ctx->regs[2];
+  output->onAllowlist = ctx->regs[3];
+  bpf_output_buf_submit(output);
+  return 0;
+}
+
+LICENSE("GPL");
diff --git a/src/test/test_setUidTempAllowlistStateLSP.textproto b/src/test/test_setUidTempAllowlistStateLSP.textproto
new file mode 100644
index 0000000..d7b7b6e
--- /dev/null
+++ b/src/test/test_setUidTempAllowlistStateLSP.textproto
@@ -0,0 +1,16 @@
+# proto-file: config.proto
+# proto-message: UprobestatsConfig
+
+tasks {
+    probe_configs: {
+        bpf_name: "prog_ProcessManagement_uprobe_set_uid_temp_allowlist_state"
+        file_paths: "/system/framework/oat/arm64/services.odex"
+        method_signature: "void com.android.server.am.OomAdjuster.setUidTempAllowlistStateLSP(int, boolean)"
+    }
+    bpf_maps: "map_ProcessManagement_output_buf"
+    target_process_name: "system_server"
+    duration_seconds: 180
+    statsd_logging_config {
+      atom_id: 926
+    }
+}