fix(net): potential deadlock

Author: robinhundt

.github/workflows/pull_request.yml

@@ -24,6 +24,7 @@ jobs:
             rustflags: ""
     env:
       RUSTFLAGS: ${{ matrix.rustflags }}
+      RUST_BACKTRACE: 1
     steps:
       - &checkout
         name: Check out repository
@@ -45,10 +46,10 @@ jobs:
         uses: ilammy/setup-nasm@72793074d3c8cdda771dba85f6deafe00623038b # v1.5.2
 
       - name: Run tests (all features)
-        run: cargo test --workspace --verbose --all-features --no-fail-fast ${{ runner.os == 'macOS' && '-- --test-threads=1' || '' }}
+        run: cargo test --workspace --verbose --all-features --no-fail-fast
 
       - name: Run tests (no features)
-        run: cargo test --workspace --verbose --no-fail-fast ${{ runner.os == 'macOS' && '-- --test-threads=1' || '' }}
+        run: cargo test --workspace --verbose --no-fail-fast
 
   miri:
     name: Miri

cryprot-net/src/lib.rs

@@ -68,6 +68,11 @@ pub struct StreamManager {
     acceptor: QuicStreamAcceptor,
     cmd_send: mpsc::UnboundedSender<Cmd>,
     cmd_recv: mpsc::UnboundedReceiver<Cmd>,
+    maps: StreamMaps,
+}
+
+#[derive(Default)]
+struct StreamMaps {
     pending: HashMap<UniqueId, StreamSend>,
     accepted: HashMap<UniqueId, (QuicRecvStream, usize)>,
 }
@@ -150,8 +155,7 @@ impl StreamManager {
             acceptor,
             cmd_send,
             cmd_recv,
-            pending: Default::default(),
-            accepted: Default::default(),
+            maps: Default::default(),
         }
     }
 
@@ -171,47 +175,32 @@ impl StreamManager {
                             Self::accepted(stream, self.cmd_send.clone());
                         }
                         Ok(None) => {
-                            debug!("remote closed");
-                            return;
+                            debug!("remote closed, draining pending commands");
+                            break;
                         }
                         Err(err) => {
-                            error!(%err, "unable to accept stream");
-                            return;
+                            error!(%err, "unable to accept stream, draining pending commands");
+                            break;
                         }
                     }
                 }
                 Some(cmd) = self.cmd_recv.recv() => {   // recv() is cancel safe
                     debug!(?cmd, "received cmd");
-                    match cmd {
-                        Cmd::NewStream {uid, stream_return} => {
-                            if let Some(accepted) = self.accepted.remove(&uid) {
-                                if stream_return.send(accepted).is_err() {
-                                    debug!("accepted remote stream but local receiver is closed");
-                                }
-                                debug!("sending new stream to receiver");
-                                continue;
-                            }
-                            match self.pending.entry(uid) {
-                                Entry::Occupied(occupied_entry) => {
-                                    panic!("Duplicate unique id: {:?}", occupied_entry.key())
-                                },
-                                Entry::Vacant(vacant_entry) => {vacant_entry.insert(stream_return);},
-                            }
-                        }
-                        Cmd::AcceptedStream {uid, stream, bytes_read} => {
-                            if let Some(stream_ret) = self.pending.remove(&uid) {
-                               if stream_ret.send((stream, bytes_read)).is_err() {
-                                debug!("accepted remote stream but local receiver is closed");
-                               }
-                            } else {
-                                debug!("accepted stream but no pending");
-                                self.accepted.insert(uid, (stream, bytes_read));
-                            }
-                        }
-                    }
+                    self.maps.handle_cmd(cmd);
                 }
             }
         }
+        // The QUIC acceptor is done (remote closed or error), but there may be
+        // in-flight `Self::accepted` tasks that already received a stream and
+        // are reading the UniqueId. Drain remaining commands so those streams
+        // are matched with pending requests.
+        while let Ok(cmd) = self.cmd_recv.try_recv() {
+            debug!(?cmd, "received cmd (draining)");
+            self.maps.handle_cmd(cmd);
+            if self.maps.pending.is_empty() {
+                break;
+            }
+        }
     }
 
     // not taking &self to work around borrow issue
@@ -224,17 +213,54 @@ impl StreamManager {
                     return;
                 }
             };
-            cmd_send
-                .send(Cmd::AcceptedStream {
-                    uid,
-                    stream,
-                    bytes_read,
-                })
-                .expect("cmd_rcv is owned by StreamManager")
+            // StreamManager may have already exited if the connection closed
+            let _ = cmd_send.send(Cmd::AcceptedStream {
+                uid,
+                stream,
+                bytes_read,
+            });
         });
     }
 }
 
+impl StreamMaps {
+    fn handle_cmd(&mut self, cmd: Cmd) {
+        match cmd {
+            Cmd::NewStream { uid, stream_return } => {
+                if let Some(accepted) = self.accepted.remove(&uid) {
+                    if stream_return.send(accepted).is_err() {
+                        debug!("accepted remote stream but local receiver is closed");
+                    }
+                    debug!("sending new stream to receiver");
+                    return;
+                }
+                match self.pending.entry(uid) {
+                    Entry::Occupied(occupied_entry) => {
+                        panic!("Duplicate unique id: {:?}", occupied_entry.key())
+                    }
+                    Entry::Vacant(vacant_entry) => {
+                        vacant_entry.insert(stream_return);
+                    }
+                }
+            }
+            Cmd::AcceptedStream {
+                uid,
+                stream,
+                bytes_read,
+            } => {
+                if let Some(stream_ret) = self.pending.remove(&uid) {
+                    if stream_ret.send((stream, bytes_read)).is_err() {
+                        debug!("accepted remote stream but local receiver is closed");
+                    }
+                } else {
+                    debug!("accepted stream but no pending");
+                    self.accepted.insert(uid, (stream, bytes_read));
+                }
+            }
+        }
+    }
+}
+
 /// Possible connection errors.
 #[derive(thiserror::Error, Debug)]
 pub enum ConnectionError {

cryprot-net/src/testing.rs

@@ -1,4 +1,6 @@
-use std::net::{Ipv4Addr, SocketAddr};
+use std::{
+    net::{Ipv4Addr, SocketAddr},
+};
 
 use anyhow::Context;
 use s2n_quic::{Client, Server, client::Connect, provider::limits::Limits};

cryprot-ot/src/silent_ot.rs

@@ -652,7 +652,7 @@ mod tests {
         }
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn correlated_silent_ot() {
         let _g = init_tracing();
         let (c1, c2) = local_conn().await.unwrap();
@@ -674,7 +674,7 @@ mod tests {
         check_correlated(&r_ot, &s_ot, Some(&choices), delta);
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn random_silent_ot() {
         let _g = init_tracing();
         let (c1, c2) = local_conn().await.unwrap();
@@ -689,7 +689,7 @@ mod tests {
         check_random(count, &s_ot, &r_ot[..], &choices);
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn test_rot_trait_for_silent_ot() {
         let _g = init_tracing();
         let (c1, c2) = local_conn().await.unwrap();
@@ -704,7 +704,7 @@ mod tests {
         check_random(count, &s_ot, &r_ot, &c);
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn test_malicious_silent_ot() {
         let _g = init_tracing();
         let (c1, c2) = local_conn().await.unwrap();