ART: Add CheckJNI lock checking

JNI MonitorEnter and MonitorExit have similar rules to structured
locking. Count locks in CheckJNI mode.

Bug: 23502994
Change-Id: Ie3f53d3aa669a6bd0c7153c50c168116b43764d9
diff --git a/runtime/check_jni.cc b/runtime/check_jni.cc
index b6ad547..beabce3 100644
--- a/runtime/check_jni.cc
+++ b/runtime/check_jni.cc
@@ -2463,6 +2463,9 @@
     ScopedCheck sc(kFlag_Default, __FUNCTION__);
     JniValueType args[2] = {{.E = env}, {.L = obj}};
     if (sc.Check(soa, true, "EL", args)) {
+      if (obj != nullptr) {
+        down_cast<JNIEnvExt*>(env)->RecordMonitorEnter(obj);
+      }
       JniValueType result;
       result.i = baseEnv(env)->MonitorEnter(env, obj);
       if (sc.Check(soa, false, "i", &result)) {
@@ -2477,6 +2480,9 @@
     ScopedCheck sc(kFlag_ExcepOkay, __FUNCTION__);
     JniValueType args[2] = {{.E = env}, {.L = obj}};
     if (sc.Check(soa, true, "EL", args)) {
+      if (obj != nullptr) {
+        down_cast<JNIEnvExt*>(env)->CheckMonitorRelease(obj);
+      }
       JniValueType result;
       result.i = baseEnv(env)->MonitorExit(env, obj);
       if (sc.Check(soa, false, "i", &result)) {
diff --git a/runtime/class_linker.h b/runtime/class_linker.h
index 739403f..7f3e938 100644
--- a/runtime/class_linker.h
+++ b/runtime/class_linker.h
@@ -883,6 +883,7 @@
   friend class ImageWriter;  // for GetClassRoots
   friend class ImageDumper;  // for FindOpenedOatFileFromOatLocation
   friend class JniCompilerTest;  // for GetRuntimeQuickGenericJniStub
+  friend class JniInternalTest;  // for GetRuntimeQuickGenericJniStub
   ART_FRIEND_TEST(mirror::DexCacheTest, Open);  // for AllocDexCache
 
   DISALLOW_COPY_AND_ASSIGN(ClassLinker);
diff --git a/runtime/entrypoints/quick/quick_jni_entrypoints.cc b/runtime/entrypoints/quick/quick_jni_entrypoints.cc
index fc5c52e..58f256a 100644
--- a/runtime/entrypoints/quick/quick_jni_entrypoints.cc
+++ b/runtime/entrypoints/quick/quick_jni_entrypoints.cc
@@ -65,6 +65,9 @@
 static void PopLocalReferences(uint32_t saved_local_ref_cookie, Thread* self)
     SHARED_REQUIRES(Locks::mutator_lock_) {
   JNIEnvExt* env = self->GetJniEnv();
+  if (UNLIKELY(env->check_jni)) {
+    env->CheckNoHeldMonitors();
+  }
   env->locals.SetSegmentState(env->local_ref_cookie);
   env->local_ref_cookie = saved_local_ref_cookie;
   self->PopHandleScope();
diff --git a/runtime/jni_env_ext.cc b/runtime/jni_env_ext.cc
index b18b430..4104d7a 100644
--- a/runtime/jni_env_ext.cc
+++ b/runtime/jni_env_ext.cc
@@ -16,10 +16,17 @@
 
 #include "jni_env_ext.h"
 
+#include <algorithm>
+#include <vector>
+
 #include "check_jni.h"
 #include "indirect_reference_table.h"
 #include "java_vm_ext.h"
 #include "jni_internal.h"
+#include "lock_word.h"
+#include "mirror/object-inl.h"
+#include "nth_caller_visitor.h"
+#include "thread-inl.h"
 
 namespace art {
 
@@ -63,14 +70,14 @@
 JNIEnvExt::~JNIEnvExt() {
 }
 
-jobject JNIEnvExt::NewLocalRef(mirror::Object* obj) SHARED_REQUIRES(Locks::mutator_lock_) {
+jobject JNIEnvExt::NewLocalRef(mirror::Object* obj) {
   if (obj == nullptr) {
     return nullptr;
   }
   return reinterpret_cast<jobject>(locals.Add(local_ref_cookie, obj));
 }
 
-void JNIEnvExt::DeleteLocalRef(jobject obj) SHARED_REQUIRES(Locks::mutator_lock_) {
+void JNIEnvExt::DeleteLocalRef(jobject obj) {
   if (obj != nullptr) {
     locals.Remove(local_ref_cookie, reinterpret_cast<IndirectRef>(obj));
   }
@@ -86,14 +93,14 @@
   monitors.Dump(os);
 }
 
-void JNIEnvExt::PushFrame(int capacity) SHARED_REQUIRES(Locks::mutator_lock_) {
+void JNIEnvExt::PushFrame(int capacity) {
   UNUSED(capacity);  // cpplint gets confused with (int) and thinks its a cast.
   // TODO: take 'capacity' into account.
   stacked_local_ref_cookies.push_back(local_ref_cookie);
   local_ref_cookie = locals.GetSegmentState();
 }
 
-void JNIEnvExt::PopFrame() SHARED_REQUIRES(Locks::mutator_lock_) {
+void JNIEnvExt::PopFrame() {
   locals.SetSegmentState(local_ref_cookie);
   local_ref_cookie = stacked_local_ref_cookies.back();
   stacked_local_ref_cookies.pop_back();
@@ -104,4 +111,118 @@
                 IndirectReferenceTable::SegmentStateOffset().Int32Value());
 }
 
+// Use some defining part of the caller's frame as the identifying mark for the JNI segment.
+static uintptr_t GetJavaCallFrame(Thread* self) SHARED_REQUIRES(Locks::mutator_lock_) {
+  NthCallerVisitor zeroth_caller(self, 0, false);
+  zeroth_caller.WalkStack();
+  if (zeroth_caller.caller == nullptr) {
+    // No Java code, must be from pure native code.
+    return 0;
+  } else if (zeroth_caller.GetCurrentQuickFrame() == nullptr) {
+    // Shadow frame = interpreter. Use the actual shadow frame's address.
+    DCHECK(zeroth_caller.GetCurrentShadowFrame() != nullptr);
+    return reinterpret_cast<uintptr_t>(zeroth_caller.GetCurrentShadowFrame());
+  } else {
+    // Quick frame = compiled code. Use the bottom of the frame.
+    return reinterpret_cast<uintptr_t>(zeroth_caller.GetCurrentQuickFrame());
+  }
+}
+
+void JNIEnvExt::RecordMonitorEnter(jobject obj) {
+  locked_objects_.push_back(std::make_pair(GetJavaCallFrame(self), obj));
+}
+
+static std::string ComputeMonitorDescription(Thread* self,
+                                             jobject obj) SHARED_REQUIRES(Locks::mutator_lock_) {
+  mirror::Object* o = self->DecodeJObject(obj);
+  if ((o->GetLockWord(false).GetState() == LockWord::kThinLocked) &&
+      Locks::mutator_lock_->IsExclusiveHeld(self)) {
+    // Getting the identity hashcode here would result in lock inflation and suspension of the
+    // current thread, which isn't safe if this is the only runnable thread.
+    return StringPrintf("<@addr=0x%" PRIxPTR "> (a %s)",
+                        reinterpret_cast<intptr_t>(o),
+                        PrettyTypeOf(o).c_str());
+  } else {
+    // IdentityHashCode can cause thread suspension, which would invalidate o if it moved. So
+    // we get the pretty type before we call IdentityHashCode.
+    const std::string pretty_type(PrettyTypeOf(o));
+    return StringPrintf("<0x%08x> (a %s)", o->IdentityHashCode(), pretty_type.c_str());
+  }
+}
+
+static void RemoveMonitors(Thread* self,
+                           uintptr_t frame,
+                           ReferenceTable* monitors,
+                           std::vector<std::pair<uintptr_t, jobject>>* locked_objects)
+    SHARED_REQUIRES(Locks::mutator_lock_) {
+  auto kept_end = std::remove_if(
+      locked_objects->begin(),
+      locked_objects->end(),
+      [self, frame, monitors](const std::pair<uintptr_t, jobject>& pair)
+          SHARED_REQUIRES(Locks::mutator_lock_) {
+        if (frame == pair.first) {
+          mirror::Object* o = self->DecodeJObject(pair.second);
+          monitors->Remove(o);
+          return true;
+        }
+        return false;
+      });
+  locked_objects->erase(kept_end, locked_objects->end());
+}
+
+void JNIEnvExt::CheckMonitorRelease(jobject obj) {
+  uintptr_t current_frame = GetJavaCallFrame(self);
+  std::pair<uintptr_t, jobject> exact_pair = std::make_pair(current_frame, obj);
+  auto it = std::find(locked_objects_.begin(), locked_objects_.end(), exact_pair);
+  bool will_abort = false;
+  if (it != locked_objects_.end()) {
+    locked_objects_.erase(it);
+  } else {
+    // Check whether this monitor was locked in another JNI "session."
+    mirror::Object* mirror_obj = self->DecodeJObject(obj);
+    for (std::pair<uintptr_t, jobject>& pair : locked_objects_) {
+      if (self->DecodeJObject(pair.second) == mirror_obj) {
+        std::string monitor_descr = ComputeMonitorDescription(self, pair.second);
+        vm->JniAbortF("<JNI MonitorExit>",
+                      "Unlocking monitor that wasn't locked here: %s",
+                      monitor_descr.c_str());
+        will_abort = true;
+        break;
+      }
+    }
+  }
+
+  // When we abort, also make sure that any locks from the current "session" are removed from
+  // the monitors table, otherwise we may visit local objects in GC during abort (which won't be
+  // valid anymore).
+  if (will_abort) {
+    RemoveMonitors(self, current_frame, &monitors, &locked_objects_);
+  }
+}
+
+void JNIEnvExt::CheckNoHeldMonitors() {
+  uintptr_t current_frame = GetJavaCallFrame(self);
+  // The locked_objects_ are grouped by their stack frame component, as this enforces structured
+  // locking, and the groups form a stack. So the current frame entries are at the end. Check
+  // whether the vector is empty, and when there are elements, whether the last element belongs
+  // to this call - this signals that there are unlocked monitors.
+  if (!locked_objects_.empty()) {
+    std::pair<uintptr_t, jobject>& pair = locked_objects_[locked_objects_.size() - 1];
+    if (pair.first == current_frame) {
+      std::string monitor_descr = ComputeMonitorDescription(self, pair.second);
+      vm->JniAbortF("<JNI End>",
+                    "Still holding a locked object on JNI end: %s",
+                    monitor_descr.c_str());
+      // When we abort, also make sure that any locks from the current "session" are removed from
+      // the monitors table, otherwise we may visit local objects in GC during abort.
+      RemoveMonitors(self, current_frame, &monitors, &locked_objects_);
+    } else if (kIsDebugBuild) {
+      // Make sure there are really no other entries and our checking worked as expected.
+      for (std::pair<uintptr_t, jobject>& check_pair : locked_objects_) {
+        CHECK_NE(check_pair.first, current_frame);
+      }
+    }
+  }
+}
+
 }  // namespace art
diff --git a/runtime/jni_env_ext.h b/runtime/jni_env_ext.h
index 9b55536..3828ff0 100644
--- a/runtime/jni_env_ext.h
+++ b/runtime/jni_env_ext.h
@@ -43,8 +43,8 @@
 
   void SetCheckJniEnabled(bool enabled);
 
-  void PushFrame(int capacity);
-  void PopFrame();
+  void PushFrame(int capacity) SHARED_REQUIRES(Locks::mutator_lock_);
+  void PopFrame() SHARED_REQUIRES(Locks::mutator_lock_);
 
   template<typename T>
   T AddLocalReference(mirror::Object* obj)
@@ -89,10 +89,27 @@
   // Used by -Xcheck:jni.
   const JNINativeInterface* unchecked_functions;
 
+  // Functions to keep track of monitor lock and unlock operations. Used to ensure proper locking
+  // rules in CheckJNI mode.
+
+  // Record locking of a monitor.
+  void RecordMonitorEnter(jobject obj) SHARED_REQUIRES(Locks::mutator_lock_);
+
+  // Check the release, that is, that the release is performed in the same JNI "segment."
+  void CheckMonitorRelease(jobject obj) SHARED_REQUIRES(Locks::mutator_lock_);
+
+  // Check that no monitors are held that have been acquired in this JNI "segment."
+  void CheckNoHeldMonitors() SHARED_REQUIRES(Locks::mutator_lock_);
+
  private:
   // The constructor should not be called directly. It may leave the object in an erronuous state,
   // and the result needs to be checked.
   JNIEnvExt(Thread* self, JavaVMExt* vm);
+
+  // All locked objects, with the (Java caller) stack frame that locked them. Used in CheckJNI
+  // to ensure that only monitors locked in this native frame are being unlocked, and that at
+  // the end all are unlocked.
+  std::vector<std::pair<uintptr_t, jobject>> locked_objects_;
 };
 
 // Used to save and restore the JNIEnvExt state when not going through code created by the JNI
diff --git a/runtime/jni_internal_test.cc b/runtime/jni_internal_test.cc
index 2a0cb28..41b368e 100644
--- a/runtime/jni_internal_test.cc
+++ b/runtime/jni_internal_test.cc
@@ -607,11 +607,64 @@
     EXPECT_EQ(check_jni, vm_->SetCheckJniEnabled(old_check_jni));
   }
 
+  void SetUpForTest(bool direct, const char* method_name, const char* method_sig,
+                    void* native_fnptr) {
+    // Initialize class loader and set generic JNI entrypoint.
+    // Note: this code is adapted from the jni_compiler_test, and taken with minimal modifications.
+    if (!runtime_->IsStarted()) {
+      {
+        ScopedObjectAccess soa(Thread::Current());
+        class_loader_ = LoadDex("MyClassNatives");
+        StackHandleScope<1> hs(soa.Self());
+        Handle<mirror::ClassLoader> loader(
+            hs.NewHandle(soa.Decode<mirror::ClassLoader*>(class_loader_)));
+        mirror::Class* c = class_linker_->FindClass(soa.Self(), "LMyClassNatives;", loader);
+        const auto pointer_size = class_linker_->GetImagePointerSize();
+        ArtMethod* method = direct ? c->FindDirectMethod(method_name, method_sig, pointer_size) :
+            c->FindVirtualMethod(method_name, method_sig, pointer_size);
+        ASSERT_TRUE(method != nullptr) << method_name << " " << method_sig;
+        method->SetEntryPointFromQuickCompiledCode(class_linker_->GetRuntimeQuickGenericJniStub());
+      }
+      // Start runtime.
+      Thread::Current()->TransitionFromSuspendedToRunnable();
+      bool started = runtime_->Start();
+      CHECK(started);
+    }
+    // JNI operations after runtime start.
+    env_ = Thread::Current()->GetJniEnv();
+    jklass_ = env_->FindClass("MyClassNatives");
+    ASSERT_TRUE(jklass_ != nullptr) << method_name << " " << method_sig;
+
+    if (direct) {
+      jmethod_ = env_->GetStaticMethodID(jklass_, method_name, method_sig);
+    } else {
+      jmethod_ = env_->GetMethodID(jklass_, method_name, method_sig);
+    }
+    ASSERT_TRUE(jmethod_ != nullptr) << method_name << " " << method_sig;
+
+    if (native_fnptr != nullptr) {
+      JNINativeMethod methods[] = { { method_name, method_sig, native_fnptr } };
+      ASSERT_EQ(JNI_OK, env_->RegisterNatives(jklass_, methods, 1))
+          << method_name << " " << method_sig;
+    } else {
+      env_->UnregisterNatives(jklass_);
+    }
+
+    jmethodID constructor = env_->GetMethodID(jklass_, "<init>", "()V");
+    jobj_ = env_->NewObject(jklass_, constructor);
+    ASSERT_TRUE(jobj_ != nullptr) << method_name << " " << method_sig;
+  }
+
   JavaVMExt* vm_;
   JNIEnv* env_;
   jclass aioobe_;
   jclass ase_;
   jclass sioobe_;
+
+  jclass jklass_;
+  jobject jobj_;
+  jobject class_loader_;
+  jmethodID jmethod_;
 };
 
 TEST_F(JniInternalTest, AllocObject) {
@@ -2111,4 +2164,38 @@
   }
 }
 
+void Java_MyClassNatives_foo_exit(JNIEnv* env, jobject thisObj) {
+  // Release the monitor on self. This should trigger an abort.
+  env->MonitorExit(thisObj);
+}
+
+TEST_F(JniInternalTest, MonitorExitLockedInDifferentCall) {
+  SetUpForTest(false, "foo", "()V", reinterpret_cast<void*>(&Java_MyClassNatives_foo_exit));
+  ASSERT_NE(jobj_, nullptr);
+
+  env_->MonitorEnter(jobj_);
+  EXPECT_FALSE(env_->ExceptionCheck());
+
+  CheckJniAbortCatcher check_jni_abort_catcher;
+  env_->CallNonvirtualVoidMethod(jobj_, jklass_, jmethod_);
+  check_jni_abort_catcher.Check("Unlocking monitor that wasn't locked here");
+}
+
+void Java_MyClassNatives_foo_enter_no_exit(JNIEnv* env, jobject thisObj) {
+  // Acquire but don't release the monitor on self. This should trigger an abort on return.
+  env->MonitorEnter(thisObj);
+}
+
+TEST_F(JniInternalTest, MonitorExitNotAllUnlocked) {
+  SetUpForTest(false,
+               "foo",
+               "()V",
+               reinterpret_cast<void*>(&Java_MyClassNatives_foo_enter_no_exit));
+  ASSERT_NE(jobj_, nullptr);
+
+  CheckJniAbortCatcher check_jni_abort_catcher;
+  env_->CallNonvirtualVoidMethod(jobj_, jklass_, jmethod_);
+  check_jni_abort_catcher.Check("Still holding a locked object on JNI end");
+}
+
 }  // namespace art