diff --git a/src/port_forward_mux.rs b/src/port_forward_mux.rs index 8efc251d0..e71971c43 100644 --- a/src/port_forward_mux.rs +++ b/src/port_forward_mux.rs @@ -66,6 +66,16 @@ impl RecvWindow { } } + /// Widens the window without advertising: the grant travels in `opened`. + pub fn grant(&mut self, n: u32) { + self.remaining = self.remaining.saturating_add(n); + } + + #[cfg(test)] + pub fn remaining(&self) -> u32 { + self.remaining + } + /// Returns the amount to advertise in a `window_update` once enough has /// been drained; the same amount is credited back. pub fn drained(&mut self, n: usize) -> Option { @@ -781,6 +791,16 @@ mod tests { } } + #[test] + fn grant_extends_a_window_that_has_been_used_up() { + let mut w = RecvWindow::new(INITIAL_WINDOW); + assert!(w.accept(INITIAL_WINDOW as usize)); + assert!(!w.accept(1)); + w.grant(CHANNEL_WINDOW - INITIAL_WINDOW); + assert!(w.accept((CHANNEL_WINDOW - INITIAL_WINDOW) as usize)); + assert!(!w.accept(1)); + } + #[test] fn send_credit_blocks_at_zero_and_resumes_on_add() { rt().block_on(async { diff --git a/src/server/port_forward_mux.rs b/src/server/port_forward_mux.rs index fedb6ba9e..f928490b5 100644 --- a/src/server/port_forward_mux.rs +++ b/src/server/port_forward_mux.rs @@ -121,7 +121,7 @@ impl PortForwardMux { let (addr, is_rdp) = Connection::normalize_port_forward_target(&mut pf); let (inbound_tx, inbound_rx) = mpsc::unbounded_channel(); let credit = Arc::new(SendCredit::new(effective_window(open.window))); - let window = Arc::new(Mutex::new(RecvWindow::new(CHANNEL_WINDOW))); + let window = Arc::new(Mutex::new(RecvWindow::new(INITIAL_WINDOW))); self.channels.insert( id, Entry { @@ -178,6 +178,11 @@ impl PortForwardMux { self.channels.len() } + #[cfg(test)] + pub fn recv_window_remaining(&self, id: i32) -> Option { + self.channels.get(&id).map(|e| e.window.lock().unwrap().remaining()) + } + /// Dropping every sender ends every task; each drops its target socket. pub fn close_all(&mut self) { self.channels.clear(); @@ -234,6 +239,8 @@ async fn run_controlled_channel( } } }; + // Granted before `opened` leaves, so the peer can never be ahead of it. + window.lock().unwrap().grant(CHANNEL_WINDOW - INITIAL_WINDOW); if sink .send_ordered(opened_msg(id, true, "", CHANNEL_WINDOW)) .await @@ -429,6 +436,21 @@ mod tests { }); } + #[test] + fn demux_admits_only_the_initial_window_before_opened() { + rt().block_on(async { + let port = echo_target().await; + let (tx, mut rx) = mpsc::unbounded_channel(); + let mut mux = PortForwardMux::new(tx, format!("127.0.0.1:{}", port)); + mux.handle(open(1, port), || true); + // The channel task has not run yet: the demultiplexer alone + // decides what may sit in the queue before `opened`. + assert_eq!(mux.recv_window_remaining(1), Some(INITIAL_WINDOW)); + assert_eq!(opened(&next_frame(&mut rx).await), (1, true)); + assert_eq!(mux.recv_window_remaining(1), Some(CHANNEL_WINDOW)); + }); + } + #[test] fn pending_bytes_are_bounded_by_initial_window_before_opened() { // A loopback connect completes before a task can observe "connecting",