3 files changed, 91 insertions, 22 deletions

diff --git a/libs/samba/src/main/cpp/jni.cpp b/libs/samba/src/main/cpp/jni.cpp
index aac1d28..5a69dc5 100644
--- a/libs/samba/src/main/cpp/jni.cpp
+++ b/libs/samba/src/main/cpp/jni.cpp
@@ -79,6 +79,22 @@ void _abort_with_exception(const char* file, int line, JNIEnv* env) {
   }
 }
 
+JNIEnv* NonFatalAttachCurrentThread() {
+#ifdef ANDROID
+  JNIEnv* env;
+  auto ret = g_vm->AttachCurrentThread(&env, nullptr);
+  if (ret == JNI_OK)
+    return env;
+  return nullptr;
+#else
+  void* v_env;
+  auto ret = g_vm->AttachCurrentThread(&v_env, nullptr);
+  if (ret == JNI_OK)
+    return reinterpret_cast<JNIEnv*>(v_env);
+  return nullptr;
+#endif
+}
+
 }  // namespace internal
 
 JNIEnv* AttachCurrentThread() {
@@ -99,6 +115,7 @@ JNIEnv* OnLoad(JavaVM* vm) {
   void* v_env;
   auto ret = vm->GetEnv(&v_env, JNI_VERSION);
   ABORT_IF_NOT_OK(ret);
+  g_vm = vm;
   return reinterpret_cast<JNIEnv*>(v_env);
 }
 
diff --git a/libs/samba/src/main/cpp/jni.hpp b/libs/samba/src/main/cpp/jni.hpp
index 450fa71..3b5078a 100644
--- a/libs/samba/src/main/cpp/jni.hpp
+++ b/libs/samba/src/main/cpp/jni.hpp
@@ -15,28 +15,30 @@ namespace internal {
 void _abort_if_not_ok(const char *file, int line, jint ret);
 void _abort_with_exception(const char* file, int line, JNIEnv* env);
 
+JNIEnv* NonFatalAttachCurrentThread();
+
 }  // namespace internal
 
 template<class T>
 class Ref {
  public:
-  constexpr Ref() : env_(nullptr), ptr_(0) {}
+  constexpr Ref() : ptr_(nullptr) {}
+  constexpr Ref(std::nullptr_t) : ptr_(nullptr) {}
   Ref(const Ref<T>&) = delete;
 
-  [[nodiscard]] JNIEnv* env() const { return env_; }
   [[nodiscard]] T get() const { return ptr_; }
   [[nodiscard]] T release() {
     auto ret = release_to_local();
-    ptr_ = 0;
+    ptr_ = nullptr;
     return ret;
   }
 
   void reset() {
     del();
-    ptr_ = 0;
+    ptr_ = nullptr;
   }
 
-  explicit operator bool() const { return ptr_ != 0; }
+  explicit operator bool() const { return ptr_ != nullptr; }
 
  protected:
   Ref(JNIEnv* env, T ptr): env_(env), ptr_(ptr) {}
@@ -49,9 +51,17 @@ class Ref {
   T ptr_;
 };
 
+constexpr jint JNI_VERSION = JNI_VERSION_1_2;
+
+JNIEnv* AttachCurrentThread();
+
+JNIEnv* OnLoad(JavaVM* vm);
+
 template<class T>
 class LocalRef : public Ref<T> {
  public:
+  constexpr LocalRef(): Ref<T>() {}
+  constexpr LocalRef(std::nullptr_t): Ref<T>() {}
   LocalRef(JNIEnv* env, T ptr): Ref<T>(env, ptr) {}
   ~LocalRef() override { free(); }
 
@@ -59,6 +69,13 @@ class LocalRef : public Ref<T> {
 
   LocalRef& operator=(const Ref<T>& other) = delete;
 
+  LocalRef& operator=(LocalRef<T> &&ref) noexcept {
+    free();
+    this->env_ = ref.env_;
+    this->ptr_ = ref.release();
+    return *this;
+  }
+
  protected:
   T release_to_local() override { return this->ptr_; }
   void del() override { free(); }
@@ -82,7 +99,7 @@ class ParamRef : public Ref<T> {
   T release_to_local() override {
     if (this->ptr_)
       return static_cast<T>(this->env_->NewLocalRef(static_cast<jobject>(this->ptr_)));
-    return 0;
+    return nullptr;
   }
   void del() override {}
 };
@@ -90,43 +107,67 @@ class ParamRef : public Ref<T> {
 template<class T>
 class GlobalRef : public Ref<T> {
  public:
-  GlobalRef(JNIEnv* env, T ptr) : Ref<T>(env, ptr ? static_cast<T>(env->NewGlobalRef(static_cast<jobject>(ptr))) : 0) {}
-  explicit GlobalRef(const Ref<T>& other) : Ref<T>(other.env(), other ? static_cast<T>(other.env()->NewGlobalRef(static_cast<jobject>(other.get()))) : 0) {}
+  constexpr GlobalRef() : Ref<T>() {}
+  constexpr GlobalRef(std::nullptr_t) : Ref<T>() {}
+  GlobalRef(JNIEnv* env, T ptr) : Ref<T>() {
+    if (ptr) {
+      if (!env)
+        env = AttachCurrentThread();
+      this->ptr_ = static_cast<T>(env->NewGlobalRef(static_cast<jobject>(ptr)));
+    }
+  }
+  GlobalRef(const Ref<T>& other) : Ref<T>() {
+    if (other) {
+      auto* env = AttachCurrentThread();
+      this->ptr_ = static_cast<T>(env->NewGlobalRef(static_cast<jobject>(other.get())));
+    }
+  }
+  GlobalRef(GlobalRef<T>&& other) noexcept : Ref<T>(nullptr, other.ptr_) {
+    other.ptr_ = nullptr;
+    return *this;
+  }
 
   ~GlobalRef() override { free(); }
 
   GlobalRef& operator=(const Ref<T>& other) {
     free();
-    this->env_ = other.env();
-    this->ptr_ = other ? static_cast<T>(this->env_->NewGlobalRef(static_cast<jobject>(other.get()))) : 0;
+    if (other) {
+      auto* env = AttachCurrentThread();
+      this->ptr_ = other ? static_cast<T>(env->NewGlobalRef(static_cast<jobject>(other.get()))) : nullptr;
+    } else {
+      this->ptr_ = nullptr;
+    }
     return *this;
   }
 
+  GlobalRef& operator=(GlobalRef<T>&& other) {
+    free();
+    this->ptr_ = other.ptr_;
+    other.ptr_ = nullptr;
+  }
+
  protected:
   T release_to_local() override {
     if (this->ptr_) {
-      auto ret = static_cast<T>(this->env_->NewLocalRef(
-          static_cast<jobject>(this->ptr_)));
-      free();
+      auto* env = AttachCurrentThread();
+      auto ret = static_cast<T>(env->NewLocalRef(static_cast<jobject>(this->ptr_)));
+      env->DeleteGlobalRef(static_cast<jobject>(this->ptr_));
       return ret;
     }
-    return 0;
+    return nullptr;
   }
   void del() override { free(); }
 
  private:
   void free() {
-    if (this->ptr_)
-      this->env_->DeleteGlobalRef(this->ptr_);
+    if (this->ptr_) {
+      auto* env = internal::NonFatalAttachCurrentThread();
+      if (env)
+        env->DeleteGlobalRef(static_cast<jobject>(this->ptr_));
+    }
   }
 };
 
-constexpr jint JNI_VERSION = JNI_VERSION_1_2;
-
-JNIEnv* AttachCurrentThread();
-
-JNIEnv* OnLoad(JavaVM* vm);
-
 LocalRef<jclass> FindClass(JNIEnv *env, const char *name);
 
 LocalRef<jthrowable> ExceptionOccurred(JNIEnv* env);
diff --git a/libs/samba/src/main/cpp/samba.cpp b/libs/samba/src/main/cpp/samba.cpp
index 1d1d9d2..1e6b425 100644
--- a/libs/samba/src/main/cpp/samba.cpp
+++ b/libs/samba/src/main/cpp/samba.cpp
@@ -260,6 +260,12 @@ void RegisterSamba(JNIEnv* env) {
   g_DirEntryClass = dir_entry_clazz;
 }
 
+void UnregisterSamba() {
+  g_CreateDirEntry = 0;
+  g_NativeSambaClass.reset();
+  g_DirEntryClass.reset();
+}
+
 jni::LocalRef<jobject> CreateDirEntry(JNIEnv* env, const std::string& name, const smb2_stat_64& stat) {
   auto j_name = jni::UTF8ToString(env, name);
   // Kotlin size casts Long to ULong
@@ -292,3 +298,8 @@ jint JNI_OnLoad(JavaVM *vm, void *reserved) {
 
   return jni::JNI_VERSION;
 }
+
+void JNI_OnUnload(JavaVM *vm, void *reserved) {
+  // Not called on Android (or in general), but if it where, this would be the place to unregister.
+  UnregisterSamba();
+}