rws_gateway_api/gateway_api/
l4.rs

1//! TLS, TCP and UDP listeners: TLSRoutes (passthrough by SNI, or
2//! terminated with the listener's certificates), TCPRoutes and UDPRoutes.
3//! A connection (or a UDP client's datagrams) goes to one route — the one
4//! whose hostname best matches the SNI, else the oldest — and to a backend
5//! of its first rule by weight; bytes are relayed both ways.
6
7use std::collections::HashMap;
8use std::net::SocketAddr;
9use std::sync::atomic::{AtomicUsize, Ordering};
10use std::sync::{Arc, Mutex};
11use std::time::{Duration, Instant};
12
13use tokio::io::AsyncWriteExt;
14use tokio::net::{TcpStream, UdpSocket};
15
16use super::router::{host_score, pick_listener, PortListeners};
17use super::routes::CompiledHttpRoute;
18use super::translate::ListenerConfig;
19
20/// The most a ClientHello may take before passthrough gives up on it.
21const MAX_CLIENT_HELLO: usize = 16 * 1024;
22/// How long a UDP client's session lives without traffic.
23pub const UDP_IDLE: Duration = Duration::from_secs(60);
24/// Connect timeout for backends.
25const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
26
27/// The route a connection goes to: the best hostname match for `sni`
28/// (a route without hostnames matches anything, below any hostname match),
29/// ties to the oldest route, then `namespace/name`.
30pub fn pick_route<'a>(routes: &'a [CompiledHttpRoute], sni: Option<&str>) -> Option<&'a CompiledHttpRoute> {
31    let sni = sni.unwrap_or("").to_ascii_lowercase();
32    routes
33        .iter()
34        .filter_map(|route| {
35            let score = if route.hostnames.is_empty() { Some((0, 0)) } else { route.hostnames.iter().filter_map(|h| host_score(h, &sni)).max() };
36            score.map(|s| (s, route))
37        })
38        .max_by(|(a, ra), (b, rb)| a.cmp(b).then_with(|| (rb.created.as_str(), &rb.key).cmp(&(ra.created.as_str(), &ra.key))))
39        .map(|(_, route)| route)
40}
41
42/// An endpoint of the route's first rule with backends, by weight; `None`
43/// when there is nothing to send to.
44pub fn pick_endpoint(route: &CompiledHttpRoute, counter: &AtomicUsize) -> Option<String> {
45    let rule = route.rules.iter().find(|r| !r.backends.is_empty())?;
46    let total: u32 = rule.backends.iter().map(|b| b.weight).sum();
47    if total == 0 {
48        return None;
49    }
50    let n = counter.fetch_add(1, Ordering::Relaxed);
51    let mut pick = ((n as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 32) as u32 % total;
52    let backend = rule.backends.iter().find(|b| {
53        if pick < b.weight {
54            true
55        } else {
56            pick -= b.weight;
57            false
58        }
59    })?;
60    let endpoints = backend.endpoints.as_ref().ok().filter(|e| !e.is_empty())?;
61    Some(endpoints[n % endpoints.len()].clone())
62}
63
64/// The listener and endpoint for a connection with this SNI.
65fn target(listeners: &PortListeners, sni: Option<&str>, counter: &AtomicUsize) -> Option<String> {
66    let listeners: Arc<Vec<ListenerConfig>> = listeners.read().unwrap_or_else(|e| e.into_inner()).clone();
67    let listener = pick_listener(&listeners, &sni.unwrap_or("").to_ascii_lowercase())?;
68    pick_endpoint(pick_route(&listener.http_routes, sni)?, counter)
69}
70
71async fn connect(endpoint: &str) -> Option<TcpStream> {
72    let address: SocketAddr = endpoint.parse().ok()?;
73    tokio::time::timeout(CONNECT_TIMEOUT, TcpStream::connect(address)).await.ok()?.ok()
74}
75
76/// A TCPRoute connection: straight to a backend.
77pub async fn serve_tcp(mut client: TcpStream, listeners: PortListeners, counter: Arc<AtomicUsize>) {
78    let Some(endpoint) = target(&listeners, None, &counter) else { return };
79    let Some(mut backend) = connect(&endpoint).await else { return };
80    let _ = tokio::io::copy_bidirectional(&mut client, &mut backend).await;
81}
82
83/// What the start of a TLS connection says.
84#[derive(Debug, PartialEq)]
85pub enum Hello {
86    /// Not a whole ClientHello yet.
87    NeedMore,
88    /// A ClientHello, with its SNI if it has one.
89    Sni(Option<String>),
90    /// Not a ClientHello rustls accepts.
91    NotTls,
92}
93
94/// Read the ClientHello at the start of `bytes`.
95pub fn client_hello(bytes: &[u8]) -> Hello {
96    let mut acceptor = rustls::server::Acceptor::default();
97    let mut input = bytes;
98    while !input.is_empty() {
99        if acceptor.read_tls(&mut input).is_err() {
100            return Hello::NotTls;
101        }
102    }
103    match acceptor.accept() {
104        Ok(Some(accepted)) => Hello::Sni(accepted.client_hello().server_name().map(str::to_string)),
105        Ok(None) => Hello::NeedMore,
106        Err(_) => Hello::NotTls,
107    }
108}
109
110/// A passthrough TLSRoute connection: read the ClientHello for its SNI,
111/// then replay it to the backend and relay the encrypted stream.
112pub async fn serve_tls_passthrough(mut client: TcpStream, listeners: PortListeners, counter: Arc<AtomicUsize>) {
113    use tokio::io::AsyncReadExt;
114    let mut hello = Vec::new();
115    let sni = loop {
116        let mut buffer = [0u8; 4096];
117        let n = match tokio::time::timeout(Duration::from_secs(10), client.read(&mut buffer)).await {
118            Ok(Ok(n)) if n > 0 => n,
119            _ => return,
120        };
121        hello.extend_from_slice(&buffer[..n]);
122        match client_hello(&hello) {
123            Hello::Sni(sni) => break sni,
124            Hello::NeedMore if hello.len() < MAX_CLIENT_HELLO => continue,
125            _ => return,
126        }
127    };
128    let Some(endpoint) = target(&listeners, sni.as_deref(), &counter) else { return };
129    let Some(mut backend) = connect(&endpoint).await else { return };
130    if backend.write_all(&hello).await.is_err() {
131        return;
132    }
133    let _ = tokio::io::copy_bidirectional(&mut client, &mut backend).await;
134}
135
136/// A terminated TLSRoute connection: handshake with the listener's
137/// certificates, then relay the plaintext.
138pub async fn serve_tls_terminate(client: TcpStream, acceptor: tokio_rustls::TlsAcceptor, listeners: PortListeners, counter: Arc<AtomicUsize>) {
139    let Ok(Ok(mut tls)) = tokio::time::timeout(Duration::from_secs(10), acceptor.accept(client)).await else { return };
140    let sni = tls.get_ref().1.server_name().map(str::to_string);
141    let Some(endpoint) = target(&listeners, sni.as_deref(), &counter) else { return };
142    let Some(mut backend) = connect(&endpoint).await else { return };
143    let _ = tokio::io::copy_bidirectional(&mut tls, &mut backend).await;
144}
145
146/// A UDP client's session: its socket to the backend and when it was last used.
147struct Session {
148    backend: Arc<UdpSocket>,
149    seen: Arc<Mutex<Instant>>,
150}
151
152/// A UDPRoute listener: each client address gets its own socket to a
153/// backend (picked once), replies go back from the listener's socket;
154/// sessions end after [`UDP_IDLE`] without traffic — or with the port:
155/// they're in a `JoinSet` this task owns, so dropping it (the port given up)
156/// aborts them and releases the socket at once.
157pub async fn serve_udp(socket: Arc<UdpSocket>, listeners: PortListeners) {
158    let counter = Arc::new(AtomicUsize::new(0));
159    let mut session_tasks = tokio::task::JoinSet::new();
160    let sessions: Arc<Mutex<HashMap<SocketAddr, Session>>> = Arc::default();
161    let mut buffer = vec![0u8; 65_535];
162    loop {
163        // Reap finished sessions so the set doesn't grow.
164        while session_tasks.try_join_next().is_some() {}
165        let Ok((n, client)) = socket.recv_from(&mut buffer).await else {
166            tokio::time::sleep(Duration::from_millis(50)).await;
167            continue;
168        };
169        let existing = sessions.lock().unwrap_or_else(|e| e.into_inner()).get(&client).map(|s| (Arc::clone(&s.backend), Arc::clone(&s.seen)));
170        let (backend, seen) = match existing {
171            Some(session) => session,
172            None => {
173                let Some(endpoint) = target(&listeners, None, &counter) else { continue };
174                let Some(backend) = udp_connect(&endpoint).await else { continue };
175                let backend = Arc::new(backend);
176                let seen = Arc::new(Mutex::new(Instant::now()));
177                sessions.lock().unwrap_or_else(|e| e.into_inner()).insert(client, Session { backend: Arc::clone(&backend), seen: Arc::clone(&seen) });
178                session_tasks.spawn(udp_replies(Arc::clone(&socket), Arc::clone(&backend), client, Arc::clone(&seen), Arc::clone(&sessions)));
179                (backend, seen)
180            }
181        };
182        *seen.lock().unwrap_or_else(|e| e.into_inner()) = Instant::now();
183        let _ = backend.send(&buffer[..n]).await;
184    }
185}
186
187async fn udp_connect(endpoint: &str) -> Option<UdpSocket> {
188    let address: SocketAddr = endpoint.parse().ok()?;
189    let local = if address.is_ipv6() { "[::]:0" } else { "0.0.0.0:0" };
190    let socket = UdpSocket::bind(local).await.ok()?;
191    socket.connect(address).await.ok()?;
192    Some(socket)
193}
194
195/// Backend → client for one session, until it's idle for [`UDP_IDLE`].
196async fn udp_replies(listener: Arc<UdpSocket>, backend: Arc<UdpSocket>, client: SocketAddr, seen: Arc<Mutex<Instant>>, sessions: Arc<Mutex<HashMap<SocketAddr, Session>>>) {
197    let mut buffer = vec![0u8; 65_535];
198    loop {
199        let idle = seen.lock().unwrap_or_else(|e| e.into_inner()).elapsed();
200        if idle >= UDP_IDLE {
201            break;
202        }
203        match tokio::time::timeout(UDP_IDLE - idle, backend.recv(&mut buffer)).await {
204            Ok(Ok(n)) => {
205                *seen.lock().unwrap_or_else(|e| e.into_inner()) = Instant::now();
206                let _ = listener.send_to(&buffer[..n], client).await;
207            }
208            Ok(Err(_)) => break,
209            Err(_) => continue, // re-check the idle time
210        }
211    }
212    sessions.lock().unwrap_or_else(|e| e.into_inner()).remove(&client);
213}
214
215#[cfg(test)]
216mod tests;
rws_gateway_api/gateway_api/l4.rs

rws_gateway_api/gateway_api/
l4.rs
