fix(yamux): protect from overflow vulnerability

Author: vlad9486

p2p/src/network/scheduler/p2p_network_scheduler_actions.rs

@@ -72,6 +72,7 @@ pub enum P2pNetworkSchedulerAction {
         addr: ConnectionAddr,
         peer_id: PeerId,
         message_size_limit: Limit<usize>,
+        pending_outgoing_limit: Limit<usize>,
     },
 
     /// Action that initiate the specified peer disconnection.

p2p/src/network/scheduler/p2p_network_scheduler_reducer.rs

@@ -249,8 +249,9 @@ impl P2pNetworkSchedulerState {
             }
             P2pNetworkSchedulerAction::YamuxDidInit {
                 addr,
-                message_size_limit,
                 peer_id,
+                message_size_limit,
+                pending_outgoing_limit,
             } => {
                 let Some(cn) = scheduler_state.connections.get_mut(&addr) else {
                     bug_condition!(
@@ -261,6 +262,7 @@ impl P2pNetworkSchedulerState {
                 if let Some(P2pNetworkConnectionMuxState::Yamux(yamux)) = &mut cn.mux {
                     yamux.init = true;
                     yamux.message_size_limit = message_size_limit;
+                    yamux.pending_outgoing_limit = pending_outgoing_limit;
                 }
 
                 let incoming = cn.incoming;
@@ -503,10 +505,13 @@ impl P2pNetworkSchedulerState {
                     return;
                 };
                 let message_size_limit = p2p_state.config.limits.yamux_message_size();
+                let pending_outgoing_limit =
+                    p2p_state.config.limits.yamux_pending_outgoing_per_peer();
                 dispatcher.push(P2pNetworkSchedulerAction::YamuxDidInit {
                     addr,
                     peer_id,
                     message_size_limit,
+                    pending_outgoing_limit,
                 });
             }
             Some(Protocol::Stream(kind)) => {

p2p/src/network/scheduler/p2p_network_scheduler_state.rs

@@ -196,6 +196,8 @@ pub enum P2pNetworkConnectionError {
     StreamReset(StreamId),
     #[error("pubsub error: {0}")]
     PubSubError(String),
+    #[error("peer make us keep too much data")]
+    YamuxOverflow(StreamId),
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]

p2p/src/network/yamux/p2p_network_yamux_reducer.rs

@@ -222,20 +222,15 @@ impl P2pNetworkYamuxState {
                             if difference > 0 {
                                 // have some fresh space in the window
                                 // try send as many frames as can
-                                // last frame
                                 let mut window = stream.window_theirs;
                                 while let Some(mut frame) = stream.pending.pop_front() {
-                                    let size = if let YamuxFrameInner::Data(data) = &frame.inner {
-                                        data.len()
-                                    } else {
-                                        0
-                                    } as u32;
-                                    if let Some(new_window) = window.checked_sub(size) {
+                                    let len = frame.len() as u32;
+                                    if let Some(new_window) = window.checked_sub(len) {
                                         pending_outgoing.push_back(frame);
                                         window = new_window;
                                     } else {
                                         if let Some(remaining) =
-                                            frame.split_at((size - window) as usize)
+                                            frame.split_at((len - window) as usize)
                                         {
                                             stream.pending.push_front(remaining);
                                         }
@@ -360,7 +355,8 @@ impl P2pNetworkYamuxState {
                 Ok(())
             }
             P2pNetworkYamuxAction::OutgoingFrame { mut frame, addr } => {
-                let Some(stream) = yamux_state.streams.get_mut(&frame.stream_id) else {
+                let stream_id = frame.stream_id;
+                let Some(stream) = yamux_state.streams.get_mut(&stream_id) else {
                     return Ok(());
                 };
                 match &mut frame.inner {
@@ -386,6 +382,14 @@ impl P2pNetworkYamuxState {
                             // or the queue is already not empty
                             // in both cases the whole frame goes in the queue and nothing to send
                             stream.pending.push_back(frame);
+                            if stream.pending.iter().map(YamuxFrame::len).sum::<usize>()
+                                > yamux_state.pending_outgoing_limit
+                            {
+                                let dispatcher = state_context.into_dispatcher();
+                                let error = P2pNetworkConnectionError::YamuxOverflow(stream_id);
+                                dispatcher.push(P2pNetworkSchedulerAction::Error { addr, error });
+                            }
+
                             return Ok(());
                         }
                     }

p2p/src/network/yamux/p2p_network_yamux_state.rs

@@ -7,6 +7,7 @@ use super::super::*;
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
 pub struct P2pNetworkYamuxState {
     pub message_size_limit: Limit<usize>,
+    pub pending_outgoing_limit: Limit<usize>,
     pub buffer: Vec<u8>,
     pub incoming: VecDeque<YamuxFrame>,
     pub streams: BTreeMap<StreamId, YamuxStreamState>,
@@ -185,13 +186,21 @@ impl YamuxFrame {
         vec
     }
 
+    pub fn len(&self) -> usize {
+        if let YamuxFrameInner::Data(data) = &self.inner {
+            data.len()
+        } else {
+            0
+        }
+    }
+
     /// If this data is bigger then `pos`, keep only first `pos` bytes and return some remaining
     /// otherwise return none
     pub fn split_at(&mut self, pos: usize) -> Option<Self> {
         use std::ops::Sub;
 
         if let YamuxFrameInner::Data(data) = &mut self.inner {
-            if data.len() == pos {
+            if data.len() <= pos {
                 return None;
             }
             let (keep, rest) = data.split_at(pos);

p2p/src/p2p_config.rs

@@ -252,6 +252,7 @@ pub struct P2pLimits {
     max_peers_in_state: Limit<usize>,
     max_streams: Limit<usize>,
     yamux_message_size: Limit<usize>,
+    yamux_pending_outgoing_per_peer: Limit<usize>,
 
     identify_message: Limit<usize>,
     kademlia_request: Limit<usize>,
@@ -319,6 +320,12 @@ impl P2pLimits {
         /// Sets the maximum number of streams that a peer is allowed to open simultaneously.
         with_yamux_message_size
     );
+    limit!(
+        /// Maximum number of streams from a peer.
+        yamux_pending_outgoing_per_peer,
+        /// Sets the maximum number of streams that a peer is allowed to open simultaneously.
+        with_yamux_pending_outgoing_per_peer
+    );
 
     limit!(
         /// Minimum number of peers.
@@ -400,6 +407,7 @@ impl Default for P2pLimits {
             max_peers_in_state,
             max_streams,
             yamux_message_size,
+            yamux_pending_outgoing_per_peer: rpc_get_staged_ledger,
 
             identify_message,
             kademlia_request,