tweak(yamux): Simplify reducer a bit more

Author: tizoc

p2p/src/network/yamux/p2p_network_yamux_reducer.rs

@@ -107,69 +107,21 @@ impl P2pNetworkYamuxState {
                 Ok(())
             }
             P2pNetworkYamuxAction::IncomingFrame { addr } => {
-                let mut pending_outgoing = VecDeque::default();
                 let Some(frame) = yamux_state.incoming.pop_front() else {
                     bug_condition!(
                         "Frame not found for action `P2pNetworkYamuxAction::IncomingFrame`"
                     );
                     return Ok(());
                 };
 
-                if frame.flags.contains(YamuxFlags::SYN) {
-                    yamux_state
-                        .streams
-                        .insert(frame.stream_id, YamuxStreamState::incoming());
-
-                    if frame.stream_id != 0 {
-                        connection_state.streams.insert(
-                            frame.stream_id,
-                            P2pNetworkStreamState::new_incoming(meta.time()),
-                        );
-                    }
-                }
-                if frame.flags.contains(YamuxFlags::ACK) {
-                    yamux_state
-                        .streams
-                        .entry(frame.stream_id)
-                        .or_default()
-                        .established = true;
-                }
+                YamuxStreamState::handle_frame_syn_ack_flags(
+                    &mut yamux_state.streams,
+                    &mut connection_state.streams,
+                    &frame,
+                    meta.time(),
+                );
 
-                match &frame.inner {
-                    YamuxFrameInner::Data(_) => {
-                        if let Some(stream) = yamux_state.streams.get_mut(&frame.stream_id) {
-                            // must not underflow
-                            // TODO: check it and disconnect peer that violates flow rules
-                            stream.window_ours =
-                                stream.window_ours.saturating_sub(frame.len_as_u32());
-                        }
-                    }
-                    YamuxFrameInner::WindowUpdate { difference } => {
-                        let stream = yamux_state
-                            .streams
-                            .entry(frame.stream_id)
-                            .or_insert_with(YamuxStreamState::incoming);
-
-                        stream.window_theirs = stream.window_theirs.saturating_add(*difference);
-
-                        if *difference > 0 {
-                            // have some fresh space in the window
-                            // try send as many frames as can
-                            let mut window = stream.window_theirs;
-                            while let Some(frame) = stream.pending.pop_front() {
-                                let len = frame.len_as_u32();
-                                pending_outgoing.push_back(frame);
-                                if let Some(new_window) = window.checked_sub(len) {
-                                    window = new_window;
-                                } else {
-                                    break;
-                                }
-                            }
-                        }
-                    }
-                    YamuxFrameInner::Ping { .. } => {}
-                    YamuxFrameInner::GoAway(res) => yamux_state.set_res(*res),
-                }
+                let mut pending_outgoing = yamux_state.handle_frame_message(&frame);
 
                 let (dispatcher, state) = state_context.into_dispatcher_and_state();
                 let limits: &P2pLimits = state.substate()?;
@@ -368,4 +320,45 @@ impl P2pNetworkYamuxState {
             }
         }
     }
+
+    fn handle_frame_message(&mut self, frame: &YamuxFrame) -> VecDeque<YamuxFrame> {
+        let mut pending_outgoing = VecDeque::default();
+
+        match &frame.inner {
+            YamuxFrameInner::Data(_) => {
+                if let Some(stream) = self.streams.get_mut(&frame.stream_id) {
+                    // must not underflow
+                    // TODO: check it and disconnect peer that violates flow rules
+                    stream.window_ours = stream.window_ours.saturating_sub(frame.len_as_u32());
+                }
+            }
+            YamuxFrameInner::WindowUpdate { difference } => {
+                let stream = self
+                    .streams
+                    .entry(frame.stream_id)
+                    .or_insert_with(YamuxStreamState::incoming);
+
+                stream.window_theirs = stream.window_theirs.saturating_add(*difference);
+
+                if *difference > 0 {
+                    // have some fresh space in the window
+                    // try send as many frames as can
+                    let mut window = stream.window_theirs;
+                    while let Some(frame) = stream.pending.pop_front() {
+                        let len = frame.len_as_u32();
+                        pending_outgoing.push_back(frame);
+                        if let Some(new_window) = window.checked_sub(len) {
+                            window = new_window;
+                        } else {
+                            break;
+                        }
+                    }
+                }
+            }
+            YamuxFrameInner::Ping { .. } => {}
+            YamuxFrameInner::GoAway(res) => self.set_res(*res),
+        }
+
+        pending_outgoing
+    }
 }

p2p/src/network/yamux/p2p_network_yamux_state.rs

@@ -228,6 +228,30 @@ impl YamuxStreamState {
             ..Default::default()
         }
     }
+
+    /// Processes SYN and ACK flags from an incoming Yamux frame and updates stream states accordingly.
+    pub fn handle_frame_syn_ack_flags(
+        yamux_streams: &mut BTreeMap<StreamId, YamuxStreamState>,
+        connection_streams: &mut BTreeMap<u32, P2pNetworkStreamState>,
+        frame: &YamuxFrame,
+        time: redux::Timestamp,
+    ) {
+        if frame.flags.contains(YamuxFlags::SYN) {
+            yamux_streams.insert(frame.stream_id, YamuxStreamState::incoming());
+
+            // TODO: when is stream 0 used? why is it ignored?
+            if frame.stream_id != 0 {
+                connection_streams
+                    .insert(frame.stream_id, P2pNetworkStreamState::new_incoming(time));
+            }
+        }
+
+        if frame.flags.contains(YamuxFlags::ACK) {
+            if let Some(stream) = yamux_streams.get_mut(&frame.stream_id) {
+                stream.established = true;
+            }
+        }
+    }
 }
 
 bitflags::bitflags! {