xref: /aosp_15_r20/external/federated-compute/fcp/jni/jni_util.h (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2022 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef FCP_JNI_JNI_UTIL_H_
18 #define FCP_JNI_JNI_UTIL_H_
19 
20 #include <jni.h>
21 
22 #include "absl/cleanup/cleanup.h"
23 #include "absl/container/fixed_array.h"
24 #include "fcp/base/monitoring.h"
25 
26 namespace fcp {
27 namespace jni {
28 
29 // Creates a JNIEnv via the passed JavaVM*, attaching the current thread if it
30 // is not already. If an attach was needed, detaches when this is destroyed.
31 //
32 // ScopedJniEnv must not be shared among threads and destructs on the same
33 // thread.
34 class ScopedJniEnv final {
35  public:
ScopedJniEnv(JavaVM * jvm)36   explicit ScopedJniEnv(JavaVM* jvm)
37       : jvm_(jvm), env_(nullptr), is_attached_(false) {
38     // We don't make any assumptions about the state of the current thread, and
39     // we want to leave it in the state we received it with respect to the
40     // JavaVm. So we only attach and detach when needed, and we always delete
41     // local references.
42     jint error = jvm_->GetEnv(reinterpret_cast<void**>(&env_), JNI_VERSION_1_2);
43     if (error != JNI_OK) {
44       error = AttachCurrentThread(jvm_, &env_);
45       FCP_CHECK(error == JNI_OK);
46       is_attached_ = true;
47     }
48   }
49 
~ScopedJniEnv()50   virtual ~ScopedJniEnv() {
51     if (is_attached_) {
52       (void)jvm_->DetachCurrentThread();
53     }
54   }
55 
env()56   JNIEnv* env() { return env_; }
57 
58  private:
59   template <typename JNIEnvArgType>
AttachCurrentThreadImpl(JavaVM * vm,jint (JavaVM::* fn)(JNIEnvArgType,void *),JNIEnv ** env)60   static jint AttachCurrentThreadImpl(JavaVM* vm,
61                                       jint (JavaVM::*fn)(JNIEnvArgType, void*),
62                                       JNIEnv** env) {
63     static_assert(std::is_same_v<JNIEnvArgType, void**> ||
64                   std::is_same_v<JNIEnvArgType, JNIEnv**>);
65     return (vm->*fn)(reinterpret_cast<JNIEnvArgType>(env), nullptr);
66   }
67 
AttachCurrentThread(JavaVM * vm,JNIEnv ** env)68   static jint AttachCurrentThread(JavaVM* vm, JNIEnv** env) {
69     // The NDK and JDK versions of jni.h disagree on the signatures for the
70     // JavaVM::AttachCurrentThread member function (the former uses 'JavaVM*'
71     // and the latter uses 'void**'). To avoid causing linker errors when the
72     // JDK's jni.h is accidentally put on the include path during an Android
73     // build, we use the indirection below when calling the function. It's not
74     // sufficient to #ifdef around __ANDROID__, because whatever is including
75     // this header file might put the JDK jni.h version on the include path.
76     return AttachCurrentThreadImpl(vm, &JavaVM::AttachCurrentThread, env);
77   }
78 
79   ScopedJniEnv(const ScopedJniEnv&) = delete;
80   void operator=(const ScopedJniEnv&) = delete;
81 
82   JavaVM* jvm_;
83   JNIEnv* env_;
84   bool is_attached_;
85 };
86 
87 // Parses a proto from a Java byte array.
88 //
89 // If any JNI calls fail, or if the parsing of the proto fails, then this
90 // FCP_CHECK-fails.
91 //
92 // This method does not call `JNIEnv::DeleteLocalRef` on the given `jbyteArray`.
93 //
94 // This is meant to be used as a convenient way to use serialized protobufs as
95 // part of a JNI API contract, since in such cases we can safely assume that the
96 // input argument will always be a valid proto (and anything else would be a
97 // programmer error).
98 template <typename MessageT>
ParseProtoFromJByteArray(JNIEnv * env,jbyteArray byte_array)99 static MessageT ParseProtoFromJByteArray(JNIEnv* env, jbyteArray byte_array) {
100   MessageT out_message;
101 
102   jsize length = env->GetArrayLength(byte_array);
103   FCP_CHECK(!env->ExceptionCheck());
104 
105   if (length == 0) {
106     return std::move(out_message);
107   }
108   // This will make a copy of the data into buffer, but generally the proto data
109   // will small enough that this shouldn't matter.
110   absl::FixedArray<jbyte> buffer(length);
111   env->GetByteArrayRegion(byte_array, 0, length, buffer.data());
112   FCP_CHECK(!env->ExceptionCheck());
113 
114   FCP_CHECK(out_message.ParseFromArray(buffer.data(), length));
115 
116   return std::move(out_message);
117 }
118 
119 // Serializes a proto to a `jbyteArray`.
120 //
121 // The caller must call `JNIEnv::DeleteLocalRef` on the returned `jbyteArray`
122 // once it is done with it.
123 //
124 // If any JNI calls fail, then this FCP_CHECK-fails.
125 template <typename MessageT>
SerializeProtoToJByteArray(JNIEnv * env,const MessageT & proto)126 static jbyteArray SerializeProtoToJByteArray(JNIEnv* env,
127                                              const MessageT& proto) {
128   int length = static_cast<int>(proto.ByteSizeLong());
129 
130   jbyteArray byte_array = env->NewByteArray(length);
131   FCP_CHECK(byte_array != nullptr);
132   FCP_CHECK(!env->ExceptionCheck());
133 
134   // This serializes into a buffer and then copies that buffer to the Java byte
135   // array. The proto data is generally small enough that this extra copy
136   // shouldn't matter.
137   absl::FixedArray<jbyte> buffer(length);
138   proto.SerializeToArray(buffer.data(), length);
139 
140   env->SetByteArrayRegion(byte_array, 0, length, buffer.data());
141   FCP_CHECK(!env->ExceptionCheck());
142 
143   return byte_array;
144 }
145 
146 // Describes the method name and JNI method signature of a Java callback.
147 struct JavaMethodSig {
148   char const* name;
149   char const* signature;
150 };
151 // Describes the field name and JNI type signature of a Java field.
152 struct JavaFieldSig {
153   char const* name;
154   char const* signature;
155 };
156 
157 // A utility for ensuring that a local JNI reference is deleted once the object
158 // goes out of scope. This class is only intended to be used inside a function
159 // body (and not to be returned or passed as an argument).
160 class LocalRefDeleter {
161  public:
LocalRefDeleter(JNIEnv * env,jobject local_ref)162   LocalRefDeleter(JNIEnv* env, jobject local_ref)
163       : env_(env), local_ref_(local_ref) {}
164   // Prevent copies & moves, to make it harder to accidentally have this object
165   // be passed as a parameter or return type.
166   LocalRefDeleter(LocalRefDeleter& other) = delete;
167   LocalRefDeleter(LocalRefDeleter&& other) = delete;
~LocalRefDeleter()168   ~LocalRefDeleter() { env_->DeleteLocalRef(local_ref_); }
169 
170  private:
171   JNIEnv* env_;
172   jobject local_ref_;
173 };
174 
175 }  // namespace jni
176 }  // namespace fcp
177 
178 #endif  // FCP_JNI_JNI_UTIL_H_
179