featherbit/stream/
mod.rs

1//! L4 (TCP/UDP) stream proxying — a data path independent of the HTTP engine.
2//!
3//! Each configured [`StreamListenerConfig`] binds a port at startup and relays
4//! raw bytes to an upstream pool selected by the shared [`Balancer`]. There is
5//! no `Context`, no node-graph, and no per-request matching: a listener maps a
6//! port straight to a backend pool. TCP uses `copy_bidirectional`; UDP tracks
7//! per-client sessions (see [`tcp`] and [`udp`]).
8
9pub(crate) mod sni;
10mod tcp;
11mod udp;
12
13use std::sync::Arc;
14use std::time::Duration;
15
16use tokio::sync::watch;
17use tracing::{info, warn};
18
19use crate::balancer::{Balancer, Strategy};
20use crate::config::{StreamListenerConfig, StreamProtocol, TimeoutConfig};
21use crate::stream::sni::SniRouter;
22
23/// Binds every configured stream listener and spawns its accept/receive loop.
24///
25/// Binding is fail-fast: the first bind error (or invalid config) is returned
26/// so startup can abort with a clear message. On success every listener is
27/// bound and running in its own detached task, and this returns immediately.
28/// Each loop stops accepting when `shutdown_rx` flips to `true` (in-flight
29/// relays/sessions are then dropped at process exit).
30pub async fn start_all(
31    streams: &[StreamListenerConfig],
32    timeouts: &TimeoutConfig,
33    shutdown_rx: watch::Receiver<bool>,
34) -> Result<(), String> {
35    let connect_timeout = Duration::from_secs(timeouts.connection_seconds);
36    let idle = Duration::from_secs(timeouts.idle_seconds);
37
38    for cfg in streams {
39        let strategy = match &cfg.upstream.load_balancing {
40            Some(s) => Strategy::parse(s)?,
41            None => Strategy::default(),
42        };
43        let balancer = Arc::new(Balancer::new(cfg.upstream.targets.clone(), strategy)?);
44
45        let addr = match cfg.protocol {
46            StreamProtocol::Tcp => {
47                // Build one pool per SNI route; `balancer` is the default.
48                let mut routes = Vec::with_capacity(cfg.sni_routes.len());
49                for route in &cfg.sni_routes {
50                    let route_strategy = match &route.upstream.load_balancing {
51                        Some(s) => Strategy::parse(s)?,
52                        None => Strategy::default(),
53                    };
54                    let pool = Arc::new(Balancer::new(
55                        route.upstream.targets.clone(),
56                        route_strategy,
57                    )?);
58                    routes.push((route.server_name.clone(), pool));
59                }
60                let router = Arc::new(SniRouter::new(routes, balancer));
61                tcp::spawn(cfg, router, connect_timeout, shutdown_rx.clone())
62                    .await
63                    .map_err(|e| {
64                        format!("failed to bind tcp stream {}:{}: {}", cfg.bind, cfg.port, e)
65                    })?
66            }
67            StreamProtocol::Udp => {
68                if !cfg.sni_routes.is_empty() {
69                    warn!(
70                        "sni_routes set on UDP stream {}:{} — ignored (TCP only)",
71                        cfg.bind, cfg.port
72                    );
73                }
74                udp::spawn(cfg, balancer, idle, shutdown_rx.clone())
75                    .await
76                    .map_err(|e| {
77                        format!("failed to bind udp stream {}:{}: {}", cfg.bind, cfg.port, e)
78                    })?
79            }
80        };
81
82        let proto = match cfg.protocol {
83            StreamProtocol::Tcp => "tcp",
84            StreamProtocol::Udp => "udp",
85        };
86        info!("Stream ({}) listening on {}", proto, addr);
87    }
88
89    Ok(())
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use crate::balancer::Target;
96    use crate::config::{StreamProtocol, StreamUpstreamConfig};
97    use std::net::SocketAddr;
98    use tokio::io::{AsyncReadExt, AsyncWriteExt};
99    use tokio::net::{TcpListener, TcpStream, UdpSocket};
100
101    fn stream_cfg(protocol: StreamProtocol, target: SocketAddr) -> StreamListenerConfig {
102        StreamListenerConfig {
103            protocol,
104            bind: "127.0.0.1".to_string(),
105            port: 0,
106            upstream: StreamUpstreamConfig {
107                targets: vec![Target {
108                    host: target.ip().to_string(),
109                    port: target.port(),
110                }],
111                load_balancing: None,
112            },
113            sni_routes: Vec::new(),
114        }
115    }
116
117    fn balancer_for(cfg: &StreamListenerConfig) -> Arc<Balancer> {
118        Arc::new(Balancer::new(cfg.upstream.targets.clone(), Strategy::RoundRobin).unwrap())
119    }
120
121    /// A router with no SNI routes (default pool only), for the plain TCP tests.
122    fn router_for(cfg: &StreamListenerConfig) -> Arc<SniRouter> {
123        Arc::new(SniRouter::new(Vec::new(), balancer_for(cfg)))
124    }
125
126    /// A TCP echo server; returns its bound address.
127    async fn tcp_echo() -> SocketAddr {
128        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
129        let addr = listener.local_addr().unwrap();
130        tokio::spawn(async move {
131            while let Ok((mut sock, _)) = listener.accept().await {
132                tokio::spawn(async move {
133                    let (mut r, mut w) = sock.split();
134                    let _ = tokio::io::copy(&mut r, &mut w).await;
135                });
136            }
137        });
138        addr
139    }
140
141    /// A UDP echo server; returns its bound address.
142    async fn udp_echo() -> SocketAddr {
143        let sock = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
144        let addr = sock.local_addr().unwrap();
145        tokio::spawn(async move {
146            let mut buf = [0u8; 2048];
147            while let Ok((n, peer)) = sock.recv_from(&mut buf).await {
148                let _ = sock.send_to(&buf[..n], peer).await;
149            }
150        });
151        addr
152    }
153
154    #[tokio::test]
155    async fn test_tcp_proxy_round_trip() {
156        let echo = tcp_echo().await;
157        let cfg = stream_cfg(StreamProtocol::Tcp, echo);
158        let (_tx, rx) = watch::channel(false);
159        let proxy = tcp::spawn(&cfg, router_for(&cfg), Duration::from_secs(5), rx)
160            .await
161            .unwrap();
162
163        let mut client = TcpStream::connect(proxy).await.unwrap();
164        client.write_all(b"hello").await.unwrap();
165        let mut buf = [0u8; 5];
166        client.read_exact(&mut buf).await.unwrap();
167        assert_eq!(&buf, b"hello");
168
169        // A second frame confirms the relay stays open.
170        client.write_all(b"world").await.unwrap();
171        client.read_exact(&mut buf).await.unwrap();
172        assert_eq!(&buf, b"world");
173    }
174
175    #[tokio::test]
176    async fn test_udp_proxy_round_trip_and_session_reuse() {
177        let echo = udp_echo().await;
178        let cfg = stream_cfg(StreamProtocol::Udp, echo);
179        let (_tx, rx) = watch::channel(false);
180        let proxy = udp::spawn(&cfg, balancer_for(&cfg), Duration::from_secs(2), rx)
181            .await
182            .unwrap();
183
184        let client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
185        client.send_to(b"ping", proxy).await.unwrap();
186        let mut buf = [0u8; 16];
187        let (n, _) = tokio::time::timeout(Duration::from_secs(1), client.recv_from(&mut buf))
188            .await
189            .unwrap()
190            .unwrap();
191        assert_eq!(&buf[..n], b"ping");
192
193        // A second datagram from the same client reuses the session.
194        client.send_to(b"pong", proxy).await.unwrap();
195        let (n, _) = tokio::time::timeout(Duration::from_secs(1), client.recv_from(&mut buf))
196            .await
197            .unwrap()
198            .unwrap();
199        assert_eq!(&buf[..n], b"pong");
200    }
201
202    /// A backend that reads its first bytes (the replayed ClientHello — proving
203    /// passthrough) then writes back a one-byte id and echoes.
204    async fn id_backend(id: u8) -> SocketAddr {
205        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
206        let addr = listener.local_addr().unwrap();
207        tokio::spawn(async move {
208            while let Ok((mut sock, _)) = listener.accept().await {
209                tokio::spawn(async move {
210                    let mut buf = [0u8; 512];
211                    if sock.read(&mut buf).await.unwrap_or(0) > 0 {
212                        let _ = sock.write_all(&[id]).await;
213                    }
214                });
215            }
216        });
217        addr
218    }
219
220    /// Builds a minimal, well-formed TLS ClientHello record carrying `sni`.
221    fn client_hello(sni: &str) -> Vec<u8> {
222        let name = sni.as_bytes();
223        let mut sni_body = Vec::new();
224        sni_body.extend_from_slice(&((1 + 2 + name.len()) as u16).to_be_bytes());
225        sni_body.push(0x00);
226        sni_body.extend_from_slice(&(name.len() as u16).to_be_bytes());
227        sni_body.extend_from_slice(name);
228        let mut ext = Vec::new();
229        ext.extend_from_slice(&0x0000u16.to_be_bytes());
230        ext.extend_from_slice(&(sni_body.len() as u16).to_be_bytes());
231        ext.extend_from_slice(&sni_body);
232        let mut body = Vec::new();
233        body.extend_from_slice(&[0x03, 0x03]);
234        body.extend_from_slice(&[0u8; 32]);
235        body.push(0x00);
236        body.extend_from_slice(&2u16.to_be_bytes());
237        body.extend_from_slice(&[0x00, 0x2f]);
238        body.push(0x01);
239        body.push(0x00);
240        body.extend_from_slice(&(ext.len() as u16).to_be_bytes());
241        body.extend_from_slice(&ext);
242        let mut hs = vec![0x01];
243        hs.extend_from_slice(&(body.len() as u32).to_be_bytes()[1..]);
244        hs.extend_from_slice(&body);
245        let mut rec = vec![0x16, 0x03, 0x01];
246        rec.extend_from_slice(&(hs.len() as u16).to_be_bytes());
247        rec.extend_from_slice(&hs);
248        rec
249    }
250
251    async fn probe(proxy: SocketAddr, first: &[u8]) -> u8 {
252        let mut c = TcpStream::connect(proxy).await.unwrap();
253        c.write_all(first).await.unwrap();
254        let mut id = [0u8; 1];
255        c.read_exact(&mut id).await.unwrap();
256        id[0]
257    }
258
259    #[tokio::test]
260    async fn test_sni_routing_selects_backend_and_replays() {
261        use crate::config::SniRoute;
262
263        let backend_a = id_backend(b'A').await;
264        let backend_b = id_backend(b'B').await;
265
266        // Default -> B; a.example.com -> A.
267        let mut cfg = stream_cfg(StreamProtocol::Tcp, backend_b);
268        cfg.sni_routes = vec![SniRoute {
269            server_name: "a.example.com".to_string(),
270            upstream: StreamUpstreamConfig {
271                targets: vec![Target {
272                    host: backend_a.ip().to_string(),
273                    port: backend_a.port(),
274                }],
275                load_balancing: None,
276            },
277        }];
278
279        // Build the router exactly as `start_all` does.
280        let default = balancer_for(&cfg);
281        let routes = cfg
282            .sni_routes
283            .iter()
284            .map(|r| {
285                (
286                    r.server_name.clone(),
287                    Arc::new(
288                        Balancer::new(r.upstream.targets.clone(), Strategy::RoundRobin).unwrap(),
289                    ),
290                )
291            })
292            .collect();
293        let router = Arc::new(SniRouter::new(routes, default));
294
295        let (_tx, rx) = watch::channel(false);
296        let proxy = tcp::spawn(&cfg, router, Duration::from_secs(2), rx)
297            .await
298            .unwrap();
299
300        // Matched SNI -> backend A (also proves the ClientHello was replayed,
301        // since A only responds after reading its first bytes).
302        assert_eq!(probe(proxy, &client_hello("a.example.com")).await, b'A');
303        // Unmatched SNI -> default backend B.
304        assert_eq!(probe(proxy, &client_hello("other.com")).await, b'B');
305        // Non-TLS garbage -> default backend B (NotPresent, no wait).
306        assert_eq!(probe(proxy, b"not a tls hello").await, b'B');
307    }
308}
featherbit/stream/mod.rs

featherbit/stream/
mod.rs
