Separate JNI for diff ver of multicast update

This CL separates the JNI calls for different versions of dynamic update
of multicast list for controlees. The version 1 and version 2 now uses
different JNI function with parameters aligned with the commands.

Bug: 256092554
Test: make com.gooogle.android.uwb
Change-Id: Ic51ff24c01970338ee7842ed6d95b0789eb4dcd0

5 files changed, 212 insertions, 43 deletions

diff --git a/service/java/com/android/server/uwb/UwbSessionManager.java b/service/java/com/android/server/uwb/UwbSessionManager.java
index d31cf77c..5f33583a 100644
--- a/service/java/com/android/server/uwb/UwbSessionManager.java
+++ b/service/java/com/android/server/uwb/UwbSessionManager.java
@@ -1065,21 +1065,31 @@ public class UwbSessionManager implements INativeUwbManager.SessionNotification
                                     // Set to 0's for the UCI stack.
                                     subSessionIdList = new int[dstAddressListSize];
                                 }
-                                int messageControl =
-                                        rangingReconfigureParams.getMessageControl() == null
-                                        ? -1 : rangingReconfigureParams.getMessageControl();
-                                int[] subsessionKeyList =
-                                        rangingReconfigureParams.getSubSessionKeyList();
-
-                                status = mNativeUwbManager.controllerMulticastListUpdate(
-                                        uwbSession.getSessionId(),
-                                        action,
-                                        subSessionIdList.length,
-                                        ArrayUtils.toPrimitive(dstAddressList),
-                                        subSessionIdList,
-                                        messageControl,
-                                        subsessionKeyList,
-                                        uwbSession.getChipId());
+                                boolean isV2 = rangingReconfigureParams.getMessageControl() != null;
+                                if (isV2) {
+                                    int messageControl =
+                                            rangingReconfigureParams.getMessageControl();
+                                    int[] subsessionKeyList =
+                                            rangingReconfigureParams.getSubSessionKeyList();
+
+                                    status = mNativeUwbManager.controllerMulticastListUpdateV2(
+                                            uwbSession.getSessionId(),
+                                            action,
+                                            subSessionIdList.length,
+                                            ArrayUtils.toPrimitive(dstAddressList),
+                                            subSessionIdList,
+                                            messageControl,
+                                            subsessionKeyList,
+                                            uwbSession.getChipId());
+                                } else {
+                                    status = mNativeUwbManager.controllerMulticastListUpdateV1(
+                                            uwbSession.getSessionId(),
+                                            action,
+                                            subSessionIdList.length,
+                                            ArrayUtils.toPrimitive(dstAddressList),
+                                            subSessionIdList,
+                                            uwbSession.getChipId());
+                                }
                                 if (status != UwbUciConstants.STATUS_CODE_OK) {
                                     Log.e(TAG, "Unable to update controller multicast list.");
                                     if (action == MULTICAST_LIST_UPDATE_ACTION_ADD) {
diff --git a/service/java/com/android/server/uwb/jni/NativeUwbManager.java b/service/java/com/android/server/uwb/jni/NativeUwbManager.java
index f3aa65fb..d9866a42 100644
--- a/service/java/com/android/server/uwb/jni/NativeUwbManager.java
+++ b/service/java/com/android/server/uwb/jni/NativeUwbManager.java
@@ -304,7 +304,28 @@ public class NativeUwbManager {
     }
 
     /**
-     * Update Multicast list for the requested UWB session
+     * Update Multicast list for the requested UWB session using V1 command.
+     *
+     * @param sessionId         : Session ID to which multicast list to be updated
+     * @param action            : Update the multicast list by adding or removing
+     *                          0x00 - Adding
+     *                          0x01 - removing
+     * @param noOfControlee     : The number(n) of Controlees
+     * @param addresses         : address list of Controlees
+     * @param subSessionIds     : Specific sub-session ID list of Controlees
+     * @return : refer to SESSION_SET_APP_CONFIG_RSP
+     * in the Table 16: Control messages to set Application configurations
+     */
+    public byte controllerMulticastListUpdateV1(int sessionId, int action, int noOfControlee,
+            short[] addresses, int[] subSessionIds, String chipId) {
+        synchronized (mSessionFnLock) {
+            return nativeControllerMulticastListUpdateV1(sessionId, (byte) action,
+                    (byte) noOfControlee, addresses, subSessionIds, chipId);
+        }
+    }
+
+    /**
+     * Update Multicast list for the requested UWB session using V2 command.
      *
      * @param sessionId         : Session ID to which multicast list to be updated
      * @param action            : Update the multicast list by adding or removing
@@ -325,11 +346,11 @@ public class NativeUwbManager {
      * @return : refer to SESSION_SET_APP_CONFIG_RSP
      * in the Table 16: Control messages to set Application configurations
      */
-    public byte controllerMulticastListUpdate(int sessionId, int action, int noOfControlee,
+    public byte controllerMulticastListUpdateV2(int sessionId, int action, int noOfControlee,
             short[] addresses, int[] subSessionIds, int messageControl,
             int[] subSessionKeyList, String chipId) {
         synchronized (mSessionFnLock) {
-            return nativeControllerMulticastListUpdate(sessionId, (byte) action,
+            return nativeControllerMulticastListUpdateV2(sessionId, (byte) action,
                     (byte) noOfControlee, addresses, subSessionIds, messageControl,
                     subSessionKeyList, chipId);
         }
@@ -435,7 +456,10 @@ public class NativeUwbManager {
 
     private native UwbTlvData nativeGetCapsInfo(String chipId);
 
-    private native byte nativeControllerMulticastListUpdate(int sessionId, byte action,
+    private native byte nativeControllerMulticastListUpdateV1(int sessionId, byte action,
+            byte noOfControlee, short[] address, int[] subSessionId, String chipId);
+
+    private native byte nativeControllerMulticastListUpdateV2(int sessionId, byte action,
             byte noOfControlee, short[] address, int[] subSessionId, int messageControl,
             int[] subSessionKeyList, String chipId);
 
diff --git a/service/tests/src/com/android/server/uwb/UwbSessionManagerTest.java b/service/tests/src/com/android/server/uwb/UwbSessionManagerTest.java
index 657fa056..7322b31f 100644
--- a/service/tests/src/com/android/server/uwb/UwbSessionManagerTest.java
+++ b/service/tests/src/com/android/server/uwb/UwbSessionManagerTest.java
@@ -1488,11 +1488,27 @@ public class UwbSessionManagerTest {
         assertThat(status).isEqualTo(UwbUciConstants.STATUS_CODE_ERROR_SESSION_NOT_EXIST);
     }
 
-    private FiraRangingReconfigureParams buildReconfigureParams() {
-        return buildReconfigureParams(FiraParams.MULTICAST_LIST_UPDATE_ACTION_ADD);
+    private FiraRangingReconfigureParams buildReconfigureParamsV1() {
+        return buildReconfigureParamsV1(FiraParams.MULTICAST_LIST_UPDATE_ACTION_ADD);
     }
 
-    private FiraRangingReconfigureParams buildReconfigureParams(int action) {
+    private FiraRangingReconfigureParams buildReconfigureParamsV1(int action) {
+        FiraRangingReconfigureParams reconfigureParams =
+                new FiraRangingReconfigureParams.Builder()
+                        .setAddressList(new UwbAddress[] {
+                                UwbAddress.fromBytes(new byte[] { (byte) 0x01, (byte) 0x02 }) })
+                        .setAction(action)
+                        .setSubSessionIdList(new int[] { 2 })
+                        .build();
+
+        return spy(reconfigureParams);
+    }
+
+    private FiraRangingReconfigureParams buildReconfigureParamsV2() {
+        return buildReconfigureParamsV2(FiraParams.MULTICAST_LIST_UPDATE_ACTION_ADD);
+    }
+
+    private FiraRangingReconfigureParams buildReconfigureParamsV2(int action) {
         FiraRangingReconfigureParams reconfigureParams =
                 new FiraRangingReconfigureParams.Builder()
                         .setAddressList(new UwbAddress[] {
@@ -1510,19 +1526,106 @@ public class UwbSessionManagerTest {
         UwbSession uwbSession = prepareExistingUwbSession();
 
         int status = mUwbSessionManager.reconfigure(
-                uwbSession.getSessionHandle(), buildReconfigureParams());
+                uwbSession.getSessionHandle(), buildReconfigureParamsV2());
 
         assertThat(status).isEqualTo(0);
         assertThat(mTestLooper.nextMessage().what).isEqualTo(4); // SESSION_RECONFIGURE_RANGING
     }
 
     @Test
-    public void execReconfigureAddControlee_success() throws Exception {
+    public void execReconfigureAddControleeV1_success() throws Exception {
+        UwbSession uwbSession = prepareExistingUwbSession();
+        FiraRangingReconfigureParams reconfigureParams =
+                buildReconfigureParamsV1();
+        when(mNativeUwbManager
+                .controllerMulticastListUpdateV1(anyInt(), anyInt(), anyInt(), any(), any(),
+                        anyString()))
+                .thenReturn((byte) UwbUciConstants.STATUS_CODE_OK);
+        UwbMulticastListUpdateStatus uwbMulticastListUpdateStatus =
+                mock(UwbMulticastListUpdateStatus.class);
+        when(uwbMulticastListUpdateStatus.getNumOfControlee()).thenReturn(1);
+        when(uwbMulticastListUpdateStatus.getControleeUwbAddresses())
+                .thenReturn(new UwbAddress[] {UWB_DEST_ADDRESS_2});
+        when(uwbMulticastListUpdateStatus.getStatus()).thenReturn(
+                new int[] { UwbUciConstants.STATUS_CODE_OK });
+        doReturn(uwbMulticastListUpdateStatus).when(uwbSession).getMulticastListUpdateStatus();
+        when(mUwbConfigurationManager.setAppConfigurations(anyInt(), any(), anyString()))
+                .thenReturn(UwbUciConstants.STATUS_CODE_OK);
+
+        mUwbSessionManager.reconfigure(uwbSession.getSessionHandle(), reconfigureParams);
+        mTestLooper.dispatchNext();
+
+        // Make sure the original address is still there.
+        assertThat(uwbSession.getControleeList().stream()
+                .anyMatch(e -> e.getUwbAddress().equals(UWB_DEST_ADDRESS)))
+                .isTrue();
+
+        // Make sure this new address was added.
+        assertThat(uwbSession.getControleeList().stream()
+                .anyMatch(e -> e.getUwbAddress().equals(UWB_DEST_ADDRESS_2)))
+                .isTrue();
+
+        short dstAddress =
+                ByteBuffer.wrap(reconfigureParams.getAddressList()[0].toBytes()).getShort(0);
+        verify(mNativeUwbManager).controllerMulticastListUpdateV1(
+                uwbSession.getSessionId(), reconfigureParams.getAction(), 1,
+                new short[] {dstAddress}, reconfigureParams.getSubSessionIdList(),
+                uwbSession.getChipId());
+        verify(mUwbSessionNotificationManager).onControleeAdded(eq(uwbSession));
+        verify(mUwbSessionNotificationManager).onRangingReconfigured(eq(uwbSession));
+    }
+
+    @Test
+    public void execReconfigureRemoveControleeV1_success() throws Exception {
+        UwbSession uwbSession = prepareExistingUwbSession();
+        FiraRangingReconfigureParams reconfigureParams =
+                buildReconfigureParamsV1(FiraParams.MULTICAST_LIST_UPDATE_ACTION_DELETE);
+        when(mNativeUwbManager
+                .controllerMulticastListUpdateV1(anyInt(), anyInt(), anyInt(), any(), any(),
+                        anyString()))
+                .thenReturn((byte) UwbUciConstants.STATUS_CODE_OK);
+        UwbMulticastListUpdateStatus uwbMulticastListUpdateStatus =
+                mock(UwbMulticastListUpdateStatus.class);
+        when(uwbMulticastListUpdateStatus.getNumOfControlee()).thenReturn(1);
+        when(uwbMulticastListUpdateStatus.getControleeUwbAddresses())
+                .thenReturn(new UwbAddress[] {UWB_DEST_ADDRESS});
+        when(uwbMulticastListUpdateStatus.getStatus()).thenReturn(
+                new int[] { UwbUciConstants.STATUS_CODE_OK });
+        doReturn(uwbMulticastListUpdateStatus).when(uwbSession).getMulticastListUpdateStatus();
+        when(mUwbConfigurationManager.setAppConfigurations(anyInt(), any(), anyString()))
+                .thenReturn(UwbUciConstants.STATUS_CODE_OK);
+
+        // Make sure the address exists in the first place. This should have been set up by
+        //  prepareExistingUwbSession
+        assertThat(uwbSession.getControleeList().stream()
+                .anyMatch(e -> e.getUwbAddress().equals(UWB_DEST_ADDRESS)))
+                .isTrue();
+
+        mUwbSessionManager.reconfigure(uwbSession.getSessionHandle(), reconfigureParams);
+        mTestLooper.dispatchNext();
+
+        // Make sure the address was removed.
+        assertThat(uwbSession.getControleeList().stream()
+                .anyMatch(e -> e.getUwbAddress().equals(UWB_DEST_ADDRESS)))
+                .isFalse();
+
+        short dstAddress =
+                ByteBuffer.wrap(reconfigureParams.getAddressList()[0].toBytes()).getShort(0);
+        verify(mNativeUwbManager).controllerMulticastListUpdateV1(
+                uwbSession.getSessionId(), reconfigureParams.getAction(), 1,
+                new short[] {dstAddress}, reconfigureParams.getSubSessionIdList(),
+                uwbSession.getChipId());
+        verify(mUwbSessionNotificationManager).onControleeRemoved(eq(uwbSession));
+        verify(mUwbSessionNotificationManager).onRangingReconfigured(eq(uwbSession));
+    }
+
+    @Test
+    public void execReconfigureAddControleeV2_success() throws Exception {
         UwbSession uwbSession = prepareExistingUwbSession();
         FiraRangingReconfigureParams reconfigureParams =
-                buildReconfigureParams();
+                buildReconfigureParamsV2();
         when(mNativeUwbManager
-                .controllerMulticastListUpdate(anyInt(), anyInt(), anyInt(), any(), any(),
+                .controllerMulticastListUpdateV2(anyInt(), anyInt(), anyInt(), any(), any(),
                                 anyInt(), any(), anyString()))
                 .thenReturn((byte) UwbUciConstants.STATUS_CODE_OK);
         UwbMulticastListUpdateStatus uwbMulticastListUpdateStatus =
@@ -1551,7 +1654,7 @@ public class UwbSessionManagerTest {
 
         short dstAddress =
                 ByteBuffer.wrap(reconfigureParams.getAddressList()[0].toBytes()).getShort(0);
-        verify(mNativeUwbManager).controllerMulticastListUpdate(
+        verify(mNativeUwbManager).controllerMulticastListUpdateV2(
                 uwbSession.getSessionId(), reconfigureParams.getAction(), 1,
                 new short[] {dstAddress}, reconfigureParams.getSubSessionIdList(),
                 reconfigureParams.getMessageControl(), reconfigureParams.getSubSessionKeyList(),
@@ -1561,12 +1664,12 @@ public class UwbSessionManagerTest {
     }
 
     @Test
-    public void execReconfigureRemoveControlee_success() throws Exception {
+    public void execReconfigureRemoveControleeV2_success() throws Exception {
         UwbSession uwbSession = prepareExistingUwbSession();
         FiraRangingReconfigureParams reconfigureParams =
-                buildReconfigureParams(FiraParams.MULTICAST_LIST_UPDATE_ACTION_DELETE);
+                buildReconfigureParamsV2(FiraParams.MULTICAST_LIST_UPDATE_ACTION_DELETE);
         when(mNativeUwbManager
-                .controllerMulticastListUpdate(anyInt(), anyInt(), anyInt(), any(), any(),
+                .controllerMulticastListUpdateV2(anyInt(), anyInt(), anyInt(), any(), any(),
                                 anyInt(), any(), anyString()))
                 .thenReturn((byte) UwbUciConstants.STATUS_CODE_OK);
         UwbMulticastListUpdateStatus uwbMulticastListUpdateStatus =
@@ -1596,7 +1699,7 @@ public class UwbSessionManagerTest {
 
         short dstAddress =
                 ByteBuffer.wrap(reconfigureParams.getAddressList()[0].toBytes()).getShort(0);
-        verify(mNativeUwbManager).controllerMulticastListUpdate(
+        verify(mNativeUwbManager).controllerMulticastListUpdateV2(
                 uwbSession.getSessionId(), reconfigureParams.getAction(), 1,
                 new short[] {dstAddress}, reconfigureParams.getSubSessionIdList(),
                 reconfigureParams.getMessageControl(), reconfigureParams.getSubSessionKeyList(),
@@ -1609,9 +1712,9 @@ public class UwbSessionManagerTest {
     public void execReconfigure_nativeUpdateFailed() throws Exception {
         UwbSession uwbSession = prepareExistingUwbSession();
         FiraRangingReconfigureParams reconfigureParams =
-                buildReconfigureParams();
+                buildReconfigureParamsV2();
         when(mNativeUwbManager
-                .controllerMulticastListUpdate(anyInt(), anyInt(), anyInt(), any(), any(),
+                .controllerMulticastListUpdateV2(anyInt(), anyInt(), anyInt(), any(), any(),
                                 anyInt(), any(), anyString()))
                 .thenReturn((byte) UwbUciConstants.STATUS_CODE_FAILED);
 
@@ -1628,9 +1731,9 @@ public class UwbSessionManagerTest {
     public void execReconfigure_uwbSessionUpdateMixedSuccess() throws Exception {
         UwbSession uwbSession = prepareExistingUwbSession();
         FiraRangingReconfigureParams reconfigureParams =
-                buildReconfigureParams();
+                buildReconfigureParamsV2();
         when(mNativeUwbManager
-                .controllerMulticastListUpdate(anyInt(), anyInt(), anyInt(), any(), any(),
+                .controllerMulticastListUpdateV2(anyInt(), anyInt(), anyInt(), any(), any(),
                         anyInt(), any(), anyString()))
                 .thenReturn((byte) UwbUciConstants.STATUS_CODE_OK);
         UwbMulticastListUpdateStatus uwbMulticastListUpdateStatus =
@@ -1666,9 +1769,9 @@ public class UwbSessionManagerTest {
     public void execReconfigure_uwbSessionUpdateFailed() throws Exception {
         UwbSession uwbSession = prepareExistingUwbSession();
         FiraRangingReconfigureParams reconfigureParams =
-                buildReconfigureParams();
+                buildReconfigureParamsV2();
         when(mNativeUwbManager
-                .controllerMulticastListUpdate(anyInt(), anyInt(), anyInt(), any(), any(),
+                .controllerMulticastListUpdateV2(anyInt(), anyInt(), anyInt(), any(), any(),
                                 anyInt(), any(), anyString()))
                 .thenReturn((byte) UwbUciConstants.STATUS_CODE_OK);
         UwbMulticastListUpdateStatus uwbMulticastListUpdateStatus =
@@ -1722,9 +1825,9 @@ public class UwbSessionManagerTest {
     public void execReconfigure_setAppConfigurationsFailed() throws Exception {
         UwbSession uwbSession = prepareExistingUwbSession();
         FiraRangingReconfigureParams reconfigureParams =
-                buildReconfigureParams();
+                buildReconfigureParamsV2();
         when(mNativeUwbManager
-                .controllerMulticastListUpdate(anyInt(), anyInt(), anyInt(), any(), any(),
+                .controllerMulticastListUpdateV2(anyInt(), anyInt(), anyInt(), any(), any(),
                                 anyInt(), any(), anyString()))
                 .thenReturn((byte) UwbUciConstants.STATUS_CODE_FAILED);
         UwbMulticastListUpdateStatus uwbMulticastListUpdateStatus =
diff --git a/service/uci/jni/rust/lib.rs b/service/uci/jni/rust/lib.rs
index a75df3ea..0287e35b 100644
--- a/service/uci/jni/rust/lib.rs
+++ b/service/uci/jni/rust/lib.rs
@@ -428,9 +428,41 @@ pub extern "system" fn Java_com_android_server_uwb_jni_NativeUwbManager_nativeGe
     }
 }
 
-/// update multicast list
+/// update multicast list by SESSION_UPDATE_CONTROLLER_MULTICAST_LIST_CMD
 #[no_mangle]
-pub extern "system" fn Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdate(
+pub extern "system" fn Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdateV1(
+    env: JNIEnv,
+    obj: JObject,
+    session_id: jint,
+    action: jbyte,
+    no_of_controlee: jbyte,
+    addresses: jshortArray,
+    sub_session_ids: jintArray,
+    chip_id: JString,
+) -> jbyte {
+    info!("Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdateV1: enter");
+    let controlee_data = ControleeData {
+        addresses,
+        sub_session_ids,
+        message_control: -1,
+        sub_session_keys: *JObject::null(),
+    };
+    byte_result_helper(
+        multicast_list_update(
+            &JniContext::new(env, obj),
+            session_id as u32,
+            action as u8,
+            no_of_controlee as u8,
+            controlee_data,
+            env.get_string(chip_id).unwrap().into(),
+        ),
+        "ControllerMulticastListUpdate",
+    )
+}
+
+/// update multicast list by SESSION_UPDATE_CONTROLLER_MULTICAST_LIST_V2_CMD
+#[no_mangle]
+pub extern "system" fn Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdateV2(
     env: JNIEnv,
     obj: JObject,
     session_id: jint,
@@ -442,7 +474,7 @@ pub extern "system" fn Java_com_android_server_uwb_jni_NativeUwbManager_nativeCo
     sub_session_keys: jbyteArray,
     chip_id: JString,
 ) -> jbyte {
-    info!("Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdate: enter");
+    info!("Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdateV2: enter");
     let controlee_data =
         ControleeData { addresses, sub_session_ids, message_control, sub_session_keys };
     byte_result_helper(
diff --git a/service/uci/jni_new/src/uci_jni_android_new.rs b/service/uci/jni_new/src/uci_jni_android_new.rs
index 034fdd7d..251a25fb 100644
--- a/service/uci/jni_new/src/uci_jni_android_new.rs
+++ b/service/uci/jni_new/src/uci_jni_android_new.rs
@@ -563,7 +563,7 @@ fn native_get_caps_info(env: JNIEnv, obj: JObject, chip_id: JString) -> Result<V
 
 /// Update multicast list on a single UWB device. Return value defined by uci_packets.pdl
 #[no_mangle]
-pub extern "system" fn Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdate(
+pub extern "system" fn Java_com_android_server_uwb_jni_NativeUwbManager_nativeControllerMulticastListUpdateV1(
     env: JNIEnv,
     obj: JObject,
     session_id: jint,