featherbit/stream/
tcp.rs

1//! TCP stream proxy: accept a client, pick a load-balanced upstream (optionally
2//! by TLS SNI), and relay bytes in both directions until either side closes.
3
4use std::io;
5use std::net::SocketAddr;
6use std::sync::Arc;
7use std::time::Duration;
8
9use tokio::io::{AsyncReadExt, AsyncWriteExt};
10use tokio::net::{TcpListener, TcpStream};
11use tokio::sync::watch;
12use tracing::{debug, warn};
13
14use crate::stream::sni::{extract_sni, SniResult, SniRouter};
15
16/// Cap on how much of a ClientHello is buffered while looking for the SNI.
17const MAX_CLIENT_HELLO: usize = 8 * 1024;
18
19/// Binds the TCP listener (fail-fast) and spawns its accept loop, returning the
20/// bound address (the OS-assigned port when `cfg.port == 0`). The loop stops
21/// accepting when `shutdown_rx` flips to `true`. When `router` has SNI routes,
22/// each connection's TLS ClientHello is peeked to pick a backend by hostname.
23pub async fn spawn(
24    cfg: &crate::config::StreamListenerConfig,
25    router: Arc<SniRouter>,
26    connect_timeout: Duration,
27    mut shutdown_rx: watch::Receiver<bool>,
28) -> io::Result<SocketAddr> {
29    let ip = cfg
30        .bind
31        .parse()
32        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("invalid bind: {}", e)))?;
33    let addr = SocketAddr::new(ip, cfg.port);
34    let listener = TcpListener::bind(addr).await?;
35    let local = listener.local_addr()?;
36
37    tokio::spawn(async move {
38        loop {
39            let (mut client, peer) = tokio::select! {
40                accepted = listener.accept() => match accepted {
41                    Ok(v) => v,
42                    Err(e) => {
43                        warn!("tcp stream accept error: {}", e);
44                        continue;
45                    }
46                },
47                _ = shutdown_rx.changed() => break,
48            };
49            let router = router.clone();
50            tokio::spawn(async move {
51                // Pick the backend pool. With SNI routes, peek the ClientHello
52                // (consuming its bytes into `prebuffer`, which must be replayed
53                // to the upstream to keep the passthrough handshake intact).
54                let (balancer, prebuffer) = if router.has_sni_routes() {
55                    let (sni, buf) = peek_sni(&mut client, connect_timeout).await;
56                    (router.select(sni.as_deref()).clone(), buf)
57                } else {
58                    (router.select(None).clone(), Vec::new())
59                };
60
61                let idx = balancer.select(&peer.to_string());
62                // Hold the in-flight guard for the connection's whole lifetime
63                // so least-connections balancing reflects live streams.
64                let _guard = balancer.owned_acquire(idx);
65                let target = balancer.target(idx);
66                let dest = (target.host.as_str(), target.port);
67
68                match tokio::time::timeout(connect_timeout, TcpStream::connect(dest)).await {
69                    Ok(Ok(mut upstream)) => {
70                        // Replay the consumed ClientHello before relaying.
71                        if !prebuffer.is_empty() {
72                            if let Err(e) = upstream.write_all(&prebuffer).await {
73                                warn!(
74                                    "tcp stream replay to {}:{} failed: {}",
75                                    target.host, target.port, e
76                                );
77                                return;
78                            }
79                        }
80                        if let Err(e) =
81                            tokio::io::copy_bidirectional(&mut client, &mut upstream).await
82                        {
83                            debug!("tcp stream relay closed: {}", e);
84                        }
85                    }
86                    Ok(Err(e)) => {
87                        warn!(
88                            "tcp stream connect to {}:{} failed: {}",
89                            target.host, target.port, e
90                        )
91                    }
92                    Err(_) => warn!(
93                        "tcp stream connect to {}:{} timed out",
94                        target.host, target.port
95                    ),
96                }
97            });
98        }
99    });
100
101    Ok(local)
102}
103
104/// Reads (consuming) the start of the connection until the SNI is resolved,
105/// EOF, the 8 KiB cap, or the timeout. Returns the hostname (if any) and the
106/// bytes consumed so they can be replayed to the upstream. Any short / garbage
107/// / timeout path yields `(None, buf)` so the caller uses the default pool.
108async fn peek_sni(client: &mut TcpStream, timeout: Duration) -> (Option<String>, Vec<u8>) {
109    let mut buf = Vec::with_capacity(1024);
110    let deadline = tokio::time::Instant::now() + timeout;
111    let mut tmp = [0u8; 4096];
112    loop {
113        match extract_sni(&buf) {
114            SniResult::Found(host) => return (Some(host), buf),
115            SniResult::NotPresent => return (None, buf),
116            SniResult::Incomplete => {}
117        }
118        if buf.len() >= MAX_CLIENT_HELLO {
119            return (None, buf);
120        }
121        match tokio::time::timeout_at(deadline, client.read(&mut tmp)).await {
122            Ok(Ok(0)) => return (None, buf), // EOF
123            Ok(Ok(n)) => {
124                let take = n.min(MAX_CLIENT_HELLO - buf.len());
125                buf.extend_from_slice(&tmp[..take]);
126            }
127            Ok(Err(_)) | Err(_) => return (None, buf), // read error or timeout
128        }
129    }
130}
featherbit/stream/tcp.rs

featherbit/stream/
tcp.rs
