1use std::net::{Ipv4Addr, Ipv6Addr};
7use std::sync::Arc;
8use std::time::Duration;
9
10use tokio::io::{AsyncReadExt, AsyncWriteExt};
11use tokio::net::{TcpListener, TcpStream};
12use tokio::sync::mpsc;
13use tokio_stream::wrappers::ReceiverStream;
14
15use crate::error::SailError;
16use crate::pb::workerproxy::v1 as pb;
17use crate::worker::WorkerProxy;
18
19const TUNNEL_CHANNEL_CAP: usize = 32;
22const TUNNEL_READ_CHUNK: usize = 32 * 1024;
24
25pub struct PortForward {
28 local_port: u16,
29 accept: tokio::task::JoinHandle<()>,
30}
31
32impl Drop for PortForward {
33 fn drop(&mut self) {
34 self.accept.abort();
35 }
36}
37
38impl PortForward {
39 pub fn local_port(&self) -> u16 {
43 self.local_port
44 }
45}
46
47pub async fn forward_port(
52 worker: Arc<WorkerProxy>,
53 endpoint: String,
54 sailbox_id: String,
55 local_port: u16,
56 remote_port: u16,
57) -> Result<PortForward, SailError> {
58 let v4 = TcpListener::bind((Ipv4Addr::LOCALHOST, local_port))
59 .await
60 .map_err(|err| SailError::Internal {
61 message: format!("bind 127.0.0.1:{local_port}: {err}"),
62 })?;
63 let bound = v4.local_addr().map_or(local_port, |addr| addr.port());
64 let v6 = match TcpListener::bind((Ipv6Addr::LOCALHOST, bound)).await {
68 Ok(v6) => Some(v6),
69 Err(err) if err.kind() == std::io::ErrorKind::AddrInUse => {
74 return Err(SailError::Internal {
75 message: format!("bind [::1]:{bound}: {err}"),
76 });
77 }
78 Err(_) => None,
81 };
82 let accept = tokio::spawn(async move {
83 loop {
84 let accepted = match &v6 {
85 Some(v6) => tokio::select! {
86 conn = v4.accept() => conn,
87 conn = v6.accept() => conn,
88 },
89 None => v4.accept().await,
90 };
91 match accepted {
92 Ok((conn, _)) => {
93 let worker = Arc::clone(&worker);
94 let endpoint = endpoint.clone();
95 let sailbox_id = sailbox_id.clone();
96 tokio::spawn(async move {
97 let _ =
98 tunnel_connection(worker, &endpoint, &sailbox_id, remote_port, conn)
99 .await;
100 });
101 }
102 Err(_) => tokio::time::sleep(Duration::from_millis(50)).await,
105 }
106 }
107 });
108 Ok(PortForward {
109 local_port: bound,
110 accept,
111 })
112}
113
114async fn tunnel_connection(
118 worker: Arc<WorkerProxy>,
119 endpoint: &str,
120 sailbox_id: &str,
121 remote_port: u16,
122 conn: TcpStream,
123) -> Result<(), SailError> {
124 let (tx, rx) = mpsc::channel::<pb::TunnelSailboxPortRequest>(TUNNEL_CHANNEL_CAP);
125 let handshake = pb::TunnelSailboxPortRequest {
127 sailbox_id: sailbox_id.to_string(),
128 port: u32::from(remote_port),
129 data: Vec::new(),
130 eof: false,
131 };
132 if tx.send(handshake).await.is_err() {
133 return Ok(());
134 }
135 let (mut read_half, mut write_half) = conn.into_split();
136
137 let uplink = tokio::spawn(async move {
144 let mut buf = vec![0u8; TUNNEL_READ_CHUNK];
145 loop {
146 match read_half.read(&mut buf).await {
147 Ok(n) if n > 0 => {
148 if tx.send(data_frame(buf[..n].to_vec())).await.is_err() {
149 break;
150 }
151 }
152 _ => {
159 let _ = tx.send(eof_frame()).await;
160 break;
161 }
162 }
163 }
164 });
165
166 let response = match open_tunnel_stream(&worker, endpoint, rx).await {
167 Ok(response) => response,
168 Err(err) => {
169 uplink.abort();
170 return Err(err);
171 }
172 };
173 let mut server = response.into_inner();
174
175 while let Ok(Some(frame)) = server.message().await {
177 if !frame.data.is_empty() && write_half.write_all(&frame.data).await.is_err() {
178 break;
179 }
180 if frame.eof {
181 let _ = write_half.shutdown().await;
182 }
183 }
184 uplink.abort();
185 Ok(())
186}
187
188pub fn forwardable_local_port(url: &str) -> Option<u16> {
193 let parsed = url::Url::parse(url).ok()?;
194 if is_forwardable_local_host(&parsed) {
195 return parsed.port_or_known_default();
196 }
197 None
198}
199
200fn is_forwardable_local_host(url: &url::Url) -> bool {
206 match url.host() {
207 Some(url::Host::Domain(domain)) => domain == "localhost",
208 Some(url::Host::Ipv4(ip)) => ip == Ipv4Addr::LOCALHOST || ip.is_unspecified(),
209 Some(url::Host::Ipv6(ip)) => ip.is_unspecified(),
210 _ => false,
211 }
212}
213
214pub fn is_unforwardable_loopback_url(url: &str) -> bool {
220 let Ok(parsed) = url::Url::parse(url) else {
221 return false;
222 };
223 match parsed.host() {
224 Some(url::Host::Ipv4(ip)) => ip.is_loopback() && ip != Ipv4Addr::LOCALHOST,
225 Some(url::Host::Ipv6(ip)) => match ip.to_ipv4_mapped() {
229 Some(v4) => v4.is_loopback(),
230 None => ip.is_loopback(),
231 },
232 _ => false,
233 }
234}
235
236pub fn rewrite_loopback_url(url: &str, local_port_for: impl Fn(u16) -> Option<u16>) -> String {
245 let Ok(mut parsed) = url::Url::parse(url) else {
246 return url.to_string();
247 };
248 if !is_forwardable_local_host(&parsed) {
249 return url.to_string();
250 }
251 let Some(remote) = parsed.port_or_known_default() else {
252 return url.to_string();
253 };
254 let Some(local) = local_port_for(remote) else {
255 return url.to_string();
256 };
257 let wildcard = matches!(parsed.host(), Some(url::Host::Ipv4(ip)) if ip.is_unspecified())
258 || matches!(parsed.host(), Some(url::Host::Ipv6(ip)) if ip.is_unspecified());
259 if !wildcard {
260 return url.to_string();
261 }
262 if parsed.set_host(Some("localhost")).is_ok() && parsed.set_port(Some(local)).is_ok() {
263 parsed.to_string()
264 } else {
265 url.to_string()
266 }
267}
268
269pub enum RedirectCallback {
272 Forwardable(u16),
275 Unreachable,
279}
280
281pub fn redirect_callback(url: &str) -> Option<RedirectCallback> {
285 let parsed = url::Url::parse(url).ok()?;
286 for (key, value) in parsed.query_pairs() {
287 if key == "redirect_uri" || key == "redirect_url" {
288 if let Some(port) = forwardable_local_port(&value) {
289 return Some(RedirectCallback::Forwardable(port));
290 }
291 if is_unforwardable_loopback_url(&value) {
292 return Some(RedirectCallback::Unreachable);
293 }
294 }
295 }
296 None
297}
298
299type TunnelResponse = tonic::Response<tonic::Streaming<pb::TunnelSailboxPortResponse>>;
301
302async fn open_tunnel_stream(
303 worker: &Arc<WorkerProxy>,
304 endpoint: &str,
305 rx: mpsc::Receiver<pb::TunnelSailboxPortRequest>,
306) -> Result<TunnelResponse, SailError> {
307 let request = worker.request_for(ReceiverStream::new(rx), &[], None)?;
308 worker
309 .client_for(endpoint)?
310 .tunnel_sailbox_port(request)
311 .await
312 .map_err(|status| SailError::from_file_rpc_status(&status))
313}
314
315fn data_frame(data: Vec<u8>) -> pb::TunnelSailboxPortRequest {
316 pb::TunnelSailboxPortRequest {
317 sailbox_id: String::new(),
318 port: 0,
319 data,
320 eof: false,
321 }
322}
323
324fn eof_frame() -> pb::TunnelSailboxPortRequest {
325 pb::TunnelSailboxPortRequest {
326 sailbox_id: String::new(),
327 port: 0,
328 data: Vec::new(),
329 eof: true,
330 }
331}
332
333#[cfg(test)]
334mod tests {
335 use super::*;
336 use std::collections::HashMap;
337
338 #[test]
339 fn forwardable_local_port_covers_loopback_wildcard_and_defaults() {
340 assert_eq!(forwardable_local_port("http://localhost:5173/"), Some(5173));
341 assert_eq!(
342 forwardable_local_port("http://127.0.0.1:9000/x"),
343 Some(9000)
344 );
345 assert_eq!(forwardable_local_port("http://0.0.0.0:3000/"), Some(3000));
348 assert_eq!(forwardable_local_port("http://[::]:3000/"), Some(3000));
349 assert_eq!(forwardable_local_port("http://localhost/cb"), Some(80));
351 assert_eq!(forwardable_local_port("https://localhost/cb"), Some(443));
352 assert_eq!(forwardable_local_port("https://example.com:8080/"), None);
355 assert_eq!(forwardable_local_port("http://[::1]:5173/"), None);
356 assert_eq!(
357 forwardable_local_port("http://[::ffff:127.0.0.1]:5173/"),
358 None
359 );
360 assert_eq!(forwardable_local_port("http://127.0.0.2:3000/"), None);
361 }
362
363 #[test]
364 fn is_unforwardable_loopback_url_matches_undialed_loopbacks() {
365 assert!(is_unforwardable_loopback_url("http://[::1]:5173/"));
368 assert!(is_unforwardable_loopback_url("http://127.0.0.2:3000/"));
369 assert!(is_unforwardable_loopback_url(
370 "http://[::ffff:127.0.0.1]:3000/"
371 ));
372 assert!(!is_unforwardable_loopback_url("http://localhost:8080/"));
374 assert!(!is_unforwardable_loopback_url("http://127.0.0.1:8080/"));
375 assert!(!is_unforwardable_loopback_url("http://0.0.0.0:8080/"));
376 assert!(!is_unforwardable_loopback_url("http://[::]:8080/"));
377 assert!(!is_unforwardable_loopback_url("https://example.com/"));
378 }
379
380 #[test]
381 fn rewrite_loopback_url_only_rewrites_wildcard_hosts() {
382 let mut forwards = HashMap::new();
383 forwards.insert(3000u16, 3000u16); forwards.insert(8080u16, 8080u16);
385 let lookup = |remote: u16| forwards.get(&remote).copied();
386
387 assert_eq!(
390 rewrite_loopback_url("https://localhost:8080/app", lookup),
391 "https://localhost:8080/app"
392 );
393 assert_eq!(
394 rewrite_loopback_url("http://127.0.0.1:3000/x", lookup),
395 "http://127.0.0.1:3000/x"
396 );
397 assert_eq!(
399 rewrite_loopback_url("http://0.0.0.0:3000/x", lookup),
400 "http://localhost:3000/x"
401 );
402 assert_eq!(
403 rewrite_loopback_url("http://[::]:3000/x", lookup),
404 "http://localhost:3000/x"
405 );
406 assert_eq!(
408 rewrite_loopback_url("http://localhost:9999/x", lookup),
409 "http://localhost:9999/x"
410 );
411 assert_eq!(
412 rewrite_loopback_url("https://example.com:8080/x", lookup),
413 "https://example.com:8080/x"
414 );
415 }
416
417 #[test]
418 fn redirect_callback_classifies_the_redirect_uri() {
419 use RedirectCallback::{Forwardable, Unreachable};
420 assert!(matches!(
423 redirect_callback(
424 "https://example.com/oauth?client_id=x&redirect_uri=http%3A%2F%2Flocalhost%3A43222%2Fcallback"
425 ),
426 Some(Forwardable(43222))
427 ));
428 assert!(matches!(
429 redirect_callback("https://id.example/a?redirect_url=http://127.0.0.1:5000/cb"),
430 Some(Forwardable(5000))
431 ));
432 assert!(matches!(
434 redirect_callback("https://example.com/oauth?redirect_uri=http://[::1]:5173/cb"),
435 Some(Unreachable)
436 ));
437 assert!(matches!(
438 redirect_callback("https://example.com/oauth?redirect_uri=http://127.0.0.2:5173/cb"),
439 Some(Unreachable)
440 ));
441 assert!(
443 redirect_callback("https://example.com/oauth?redirect_uri=https://example.com/cb")
444 .is_none()
445 );
446 assert!(redirect_callback("https://example.com/login").is_none());
447 }
448}