Merge changes Ia6007bef,Ifba9819a into tm-dev am: 0e94af61b7

Original change: https://googleplex-android-review.googlesource.com/c/platform/external/uwb/+/17810512

Change-Id: Ie4aeb0df5a63af3c456cd0737e9030de4fedda66
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

6 files changed, 331 insertions, 45 deletions

diff --git a/src/rust/uwb_core/src/session/session_manager.rs b/src/rust/uwb_core/src/session/session_manager.rs
index 352b2f8..592affe 100644
--- a/src/rust/uwb_core/src/session/session_manager.rs
+++ b/src/rust/uwb_core/src/session/session_manager.rs
@@ -21,7 +21,9 @@ use crate::session::error::{Error, Result};
 use crate::session::params::AppConfigParams;
 use crate::session::uwb_session::UwbSession;
 use crate::uci::notification::UciNotification;
-use crate::uci::params::{SessionId, SessionState, SessionType};
+use crate::uci::params::{
+    Controlee, SessionId, SessionState, SessionType, UpdateMulticastListAction,
+};
 use crate::uci::uci_manager::UciManager;
 
 const MAX_SESSION_COUNT: usize = 5;
@@ -76,6 +78,20 @@ impl SessionManager {
         self.send_cmd(SessionCommand::Reconfigure { session_id, params }).await
     }
 
+    async fn update_controller_multicast_list(
+        &mut self,
+        session_id: SessionId,
+        action: UpdateMulticastListAction,
+        controlees: Vec<Controlee>,
+    ) -> Result<()> {
+        self.send_cmd(SessionCommand::UpdateControllerMulticastList {
+            session_id,
+            action,
+            controlees,
+        })
+        .await
+    }
+
     // Send the |cmd| to the SessionManagerActor.
     async fn send_cmd(&self, cmd: SessionCommand) -> Result<()> {
         let (result_sender, result_receiver) = oneshot::channel();
@@ -166,7 +182,6 @@ impl<T: UciManager> SessionManagerActor<T> {
                     }
                 }
             }
