uwb(jni-rust): Cache classloader for creating uwb class objects

This is needed for invoking jni callbacks from the native thread
which needs to pass uwb class objects.

Use the infra to implement the range data notifications for short
mac address and extended mac address.

Bug: 197341298
Test: Manual verification by invoking the callback (ag/16659743).
Change-Id: I6df679c09391bde43d1aba2c5ab9e38e995f5c92

1 files changed, 467 insertions, 34 deletions

diff --git a/src/rust/event_manager/mod.rs b/src/rust/event_manager/mod.rs
index c83fb68..2499f45 100644
--- a/src/rust/event_manager/mod.rs
+++ b/src/rust/event_manager/mod.rs
@@ -15,17 +15,26 @@
  */
 
 use crate::uci::uci_hrcv::UciNotification;
-use jni::errors::Result;
-use jni::objects::{GlobalRef, JObject, JValue};
+use jni::errors::{Error, Result};
+use jni::objects::{GlobalRef, JClass, JMethodID, JObject, JValue, JValue::Void};
+use jni::signature::JavaType;
+use jni::sys::jobjectArray;
 use jni::{AttachGuard, JNIEnv, JavaVM};
 use log::error;
 use num_traits::ToPrimitive;
 use std::convert::TryInto;
+use std::vec::Vec;
 use uwb_uci_packets::{
-    DeviceStatusNtfPacket, GenericErrorPacket, SessionStatusNtfPacket,
-    ShortMacTwoWayRangeDataNtfPacket,
+    DeviceStatusNtfPacket, ExtendedAddressTwoWayRangingMeasurement,
+    ExtendedMacTwoWayRangeDataNtfPacket, GenericErrorPacket, RangeDataNtfPacket,
+    SessionStatusNtfPacket, ShortAddressTwoWayRangingMeasurement, ShortMacTwoWayRangeDataNtfPacket,
 };
 
+const UWB_RANGING_DATA_CLASS: &str = "com/android/uwb/data/UwbRangingData";
+const UWB_TWO_WAY_MEASUREMENT_CLASS: &str = "com/android/uwb/data/UwbTwoWayMeasurement";
+const SHORT_MAC_ADDRESS_LEN: usize = 2;
+const EXTENDED_MAC_ADDRESS_LEN: usize = 8;
+
 // TODO: Reconsider the best way to cache the JNIEnv.  We currently attach and detach for every
 // call, which the documentation warns could be expensive.  We could attach the thread permanently,
 // but that would not allow us to detach when we drop this structure.  We could cache the
@@ -40,6 +49,8 @@ use uwb_uci_packets::{
 pub struct EventManager {
     jvm: JavaVM,
     obj: GlobalRef,
+    // cache used to lookup uwb classes in callback.
+    class_loader_obj: GlobalRef,
 }
 
 impl EventManager {
@@ -47,58 +58,480 @@ impl EventManager {
     pub fn new(env: JNIEnv, obj: JObject) -> Result<Self> {
         let jvm = env.get_java_vm()?;
         let obj = env.new_global_ref(obj)?;
-        Ok(EventManager { jvm, obj })
+        let class_loader_obj = EventManager::get_classloader_obj(&env)?;
+        let class_loader_obj = env.new_global_ref(class_loader_obj)?;
+        Ok(EventManager { jvm, obj, class_loader_obj })
     }
 
-    pub fn device_status_notification_received(&self, data: DeviceStatusNtfPacket) -> Result<()> {
-        let state = data.get_device_state().to_u8().expect("Failed converting device_state to u8");
-        let env = self.jvm.attach_current_thread()?;
-        let result = env.call_method(
+    fn get_classloader_obj<'a>(env: &'a JNIEnv) -> Result<JObject<'a>> {
+        // Use UwbRangingData class to find the classloader used by the java service.
+        let ranging_data_class = env.find_class(&UWB_RANGING_DATA_CLASS)?;
+        let ranging_data_class_class = env.get_object_class(ranging_data_class)?;
+        let get_class_loader_method = env.get_method_id(
+            ranging_data_class_class,
+            "getClassLoader",
+            "()Ljava/lang/ClassLoader;",
+        )?;
+        let class_loader = env.call_method_unchecked(
+            ranging_data_class,
+            get_class_loader_method,
+            JavaType::Object("java/lang/ClassLoader".into()),
+            &[Void],
+        )?;
+        class_loader.l()
+    }
+
+    fn find_class<'a>(&'a self, env: &'a JNIEnv, class_name: &'a str) -> Result<JClass<'a>> {
+        let class_value = env.call_method(
+            self.class_loader_obj.as_obj(),
+            "findClass",
+            "(Ljava/lang/String;)Ljava/lang/Class;",
+            &[JValue::Object(JObject::from(env.new_string(class_name)?))],
+        )?;
+        class_value.l().map(|value| JClass::from(value))
+    }
+
+    fn handle_device_status_notification_received(
+        &self,
+        env: &JNIEnv,
+        data: DeviceStatusNtfPacket,
+    ) -> Result<()> {
+        let state =
+            data.get_device_state().to_i32().expect("Failed converting device_state to i32");
+        env.call_method(
             self.obj.as_obj(),
             "onDeviceStatusNotificationReceived",
             "(I)V",
-            &[JValue::Int(state.try_into().expect("Could not convert device_state"))],
-        );
-        self.cleanup_and_return(env, result)
+            &[JValue::Int(state)],
+        )
+        .map(|_| ()) // drop void method return
     }
 
