Allow duplicate pending APS ACKs and take ownership over APS seq

crates/ziggurat-driver/src/rng.rs

@@ -8,6 +8,13 @@ pub fn random_f32() -> f32 {
     (u32::from_le_bytes(bytes) >> 8) as f32 / (1u32 << 24) as f32
 }
 
+/// A uniform `u8`, for the initial APS counter.
+pub fn random_u8() -> u8 {
+    let mut bytes = [0u8; 1];
+    fill_bytes(&mut bytes);
+    bytes[0]
+}
+
 /// A uniform `u16`, for stochastic address allocation.
 pub fn random_u16() -> u16 {
     let mut bytes = [0u8; 2];

crates/ziggurat-driver/src/zigbee_stack.rs

@@ -222,7 +222,8 @@ pub enum HostRoute {
 }
 
 /// Whether a unicast APS data frame requests an end-to-end acknowledgement. When it
-/// does, [`ZigbeeStack::send_aps_command`] returns an [`ApsAckWaiter`] to await it.
+/// does, the [`SendHandle`]'s `delivered` stage resolves on that ack rather than on
+/// next-hop acceptance.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum ApsAck {
     Request,
@@ -400,6 +401,7 @@ pub struct Broadcast {
 /// the held slot.
 #[derive(Debug)]
 pub struct PendingApsAck {
+    pub(crate) ack_data: ApsAckData,
     pub(crate) slot: Arc<SendSlot>,
     pub(crate) deadline: CoreInstant,
 }
@@ -572,6 +574,10 @@ pub struct Nib {
 /// link-key store and APS-layer counter.
 #[derive(Debug)]
 pub struct Aib {
+    /// The APS counter stamped on every outgoing APS frame, stack-owned so that one
+    /// counter space covers both stack-originated frames (ZDP, APS commands) and host
+    /// sends. Randomly seeded, since a restart must not reuse the counters a peer still
+    /// holds in its duplicate-rejection table.
     pub aps_counter: u8,
     /// APS-layer security material and operations (`apsDeviceKeyPairSet`, link-key
     /// derivation, command encryption). Holds the non-spec TCLK seed used to derive
@@ -678,7 +684,11 @@ pub struct State {
     /// All mutable protocol state, behind one lock
     pub core: Mutex<ZigbeeCore>,
 
-    pub pending_aps_acks: Mutex<FlatMap<ApsAckData, PendingApsAck>>,
+    /// Sends awaiting an end-to-end APS ack. Unkeyed, like
+    /// [`Self::pending_unicast_retries`]: an ack key does not identify an entry, so
+    /// several in-flight frames can share one. Insertion order is load-bearing — an ack
+    /// resolves the oldest match — so every mutation here must preserve it.
+    pub pending_aps_acks: Mutex<Vec<PendingApsAck>>,
     pub pending_routes: Mutex<FlatMap<Nwk, PendingRoute>>,
     /// Broadcasts awaiting retransmission, keyed by (source, sequence number).
     pub pending_broadcasts: Mutex<FlatMap<(Nwk, u8), PendingBroadcast>>,
@@ -763,7 +773,7 @@ impl State {
                     address_map: AddressMap::new(config.network_address, config.ieee_address),
                 },
                 aib: Aib {
-                    aps_counter: 0,
+                    aps_counter: crate::rng::random_u8(),
                     aps_security: ApsSecurity::new(
                         config.tc_link_key.clone(),
                         config.ieee_address,
@@ -780,7 +790,7 @@ impl State {
                 trust_center_joins_until: None,
                 beacon_spam_until: None,
             }),
-            pending_aps_acks: Mutex::new(FlatMap::new()),
+            pending_aps_acks: Mutex::new(Vec::new()),
             pending_routes: Mutex::new(FlatMap::new()),
             pending_broadcasts: Mutex::new(FlatMap::new()),
             pending_unicast_retries: Mutex::new(Vec::new()),

crates/ziggurat-driver/src/zigbee_stack/aps.rs

@@ -85,8 +85,15 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
         let ack_data = ApsAckData::from_aps_ack(nwk_frame.nwk_header.source, ack);
         tracing::trace!("Received APS ack: {ack_data:?}");
 
-        let pending = self.state.pending_aps_acks.lock().remove(&ack_data);
-        if let Some(PendingApsAck { slot, .. }) = pending {
+        // The oldest match: the ack carries no way to tell two frames sharing a key apart
+        let mut pending = self.state.pending_aps_acks.lock();
+        let matched = pending
+            .iter()
+            .position(|entry| entry.ack_data == ack_data)
+            .map(|index| pending.remove(index));
+        drop(pending);
+
+        if let Some(PendingApsAck { slot, .. }) = matched {
             slot.resolve(TrackStage::Delivery, Ok(()));
         }
     }
@@ -210,14 +217,14 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
         dst_ep: u8,
         aps_ack: ApsAck,
         radius: u8,
-        aps_seq: u8,
         data: Vec<u8>,
         aps_security: Option<Eui64>,
         sleepy_destination: bool,
         priority: TxPriority,
         route: RouteDirective,
     ) -> Result<SendHandle, EnqueueError> {
         let asdu = FrameBytes::from_slice(&data).map_err(|_| EnqueueError::PayloadTooLong)?;
+        let aps_seq = self.next_aps_counter();
 
         let aps_frame = ApsDataFrame {
             frame_control: ApsFrameControl {
@@ -286,15 +293,14 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
 
         // An APS-ack send registers its pending ack (with the deadline the timeout
         // reactor uses) before enqueueing so a fast reply is caught.
-        if let Some(ack_data) = &ack_data {
+        let registered_slot = ack_data.is_some().then(|| slot.clone());
+        if let Some(ack_data) = ack_data {
             let deadline = self.core_now() + self.aps_ack_timeout(destination, sleepy_destination);
-            self.state.pending_aps_acks.lock().insert(
-                ack_data.clone(),
-                PendingApsAck {
-                    slot: slot.clone(),
-                    deadline,
-                },
-            );
+            self.state.pending_aps_acks.lock().push(PendingApsAck {
+                ack_data,
+                slot: slot.clone(),
+                deadline,
+            });
             self.aps_ack_wake.notify_one();
         }
 
@@ -312,8 +318,11 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
         // A rejected frame gets no confirmation: unregister its pending ack and drop the
         // handle by returning the admission error.
         if let Err(err) = accepted {
-            if let Some(ack_data) = ack_data {
-                self.state.pending_aps_acks.lock().remove(&ack_data);
+            if let Some(registered_slot) = registered_slot {
+                self.state
+                    .pending_aps_acks
+                    .lock()
+                    .retain(|entry| !Arc::ptr_eq(&entry.slot, &registered_slot));
             }
             return Err(err);
         }
@@ -334,7 +343,6 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
         src_ep: u8,
         dst_ep: u8,
         radius: u8,
-        aps_seq: u8,
         data: Vec<u8>,
         priority: TxPriority,
     ) -> Result<SendHandle, EnqueueError> {
@@ -354,7 +362,7 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
             cluster_id,
             profile_id,
             source_endpoint: src_ep,
-            counter: aps_seq,
+            counter: self.next_aps_counter(),
             asdu,
         };
 
@@ -392,7 +400,6 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
         cluster_id: u16,
         src_ep: u8,
         radius: u8,
-        aps_seq: u8,
         data: Vec<u8>,
         priority: TxPriority,
     ) -> Result<SendHandle, EnqueueError> {
@@ -412,7 +419,7 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
             cluster_id,
             profile_id,
             source_endpoint: src_ep,
-            counter: aps_seq,
+            counter: self.next_aps_counter(),
             asdu,
         };
 
@@ -458,7 +465,7 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
         self.state
             .pending_aps_acks
             .lock()
-            .values()
+            .iter()
             .map(|pending| pending.deadline)
             .min()
     }
@@ -470,21 +477,18 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
     fn expire_aps_acks(&self) {
         let now = self.core_now();
 
-        let due: Vec<(Arc<SendSlot>, bool)> = {
-            let mut pending = self.state.pending_aps_acks.lock();
-            let due: Vec<(ApsAckData, Arc<SendSlot>, bool)> = pending
-                .iter()
-                .filter(|(_, p)| p.deadline <= now || p.slot.is_cancelled())
-                .map(|(key, p)| (key.clone(), p.slot.clone(), p.slot.is_cancelled()))
-                .collect();
-            for (key, _, _) in &due {
-                pending.remove(key);
-            }
-            drop(pending);
-            due.into_iter()
-                .map(|(_, slot, cancelled)| (slot, cancelled))
-                .collect()
-        };
+        let due: Vec<(Arc<SendSlot>, bool)> = self
+            .state
+            .pending_aps_acks
+            .lock()
+            .extract_if(.., |entry| {
+                entry.deadline <= now || entry.slot.is_cancelled()
+            })
+            .map(|entry| {
+                let cancelled = entry.slot.is_cancelled();
+                (entry.slot, cancelled)
+            })
+            .collect();
 
         for (slot, cancelled) in due {
             let result = if cancelled {

crates/ziggurat-protocol/src/bridge.rs

@@ -293,7 +293,6 @@ pub fn send_unicast<P: RadioPhy, R: Runtime>(
             payload.dst_ep,
             aps_ack,
             payload.radius,
-            payload.aps_seq,
             payload.asdu,
             aps_security,
             payload.flags.sleepy_destination,
@@ -318,7 +317,6 @@ pub fn send_broadcast<P: RadioPhy, R: Runtime>(
             payload.src_ep,
             payload.dst_ep,
             payload.radius,
-            payload.aps_seq,
             payload.asdu,
             TxPriority::from_host(payload.priority as i8),
         )
@@ -339,7 +337,6 @@ pub fn send_groupcast<P: RadioPhy, R: Runtime>(
             payload.cluster_id,
             payload.src_ep,
             payload.radius,
-            payload.aps_seq,
             payload.asdu,
             TxPriority::from_host(payload.priority as i8),
         )

crates/ziggurat-protocol/src/wire.rs

@@ -455,7 +455,9 @@ pub struct SendUnicastPayload {
     pub cluster_id: u16,
     pub src_ep: u8,
     pub dst_ep: u8,
-    pub aps_seq: u8,
+    /// Ignored: the stack owns the APS counter space so that host sends cannot collide
+    /// with stack-originated frames (ZDP, APS commands). Kept for wire stability.
+    pub _aps_seq: u8,
     pub radius: u8,
     pub priority: u8, // i8 two's complement
     pub route: RouteControl,
@@ -489,7 +491,8 @@ pub struct SendBroadcastPayload {
     pub cluster_id: u16,
     pub src_ep: u8,
     pub dst_ep: u8,
-    pub aps_seq: u8,
+    /// Ignored, as in [`SendUnicastPayload::_aps_seq`].
+    pub _aps_seq: u8,
     pub radius: u8,
     pub priority: u8, // i8 two's complement
     pub asdu_len: u16,
@@ -513,7 +516,8 @@ pub struct SendGroupcastPayload {
     pub profile_id: u16,
     pub cluster_id: u16,
     pub src_ep: u8,
-    pub aps_seq: u8,
+    /// Ignored, as in [`SendUnicastPayload::_aps_seq`].
+    pub _aps_seq: u8,
     pub radius: u8,
     pub priority: u8, // i8 two's complement
     pub asdu_len: u16,