ART: Add JNI function table manipulation

Add support for a function table override. This will override the
decision between the regular and the CheckJNI function tables, if
set.

Bug: 34343708
Test: m test-art-host-gtest-jni_internal_test
Change-Id: I0e95b0cbd21f4efdcd8c3d312781d9aeeff54a1e
diff --git a/runtime/jni_env_ext.cc b/runtime/jni_env_ext.cc
index 5a3fafa..0148a1c 100644
--- a/runtime/jni_env_ext.cc
+++ b/runtime/jni_env_ext.cc
@@ -29,6 +29,7 @@
 #include "mirror/object-inl.h"
 #include "nth_caller_visitor.h"
 #include "thread-inl.h"
+#include "thread_list.h"
 
 namespace art {
 
@@ -37,6 +38,8 @@
 static constexpr size_t kMonitorsInitial = 32;  // Arbitrary.
 static constexpr size_t kMonitorsMax = 4096;  // Arbitrary sanity check.
 
+const JNINativeInterface* JNIEnvExt::table_override_ = nullptr;
+
 // Checking "locals" requires the mutator lock, but at creation time we're really only interested
 // in validity, which isn't changing. To avoid grabbing the mutator lock, factored out and tagged
 // with NO_THREAD_SAFETY_ANALYSIS.
@@ -78,10 +81,10 @@
       runtime_deleted(false),
       critical(0),
       monitors("monitors", kMonitorsInitial, kMonitorsMax) {
-  functions = unchecked_functions = GetJniNativeInterface();
-  if (vm->IsCheckJniEnabled()) {
-    SetCheckJniEnabled(true);
-  }
+  MutexLock mu(Thread::Current(), *Locks::jni_function_table_lock_);
+  check_jni = vm->IsCheckJniEnabled();
+  functions = GetFunctionTable(check_jni);
+  unchecked_functions = GetJniNativeInterface();
 }
 
 void JNIEnvExt::SetFunctionsToRuntimeShutdownFunctions() {
@@ -107,7 +110,12 @@
 
 void JNIEnvExt::SetCheckJniEnabled(bool enabled) {
   check_jni = enabled;
-  functions = enabled ? GetCheckJniNativeInterface() : GetJniNativeInterface();
+  MutexLock mu(Thread::Current(), *Locks::jni_function_table_lock_);
+  functions = GetFunctionTable(enabled);
+  // Check whether this is a no-op because of override.
+  if (enabled && JNIEnvExt::table_override_ != nullptr) {
+    LOG(WARNING) << "Enabling CheckJNI after a JNIEnv function table override is not functional.";
+  }
 }
 
 void JNIEnvExt::DumpReferenceTables(std::ostream& os) {
@@ -269,4 +277,33 @@
   }
 }
 
+static void ThreadResetFunctionTable(Thread* thread, void* arg ATTRIBUTE_UNUSED)
+    REQUIRES(Locks::jni_function_table_lock_) {
+  JNIEnvExt* env = thread->GetJniEnv();
+  bool check_jni = env->check_jni;
+  env->functions = JNIEnvExt::GetFunctionTable(check_jni);
+}
+
+void JNIEnvExt::SetTableOverride(const JNINativeInterface* table_override) {
+  MutexLock mu(Thread::Current(), *Locks::thread_list_lock_);
+  MutexLock mu2(Thread::Current(), *Locks::jni_function_table_lock_);
+
+  JNIEnvExt::table_override_ = table_override;
+
+  // See if we have a runtime. Note: we cannot run other code (like JavaVMExt's CheckJNI install
+  // code), as we'd have to recursively lock the mutex.
+  Runtime* runtime = Runtime::Current();
+  if (runtime != nullptr) {
+    runtime->GetThreadList()->ForEach(ThreadResetFunctionTable, nullptr);
+  }
+}
+
+const JNINativeInterface* JNIEnvExt::GetFunctionTable(bool check_jni) {
+  const JNINativeInterface* override = JNIEnvExt::table_override_;
+  if (override != nullptr) {
+    return override;
+  }
+  return check_jni ? GetCheckJniNativeInterface() : GetJniNativeInterface();
+}
+
 }  // namespace art