Fix endianness conversion for multicast

Bug: 271901693
Test: make com.google.android.uwb
(cherry picked from https://android-review.googlesource.com/q/commit:705a42a496f0c4e80893a6a6dd2f445454bcf968)
Merged-In: I4a4928bd579d599cc960968c3ce8a69667e9cedd
Change-Id: I4a4928bd579d599cc960968c3ce8a69667e9cedd

6 files changed, 22 insertions, 19 deletions

diff --git a/src/rust/uwb_core/src/proto/mappings.rs b/src/rust/uwb_core/src/proto/mappings.rs
index 6444326..cb0e078 100644
--- a/src/rust/uwb_core/src/proto/mappings.rs
+++ b/src/rust/uwb_core/src/proto/mappings.rs
@@ -958,8 +958,7 @@ impl TryFrom<ProtoControlee> for Controlee {
     type Error = String;
     fn try_from(item: ProtoControlee) -> std::result::Result<Self, Self::Error> {
         Ok(Self {
-            short_address: item
-                .short_address
+            short_address: item.short_address.to_ne_bytes()[0..2]
                 .try_into()
                 .map_err(|_| "Failed to convert short_address")?,
             subsession_id: item.subsession_id,
diff --git a/src/rust/uwb_core/src/service/uwb_service.rs b/src/rust/uwb_core/src/service/uwb_service.rs
index a2ecb85..ce00f95 100644
--- a/src/rust/uwb_core/src/service/uwb_service.rs
+++ b/src/rust/uwb_core/src/service/uwb_service.rs
@@ -709,7 +709,8 @@ mod tests {
         let session_type = SessionType::FiraRangingSession;
         let params = generate_params();
         let action = UpdateMulticastListAction::AddControlee;
-        let controlees = vec![Controlee { short_address: 0x13, subsession_id: 0x24 }];
+        let short_address: [u8; 2] = [0x12, 0x34];
+        let controlees = vec![Controlee { short_address, subsession_id: 0x24 }];
 
         let uci_manager = MockUciManager::new();
         let (service, _, _runtime) = setup_uwb_service(uci_manager);
diff --git a/src/rust/uwb_core/src/session/session_manager.rs b/src/rust/uwb_core/src/session/session_manager.rs
index d53ecc7..b1287f3 100644
--- a/src/rust/uwb_core/src/session/session_manager.rs
+++ b/src/rust/uwb_core/src/session/session_manager.rs
@@ -773,7 +773,8 @@ mod tests {
         let params = generate_params();
         let tlvs = params.generate_tlvs();
         let action = UpdateMulticastListAction::AddControlee;
-        let controlees = vec![Controlee { short_address: 0x13, subsession_id: 0x24 }];
+        let short_address: [u8; 2] = [0x12, 0x34];
+        let controlees = vec![Controlee { short_address, subsession_id: 0x24 }];
 
         let controlees_clone = controlees.clone();
         let (mut session_manager, mut mock_uci_manager, _) =
@@ -830,7 +831,8 @@ mod tests {
         let params = generate_ccc_params();
         let tlvs = params.generate_tlvs();
         let action = UpdateMulticastListAction::AddControlee;
-        let controlees = vec![Controlee { short_address: 0x13, subsession_id: 0x24 }];
+        let short_address: [u8; 2] = [0x12, 0x34];
+        let controlees = vec![Controlee { short_address, subsession_id: 0x24 }];
 
         let (mut session_manager, mut mock_uci_manager, _) =
             setup_session_manager(move |uci_manager| {
@@ -869,7 +871,8 @@ mod tests {
         let params = generate_params();
         let tlvs = params.generate_tlvs();
         let action = UpdateMulticastListAction::AddControlee;
-        let controlees = vec![Controlee { short_address: 0x13, subsession_id: 0x24 }];
+        let short_address: [u8; 2] = [0x12, 0x34];
+        let controlees = vec![Controlee { short_address, subsession_id: 0x24 }];
 
         let controlees_clone = controlees.clone();
         let (mut session_manager, mut mock_uci_manager, _) =
diff --git a/src/rust/uwb_core/src/uci/uci_manager.rs b/src/rust/uwb_core/src/uci/uci_manager.rs
index 7dfd878..11213bf 100644
--- a/src/rust/uwb_core/src/uci/uci_manager.rs
+++ b/src/rust/uwb_core/src/uci/uci_manager.rs
@@ -1960,7 +1960,8 @@ mod tests {
         let session_id = 0x123;
         let session_token = 0x123;
         let action = UpdateMulticastListAction::AddControlee;
-        let controlee = Controlee { short_address: 0x4567, subsession_id: 0x90ab };
+        let short_address: [u8; 2] = [0x45, 0x67];
+        let controlee = Controlee { short_address, subsession_id: 0x90ab };
         let controlee_clone = controlee.clone();
 
         let (uci_manager, mut mock_hal) = setup_uci_manager_with_session_initialized(
diff --git a/src/rust/uwb_uci_packets/src/lib.rs b/src/rust/uwb_uci_packets/src/lib.rs
index ce63bfb..38b0910 100644
--- a/src/rust/uwb_uci_packets/src/lib.rs
+++ b/src/rust/uwb_uci_packets/src/lib.rs
@@ -789,8 +789,7 @@ pub enum Controlees {
 // TODO(ziyiw): Replace these functions after making uwb_uci_packets::Controlee::write_to() public.
 pub fn write_controlee(controlee: &Controlee) -> BytesMut {
     let mut buffer = BytesMut::new();
-    let short_address = controlee.short_address;
-    buffer.extend_from_slice(&short_address.to_le_bytes()[0..2]);
+    buffer.extend_from_slice(&controlee.short_address);
     let subsession_id = controlee.subsession_id;
     buffer.extend_from_slice(&subsession_id.to_le_bytes()[0..4]);
     buffer
@@ -798,8 +797,7 @@ pub fn write_controlee(controlee: &Controlee) -> BytesMut {
 
 pub fn write_controlee_2_0_16byte(controlee: &Controlee_V2_0_16_Byte_Version) -> BytesMut {
     let mut buffer = BytesMut::new();
-    let short_address = controlee.short_address;
-    buffer.extend_from_slice(&short_address.to_le_bytes()[0..2]);
+    buffer.extend_from_slice(&controlee.short_address);
     let subsession_id = controlee.subsession_id;
     buffer.extend_from_slice(&subsession_id.to_le_bytes()[0..4]);
     buffer.extend_from_slice(&controlee.subsession_key);
@@ -808,8 +806,7 @@ pub fn write_controlee_2_0_16byte(controlee: &Controlee_V2_0_16_Byte_Version) ->
 
 pub fn write_controlee_2_0_32byte(controlee: &Controlee_V2_0_32_Byte_Version) -> BytesMut {
     let mut buffer = BytesMut::new();
-    let short_address = controlee.short_address;
-    buffer.extend_from_slice(&short_address.to_le_bytes()[0..2]);
+    buffer.extend_from_slice(&controlee.short_address);
     let subsession_id = controlee.subsession_id;
     buffer.extend_from_slice(&subsession_id.to_le_bytes()[0..4]);
     buffer.extend_from_slice(&controlee.subsession_key);
@@ -919,7 +916,8 @@ mod tests {
 
     #[test]
     fn test_write_controlee() {
-        let controlee: Controlee = Controlee { short_address: 2, subsession_id: 3 };
+        let short_address: [u8; 2] = [2, 3];
+        let controlee: Controlee = Controlee { short_address, subsession_id: 3 };
         let bytes = write_controlee(&controlee);
         let parsed_controlee = Controlee::parse(&bytes).unwrap();
         assert_eq!(controlee, parsed_controlee);
@@ -927,7 +925,8 @@ mod tests {
 
     #[test]
     fn test_build_multicast_update_packet() {
-        let controlee = Controlee { short_address: 0x1234, subsession_id: 0x1324_3546 };
+        let short_address: [u8; 2] = [0x12, 0x34];
+        let controlee = Controlee { short_address, subsession_id: 0x1324_3546 };
         let packet: UciControlPacket = build_session_update_controller_multicast_list_cmd(
             0x1425_3647,
             UpdateMulticastListAction::AddControlee,
@@ -942,7 +941,7 @@ mod tests {
             vec![
                 0x21, 0x07, 0x00, 0x0c, // 2(packet info), RFU, payload length(12)
                 0x47, 0x36, 0x25, 0x14, // 4(session id (LE))
-                0x00, 0x01, 0x34, 0x12, // action, # controlee, 2(short address (LE))
+                0x00, 0x01, 0x12, 0x34, // action, # controlee, 2(short address (LE))
                 0x46, 0x35, 0x24, 0x13, // 4(subsession id (LE))
             ]
         );
diff --git a/src/rust/uwb_uci_packets/uci_packets.pdl b/src/rust/uwb_uci_packets/uci_packets.pdl
index b40ca1e..d764a49 100644
--- a/src/rust/uwb_uci_packets/uci_packets.pdl
+++ b/src/rust/uwb_uci_packets/uci_packets.pdl
@@ -879,18 +879,18 @@ test SessionUpdateDtTagRangingRoundsRsp {
 }
 
 struct Controlee {
-    short_address: 16,
+    short_address: 8[2],
     subsession_id: 32,
 }
 
 struct Controlee_V2_0_16_Byte_Version {
-    short_address: 16,
+    short_address: 8[2],
     subsession_id: 32,
     subsession_key: 8[16],
 }
 
 struct Controlee_V2_0_32_Byte_Version {
-    short_address: 16,
+    short_address: 8[2],
     subsession_id: 32,
     subsession_key: 8[32],
 }