-
             SessionCommand::StartRanging { session_id } => {
                 match self.active_sessions.get_mut(&session_id) {
                     None => {
@@ -197,12 +212,20 @@ impl<T: UciManager> SessionManagerActor<T> {
                     }
                 }
             }
+            SessionCommand::UpdateControllerMulticastList { session_id, action, controlees } => {
+                match self.active_sessions.get_mut(&session_id) {
+                    None => {
+                        let _ = result_sender.send(Err(Error::UnknownSessionId(session_id)));
+                    }
+                    Some(session) => {
+                        session.update_controller_multicast_list(action, controlees, result_sender);
+                    }
+                }
+            }
         }
     }
 
     fn handle_uci_notification(&mut self, notf: UciNotification) {
-        // TODO(akahuang): Remove this after handling multiple kind of notifications.
-        #[allow(clippy::single_match)]
         match notf {
             UciNotification::SessionStatus { session_id, session_state, reason_code } => {
                 if session_state == SessionState::SessionStateDeinit {
@@ -212,7 +235,7 @@ impl<T: UciManager> SessionManagerActor<T> {
                 }
 
                 match self.active_sessions.get_mut(&session_id) {
-                    Some(session) => session.set_state(session_state),
+                    Some(session) => session.on_session_status_changed(session_state),
                     None => {
                         warn!(
                             "Received notification of the unknown Session {:?}: {:?}, {:?}",
@@ -221,6 +244,19 @@ impl<T: UciManager> SessionManagerActor<T> {
                     }
                 }
             }
+            UciNotification::SessionUpdateControllerMulticastList {
+                session_id,
+                remaining_multicast_list_size: _,
+                status_list,
+            } => match self.active_sessions.get_mut(&session_id) {
+                Some(session) => session.on_controller_multicast_list_udpated(status_list),
+                None => {
+                    warn!(
+                        "Received the notification of the unknown Session {}: {:?}",
+                        session_id, status_list
+                    );
+                }
+            },
             _ => {}
         }
     }
@@ -228,11 +264,29 @@ impl<T: UciManager> SessionManagerActor<T> {
 
 #[derive(Debug)]
 enum SessionCommand {
-    InitSession { session_id: SessionId, session_type: SessionType, params: AppConfigParams },
-    DeinitSession { session_id: SessionId },
-    StartRanging { session_id: SessionId },
-    StopRanging { session_id: SessionId },
-    Reconfigure { session_id: SessionId, params: AppConfigParams },
+    InitSession {
+        session_id: SessionId,
+        session_type: SessionType,
+        params: AppConfigParams,
+    },
+    DeinitSession {
+        session_id: SessionId,
+    },
+    StartRanging {
+        session_id: SessionId,
+    },
+    StopRanging {
+        session_id: SessionId,
+    },
+    Reconfigure {
+        session_id: SessionId,
+        params: AppConfigParams,
+    },
+    UpdateControllerMulticastList {
+        session_id: SessionId,
+        action: UpdateMulticastListAction,
+        controlees: Vec<Controlee>,
+    },
 }
 
 #[cfg(test)]
@@ -242,7 +296,9 @@ mod tests {
     use crate::session::params::fira_app_config_params::*;
     use crate::uci::error::StatusCode;
     use crate::uci::mock_uci_manager::MockUciManager;
-    use crate::uci::params::{ReasonCode, SetAppConfigResponse};
+    use crate::uci::params::{
+        ControleeStatus, MulticastUpdateStatusCode, ReasonCode, SetAppConfigResponse,
+    };
     use crate::utils::init_test_logging;
 
     fn generate_params() -> AppConfigParams {
@@ -396,4 +452,117 @@ mod tests {
 
         assert!(mock_uci_manager.wait_expected_calls_done().await);
     }
+
+    #[tokio::test]
+    async fn test_update_controller_multicast_list() {
+        let session_id = 0x123;
+        let session_type = SessionType::FiraRangingSession;
+        let params = generate_params();
+        let tlvs = params.generate_tlvs();
+        let action = UpdateMulticastListAction::AddControlee;
+        let controlees = vec![Controlee { short_address: 0x13, subsession_id: 0x24 }];
+
+        let controlees_clone = controlees.clone();
+        let (mut session_manager, mut mock_uci_manager) =
+            setup_session_manager(move |uci_manager| {
+                let state_init_notf = vec![UciNotification::SessionStatus {
+                    session_id,
+                    session_state: SessionState::SessionStateInit,
+                    reason_code: ReasonCode::StateChangeWithSessionManagementCommands,
+                }];
+                let state_idle_notf = vec![UciNotification::SessionStatus {
+                    session_id,
+                    session_state: SessionState::SessionStateIdle,
+                    reason_code: ReasonCode::StateChangeWithSessionManagementCommands,
+                }];
+                let multicast_list_notf =
+                    vec![UciNotification::SessionUpdateControllerMulticastList {
+                        session_id,
+                        remaining_multicast_list_size: 1,
+                        status_list: vec![ControleeStatus {
+                            mac_address: 0x13,
+                            subsession_id: 0x24,
+                            status: MulticastUpdateStatusCode::StatusOkMulticastListUpdate,
+                        }],
+                    }];
+                uci_manager.expect_session_init(session_id, session_type, state_init_notf, Ok(()));
+                uci_manager.expect_session_set_app_config(
+                    session_id,
+                    tlvs,
+                    state_idle_notf,
+                    Ok(SetAppConfigResponse {
+                        status: StatusCode::UciStatusOk,
+                        config_status: vec![],
+                    }),
+                );
+                uci_manager.expect_session_update_controller_multicast_list(
+                    session_id,
+                    action,
+                    controlees_clone,
+                    multicast_list_notf,
+                    Ok(()),
+                );
+            })
+            .await;
+
+        let result = session_manager.init_session(session_id, session_type, params).await;
+        assert_eq!(result, Ok(()));
+        let result =
+            session_manager.update_controller_multicast_list(session_id, action, controlees).await;
+        assert_eq!(result, Ok(()));
+
+        assert!(mock_uci_manager.wait_expected_calls_done().await);
+    }
+
+    #[tokio::test]
+    async fn test_update_controller_multicast_list_without_notification() {
+        let session_id = 0x123;
+        let session_type = SessionType::FiraRangingSession;
+        let params = generate_params();
+        let tlvs = params.generate_tlvs();
+        let action = UpdateMulticastListAction::AddControlee;
+        let controlees = vec![Controlee { short_address: 0x13, subsession_id: 0x24 }];
+
+        let controlees_clone = controlees.clone();
+        let (mut session_manager, mut mock_uci_manager) =
+            setup_session_manager(move |uci_manager| {
+                let state_init_notf = vec![UciNotification::SessionStatus {
+                    session_id,
+                    session_state: SessionState::SessionStateInit,
+                    reason_code: ReasonCode::StateChangeWithSessionManagementCommands,
+                }];
+                let state_idle_notf = vec![UciNotification::SessionStatus {
+                    session_id,
+                    session_state: SessionState::SessionStateIdle,
+                    reason_code: ReasonCode::StateChangeWithSessionManagementCommands,
+                }];
+                uci_manager.expect_session_init(session_id, session_type, state_init_notf, Ok(()));
+                uci_manager.expect_session_set_app_config(
+                    session_id,
+                    tlvs,
+                    state_idle_notf,
+                    Ok(SetAppConfigResponse {
+                        status: StatusCode::UciStatusOk,
+                        config_status: vec![],
+                    }),
+                );
+                uci_manager.expect_session_update_controller_multicast_list(
+                    session_id,
+                    action,
+                    controlees_clone,
+                    vec![], // Not sending notification.
+                    Ok(()),
+                );
+            })
+            .await;
+
+        let result = session_manager.init_session(session_id, session_type, params).await;
+        assert_eq!(result, Ok(()));
+        // This method should timeout waiting for the notification.
+        let result =
+            session_manager.update_controller_multicast_list(session_id, action, controlees).await;
+        assert_eq!(result, Err(Error::Timeout));
+
+        assert!(mock_uci_manager.wait_expected_calls_done().await);
+    }
 }
diff --git a/src/rust/uwb_core/src/session/uwb_session.rs b/src/rust/uwb_core/src/session/uwb_session.rs
index def964f..3717348 100644
--- a/src/rust/uwb_core/src/session/uwb_session.rs
+++ b/src/rust/uwb_core/src/session/uwb_session.rs
@@ -21,12 +21,18 @@ use tokio::time::timeout;
 use crate::session::error::{Error, Result};
 use crate::session::params::AppConfigParams;
 use crate::uci::error::StatusCode;
-use crate::uci::params::{SessionId, SessionState, SessionType};
+use crate::uci::params::{
+    Controlee, ControleeStatus, MulticastUpdateStatusCode, SessionId, SessionState, SessionType,
+    UpdateMulticastListAction,
+};
 use crate::uci::uci_manager::UciManager;
 
+const NOTIFICATION_TIMEOUT_MS: u64 = 1000;
+
 pub(crate) struct UwbSession {
     cmd_sender: mpsc::UnboundedSender<(Command, oneshot::Sender<Result<()>>)>,
     state_sender: watch::Sender<SessionState>,
+    controlee_status_notf_sender: Option<oneshot::Sender<Vec<ControleeStatus>>>,
 }
 
 impl UwbSession {
@@ -49,7 +55,7 @@ impl UwbSession {
         );
         tokio::spawn(async move { actor.run().await });
 
-        Self { cmd_sender, state_sender }
+        Self { cmd_sender, state_sender, controlee_status_notf_sender: None }
     }
 
     pub fn initialize(
@@ -80,9 +86,29 @@ impl UwbSession {
         let _ = self.cmd_sender.send((Command::Reconfigure { params }, result_sender));
     }
 
-    pub fn set_state(&mut self, state: SessionState) {
+    pub fn update_controller_multicast_list(
+        &mut self,
+        action: UpdateMulticastListAction,
+        controlees: Vec<Controlee>,
+        result_sender: oneshot::Sender<Result<()>>,
+    ) {
+        let (notf_sender, notf_receiver) = oneshot::channel();
+        self.controlee_status_notf_sender = Some(notf_sender);
+        let _ = self.cmd_sender.send((
+            Command::UpdateControllerMulticastList { action, controlees, notf_receiver },
+            result_sender,
+        ));
+    }
+
+    pub fn on_session_status_changed(&mut self, state: SessionState) {
         let _ = self.state_sender.send(state);
     }
+
+    pub fn on_controller_multicast_list_udpated(&mut self, status_list: Vec<ControleeStatus>) {
+        if let Some(sender) = self.controlee_status_notf_sender.take() {
+            let _ = sender.send(status_list);
+        }
+    }
 }
 
 struct UwbSessionActor<T: UciManager> {
@@ -113,14 +139,28 @@ impl<T: UciManager> UwbSessionActor<T> {
                         None => {
                             debug!("UwbSession is about to drop.");
                             break;
-                        },
+                        }
                         Some((cmd, result_sender)) => {
                             let result = match cmd {
                                 Command::Initialize { params } => self.initialize(params).await,
                                 Command::Deinitialize => self.deinitialize().await,
                                 Command::StartRanging => self.start_ranging().await,
                                 Command::StopRanging => self.stop_ranging().await,
-                                Command::Reconfigure { params } => self.reconfigure(params).await,
+                                Command::Reconfigure { params } => {
+                                    self.reconfigure(params).await
+                                }
+                                Command::UpdateControllerMulticastList {
+                                    action,
+                                    controlees,
+                                    notf_receiver,
+                                } => {
+                                    self.update_controller_multicast_list(
+                                        action,
+                                        controlees,
+                                        notf_receiver,
+                                    )
+                                    .await
+                                }
                             };
                             let _ = result_sender.send(result);
                         }
@@ -230,23 +270,65 @@ impl<T: UciManager> UwbSessionActor<T> {
         Ok(())
     }
 
-    async fn wait_state(&mut self, expected_state: SessionState) -> Result<()> {
-        const WAIT_STATE_TIMEOUT_MS: u64 = 1000;
-        match timeout(Duration::from_millis(WAIT_STATE_TIMEOUT_MS), self.state_receiver.changed())
+    async fn update_controller_multicast_list(
+        &mut self,
+        action: UpdateMulticastListAction,
+        controlees: Vec<Controlee>,
+        notf_receiver: oneshot::Receiver<Vec<ControleeStatus>>,
+    ) -> Result<()> {
+        let state = *self.state_receiver.borrow();
+        if !matches!(state, SessionState::SessionStateIdle | SessionState::SessionStateActive) {
+            error!("Cannot update multicast list at state {:?}", state);
+            return Err(Error::WrongState(state));
+        }
+
+        self.uci_manager
+            .session_update_controller_multicast_list(self.session_id, action, controlees)
             .await
-        {
-            Ok(result) => {
-                if result.is_err() {
-                    debug!("UwbSession is about to drop.");
-                    return Err(Error::TokioFailure);
+            .map_err(|e| {
+                error!("Failed to update multicast list: {:?}", e);
+                Error::Uci
+            })?;
+
+        // Wait for the notification of the update status.
+        let results = timeout(Duration::from_millis(NOTIFICATION_TIMEOUT_MS), notf_receiver)
+            .await
+            .map_err(|_| {
+                error!("Timeout waiting for the multicast list notification");
+                Error::Timeout
+            })?
+            .map_err(|_| {
+                error!("oneshot sender is dropped.");
+                Error::TokioFailure
+            })?;
+
+        // Check the update status for adding new controlees.
+        if action == UpdateMulticastListAction::AddControlee {
+            for result in results.iter() {
+                if result.status != MulticastUpdateStatusCode::StatusOkMulticastListUpdate {
+                    error!("Failed to update multicast list: {:?}", result);
+                    return Err(Error::Uci);
                 }
             }
-            Err(_) => {
-                error!("Timeout waiting for the session status notification");
-                return Err(Error::Timeout);
-            }
         }
 
+        Ok(())
+    }
+
+    async fn wait_state(&mut self, expected_state: SessionState) -> Result<()> {
+        // Wait for the notification of the session status.
+        timeout(Duration::from_millis(NOTIFICATION_TIMEOUT_MS), self.state_receiver.changed())
+            .await
+            .map_err(|_| {
+                error!("Timeout waiting for the session status notification");
+                Error::Timeout
+            })?
+            .map_err(|_| {
+                debug!("UwbSession is about to drop.");
+                Error::TokioFailure
+            })?;
+
+        // Check if the latest session status is expected or not.
         let state = *self.state_receiver.borrow();
         if state != expected_state {
             error!(
@@ -261,9 +343,18 @@ impl<T: UciManager> UwbSessionActor<T> {
 }
 
 enum Command {
-    Initialize { params: AppConfigParams },
+    Initialize {
+        params: AppConfigParams,
+    },
     Deinitialize,
     StartRanging,
     StopRanging,
-    Reconfigure { params: AppConfigParams },
+    Reconfigure {
+        params: AppConfigParams,
+    },
+    UpdateControllerMulticastList {
+        action: UpdateMulticastListAction,
+        controlees: Vec<Controlee>,
+        notf_receiver: oneshot::Receiver<Vec<ControleeStatus>>,
+    },
 }
diff --git a/src/rust/uwb_core/src/uci/mock_uci_manager.rs b/src/rust/uwb_core/src/uci/mock_uci_manager.rs
index f3aa425..47010a0 100644
--- a/src/rust/uwb_core/src/uci/mock_uci_manager.rs
+++ b/src/rust/uwb_core/src/uci/mock_uci_manager.rs
@@ -168,6 +168,7 @@ impl MockUciManager {
         expected_session_id: SessionId,
         expected_action: UpdateMulticastListAction,
         expected_controlees: Vec<Controlee>,
+        notfs: Vec<UciNotification>,
         out: Result<()>,
     ) {
         self.expected_calls.lock().unwrap().push_back(
@@ -175,6 +176,7 @@ impl MockUciManager {
                 expected_session_id,
                 expected_action,
                 expected_controlees,
+                notfs,
                 out,
             },
         );
@@ -514,6 +516,7 @@ impl UciManager for MockUciManager {
                 expected_session_id,
                 expected_action,
                 expected_controlees,
+                notfs,
                 out,
             }) if expected_session_id == session_id
                 && expected_action == action
@@ -522,6 +525,9 @@ impl UciManager for MockUciManager {
                 }) =>
             {
                 self.expect_call_consumed.notify_one();
+                for notf in notfs.into_iter() {
+                    let _ = self.notf_sender.as_mut().unwrap().send(notf);
+                }
                 out
             }
             Some(call) => {
@@ -706,6 +712,7 @@ enum ExpectedCall {
         expected_session_id: SessionId,
         expected_action: UpdateMulticastListAction,
         expected_controlees: Vec<Controlee>,
+        notfs: Vec<UciNotification>,
         out: Result<()>,
     },
     RangeStart {
diff --git a/src/rust/uwb_core/src/uci/notification.rs b/src/rust/uwb_core/src/uci/notification.rs
index 20debb6..89b4bb1 100644
--- a/src/rust/uwb_core/src/uci/notification.rs
+++ b/src/rust/uwb_core/src/uci/notification.rs
@@ -19,8 +19,8 @@ use uwb_uci_packets::Packet;
 
 use crate::uci::error::{Error, Result as UciResult, StatusCode};
 use crate::uci::params::{
-    ControleeStatus, DeviceState, ExtendedAddressTwoWayRangingMeasurement, RawVendorMessage,
-    ReasonCode, SessionId, SessionState, ShortAddressTwoWayRangingMeasurement,
+    ControleeStatus, DeviceState, ExtendedAddressTwoWayRangingMeasurement, RangingMeasurementType,
+    RawVendorMessage, ReasonCode, SessionId, SessionState, ShortAddressTwoWayRangingMeasurement,
 };
 
 #[derive(Debug, Clone)]
@@ -37,11 +37,22 @@ pub(crate) enum UciNotification {
         remaining_multicast_list_size: usize,
         status_list: Vec<ControleeStatus>,
     },
-    ShortMacTwoWayRangeData(Vec<ShortAddressTwoWayRangingMeasurement>),
-    ExtendedMacTwoWayRangeData(Vec<ExtendedAddressTwoWayRangingMeasurement>),
+    RangeData {
+        sequence_number: u32,
+        session_id: SessionId,
+        current_ranging_interval_ms: u32,
+        ranging_measurement_type: RangingMeasurementType,
+        ranging_measurements: RangingMeasurements,
+    },
     RawVendor(RawVendorMessage),
 }
 
+#[derive(Debug, Clone)]
+pub(crate) enum RangingMeasurements {
+    Short(Vec<ShortAddressTwoWayRangingMeasurement>),
+    Extended(Vec<ExtendedAddressTwoWayRangingMeasurement>),
+}
+
 impl UciNotification {
     pub fn need_retry(&self) -> bool {
         matches!(self, Self::CoreGenericError(StatusCode::UciStatusCommandRetry))
@@ -120,19 +131,22 @@ impl TryFrom<uwb_uci_packets::RangeDataNtfPacket> for UciNotification {
     type Error = Error;
     fn try_from(evt: uwb_uci_packets::RangeDataNtfPacket) -> Result<Self, Self::Error> {
         use uwb_uci_packets::RangeDataNtfChild;
-        match evt.specialize() {
+        let ranging_measurements = match evt.specialize() {
             RangeDataNtfChild::ShortMacTwoWayRangeDataNtf(evt) => {
-                Ok(UciNotification::ShortMacTwoWayRangeData(
-                    evt.get_two_way_ranging_measurements().clone(),
-                ))
+                RangingMeasurements::Short(evt.get_two_way_ranging_measurements().clone())
             }
             RangeDataNtfChild::ExtendedMacTwoWayRangeDataNtf(evt) => {
-                Ok(UciNotification::ExtendedMacTwoWayRangeData(
-                    evt.get_two_way_ranging_measurements().clone(),
-                ))
+                RangingMeasurements::Extended(evt.get_two_way_ranging_measurements().clone())
             }
-            _ => Err(Error::Specialize(evt.to_vec())),
-        }
+            _ => return Err(Error::Specialize(evt.to_vec())),
+        };
+        Ok(UciNotification::RangeData {
+            sequence_number: evt.get_sequence_number(),
+            session_id: evt.get_session_id(),
+            current_ranging_interval_ms: evt.get_current_ranging_interval(),
+            ranging_measurement_type: evt.get_ranging_measurement_type(),
+            ranging_measurements,
+        })
     }
 }
 
diff --git a/src/rust/uwb_core/src/uci/params.rs b/src/rust/uwb_core/src/uci/params.rs
index 72c39b8..0bfd815 100644
--- a/src/rust/uwb_core/src/uci/params.rs
+++ b/src/rust/uwb_core/src/uci/params.rs
@@ -26,8 +26,9 @@ use crate::uci::error::StatusCode;
 pub use uwb_uci_packets::{
     AppConfigStatus, AppConfigTlv, AppConfigTlvType, CapTlv, CapTlvType, Controlee,
     ControleeStatus, DeviceConfigId, DeviceConfigStatus, DeviceConfigTlv, DeviceState,
-    ExtendedAddressTwoWayRangingMeasurement, PowerStats, ReasonCode, ResetConfig, SessionState,
-    SessionType, ShortAddressTwoWayRangingMeasurement, UpdateMulticastListAction,
+    ExtendedAddressTwoWayRangingMeasurement, MulticastUpdateStatusCode, PowerStats,
+    RangingMeasurementType, ReasonCode, ResetConfig, SessionState, SessionType,
+    ShortAddressTwoWayRangingMeasurement, UpdateMulticastListAction,
 };
 
 pub type SessionId = u32;
diff --git a/src/rust/uwb_core/src/uci/uci_manager.rs b/src/rust/uwb_core/src/uci/uci_manager.rs
index 1f72954..2e2a912 100644
--- a/src/rust/uwb_core/src/uci/uci_manager.rs
+++ b/src/rust/uwb_core/src/uci/uci_manager.rs
@@ -276,6 +276,10 @@ impl UciManager for UciManagerImpl {
         action: UpdateMulticastListAction,
         controlees: Vec<Controlee>,
     ) -> Result<()> {
+        if !(1..=8).contains(&controlees.len()) {
+            warn!("Number of controlees should be between 1 to 8");
+            return Err(Error::InvalidArgs);
+        }
         let cmd =
             UciCommand::SessionUpdateControllerMulticastList { session_id, action, controlees };
         match self.send_cmd(UciManagerCmd::SendUciCommand { cmd }).await {