1#![forbid(unsafe_code)]
15
16use std::collections::HashMap;
17use std::io::{BufRead, BufReader, Read, Write};
18use std::net::{IpAddr, SocketAddr, TcpListener, TcpStream, ToSocketAddrs};
19use std::path::PathBuf;
20use std::process::Command;
21use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
22use std::sync::{Arc, Mutex};
23use std::thread::JoinHandle;
24use std::time::{Duration, Instant};
25
26const MAX_HEAD: usize = 16 * 1024;
27
28pub const MAX_CONNECTIONS: usize = 256;
31
32const CHUNK: usize = 16 * 1024;
35
36#[derive(Debug, Clone, Copy)]
38pub struct Request<'a> {
39 pub host: &'a str,
41 pub port: u16,
42 pub peer: SocketAddr,
44 pub local: SocketAddr,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct Verdict {
50 pub allow: bool,
51 pub allow_private: bool,
53 pub reason: String,
55}
56
57impl Verdict {
58 pub fn allow() -> Verdict {
59 Verdict {
60 allow: true,
61 allow_private: false,
62 reason: String::new(),
63 }
64 }
65
66 pub fn deny(reason: impl Into<String>) -> Verdict {
67 Verdict {
68 allow: false,
69 allow_private: false,
70 reason: reason.into(),
71 }
72 }
73}
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub enum Outcome {
78 Refused,
80 Unresolved,
82 Private,
84 Unreachable,
86 Cut,
88 Tunneled,
90}
91
92pub trait Gate: Send + Sync + 'static {
93 fn decide(&self, request: &Request<'_>) -> Verdict;
95
96 fn finished(&self, _request: &Request<'_>, _outcome: Outcome) {}
98
99 fn carried(&self, _request: &Request<'_>, _bytes: u64) -> bool {
102 true
103 }
104
105 fn relay(&self, _request: &Request<'_>, tunnel: Tunnel) -> Option<Tunnel> {
111 Some(tunnel)
112 }
113
114 fn relay_plain(
122 &self,
123 _request: &Request<'_>,
124 head: String,
125 tunnel: Tunnel,
126 ) -> Option<(String, Tunnel)> {
127 Some((head, tunnel))
128 }
129}
130
131pub struct Tunnel {
134 pub client: TcpStream,
135 pub upstream: TcpStream,
136 pub early: Vec<u8>,
139}
140
141pub fn bind(addr: &str) -> std::io::Result<TcpListener> {
144 let l = TcpListener::bind(addr)?;
145 if !l.local_addr()?.ip().is_loopback() {
146 return Err(std::io::Error::new(
147 std::io::ErrorKind::InvalidInput,
148 "the proxy only listens on loopback",
149 ));
150 }
151 Ok(l)
152}
153
154pub fn serve<G: Gate>(listener: TcpListener, gate: Arc<G>) {
156 accept(listener, gate, &Arc::new(Shared::default()));
157}
158
159pub struct Server {
162 addr: SocketAddr,
163 shared: Arc<Shared>,
164 thread: Option<JoinHandle<()>>,
165}
166
167#[derive(Default)]
168struct Shared {
169 stopping: AtomicBool,
170 open: Mutex<HashMap<u64, Vec<TcpStream>>>,
173 next: AtomicU64,
174}
175
176impl Shared {
177 fn track(&self, id: u64, stream: &TcpStream) {
180 let mut open = self.open.lock().unwrap();
181 if self.stopping.load(Ordering::SeqCst) {
182 let _ = stream.shutdown(std::net::Shutdown::Both);
183 } else if let Ok(clone) = stream.try_clone() {
184 open.entry(id).or_default().push(clone);
185 }
186 }
187}
188
189pub fn spawn<G: Gate>(listener: TcpListener, gate: Arc<G>) -> std::io::Result<Server> {
191 let addr = listener.local_addr()?;
192 let shared = Arc::new(Shared::default());
193 let served = Arc::clone(&shared);
194 let thread = std::thread::spawn(move || accept(listener, gate, &served));
195 Ok(Server {
196 addr,
197 shared,
198 thread: Some(thread),
199 })
200}
201
202impl Server {
203 pub fn local_addr(&self) -> SocketAddr {
204 self.addr
205 }
206
207 pub fn stop(mut self) {
211 self.shut();
212 }
213
214 fn shut(&mut self) {
215 let Some(thread) = self.thread.take() else {
216 return;
217 };
218 {
219 let _open = self.shared.open.lock().unwrap();
222 self.shared.stopping.store(true, Ordering::SeqCst);
223 }
224 let _ = TcpStream::connect_timeout(&self.addr, Duration::from_secs(1));
226 let _ = thread.join();
227 for stream in self.shared.open.lock().unwrap().values().flatten() {
228 let _ = stream.shutdown(std::net::Shutdown::Both);
229 }
230 let deadline = Instant::now() + Duration::from_secs(3);
233 while !self.shared.open.lock().unwrap().is_empty() && Instant::now() < deadline {
234 std::thread::sleep(Duration::from_millis(5));
235 }
236 }
237}
238
239impl Drop for Server {
240 fn drop(&mut self) {
241 self.shut();
242 }
243}
244
245fn accept<G: Gate>(listener: TcpListener, gate: Arc<G>, shared: &Arc<Shared>) {
246 for conn in listener.incoming() {
247 if shared.stopping.load(Ordering::SeqCst) {
248 return;
249 }
250 let Ok(stream) = conn else { continue };
251 if shared.open.lock().unwrap().len() >= MAX_CONNECTIONS {
252 respond(&stream, 503, "too many connections through the proxy");
253 continue;
254 }
255 let id = shared.next.fetch_add(1, Ordering::Relaxed);
256 shared.track(id, &stream);
257 let (gate, shared) = (gate.clone(), Arc::clone(shared));
258 std::thread::spawn(move || {
259 if !shared.stopping.load(Ordering::SeqCst) {
261 let _ = handle(stream, &*gate, &|upstream| shared.track(id, upstream));
262 }
263 shared.open.lock().unwrap().remove(&id);
264 });
265 }
266}
267
268pub fn is_private(ip: IpAddr) -> bool {
273 match ip {
274 IpAddr::V4(v) => private_v4(v),
275 IpAddr::V6(v) => {
276 if let Some(v4) = v.to_ipv4_mapped() {
277 return private_v4(v4);
278 }
279 let s = v.segments();
280 if s[..6] == [0x64, 0xff9b, 0, 0, 0, 0] {
282 let [a, b] = s[6].to_be_bytes();
283 let [c, d] = s[7].to_be_bytes();
284 return private_v4(std::net::Ipv4Addr::new(a, b, c, d));
285 }
286 v.is_unspecified()
287 || v.is_loopback()
288 || v.is_multicast()
289 || (s[0] & 0xfe00) == 0xfc00 || (s[0] & 0xffc0) == 0xfe80 || (s[0] & 0xffc0) == 0xfec0 || (s[0] == 0x2001 && s[1] == 0x0db8) || s[..6] == [0, 0, 0, 0, 0, 0] }
295 }
296}
297
298pub fn is_metadata(ip: IpAddr) -> bool {
302 let v4 = |v: std::net::Ipv4Addr| {
303 matches!(
304 v.octets(),
305 [169, 254, 169, 254] | [169, 254, 170, 2] | [169, 254, 170, 23] | [100, 100, 100, 200]
306 )
307 };
308 match ip {
309 IpAddr::V4(v) => v4(v),
310 IpAddr::V6(v) => {
311 v.to_ipv4_mapped().is_some_and(v4)
312 || v.segments() == [0xfd00, 0x0ec2, 0, 0, 0, 0, 0, 0x0254]
313 }
314 }
315}
316
317fn private_v4(ip: std::net::Ipv4Addr) -> bool {
318 let [a, b, c, _] = ip.octets();
319 ip.is_unspecified()
320 || ip.is_loopback()
321 || ip.is_private()
322 || ip.is_link_local()
323 || ip.is_broadcast()
324 || ip.is_documentation()
325 || ip.is_multicast()
326 || a == 0
327 || (a == 100 && (64..128).contains(&b)) || (a == 192 && b == 0 && c == 0) || (a == 198 && (18..20).contains(&b)) || a >= 240 }
332
333fn parse_authority(a: &str, default_port: u16) -> Option<(String, u16)> {
335 let (host, port) = if let Some(rest) = a.strip_prefix('[') {
336 let (h, tail) = rest.split_once(']')?;
337 let port = tail
338 .strip_prefix(':')
339 .map(|p| p.parse().ok())
340 .unwrap_or(Some(default_port))?;
341 (h.to_string(), port)
342 } else {
343 match a.rsplit_once(':') {
344 Some((h, p)) => (h.to_string(), p.parse().ok()?),
345 None => (a.to_string(), default_port),
346 }
347 };
348 let ok = !host.is_empty()
349 && host.len() <= 253
350 && host
351 .chars()
352 .all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | ':'));
353 ok.then_some((host.to_ascii_lowercase(), port))
354}
355
356pub fn peer_cwd(peer: SocketAddr, local: SocketAddr) -> Option<PathBuf> {
359 let me = std::process::id().to_string();
360 let out = Command::new("lsof")
361 .args([
362 "-nP",
363 &format!("-iTCP@{}:{}", peer.ip(), peer.port()),
364 "-sTCP:ESTABLISHED",
365 "-Fpn",
366 ])
367 .output()
368 .ok()?;
369 let text = String::from_utf8_lossy(&out.stdout);
370 let want = format!(
373 "{}:{}->{}:{}",
374 peer.ip(),
375 peer.port(),
376 local.ip(),
377 local.port()
378 );
379 let mut pid = None;
380 let mut cur = None;
381 for line in text.lines() {
382 if let Some(p) = line.strip_prefix('p') {
383 cur = Some(p.to_string());
384 } else if let Some(n) = line.strip_prefix('n')
385 && n == want
386 && cur.as_deref() != Some(me.as_str())
387 {
388 pid = cur.clone();
389 }
390 }
391 let pid = pid?;
392 let out = Command::new("lsof")
393 .args(["-a", "-p", &pid, "-d", "cwd", "-Fn"])
394 .output()
395 .ok()?;
396 String::from_utf8_lossy(&out.stdout)
397 .lines()
398 .find_map(|l| l.strip_prefix('n').map(PathBuf::from))
399}
400
401fn connect_any(addrs: &[SocketAddr], within: Duration) -> Option<TcpStream> {
406 let deadline = Instant::now() + within;
407 addrs.iter().find_map(|a| {
408 TcpStream::connect_timeout(a, deadline.saturating_duration_since(Instant::now())).ok()
409 })
410}
411
412fn respond(mut s: &TcpStream, code: u16, reason: &str) {
413 let body = format!("proxy: {reason}\n");
414 let _ = write!(
415 s,
416 "HTTP/1.1 {code} {reason}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
417 body.len()
418 );
419}
420
421fn handle(client: TcpStream, gate: &dyn Gate, track: &dyn Fn(&TcpStream)) -> std::io::Result<()> {
424 client.set_read_timeout(Some(Duration::from_secs(30)))?;
425 let peer = client.peer_addr()?;
426 let local = client.local_addr()?;
427 let mut reader = BufReader::new(client.try_clone()?);
428 let mut head = Vec::new();
429 loop {
430 let mut line = Vec::new();
431 let n = reader
432 .by_ref()
433 .take((MAX_HEAD - head.len()) as u64)
434 .read_until(b'\n', &mut line)?;
435 if n == 0 || head.len() + n >= MAX_HEAD {
436 respond(&client, 400, "bad request");
437 return Ok(());
438 }
439 head.extend_from_slice(&line);
440 if line == b"\r\n" || line == b"\n" {
441 break;
442 }
443 }
444 let head_text = String::from_utf8_lossy(&head).into_owned();
445 let first = head_text.lines().next().unwrap_or("");
446 let mut parts = first.split_whitespace();
447 let (method, target) = (parts.next().unwrap_or(""), parts.next().unwrap_or(""));
448 let (authority, rest_of_request) = if method.eq_ignore_ascii_case("CONNECT") {
449 (parse_authority(target, 443), None)
450 } else if let Some(url) = target.strip_prefix("http://") {
451 let (auth, path) = url.split_at(url.find('/').unwrap_or(url.len()));
452 let path = if path.is_empty() { "/" } else { path };
453 let rewritten = head_text.replacen(target, path, 1);
454 (parse_authority(auth, 80), Some(rewritten))
455 } else {
456 respond(&client, 400, "only CONNECT and absolute http:// requests");
457 return Ok(());
458 };
459 let Some((host, port)) = authority else {
460 respond(&client, 400, "bad target");
461 return Ok(());
462 };
463 let request = Request {
464 host: &host,
465 port,
466 peer,
467 local,
468 };
469 let verdict = gate.decide(&request);
470 if !verdict.allow {
471 respond(&client, 403, &verdict.reason);
472 gate.finished(&request, Outcome::Refused);
473 return Ok(());
474 }
475 let addrs: Vec<SocketAddr> = match (host.as_str(), port).to_socket_addrs() {
476 Ok(a) => a.collect(),
477 Err(_) => {
478 respond(&client, 502, "cannot resolve host");
479 gate.finished(&request, Outcome::Unresolved);
480 return Ok(());
481 }
482 };
483 if !addrs.is_empty() && addrs.iter().all(|a| is_metadata(a.ip())) {
484 respond(&client, 403, "cloud metadata addresses are never reachable");
485 gate.finished(&request, Outcome::Private);
486 return Ok(());
487 }
488 let allowed: Vec<SocketAddr> = addrs
489 .into_iter()
490 .filter(|a| !is_metadata(a.ip()))
491 .filter(|a| verdict.allow_private || !is_private(a.ip()))
492 .collect();
493 if allowed.is_empty() {
494 respond(
495 &client,
496 403,
497 "private and loopback addresses are not reachable",
498 );
499 gate.finished(&request, Outcome::Private);
500 return Ok(());
501 }
502 let Some(upstream) = connect_any(&allowed, Duration::from_secs(10)) else {
503 respond(&client, 502, "cannot connect");
504 gate.finished(&request, Outcome::Unreachable);
505 return Ok(());
506 };
507 track(&upstream);
508 client.set_read_timeout(None)?;
509 let buffered = reader.buffer().to_vec();
511 drop(reader);
512 let (client, upstream, buffered) = match rest_of_request {
513 None => {
514 let mut c = &client;
515 c.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n")?;
516 let tunnel = Tunnel {
517 client,
518 upstream,
519 early: buffered,
520 };
521 match gate.relay(&request, tunnel) {
522 Some(t) => (t.client, t.upstream, t.early),
523 None => {
524 gate.finished(&request, Outcome::Tunneled);
525 return Ok(());
526 }
527 }
528 }
529 Some(h) => {
530 let tunnel = Tunnel {
531 client,
532 upstream,
533 early: buffered,
534 };
535 match gate.relay_plain(&request, h, tunnel) {
536 Some((h, t)) => {
537 (&t.upstream).write_all(h.as_bytes())?;
538 (t.client, t.upstream, t.early)
539 }
540 None => {
541 gate.finished(&request, Outcome::Tunneled);
542 return Ok(());
543 }
544 }
545 }
546 };
547 let mut up_w = upstream.try_clone()?;
548 if !buffered.is_empty() {
549 if !gate.carried(&request, buffered.len() as u64) {
550 gate.finished(&request, Outcome::Cut);
551 return Ok(());
552 }
553 up_w.write_all(&buffered)?;
554 }
555 let cut = AtomicBool::new(false);
556 let (up_r, client_w) = (upstream.try_clone()?, client.try_clone()?);
557 std::thread::scope(|scope| {
558 scope.spawn(|| {
559 if !carry(&up_r, &client_w, gate, &request) {
560 cut.store(true, Ordering::SeqCst);
561 let _ = up_r.shutdown(std::net::Shutdown::Both);
562 let _ = client_w.shutdown(std::net::Shutdown::Both);
563 }
564 let _ = client_w.shutdown(std::net::Shutdown::Write);
565 });
566 if !carry(&client, &up_w, gate, &request) {
567 cut.store(true, Ordering::SeqCst);
568 let _ = client.shutdown(std::net::Shutdown::Both);
569 let _ = up_w.shutdown(std::net::Shutdown::Both);
570 }
571 let _ = up_w.shutdown(std::net::Shutdown::Write);
572 });
573 let how = if cut.load(Ordering::SeqCst) {
574 Outcome::Cut
575 } else {
576 Outcome::Tunneled
577 };
578 gate.finished(&request, how);
579 Ok(())
580}
581
582fn carry(mut from: &TcpStream, mut to: &TcpStream, gate: &dyn Gate, request: &Request<'_>) -> bool {
585 let mut buf = vec![0u8; CHUNK];
586 loop {
587 let n = match from.read(&mut buf) {
588 Ok(0) | Err(_) => return true,
589 Ok(n) => n,
590 };
591 if !gate.carried(request, n as u64) {
592 return false;
593 }
594 if to.write_all(&buf[..n]).is_err() {
595 return true;
596 }
597 }
598}
599
600#[cfg(test)]
601mod tests {
602 use super::*;
603
604 #[test]
605 fn a_connection_goes_to_the_first_address_that_takes_it() {
606 let closed = TcpListener::bind("127.0.0.1:0")
608 .unwrap()
609 .local_addr()
610 .unwrap();
611 let open = TcpListener::bind("127.0.0.1:0").unwrap();
612 let at = open.local_addr().unwrap();
613 let budget = Duration::from_secs(5);
614 let s = connect_any(&[closed, at], budget).expect("the second address");
615 assert_eq!(s.peer_addr().unwrap(), at);
616 assert!(connect_any(&[closed], budget).is_none());
617 assert!(connect_any(&[], budget).is_none());
618 assert!(
619 connect_any(&[at], Duration::ZERO).is_none(),
620 "no budget, no attempt"
621 );
622 }
623
624 #[test]
625 fn metadata_addresses() {
626 for ip in [
627 "169.254.169.254",
628 "169.254.170.2",
629 "100.100.100.200",
630 "fd00:ec2::254",
631 "::ffff:169.254.169.254",
632 ] {
633 assert!(is_metadata(ip.parse().unwrap()), "{ip}");
634 }
635 for ip in ["169.254.169.253", "10.0.0.1", "fd00:ec2::253", "8.8.8.8"] {
636 assert!(!is_metadata(ip.parse().unwrap()), "{ip}");
637 }
638 }
639
640 #[test]
641 fn private_ranges() {
642 for ip in [
643 "127.0.0.1",
644 "10.1.2.3",
645 "192.168.0.1",
646 "172.16.5.4",
647 "169.254.1.1",
648 "100.64.0.1",
649 "0.0.0.0",
650 "::1",
651 "fd00::1",
652 "fe80::1",
653 "::ffff:127.0.0.1",
654 "192.0.0.1",
655 "198.18.0.1",
656 "192.0.2.1",
657 "224.0.0.1",
658 "240.0.0.1",
659 "255.255.255.255",
660 "fec0::1",
661 "ff02::1",
662 "2001:db8::1",
663 "64:ff9b::7f00:1",
664 "::127.0.0.1",
665 ] {
666 assert!(is_private(ip.parse().unwrap()), "{ip}");
667 }
668 for ip in [
669 "93.184.216.34",
670 "1.1.1.1",
671 "2606:4700::1111",
672 "::ffff:8.8.8.8",
673 "64:ff9b::808:808",
674 ] {
675 assert!(!is_private(ip.parse().unwrap()), "{ip}");
676 }
677 }
678
679 #[test]
680 fn authorities() {
681 assert_eq!(
682 parse_authority("Example.com:443", 1),
683 Some(("example.com".into(), 443))
684 );
685 assert_eq!(
686 parse_authority("example.com", 80),
687 Some(("example.com".into(), 80))
688 );
689 assert_eq!(parse_authority("[::1]:8080", 1), Some(("::1".into(), 8080)));
690 assert_eq!(parse_authority("a b:1", 1), None);
691 assert_eq!(parse_authority("x:notaport", 1), None);
692 }
693}