1use 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
54pub 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 egress: EgressPolicy,
67}
68
69pub struct NewTcpConnection {
75 pub destination: SocketAddr,
77 pub relay_target: RelayTarget,
79 pub from_smoltcp: Receiver<Vec<u8>>,
81 pub to_smoltcp: SyncSender<Vec<u8>>,
83 pub exit_state: RelayExitState,
85}
86
87#[derive(Debug)]
88struct TrackedConnection {
89 source: SocketAddr,
91 destination: SocketAddr,
92 to_proxy: SyncSender<Vec<u8>>,
94 from_proxy: Receiver<Vec<u8>>,
96 pending_proxy_endpoints: Option<PendingProxyEndpoints>,
98 relay_spawned: bool,
100 buffered_guest_data: Option<Vec<u8>>,
103 buffered_proxy_data: Option<(Vec<u8>, usize)>,
105 close_attempts: u16,
107 exit_state: RelayExitState,
109 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#[derive(Debug)]
122pub enum RelayTarget {
123 Connect(SocketAddr),
125 Attached(TcpStream),
127}
128
129#[derive(Clone, Debug)]
137pub struct RelayExitState {
138 inner: Arc<AtomicU8>,
139}
140
141#[derive(Clone, Copy, Debug, PartialEq, Eq)]
143#[repr(u8)]
144pub enum RelayExitMode {
145 Running = 0,
147 Graceful = 1,
149 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 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 pub fn has_socket_for(&self, source: &SocketAddr, destination: &SocketAddr) -> bool {
188 self.connection_keys.contains(&(*source, *destination))
189 }
190
191 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 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 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 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 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 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 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
495pub 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 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 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 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 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 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}