-    pub fn session_status_notification_received(&self, data: SessionStatusNtfPacket) -> Result<()> {
+    pub fn device_status_notification_received(&self, data: DeviceStatusNtfPacket) -> Result<()> {
+        let env = self.jvm.attach_current_thread()?;
+        let result = self.handle_device_status_notification_received(&env, data);
+        self.clear_exception(env);
+        result
+    }
+
+    fn handle_session_status_notification_received(
+        &self,
+        env: &JNIEnv,
+        data: SessionStatusNtfPacket,
+    ) -> Result<()> {
         let session_id =
-            data.get_session_id().to_u32().expect("Failed converting session_id to u32");
+            data.get_session_id().to_i64().expect("Failed converting session_id to i64");
         let state =
-            data.get_session_state().to_u8().expect("Failed converting session_state to u8");
+            data.get_session_state().to_i32().expect("Failed converting session_state to i32");
         let reason_code =
-            data.get_reason_code().to_u8().expect("Failed coverting reason_code to u32");
-        let env = self.jvm.attach_current_thread()?;
-        let result = env.call_method(
+            data.get_reason_code().to_i32().expect("Failed coverting reason_code to i32");
+        env.call_method(
             self.obj.as_obj(),
             "onSessionStatusNotificationReceived",
             "(JII)V",
-            &[
-                JValue::Long(session_id.try_into().expect("Could not convert session_id")),
-                JValue::Int(state.try_into().expect("Could not convert session_state")),
-                JValue::Int(reason_code.try_into().expect("Could not convert reason_code")),
-            ],
-        );
-        self.cleanup_and_return(env, result)
+            &[JValue::Long(session_id), JValue::Int(state), JValue::Int(reason_code)],
+        )
+        .map(|_| ()) // drop void method return
     }
 
-    pub fn core_generic_error_notification_received(&self, data: GenericErrorPacket) -> Result<()> {
-        let status = data.get_status().to_u8().expect("Failed converting status to u8");
+    pub fn session_status_notification_received(&self, data: SessionStatusNtfPacket) -> Result<()> {
         let env = self.jvm.attach_current_thread()?;
-        let result = env.call_method(
+        let result = self.handle_session_status_notification_received(&env, data);
+        self.clear_exception(env);
+        result
+    }
+
+    fn handle_core_generic_error_notification_received(
+        &self,
+        env: &JNIEnv,
+        data: GenericErrorPacket,
+    ) -> Result<()> {
+        let status = data.get_status().to_i32().expect("Failed converting status to i32");
+        env.call_method(
             self.obj.as_obj(),
             "onCoreGenericErrorNotificationReceived",
             "(I)V",
-            &[JValue::Int(status.try_into().expect("Could not convert status"))],
-        );
-        self.cleanup_and_return(env, result)
+            &[JValue::Int(status)],
+        )
+        .map(|_| ()) // drop void method return
+    }
+
+    pub fn core_generic_error_notification_received(&self, data: GenericErrorPacket) -> Result<()> {
+        let env = self.jvm.attach_current_thread()?;
+        let result = self.handle_core_generic_error_notification_received(&env, data);
+        self.clear_exception(env);
+        result
+    }
+
+    fn create_zeroed_two_way_measurement_java<'a>(
+        env: &'a JNIEnv,
+        two_way_measurement_class: JClass,
+        mac_address_java: jobjectArray,
+    ) -> Result<JObject<'a>> {
+        env.new_object(
+            two_way_measurement_class,
+            "([BIIIIIIIIIIII)V",
+            &[
+                JValue::Object(JObject::from(mac_address_java)),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+                JValue::Int(0),
+            ],
+        )
+    }
+
+    fn create_short_mac_two_way_measurement_java<'a>(
+        env: &'a JNIEnv,
+        two_way_measurement_class: JClass,
+        two_way_measurement: &'a ShortAddressTwoWayRangingMeasurement,
+    ) -> Result<JObject<'a>> {
+        let mac_address_arr = two_way_measurement.mac_address.to_ne_bytes();
+        let mac_address_java = env.new_byte_array(
+            SHORT_MAC_ADDRESS_LEN.to_i32().expect("Failed converting mac address len to i32"),
+        )?;
+        // Convert from [u8] to [i8] since java does not support unsigned byte.
+        let mac_address_arr_i8 = mac_address_arr.map(|x| x as i8);
+        env.set_byte_array_region(mac_address_java, 0, &mac_address_arr_i8)?;
+        env.new_object(
+            two_way_measurement_class,
+            "([BIIIIIIIIIIII)V",
+            &[
+                JValue::Object(JObject::from(mac_address_java)),
+                JValue::Int(
+                    two_way_measurement.status.to_i32().expect("Failed converting status to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement.nlos.to_i32().expect("Failed converting nlos to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .distance
+                        .to_i32()
+                        .expect("Failed converting distance to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_azimuth
+                        .to_i32()
+                        .expect("Failed converting aoa azimuth to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_azimuth_fom
+                        .to_i32()
+                        .expect("Failed converting aoa azimuth fom to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_elevation
+                        .to_i32()
+                        .expect("Failed converting aoa elevation to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_elevation_fom
+                        .to_i32()
+                        .expect("Failed converting aoa elevaion fom to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_azimuth
+                        .to_i32()
+                        .expect("Failed converting dest aoa azimuth to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_azimuth_fom
+                        .to_i32()
+                        .expect("Failed converting dest aoa azimuth fom to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_elevation
+                        .to_i32()
+                        .expect("Failed converting dest aoa elevation to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_elevation_fom
+                        .to_i32()
+                        .expect("Failed converting dest aoa elevation azimuth to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .slot_index
+                        .to_i32()
+                        .expect("Failed converting slot index to i32"),
+                ),
+            ],
+        )
+    }
+
+    fn create_extended_mac_two_way_measurement_java<'a>(
+        env: &'a JNIEnv,
+        two_way_measurement_class: JClass,
+        two_way_measurement: &'a ExtendedAddressTwoWayRangingMeasurement,
+    ) -> Result<JObject<'a>> {
+        let mac_address_arr = two_way_measurement.mac_address.to_ne_bytes();
+        let mac_address_java = env.new_byte_array(
+            EXTENDED_MAC_ADDRESS_LEN.to_i32().expect("Failed converting mac address len to i32"),
+        )?;
+        // Convert from [u8] to [i8] since java does not support unsigned byte.
+        let mac_address_arr_i8 = mac_address_arr.map(|x| x as i8);
+        env.set_byte_array_region(mac_address_java, 0, &mac_address_arr_i8)?;
+        env.new_object(
+            two_way_measurement_class,
+            "([BIIIIIIIIIIII)V",
+            &[
+                JValue::Object(JObject::from(mac_address_java)),
+                JValue::Int(
+                    two_way_measurement.status.to_i32().expect("Failed converting status to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement.nlos.to_i32().expect("Failed converting nlos to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .distance
+                        .to_i32()
+                        .expect("Failed converting distance to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_azimuth
+                        .to_i32()
+                        .expect("Failed converting aoa azimuth to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_azimuth_fom
+                        .to_i32()
+                        .expect("Failed converting aoa azimuth fom to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_elevation
+                        .to_i32()
+                        .expect("Failed converting aoa elevation to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_elevation_fom
+                        .to_i32()
+                        .expect("Failed converting aoa elevaion fom to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_azimuth
+                        .to_i32()
+                        .expect("Failed converting dest aoa azimuth to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_azimuth_fom
+                        .to_i32()
+                        .expect("Failed converting dest aoa azimuth fom to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_elevation
+                        .to_i32()
+                        .expect("Failed converting dest aoa elevation to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .aoa_destination_elevation_fom
+                        .to_i32()
+                        .expect("Failed converting dest aoa elevation azimuth to i32"),
+                ),
+                JValue::Int(
+                    two_way_measurement
+                        .slot_index
+                        .to_i32()
+                        .expect("Failed converting slot index to i32"),
+                ),
+            ],
+        )
+    }
+
+    fn create_range_data_java<'a>(
+        &'a self,
+        env: &'a JNIEnv,
+        data: RangeDataNtfPacket,
+        two_way_measurements_java: jobjectArray,
+        num_two_way_measurements: i32,
+    ) -> Result<JObject<'a>> {
+        let ranging_data_class = self.find_class(env, &UWB_RANGING_DATA_CLASS)?;
+        env.new_object(
+            ranging_data_class,
+            "(JJIJIII[Lcom/android/uwb/data/UwbTwoWayMeasurement;)V",
+            &[
+                JValue::Long(
+                    data.get_sequence_number().to_i64().expect("Failed converting seq num to i64"),
+                ),
+                JValue::Long(
+                    data.get_session_id().to_i64().expect("Failed converting session id to i64"),
+                ),
+                JValue::Int(
+                    data.get_rcr_indicator()
+                        .to_i32()
+                        .expect("Failed converting rcr indicator to i32"),
+                ),
+                JValue::Long(
+                    data.get_current_ranging_interval()
+                        .to_i64()
+                        .expect("Failed converting curr ranging interval to i32"),
+                ),
+                JValue::Int(
+                    data.get_ranging_measurement_type()
+                        .to_i32()
+                        .expect("Failed converting ranging measurement type to i32"),
+                ),
+                JValue::Int(
+                    data.get_mac_address_indicator()
+                        .to_i32()
+                        .expect("Failed converting mac address indicator to i32"),
+                ),
+                JValue::Int(num_two_way_measurements),
+                JValue::Object(JObject::from(two_way_measurements_java)),
+            ],
+        )
+    }
+
+    fn handle_short_range_data_notification(
+        &self,
+        env: &JNIEnv,
+        data: ShortMacTwoWayRangeDataNtfPacket,
+    ) -> Result<()> {
+        let two_way_measurement_class = self.find_class(&env, &UWB_TWO_WAY_MEASUREMENT_CLASS)?;
+        let two_way_measurement_initial_java =
+            EventManager::create_zeroed_two_way_measurement_java(
+                &env,
+                two_way_measurement_class,
+                env.new_byte_array(
+                    EXTENDED_MAC_ADDRESS_LEN
+                        .to_i32()
+                        .expect("Failed converting mac address len to i32"),
+                )?,
+            )?;
+        let num_two_way_measurements: i32 = data
+            .get_two_way_ranging_measurements()
+            .len()
+            .to_i32()
+            .expect("Failed converting len to i32");
+        let two_way_measurements_java = env.new_object_array(
+            num_two_way_measurements,
+            two_way_measurement_class,
+            two_way_measurement_initial_java,
+        )?;
+        for (i, two_way_measurement) in data.get_two_way_ranging_measurements().iter().enumerate() {
+            let two_way_measurement_java = EventManager::create_short_mac_two_way_measurement_java(
+                &env,
+                two_way_measurement_class,
+                two_way_measurement,
+            )?;
+            env.set_object_array_element(
+                two_way_measurements_java,
+                i.to_i32().expect("Failed converting idx to i32"),
+                two_way_measurement_java,
+            )?
+        }
+        let ranging_data_java = self.create_range_data_java(
+            &env,
+            data.into(),
+            two_way_measurements_java,
+            num_two_way_measurements,
+        )?;
+        env.call_method(
+            self.obj.as_obj(),
+            "onRangeDataNotificationReceived",
+            "(Lcom/android/uwb/data/UwbRangingData;)V",
+            &[JValue::Object(JObject::from(ranging_data_java))],
+        )
+        .map(|_| ()) // drop void method return
+    }
+
+    pub fn short_range_data_notification(
+        &self,
+        data: ShortMacTwoWayRangeDataNtfPacket,
+    ) -> Result<()> {
+        let env = self.jvm.attach_current_thread()?;
+        let result = self.handle_short_range_data_notification(&env, data);
+        self.clear_exception(env);
+        result
     }
 
-    fn cleanup_and_return<T>(&self, env: AttachGuard, result: Result<T>) -> Result<()> {
+    fn handle_extended_range_data_notification(
+        &self,
+        env: &JNIEnv,
+        data: ExtendedMacTwoWayRangeDataNtfPacket,
+    ) -> Result<()> {
+        let two_way_measurement_class = self.find_class(&env, &UWB_TWO_WAY_MEASUREMENT_CLASS)?;
+        let two_way_measurement_initial_java =
+            EventManager::create_zeroed_two_way_measurement_java(
+                &env,
+                two_way_measurement_class,
+                env.new_byte_array(
+                    EXTENDED_MAC_ADDRESS_LEN
+                        .to_i32()
+                        .expect("Failed converting mac address len to i32"),
+                )?,
+            )?;
+        let num_two_way_measurements: i32 = data
+            .get_two_way_ranging_measurements()
+            .len()
+            .to_i32()
+            .expect("Failed converting len to i32");
+        let two_way_measurements_java = env.new_object_array(
+            num_two_way_measurements,
+            two_way_measurement_class,
+            two_way_measurement_initial_java,
+        )?;
+        for (i, two_way_measurement) in data.get_two_way_ranging_measurements().iter().enumerate() {
+            let two_way_measurement_java =
+                EventManager::create_extended_mac_two_way_measurement_java(
+                    &env,
+                    two_way_measurement_class,
+                    two_way_measurement,
+                )?;
+            env.set_object_array_element(
+                two_way_measurements_java,
+                i.to_i32().expect("Failed converting idx to i32"),
+                two_way_measurement_java,
+            )?;
+        }
+        let ranging_data_java = self.create_range_data_java(
+            &env,
+            data.into(),
+            two_way_measurements_java,
+            num_two_way_measurements,
+        )?;
+        env.call_method(
+            self.obj.as_obj(),
+            "onRangeDataNotificationReceived",
+            "(Lcom/android/uwb/data/UwbRangingData;)V",
+            &[JValue::Object(JObject::from(ranging_data_java))],
+        )
+        .map(|_| ()) // drop void method return
+    }
+
+    pub fn extended_range_data_notification(
+        &self,
+        data: ExtendedMacTwoWayRangeDataNtfPacket,
+    ) -> Result<()> {
+        let env = self.jvm.attach_current_thread()?;
+        let result = self.handle_extended_range_data_notification(&env, data);
         self.clear_exception(env);
-        // Discard the value returned by the call.
-        result.map(|_| ())
+        result
     }
 
     // Attempts to clear an exception.  If we do not do this, the exception continues being thrown