Skip to main content

smolvm_network/
tcp_relay.rs

1//! TCP relay support for the virtio-net backend.
2//!
3//! Context
4//! =======
5//!
6//! In the Phase 1 virtio-net design, guest TCP does not flow directly from the
7//! guest to the outside network through the host kernel. Instead, the host-side
8//! smoltcp runtime terminates the guest-visible TCP connection in userspace and
9//! relays payloads to a normal host `TcpStream`.
10//!
11//! Conceptually:
12//!
13//! ```text
14//! guest app
15//!   -> guest kernel TCP
16//!   -> Ethernet frame
17//!   -> smoltcp TCP socket (inside smolvm)
18//!   -> channel
19//!   -> host TcpStream
20//!   -> remote server
21//! ```
22//!
23//! That means:
24//! - the host runtime can observe every guest TCP byte stream on this NIC
25//! - smoltcp owns the guest-facing TCP state machine
26//! - the relay thread owns the host-facing TCP socket
27//! - channels bridge payloads between them
28
29use crate::egress::EgressPolicy;
30use crate::queues::WakePipe;
31use crate::virtio_net_log;
32use smoltcp::iface::{Interface, SocketHandle, SocketSet};
33use smoltcp::socket::tcp;
34use smoltcp::wire::IpListenEndpoint;
35use std::collections::{HashMap, HashSet};
36use std::io::{self, Read, Write};
37use std::net::{Ipv4Addr, Shutdown, SocketAddr, TcpStream};
38use std::sync::atomic::{AtomicU8, Ordering};
39use std::sync::mpsc::{self, Receiver, SyncSender, TryRecvError, TrySendError};
40use std::sync::Arc;
41use std::thread;
42use std::time::Duration;
43
44const TCP_RX_BUFFER_BYTES: usize = 64 * 1024;
45const TCP_TX_BUFFER_BYTES: usize = 64 * 1024;
46const MAX_CONNECTIONS: usize = 256;
47const CHANNEL_CAPACITY: usize = 32;
48const RELAY_BUFFER_BYTES: usize = 16 * 1024;
49const CLOSE_RETRY_LIMIT: u16 = 64;
50const PROXY_IDLE_SLEEP: Duration = Duration::from_millis(10);
51const PUBLISHED_PORT_START: u16 = 49_152;
52const PUBLISHED_PORT_END: u16 = 65_535;
53
54/// Track all active guest TCP connections bridged through host sockets.
55///
56/// One entry corresponds to one `(guest source, destination)` tuple. The table
57/// lives in the smoltcp poll thread and owns all guest-facing socket handles.
58pub struct TcpRelayTable {
59    connections: HashMap<SocketHandle, TrackedConnection>,
60    connection_keys: HashSet<(SocketAddr, SocketAddr)>,
61    used_published_ports: HashSet<u16>,
62    next_published_port: u16,
63    max_connections: usize,
64    /// Outbound allow-list applied before opening a host connection for a
65    /// guest-initiated flow. Inbound published-port connections bypass it.
66    egress: EgressPolicy,
67}
68
69/// Newly established guest connection ready for a host relay thread.
70///
71/// The poll loop emits these once the guest-side smoltcp socket reaches
72/// `Established`. At that point we can safely create the host-side relay
73/// thread and give it channel endpoints for payload exchange.
74pub struct NewTcpConnection {
75    /// Destination originally requested by the guest.
76    pub destination: SocketAddr,
77    /// How the host-side relay should be started.
78    pub relay_target: RelayTarget,
79    /// Guest-to-host payloads read from the smoltcp socket.
80    pub from_smoltcp: Receiver<Vec<u8>>,
81    /// Host-to-guest payloads written back into the smoltcp socket.
82    pub to_smoltcp: SyncSender<Vec<u8>>,
83    /// Shared relay exit state.
84    pub exit_state: RelayExitState,
85}
86
87#[derive(Debug)]
88struct TrackedConnection {
89    // `source` and `destination` identify the guest-side flow.
90    source: SocketAddr,
91    destination: SocketAddr,
92    // guest -> host relay payloads
93    to_proxy: SyncSender<Vec<u8>>,
94    // host -> guest relay payloads
95    from_proxy: Receiver<Vec<u8>>,
96    // endpoints are held here until the guest-side handshake completes
97    pending_proxy_endpoints: Option<PendingProxyEndpoints>,
98    // once true, a dedicated host relay thread exists
99    relay_spawned: bool,
100    // partial guest->host payload already consumed from smoltcp but not yet
101    // accepted by the relay thread channel
102    buffered_guest_data: Option<Vec<u8>>,
103    // partial host->guest payload not yet fully accepted by smoltcp
104    buffered_proxy_data: Option<(Vec<u8>, usize)>,
105    // bounded retry count for closing with unsent buffered data
106    close_attempts: u16,
107    // relay thread termination mode observed by the poll loop
108    exit_state: RelayExitState,
109    // reserved local source port for published inbound connections
110    reserved_published_port: Option<u16>,
111}
112
113#[derive(Debug)]
114struct PendingProxyEndpoints {
115    from_smoltcp: Receiver<Vec<u8>>,
116    to_smoltcp: SyncSender<Vec<u8>>,
117    relay_target: RelayTarget,
118}
119
120/// How a host-side TCP relay should obtain its remote socket.
121#[derive(Debug)]
122pub enum RelayTarget {
123    /// Open a new outbound host `TcpStream` to the destination.
124    Connect(SocketAddr),
125    /// Use an already-accepted host `TcpStream` from a published port listener.
126    Attached(TcpStream),
127}
128
129/// Host relay termination state shared between the poll loop and the relay thread.
130///
131/// The relay thread cannot mutate smoltcp sockets directly because those sockets
132/// are owned by the poll loop thread. Instead it reports how it finished, and
133/// the poll loop interprets that into guest-side socket actions:
134/// - `Graceful` -> close guest socket cleanly
135/// - `Abort`    -> abort/reset guest socket
136#[derive(Clone, Debug)]
137pub struct RelayExitState {
138    inner: Arc<AtomicU8>,
139}
140
141/// How a host TCP relay thread terminated.
142#[derive(Clone, Copy, Debug, PartialEq, Eq)]
143#[repr(u8)]
144pub enum RelayExitMode {
145    /// Relay thread is still running.
146    Running = 0,
147    /// Remote side closed normally; send FIN toward the guest.
148    Graceful = 1,
149    /// Remote connect or I/O failed; abort the guest TCP socket.
150    Abort = 2,
151}
152
153impl RelayExitState {
154    fn new() -> Self {
155        Self {
156            inner: Arc::new(AtomicU8::new(RelayExitMode::Running as u8)),
157        }
158    }
159
160    fn load(&self) -> RelayExitMode {
161        match self.inner.load(Ordering::Relaxed) {
162            1 => RelayExitMode::Graceful,
163            2 => RelayExitMode::Abort,
164            _ => RelayExitMode::Running,
165        }
166    }
167
168    fn store(&self, mode: RelayExitMode) {
169        self.inner.store(mode as u8, Ordering::Relaxed);
170    }
171}
172
173impl TcpRelayTable {
174    /// Create a new relay table.
175    pub fn new(max_connections: Option<usize>, egress: EgressPolicy) -> Self {
176        Self {
177            connections: HashMap::new(),
178            connection_keys: HashSet::new(),
179            used_published_ports: HashSet::new(),
180            next_published_port: PUBLISHED_PORT_START,
181            max_connections: max_connections.unwrap_or(MAX_CONNECTIONS),
182            egress,
183        }
184    }
185
186    /// Whether a relay socket already exists for the same guest source and destination.
187    pub fn has_socket_for(&self, source: &SocketAddr, destination: &SocketAddr) -> bool {
188        self.connection_keys.contains(&(*source, *destination))
189    }
190
191    /// Create a smoltcp TCP socket for a guest SYN.
192    ///
193    /// Why this happens before full ingress processing:
194    /// - when the first guest SYN arrives, smoltcp needs a matching socket to
195    ///   receive it
196    /// - the poll loop therefore pre-creates a listening socket keyed to the
197    ///   destination the guest is trying to reach
198    /// - only after the guest-facing connection reaches `Established` do we
199    ///   spawn the host relay thread
200    ///
201    /// Data path after creation:
202    ///
203    /// ```text
204    /// smoltcp socket --to_proxy channel--> host relay thread
205    /// host relay thread --from_proxy channel--> smoltcp socket
206    /// ```
207    pub fn create_tcp_socket(
208        &mut self,
209        source: SocketAddr,
210        destination: SocketAddr,
211        sockets: &mut SocketSet<'_>,
212    ) -> bool {
213        if self.connections.len() >= self.max_connections {
214            tracing::warn!("dropping TCP connection because the relay table is full");
215            return false;
216        }
217
218        // Egress policy: drop the guest SYN before any socket is created when the
219        // destination isn't allowed, so the guest just sees the connection fail.
220        // Inbound published-port flows take a separate path and are unaffected.
221        if !self.egress.allows(destination.ip()) {
222            tracing::debug!(
223                %destination,
224                "virtio-net: blocking outbound connection by egress policy"
225            );
226            return false;
227        }
228
229        let rx_buffer = tcp::SocketBuffer::new(vec![0u8; TCP_RX_BUFFER_BYTES]);
230        let tx_buffer = tcp::SocketBuffer::new(vec![0u8; TCP_TX_BUFFER_BYTES]);
231        let mut socket = tcp::Socket::new(rx_buffer, tx_buffer);
232
233        let listen_endpoint = IpListenEndpoint {
234            addr: Some(destination.ip().into()),
235            port: destination.port(),
236        };
237        if socket.listen(listen_endpoint).is_err() {
238            return false;
239        }
240
241        let handle = sockets.add(socket);
242
243        let (to_proxy_tx, to_proxy_rx) = mpsc::sync_channel(CHANNEL_CAPACITY);
244        let (from_proxy_tx, from_proxy_rx) = mpsc::sync_channel(CHANNEL_CAPACITY);
245        let exit_state = RelayExitState::new();
246
247        self.connection_keys.insert((source, destination));
248        self.connections.insert(
249            handle,
250            TrackedConnection {
251                source,
252                destination,
253                to_proxy: to_proxy_tx,
254                from_proxy: from_proxy_rx,
255                pending_proxy_endpoints: Some(PendingProxyEndpoints {
256                    from_smoltcp: to_proxy_rx,
257                    to_smoltcp: from_proxy_tx,
258                    relay_target: RelayTarget::Connect(destination),
259                }),
260                relay_spawned: false,
261                buffered_guest_data: None,
262                buffered_proxy_data: None,
263                close_attempts: 0,
264                exit_state,
265                reserved_published_port: None,
266            },
267        );
268
269        true
270    }
271
272    /// Create a guest-facing TCP connection for a published host socket.
273    ///
274    /// This is the host->guest mirror of `create_tcp_socket`:
275    ///
276    /// ```text
277    /// host client connects to published port
278    ///   -> host listener accepts TcpStream
279    ///   -> poll loop creates smoltcp TCP socket from gateway_ip:ephemeral
280    ///      to guest_ip:guest_port
281    ///   -> guest kernel sees a normal inbound TCP connection on guest_port
282    /// ```
283    ///
284    /// The guest-visible source address is the gateway IP, not the original
285    /// host peer address. That keeps the first version simple and matches the
286    /// fact that this runtime is acting as a userspace gateway/proxy.
287    pub fn create_published_socket(
288        &mut self,
289        interface: &mut Interface,
290        gateway_ip: Ipv4Addr,
291        destination: SocketAddr,
292        host_stream: TcpStream,
293        sockets: &mut SocketSet<'_>,
294    ) -> bool {
295        if self.connections.len() >= self.max_connections {
296            tracing::warn!("dropping published TCP connection because the relay table is full");
297            return false;
298        }
299
300        let Some(local_port) = self.allocate_published_port() else {
301            tracing::warn!(
302                "dropping published TCP connection because no gateway source port is available"
303            );
304            return false;
305        };
306
307        // Inbound published connections always target the guest's IPv4 on the
308        // internal link (the host listener family is independent of this).
309        let std::net::IpAddr::V4(destination_ip) = destination.ip() else {
310            self.used_published_ports.remove(&local_port);
311            return false;
312        };
313
314        let rx_buffer = tcp::SocketBuffer::new(vec![0u8; TCP_RX_BUFFER_BYTES]);
315        let tx_buffer = tcp::SocketBuffer::new(vec![0u8; TCP_TX_BUFFER_BYTES]);
316        let mut socket = tcp::Socket::new(rx_buffer, tx_buffer);
317        let local_endpoint = IpListenEndpoint {
318            addr: Some(gateway_ip.into()),
319            port: local_port,
320        };
321        if socket
322            .connect(
323                interface.context(),
324                (destination_ip, destination.port()),
325                local_endpoint,
326            )
327            .is_err()
328        {
329            self.used_published_ports.remove(&local_port);
330            return false;
331        }
332
333        let handle = sockets.add(socket);
334        let source = SocketAddr::new(std::net::IpAddr::V4(gateway_ip), local_port);
335
336        let (to_proxy_tx, to_proxy_rx) = mpsc::sync_channel(CHANNEL_CAPACITY);
337        let (from_proxy_tx, from_proxy_rx) = mpsc::sync_channel(CHANNEL_CAPACITY);
338        let exit_state = RelayExitState::new();
339
340        self.connection_keys.insert((source, destination));
341        self.connections.insert(
342            handle,
343            TrackedConnection {
344                source,
345                destination,
346                to_proxy: to_proxy_tx,
347                from_proxy: from_proxy_rx,
348                pending_proxy_endpoints: Some(PendingProxyEndpoints {
349                    from_smoltcp: to_proxy_rx,
350                    to_smoltcp: from_proxy_tx,
351                    relay_target: RelayTarget::Attached(host_stream),
352                }),
353                relay_spawned: false,
354                buffered_guest_data: None,
355                buffered_proxy_data: None,
356                close_attempts: 0,
357                exit_state,
358                reserved_published_port: Some(local_port),
359            },
360        );
361
362        true
363    }
364
365    /// Relay TCP payloads between smoltcp sockets and host relay threads.
366    ///
367    /// This runs in the poll thread. It is responsible for:
368    /// - draining bytes received from the guest-facing smoltcp socket and
369    ///   pushing them toward the host relay thread
370    /// - draining bytes received from the host relay thread and writing them
371    ///   back into the smoltcp socket
372    /// - interpreting relay exit state into guest-side `close()` or `abort()`
373    pub fn relay_data(&mut self, sockets: &mut SocketSet<'_>) {
374        let mut read_buffer = [0u8; RELAY_BUFFER_BYTES];
375
376        for (&handle, connection) in &mut self.connections {
377            if !connection.relay_spawned {
378                continue;
379            }
380
381            let socket = sockets.get_mut::<tcp::Socket>(handle);
382
383            match connection.exit_state.load() {
384                RelayExitMode::Abort => {
385                    socket.abort();
386                    continue;
387                }
388                RelayExitMode::Graceful => {
389                    flush_proxy_data(socket, connection);
390                    if connection.buffered_proxy_data.is_none() {
391                        socket.close();
392                    } else {
393                        connection.close_attempts += 1;
394                        if connection.close_attempts >= CLOSE_RETRY_LIMIT {
395                            socket.abort();
396                        }
397                    }
398                    continue;
399                }
400                RelayExitMode::Running => {}
401            }
402
403            flush_guest_data(connection);
404            while connection.buffered_guest_data.is_none() && socket.can_recv() {
405                match socket.recv_slice(&mut read_buffer) {
406                    Ok(bytes_read) if bytes_read > 0 => {
407                        let payload = read_buffer[..bytes_read].to_vec();
408                        if !send_guest_payload(connection, payload) {
409                            break;
410                        }
411                    }
412                    _ => break,
413                }
414            }
415
416            flush_proxy_data(socket, connection);
417        }
418    }
419
420    /// Collect connections that reached ESTABLISHED and need a host relay thread.
421    ///
422    /// The separation between `create_tcp_socket` and this method is important:
423    /// the guest TCP handshake is accepted first on the smoltcp side, and only
424    /// once that succeeds do we commit to opening the host-side `TcpStream`.
425    pub fn take_new_connections(&mut self, sockets: &mut SocketSet<'_>) -> Vec<NewTcpConnection> {
426        let mut new_connections = Vec::new();
427
428        for (&handle, connection) in &mut self.connections {
429            if connection.relay_spawned {
430                continue;
431            }
432
433            let socket = sockets.get::<tcp::Socket>(handle);
434            if socket.state() == tcp::State::Established {
435                connection.relay_spawned = true;
436
437                if let Some(endpoints) = connection.pending_proxy_endpoints.take() {
438                    new_connections.push(NewTcpConnection {
439                        destination: connection.destination,
440                        relay_target: endpoints.relay_target,
441                        from_smoltcp: endpoints.from_smoltcp,
442                        to_smoltcp: endpoints.to_smoltcp,
443                        exit_state: connection.exit_state.clone(),
444                    });
445                }
446            }
447        }
448
449        new_connections
450    }
451
452    /// Remove closed sockets and drop their relay endpoints.
453    ///
454    /// This is the final ownership cleanup step for a guest TCP flow.
455    pub fn cleanup_closed(&mut self, sockets: &mut SocketSet<'_>) {
456        let keys = &mut self.connection_keys;
457        let published_ports = &mut self.used_published_ports;
458        self.connections.retain(|&handle, connection| {
459            let socket = sockets.get::<tcp::Socket>(handle);
460            if socket.state() == tcp::State::Closed {
461                keys.remove(&(connection.source, connection.destination));
462                if let Some(port) = connection.reserved_published_port {
463                    published_ports.remove(&port);
464                }
465                sockets.remove(handle);
466                false
467            } else {
468                true
469            }
470        });
471    }
472
473    fn allocate_published_port(&mut self) -> Option<u16> {
474        let start = self.next_published_port;
475
476        loop {
477            let candidate = self.next_published_port;
478            self.next_published_port = if candidate == PUBLISHED_PORT_END {
479                PUBLISHED_PORT_START
480            } else {
481                candidate + 1
482            };
483
484            if self.used_published_ports.insert(candidate) {
485                return Some(candidate);
486            }
487
488            if self.next_published_port == start {
489                return None;
490            }
491        }
492    }
493}
494
495/// Spawn one host TCP relay thread for an established guest connection.
496///
497/// Thread responsibilities:
498/// - connect a host `TcpStream` to the guest-requested destination
499/// - copy bytes guest->host from `from_smoltcp`
500/// - copy bytes host->guest into `to_smoltcp`
501/// - wake the poll loop when host->guest data arrives or guest->host backpressure eases
502/// - report termination mode through `exit_state`
503pub fn spawn_tcp_relay(
504    destination: SocketAddr,
505    relay_target: RelayTarget,
506    from_smoltcp: Receiver<Vec<u8>>,
507    to_smoltcp: SyncSender<Vec<u8>>,
508    relay_wake: Arc<WakePipe>,
509    exit_state: RelayExitState,
510) {
511    let thread_name = format!("smolvm-tcp-{}", destination.port());
512    virtio_net_log!(
513        "virtio-net: spawning host TCP relay thread destination={} thread={}",
514        destination,
515        thread_name
516    );
517    let _ = thread::Builder::new().name(thread_name).spawn(move || {
518        run_tcp_relay(
519            destination,
520            relay_target,
521            from_smoltcp,
522            to_smoltcp,
523            relay_wake,
524            exit_state,
525        )
526    });
527}
528
529fn run_tcp_relay(
530    destination: SocketAddr,
531    relay_target: RelayTarget,
532    from_smoltcp: Receiver<Vec<u8>>,
533    to_smoltcp: SyncSender<Vec<u8>>,
534    relay_wake: Arc<WakePipe>,
535    exit_state: RelayExitState,
536) {
537    // The relay thread is intentionally isolated from smoltcp internals. Its
538    // contract is just channels in, channels out, and an exit code back.
539    virtio_net_log!(
540        "virtio-net: host TCP relay thread started destination={}",
541        destination
542    );
543    match tcp_relay_loop(
544        destination,
545        relay_target,
546        from_smoltcp,
547        to_smoltcp,
548        relay_wake,
549    ) {
550        Ok(mode) => {
551            virtio_net_log!(
552                "virtio-net: host TCP relay thread exited destination={} mode={:?}",
553                destination,
554                mode
555            );
556            exit_state.store(mode)
557        }
558        Err(err) => {
559            virtio_net_log!(
560                "virtio-net: host TCP relay failed destination={} error={}",
561                destination,
562                err
563            );
564            exit_state.store(RelayExitMode::Abort);
565        }
566    }
567}
568
569fn tcp_relay_loop(
570    destination: SocketAddr,
571    relay_target: RelayTarget,
572    from_smoltcp: Receiver<Vec<u8>>,
573    to_smoltcp: SyncSender<Vec<u8>>,
574    relay_wake: Arc<WakePipe>,
575) -> io::Result<RelayExitMode> {
576    // Host-side flow:
577    //
578    // 1. Connect a normal host TcpStream to the destination.
579    // 2. Non-blockingly drain guest payloads from the channel into the socket.
580    // 3. Non-blockingly read remote payloads from the socket into the channel.
581    // 4. If neither side made progress, sleep briefly to avoid a hot spin loop.
582    let mut stream = match relay_target {
583        RelayTarget::Connect(destination) => {
584            virtio_net_log!(
585                "virtio-net: connecting host TCP relay socket destination={}",
586                destination
587            );
588            let stream = TcpStream::connect(destination)?;
589            virtio_net_log!(
590                "virtio-net: host TCP relay socket connected destination={}",
591                destination
592            );
593            stream
594        }
595        RelayTarget::Attached(stream) => {
596            virtio_net_log!(
597                "virtio-net: using accepted host TCP socket for published port guest_destination={} peer_addr={:?} local_addr={:?}",
598                destination,
599                stream.peer_addr().ok(),
600                stream.local_addr().ok()
601            );
602            stream
603        }
604    };
605    stream.set_nonblocking(true)?;
606
607    let mut guest_write_closed = false;
608    let mut guest_channel_closed = false;
609    let mut pending_guest_data: Option<(Vec<u8>, usize)> = None;
610    let mut read_buffer = [0u8; RELAY_BUFFER_BYTES];
611
612    loop {
613        let mut did_work = false;
614
615        if pending_guest_data.is_none() && !guest_channel_closed {
616            match from_smoltcp.try_recv() {
617                Ok(payload) => {
618                    pending_guest_data = Some((payload, 0));
619                    // Consuming from the bounded guest->host channel may free
620                    // capacity for a payload buffered in the smoltcp poll
621                    // thread. Wake it so backpressure clears promptly even for
622                    // one-way guest->host streams.
623                    relay_wake.wake();
624                    did_work = true;
625                }
626                Err(TryRecvError::Empty) => {}
627                Err(TryRecvError::Disconnected) => {
628                    guest_channel_closed = true;
629                }
630            }
631        }
632
633        if let Some((payload, offset)) = &mut pending_guest_data {
634            while *offset < payload.len() {
635                match stream.write(&payload[*offset..]) {
636                    Ok(0) => {
637                        return Err(io::Error::new(
638                            io::ErrorKind::WriteZero,
639                            "host TCP relay wrote zero bytes",
640                        ));
641                    }
642                    Ok(bytes_written) => {
643                        *offset += bytes_written;
644                        did_work = true;
645                    }
646                    Err(err) if err.kind() == io::ErrorKind::WouldBlock => break,
647                    Err(err) if err.kind() == io::ErrorKind::Interrupted => continue,
648                    Err(err) => return Err(err),
649                }
650            }
651
652            if *offset >= payload.len() {
653                pending_guest_data = None;
654            }
655        }
656
657        if guest_channel_closed && pending_guest_data.is_none() && !guest_write_closed {
658            // The guest side closed its write half. Mirror that toward the
659            // remote peer only after all buffered guest bytes were written.
660            let _ = stream.shutdown(Shutdown::Write);
661            guest_write_closed = true;
662        }
663
664        match stream.read(&mut read_buffer) {
665            Ok(0) => return Ok(RelayExitMode::Graceful),
666            Ok(bytes_read) => {
667                if to_smoltcp.send(read_buffer[..bytes_read].to_vec()).is_err() {
668                    return Ok(RelayExitMode::Graceful);
669                }
670                relay_wake.wake();
671                did_work = true;
672            }
673            Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
674            Err(err) => return Err(err),
675        }
676
677        if !did_work {
678            thread::sleep(PROXY_IDLE_SLEEP);
679        }
680    }
681}
682
683fn flush_guest_data(connection: &mut TrackedConnection) {
684    let Some(payload) = connection.buffered_guest_data.take() else {
685        return;
686    };
687    send_guest_payload(connection, payload);
688}
689
690fn send_guest_payload(connection: &mut TrackedConnection, payload: Vec<u8>) -> bool {
691    match connection.to_proxy.try_send(payload) {
692        Ok(()) => true,
693        Err(TrySendError::Full(payload)) => {
694            connection.buffered_guest_data = Some(payload);
695            false
696        }
697        Err(TrySendError::Disconnected(_)) => false,
698    }
699}
700
701fn flush_proxy_data(socket: &mut tcp::Socket<'_>, connection: &mut TrackedConnection) {
702    // smoltcp send windows may accept only part of an inbound host payload.
703    // `buffered_proxy_data` remembers the unwritten remainder so the next poll
704    // iteration can continue where it left off instead of dropping bytes.
705    if let Some((data, offset)) = &mut connection.buffered_proxy_data {
706        if socket.can_send() {
707            match socket.send_slice(&data[*offset..]) {
708                Ok(written) => {
709                    *offset += written;
710                    if *offset >= data.len() {
711                        connection.buffered_proxy_data = None;
712                    }
713                }
714                Err(_) => return,
715            }
716        } else {
717            return;
718        }
719    }
720
721    while connection.buffered_proxy_data.is_none() {
722        match connection.from_proxy.try_recv() {
723            Ok(payload) => {
724                if socket.can_send() {
725                    match socket.send_slice(&payload) {
726                        Ok(written) if written < payload.len() => {
727                            connection.buffered_proxy_data = Some((payload, written));
728                        }
729                        Err(_) => {
730                            connection.buffered_proxy_data = Some((payload, 0));
731                        }
732                        _ => {}
733                    }
734                } else {
735                    connection.buffered_proxy_data = Some((payload, 0));
736                }
737            }
738            Err(TryRecvError::Empty | TryRecvError::Disconnected) => break,
739        }
740    }
741}
742
743#[cfg(test)]
744mod tests {
745    use super::*;
746
747    fn test_connection(to_proxy: SyncSender<Vec<u8>>) -> TrackedConnection {
748        let (_from_proxy_tx, from_proxy) = mpsc::sync_channel(CHANNEL_CAPACITY);
749        TrackedConnection {
750            source: SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 12_345),
751            destination: SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 80),
752            to_proxy,
753            from_proxy,
754            pending_proxy_endpoints: None,
755            relay_spawned: true,
756            buffered_guest_data: None,
757            buffered_proxy_data: None,
758            close_attempts: 0,
759            exit_state: RelayExitState::new(),
760            reserved_published_port: None,
761        }
762    }
763
764    #[test]
765    fn guest_payload_is_buffered_when_relay_channel_is_full() {
766        let (to_proxy, from_smoltcp) = mpsc::sync_channel(1);
767        to_proxy.send(vec![1]).unwrap();
768        let mut connection = test_connection(to_proxy);
769
770        assert!(!send_guest_payload(&mut connection, vec![2]));
771        assert_eq!(connection.buffered_guest_data.as_deref(), Some(&[2][..]));
772
773        assert_eq!(from_smoltcp.recv().unwrap(), vec![1]);
774        flush_guest_data(&mut connection);
775
776        assert!(connection.buffered_guest_data.is_none());
777        assert_eq!(from_smoltcp.recv().unwrap(), vec![2]);
778    }
779}