Fix panic in Yamux state machine when a remote closes a substream wit…

…h an active timeout (#1122) * Fix Yamux state mismatch * Review usage of `wake_up_after` * CHANGELOG * Fix PR link
smol-dot · Sep 8, 2023 · 85f6cad · 85f6cad
1 parent 752d699
commit 85f6cad
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 13 deletions.
diff --git a/lib/src/libp2p/collection/single_stream.rs b/lib/src/libp2p/collection/single_stream.rs
@@ -788,10 +788,6 @@ where
                     return;
                 }
 
-                // `read_write()` should be called again as soon as possible after `timeout` in
-                // order for the check above to work.
-                read_write.wake_up_after(&timeout);
-
                 loop {
                     let (read_before, written_before) =
                         (read_write.read_bytes, read_write.write_bytes_queued);
@@ -818,6 +814,9 @@ where
                             if (read_before, written_before)
                                 == (read_write.read_bytes, read_write.write_bytes_queued) =>
                         {
+                            // `read_write()` should be called again as soon as possible
+                            // after `timeout`.
+                            read_write.wake_up_after(&timeout);
                             self.connection = SingleStreamConnectionTaskInner::Handshake {
                                 handshake: updated_handshake,
                                 randomness_seed,

diff --git a/lib/src/libp2p/connection/established/single_stream.rs b/lib/src/libp2p/connection/established/single_stream.rs
@@ -149,7 +149,14 @@ where
                 return Ok((self, Some(Event::PingOutFailed)));
             }
         }
-        read_write.wake_up_after(&self.inner.next_ping);
+
+        // Only wake up if the connection can be closed.
+        // TODO: review w.r.t. https://github.com/smol-dot/smoldot/issues/1121
+        if read_write.expected_incoming_bytes.is_some()
+            && read_write.write_bytes_queueable.is_some()
+        {
+            read_write.wake_up_after(&self.inner.next_ping);
+        }
 
         // If we have both sent and received a GoAway frame, that means that no new substream
         // can be opened. If in addition to this there is no substream in the connection,

diff --git a/lib/src/libp2p/connection/established/substream.rs b/lib/src/libp2p/connection/established/substream.rs
@@ -476,8 +476,6 @@ where
                     );
                 }
 
-                read_write.wake_up_after(&timeout);
-
                 if let Some(extracted_negotiation) = negotiation.take() {
                     match extracted_negotiation.read_write(read_write) {
                         Ok(multistream_select::Negotiation::InProgress(nego)) => {
@@ -590,6 +588,8 @@ where
                     }
                 }
 
+                read_write.wake_up_after(&timeout);
+
                 (
                     Some(SubstreamInner::NotificationsOutHandshakeRecv {
                         timeout,
@@ -667,7 +667,6 @@ where
                         }),
                     );
                 }
-                read_write.wake_up_after(&timeout);
 
                 if let Some(extracted_nego) = negotiation.take() {
                     match extracted_nego.read_write(read_write) {
@@ -754,6 +753,8 @@ where
                     }
                 }
 
+                read_write.wake_up_after(&timeout);
+
                 (
                     Some(SubstreamInner::RequestOut {
                         timeout,
@@ -1066,7 +1067,6 @@ where
                             return (Some(SubstreamInner::PingOutFailed { queued_pings }), None);
                         }
                         if queued_pings.remove(0).is_some() {
-                            read_write.wake_up_asap();
                             return (
                                 Some(SubstreamInner::PingOut {
                                     negotiation,

diff --git a/lib/src/libp2p/connection/yamux.rs b/lib/src/libp2p/connection/yamux.rs
@@ -1187,6 +1187,10 @@ where
                                 if fin {
                                     *remote_write_closed = true;
 
+                                    // No need to send window frames anymore if the remote has
+                                    // sent a FIN.
+                                    self.inner.window_frames_to_send.remove(&stream_id);
+
                                     // Wake up the substream.
                                     match substreams_wake_up_key {
                                         Some(Some(when)) if *when <= outer_read_write.now => {}
@@ -1960,7 +1964,7 @@ where
 
         // When to wake up the substream for reading again.
         debug_assert!(matches!(substreams_wake_up_key, None));
-        let will_wake_up_again_asap_read = match (
+        let will_wake_up_read_again = match (
             self.inner_read_write.read_bytes,
             &self.inner_read_write.wake_up_after,
         ) {
@@ -1976,7 +1980,7 @@ where
                     .substreams_wake_up
                     .insert((Some(when.clone()), self.substream_id));
                 *substreams_wake_up_key = Some(Some(when.clone()));
-                false
+                true
             }
             _ => {
                 // Non-zero bytes written or `when <= now`.
@@ -2089,10 +2093,10 @@ where
             // Substream has nothing to write.
         }
 
-        // TODO: review w.r.t. wake ups
+        // Mark the substream as dead if it won't ever wake up again.
         if matches!(local_write_close, SubstreamStateLocalWrite::FinQueued)
             && *remote_write_closed
-            && !will_wake_up_again_asap_read
+            && !will_wake_up_read_again
             && !self
                 .yamux
                 .inner
@@ -2106,6 +2110,22 @@ where
             {
                 let _was_inserted = self.yamux.inner.dead_substreams.insert(self.substream_id);
                 debug_assert!(_was_inserted);
+                debug_assert!(!self
+                    .yamux
+                    .inner
+                    .substreams_wake_up
+                    .iter()
+                    .any(|(_, s)| *s == self.substream_id));
+                debug_assert!(!self
+                    .yamux
+                    .inner
+                    .substreams_write_ready
+                    .contains(&self.substream_id));
+                debug_assert!(!self
+                    .yamux
+                    .inner
+                    .window_frames_to_send
+                    .contains_key(&self.substream_id));
             }
         }
 

diff --git a/wasm-node/CHANGELOG.md b/wasm-node/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Fix panic in Yamux state machine when a remote closes a substream with an active timeout. ([#1122](https://github.com/smol-dot/smoldot/pull/1122))
+
 ## 2.0.0 - 2023-09-07
 
 ### Remove