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    pick_backend(route, counter).map(|(endpoint, _)| endpoint)
46}
47
48/// The backend's TLS when a BackendTLSPolicy covers it (`Err`: covered but
49/// unusable — such a connection is dropped, never sent in the clear).
50type BackendTlsOf = Option<Result<super::backend_tls::BackendTls, String>>;
51
52/// [`pick_endpoint`], with the picked backend's TLS.
53fn pick_backend(route: &CompiledHttpRoute, counter: &AtomicUsize) -> Option<(String, BackendTlsOf)> {
54    let rule = route.rules.iter().find(|r| !r.backends.is_empty())?;
55    let total: u32 = rule.backends.iter().map(|b| b.weight).sum();
56    if total == 0 {
57        return None;
58    }
59    let n = counter.fetch_add(1, Ordering::Relaxed);
60    let mut pick = ((n as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 32) as u32 % total;
61    let backend = rule.backends.iter().find(|b| {
62        if pick < b.weight {
63            true
64        } else {
65            pick -= b.weight;
66            false
67        }
68    })?;
69    let endpoints = backend.endpoints.as_ref().ok().filter(|e| !e.is_empty())?;
70    Some((endpoints[n % endpoints.len()].clone(), backend.tls.clone()))
71}
72
73/// The listener and endpoint for a connection with this SNI.
74fn target(listeners: &PortListeners, sni: Option<&str>, counter: &AtomicUsize) -> Option<String> {
75    let listeners: Arc<Vec<ListenerConfig>> = listeners.read().unwrap_or_else(|e| e.into_inner()).clone();
76    let listener = pick_listener(&listeners, &sni.unwrap_or("").to_ascii_lowercase())?;
77    pick_endpoint(pick_route(&listener.http_routes, sni)?, counter)
78}
79
80async fn connect(endpoint: &str) -> Option<TcpStream> {
81    let address: SocketAddr = endpoint.parse().ok()?;
82    tokio::time::timeout(CONNECT_TIMEOUT, TcpStream::connect(address)).await.ok()?.ok()
83}
84
85/// A TCPRoute connection: straight to a backend.
86pub async fn serve_tcp(mut client: TcpStream, listeners: PortListeners, counter: Arc<AtomicUsize>) {
87    let Some(endpoint) = target(&listeners, None, &counter) else { return };
88    let Some(mut backend) = connect(&endpoint).await else { return };
89    let _ = tokio::io::copy_bidirectional(&mut client, &mut backend).await;
90}
91
92/// What the start of a TLS connection says.
93#[derive(Debug, PartialEq)]
94pub enum Hello {
95    /// Not a whole ClientHello yet.
96    NeedMore,
97    /// A ClientHello, with its SNI if it has one.
98    Sni(Option<String>),
99    /// Not a ClientHello rustls accepts.
100    NotTls,
101}
102
103/// Read the ClientHello at the start of `bytes`.
104pub fn client_hello(bytes: &[u8]) -> Hello {
105    let mut acceptor = rustls::server::Acceptor::default();
106    let mut input = bytes;
107    while !input.is_empty() {
108        if acceptor.read_tls(&mut input).is_err() {
109            return Hello::NotTls;
110        }
111    }
112    match acceptor.accept() {
113        Ok(Some(accepted)) => Hello::Sni(accepted.client_hello().server_name().map(str::to_string)),
114        Ok(None) => Hello::NeedMore,
115        Err(_) => Hello::NotTls,
116    }
117}
118
119/// A passthrough TLSRoute connection: read the ClientHello for its SNI,
120/// then replay it to the backend and relay the encrypted stream.
121pub async fn serve_tls_passthrough(mut client: TcpStream, listeners: PortListeners, counter: Arc<AtomicUsize>) {
122    use tokio::io::AsyncReadExt;
123    let mut hello = Vec::new();
124    let sni = loop {
125        let mut buffer = [0u8; 4096];
126        let n = match tokio::time::timeout(Duration::from_secs(10), client.read(&mut buffer)).await {
127            Ok(Ok(n)) if n > 0 => n,
128            _ => return,
129        };
130        hello.extend_from_slice(&buffer[..n]);
131        match client_hello(&hello) {
132            Hello::Sni(sni) => break sni,
133            Hello::NeedMore if hello.len() < MAX_CLIENT_HELLO => continue,
134            _ => return,
135        }
136    };
137    let Some(endpoint) = target(&listeners, sni.as_deref(), &counter) else { return };
138    let Some(mut backend) = connect(&endpoint).await else { return };
139    if backend.write_all(&hello).await.is_err() {
140        return;
141    }
142    let _ = tokio::io::copy_bidirectional(&mut client, &mut backend).await;
143}
144
145/// A terminated TLSRoute connection: handshake with the listener's
146/// certificates, then relay the plaintext.
147pub async fn serve_tls_terminate(client: TcpStream, acceptor: tokio_rustls::TlsAcceptor, listeners: PortListeners, counter: Arc<AtomicUsize>) {
148    let Ok(Ok(mut tls)) = tokio::time::timeout(Duration::from_secs(10), acceptor.accept(client)).await else { return };
149    let sni = tls.get_ref().1.server_name().map(str::to_string);
150    let current: Arc<Vec<ListenerConfig>> = listeners.read().unwrap_or_else(|e| e.into_inner()).clone();
151    let Some(listener) = pick_listener(&current, &sni.as_deref().unwrap_or("").to_ascii_lowercase()) else { return };
152    let Some((endpoint, backend_tls)) = pick_route(&listener.http_routes, sni.as_deref()).and_then(|route| pick_backend(route, &counter)) else { return };
153    let Some(backend) = connect(&endpoint).await else { return };
154    match backend_tls {
155        None => {
156            let mut backend = backend;
157            let _ = tokio::io::copy_bidirectional(&mut tls, &mut backend).await;
158        }
159        // BackendTLSPolicy: encrypted again, checked as the policy says.
160        Some(Ok(policy)) => {
161            let Ok(name) = rustls::pki_types::ServerName::try_from(policy.server_name.clone()) else { return };
162            let connector = tokio_rustls::TlsConnector::from(Arc::clone(&policy.config));
163            let Ok(Ok(mut backend)) = tokio::time::timeout(CONNECT_TIMEOUT, connector.connect(name, backend)).await else { return };
164            let _ = tokio::io::copy_bidirectional(&mut tls, &mut backend).await;
165        }
166        Some(Err(_)) => {}
167    }
168}
169
170/// The SNI of a connection's ClientHello, read with `peek` so the bytes
171/// stay in the socket: `Err` when the start of the stream isn't TLS or the
172/// hello doesn't come within 10 s.
173async fn peek_sni(client: &TcpStream) -> Result<Option<String>, ()> {
174    let deadline = Instant::now() + Duration::from_secs(10);
175    let mut buffer = vec![0u8; MAX_CLIENT_HELLO];
176    let mut last = 0;
177    loop {
178        let left = deadline.saturating_duration_since(Instant::now());
179        let n = match tokio::time::timeout(left, client.peek(&mut buffer)).await {
180            Ok(Ok(n)) if n > 0 => n,
181            _ => return Err(()),
182        };
183        match client_hello(&buffer[..n]) {
184            Hello::Sni(sni) => return Ok(sni),
185            Hello::NeedMore if n < MAX_CLIENT_HELLO => {
186                // `peek` returns at once while the bytes it saw are unread:
187                // wait for more to arrive before looking again.
188                if n == last {
189                    tokio::time::sleep(Duration::from_millis(5)).await;
190                }
191                last = n;
192            }
193            _ => return Err(()),
194        }
195    }
196}
197
198/// A TLS port with listeners of both modes (`TLSRouteModeMixed`): the
199/// ClientHello's SNI picks the listener; a `Terminate` one gets the
200/// handshake with its certificates, a `Passthrough` one the untouched stream.
201pub async fn serve_tls_mixed(client: TcpStream, acceptor: Option<tokio_rustls::TlsAcceptor>, listeners: PortListeners, counter: Arc<AtomicUsize>) {
202    let current: Arc<Vec<ListenerConfig>> = listeners.read().unwrap_or_else(|e| e.into_inner()).clone();
203    let all_terminate = current.iter().all(|l| l.tls.is_some());
204    if all_terminate {
205        if let Some(acceptor) = acceptor {
206            serve_tls_terminate(client, acceptor, listeners, counter).await;
207        }
208        return;
209    }
210    let Ok(sni) = peek_sni(&client).await else { return };
211    let terminate = pick_listener(&current, &sni.as_deref().unwrap_or("").to_ascii_lowercase()).is_some_and(|l| l.tls.is_some());
212    match acceptor {
213        Some(acceptor) if terminate => serve_tls_terminate(client, acceptor, listeners, counter).await,
214        _ if terminate => {}
215        _ => {
216            let mut client = client;
217            let Some(endpoint) = target(&listeners, sni.as_deref(), &counter) else { return };
218            let Some(mut backend) = connect(&endpoint).await else { return };
219            let _ = tokio::io::copy_bidirectional(&mut client, &mut backend).await;
220        }
221    }
222}
223
224/// A UDP client's session: its socket to the backend and when it was last used.
225struct Session {
226    backend: Arc<UdpSocket>,
227    seen: Arc<Mutex<Instant>>,
228}
229
230/// A UDPRoute listener: each client address gets its own socket to a
231/// backend (picked once), replies go back from the listener's socket;
232/// sessions end after [`UDP_IDLE`] without traffic — or with the port:
233/// they're in a `JoinSet` this task owns, so dropping it (the port given up)
234/// aborts them and releases the socket at once.
235pub async fn serve_udp(socket: Arc<UdpSocket>, listeners: PortListeners) {
236    let counter = Arc::new(AtomicUsize::new(0));
237    let mut session_tasks = tokio::task::JoinSet::new();
238    let sessions: Arc<Mutex<HashMap<SocketAddr, Session>>> = Arc::default();
239    let mut buffer = vec![0u8; 65_535];
240    loop {
241        // Reap finished sessions so the set doesn't grow.
242        while session_tasks.try_join_next().is_some() {}
243        let Ok((n, client)) = socket.recv_from(&mut buffer).await else {
244            tokio::time::sleep(Duration::from_millis(50)).await;
245            continue;
246        };
247        let existing = sessions.lock().unwrap_or_else(|e| e.into_inner()).get(&client).map(|s| (Arc::clone(&s.backend), Arc::clone(&s.seen)));
248        let (backend, seen) = match existing {
249            Some(session) => session,
250            None => {
251                let Some(endpoint) = target(&listeners, None, &counter) else { continue };
252                let Some(backend) = udp_connect(&endpoint).await else { continue };
253                let backend = Arc::new(backend);
254                let seen = Arc::new(Mutex::new(Instant::now()));
255                sessions.lock().unwrap_or_else(|e| e.into_inner()).insert(client, Session { backend: Arc::clone(&backend), seen: Arc::clone(&seen) });
256                session_tasks.spawn(udp_replies(Arc::clone(&socket), Arc::clone(&backend), client, Arc::clone(&seen), Arc::clone(&sessions)));
257                (backend, seen)
258            }
259        };
260        *seen.lock().unwrap_or_else(|e| e.into_inner()) = Instant::now();
261        let _ = backend.send(&buffer[..n]).await;
262    }
263}
264
265async fn udp_connect(endpoint: &str) -> Option<UdpSocket> {
266    let address: SocketAddr = endpoint.parse().ok()?;
267    let local = if address.is_ipv6() { "[::]:0" } else { "0.0.0.0:0" };
268    let socket = UdpSocket::bind(local).await.ok()?;
269    socket.connect(address).await.ok()?;
270    Some(socket)
271}
272
273/// Backend → client for one session, until it's idle for [`UDP_IDLE`].
274async fn udp_replies(listener: Arc<UdpSocket>, backend: Arc<UdpSocket>, client: SocketAddr, seen: Arc<Mutex<Instant>>, sessions: Arc<Mutex<HashMap<SocketAddr, Session>>>) {
275    let mut buffer = vec![0u8; 65_535];
276    loop {
277        let idle = seen.lock().unwrap_or_else(|e| e.into_inner()).elapsed();
278        if idle >= UDP_IDLE {
279            break;
280        }
281        match tokio::time::timeout(UDP_IDLE - idle, backend.recv(&mut buffer)).await {
282            Ok(Ok(n)) => {
283                *seen.lock().unwrap_or_else(|e| e.into_inner()) = Instant::now();
284                let _ = listener.send_to(&buffer[..n], client).await;
285            }
286            Ok(Err(_)) => break,
287            Err(_) => continue, // re-check the idle time
288        }
289    }
290    sessions.lock().unwrap_or_else(|e| e.into_inner()).remove(&client);
291}
292
293#[cfg(test)]
294mod tests;
rws_gateway_api/gateway_api/l4.rs

rws_gateway_api/gateway_api/
l4.rs
