#![forbid(unsafe_code)]
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read, Write};
use std::net::{IpAddr, SocketAddr, TcpListener, TcpStream, ToSocketAddrs};
use std::path::PathBuf;
use std::process::Command;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
const MAX_HEAD: usize = 16 * 1024;
pub const MAX_CONNECTIONS: usize = 256;
const CHUNK: usize = 16 * 1024;
#[derive(Debug, Clone, Copy)]
pub struct Request<'a> {
pub host: &'a str,
pub port: u16,
pub peer: SocketAddr,
pub local: SocketAddr,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Verdict {
pub allow: bool,
pub allow_private: bool,
pub reason: String,
}
impl Verdict {
pub fn allow() -> Verdict {
Verdict {
allow: true,
allow_private: false,
reason: String::new(),
}
}
pub fn deny(reason: impl Into<String>) -> Verdict {
Verdict {
allow: false,
allow_private: false,
reason: reason.into(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Outcome {
Refused,
Unresolved,
Private,
Unreachable,
Cut,
Tunneled,
}
pub trait Gate: Send + Sync + 'static {
fn decide(&self, request: &Request<'_>) -> Verdict;
fn finished(&self, _request: &Request<'_>, _outcome: Outcome) {}
fn carried(&self, _request: &Request<'_>, _bytes: u64) -> bool {
true
}
fn relay(&self, _request: &Request<'_>, tunnel: Tunnel) -> Option<Tunnel> {
Some(tunnel)
}
fn relay_plain(
&self,
_request: &Request<'_>,
head: String,
tunnel: Tunnel,
) -> Option<(String, Tunnel)> {
Some((head, tunnel))
}
}
pub struct Tunnel {
pub client: TcpStream,
pub upstream: TcpStream,
pub early: Vec<u8>,
}
pub fn bind(addr: &str) -> std::io::Result<TcpListener> {
let l = TcpListener::bind(addr)?;
if !l.local_addr()?.ip().is_loopback() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"the proxy only listens on loopback",
));
}
Ok(l)
}
pub fn serve<G: Gate>(listener: TcpListener, gate: Arc<G>) {
accept(listener, gate, &Arc::new(Shared::default()));
}
pub struct Server {
addr: SocketAddr,
shared: Arc<Shared>,
thread: Option<JoinHandle<()>>,
}
#[derive(Default)]
struct Shared {
stopping: AtomicBool,
open: Mutex<HashMap<u64, Vec<TcpStream>>>,
next: AtomicU64,
}
impl Shared {
fn track(&self, id: u64, stream: &TcpStream) {
let mut open = self.open.lock().unwrap();
if self.stopping.load(Ordering::SeqCst) {
let _ = stream.shutdown(std::net::Shutdown::Both);
} else if let Ok(clone) = stream.try_clone() {
open.entry(id).or_default().push(clone);
}
}
}
pub fn spawn<G: Gate>(listener: TcpListener, gate: Arc<G>) -> std::io::Result<Server> {
let addr = listener.local_addr()?;
let shared = Arc::new(Shared::default());
let served = Arc::clone(&shared);
let thread = std::thread::spawn(move || accept(listener, gate, &served));
Ok(Server {
addr,
shared,
thread: Some(thread),
})
}
impl Server {
pub fn local_addr(&self) -> SocketAddr {
self.addr
}
pub fn stop(mut self) {
self.shut();
}
fn shut(&mut self) {
let Some(thread) = self.thread.take() else {
return;
};
{
let _open = self.shared.open.lock().unwrap();
self.shared.stopping.store(true, Ordering::SeqCst);
}
let _ = TcpStream::connect_timeout(&self.addr, Duration::from_secs(1));
let _ = thread.join();
for stream in self.shared.open.lock().unwrap().values().flatten() {
let _ = stream.shutdown(std::net::Shutdown::Both);
}
let deadline = Instant::now() + Duration::from_secs(3);
while !self.shared.open.lock().unwrap().is_empty() && Instant::now() < deadline {
std::thread::sleep(Duration::from_millis(5));
}
}
}
impl Drop for Server {
fn drop(&mut self) {
self.shut();
}
}
fn accept<G: Gate>(listener: TcpListener, gate: Arc<G>, shared: &Arc<Shared>) {
for conn in listener.incoming() {
if shared.stopping.load(Ordering::SeqCst) {
return;
}
let Ok(stream) = conn else { continue };
if shared.open.lock().unwrap().len() >= MAX_CONNECTIONS {
respond(&stream, 503, "too many connections through the proxy");
continue;
}
let id = shared.next.fetch_add(1, Ordering::Relaxed);
shared.track(id, &stream);
let (gate, shared) = (gate.clone(), Arc::clone(shared));
std::thread::spawn(move || {
if !shared.stopping.load(Ordering::SeqCst) {
let _ = handle(stream, &*gate, &|upstream| shared.track(id, upstream));
}
shared.open.lock().unwrap().remove(&id);
});
}
}
pub fn is_private(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(v) => private_v4(v),
IpAddr::V6(v) => {
if let Some(v4) = v.to_ipv4_mapped() {
return private_v4(v4);
}
let s = v.segments();
if s[..6] == [0x64, 0xff9b, 0, 0, 0, 0] {
let [a, b] = s[6].to_be_bytes();
let [c, d] = s[7].to_be_bytes();
return private_v4(std::net::Ipv4Addr::new(a, b, c, d));
}
v.is_unspecified()
|| v.is_loopback()
|| v.is_multicast()
|| (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] }
}
}
pub fn is_metadata(ip: IpAddr) -> bool {
let v4 = |v: std::net::Ipv4Addr| {
matches!(
v.octets(),
[169, 254, 169, 254] | [169, 254, 170, 2] | [169, 254, 170, 23] | [100, 100, 100, 200]
)
};
match ip {
IpAddr::V4(v) => v4(v),
IpAddr::V6(v) => {
v.to_ipv4_mapped().is_some_and(v4)
|| v.segments() == [0xfd00, 0x0ec2, 0, 0, 0, 0, 0, 0x0254]
}
}
}
fn private_v4(ip: std::net::Ipv4Addr) -> bool {
let [a, b, c, _] = ip.octets();
ip.is_unspecified()
|| ip.is_loopback()
|| ip.is_private()
|| ip.is_link_local()
|| ip.is_broadcast()
|| ip.is_documentation()
|| ip.is_multicast()
|| a == 0
|| (a == 100 && (64..128).contains(&b)) || (a == 192 && b == 0 && c == 0) || (a == 198 && (18..20).contains(&b)) || a >= 240 }
fn parse_authority(a: &str, default_port: u16) -> Option<(String, u16)> {
let (host, port) = if let Some(rest) = a.strip_prefix('[') {
let (h, tail) = rest.split_once(']')?;
let port = tail
.strip_prefix(':')
.map(|p| p.parse().ok())
.unwrap_or(Some(default_port))?;
(h.to_string(), port)
} else {
match a.rsplit_once(':') {
Some((h, p)) => (h.to_string(), p.parse().ok()?),
None => (a.to_string(), default_port),
}
};
let ok = !host.is_empty()
&& host.len() <= 253
&& host
.chars()
.all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | ':'));
ok.then_some((host.to_ascii_lowercase(), port))
}
pub fn peer_cwd(peer: SocketAddr, local: SocketAddr) -> Option<PathBuf> {
let me = std::process::id().to_string();
let out = Command::new("lsof")
.args([
"-nP",
&format!("-iTCP@{}:{}", peer.ip(), peer.port()),
"-sTCP:ESTABLISHED",
"-Fpn",
])
.output()
.ok()?;
let text = String::from_utf8_lossy(&out.stdout);
let want = format!(
"{}:{}->{}:{}",
peer.ip(),
peer.port(),
local.ip(),
local.port()
);
let mut pid = None;
let mut cur = None;
for line in text.lines() {
if let Some(p) = line.strip_prefix('p') {
cur = Some(p.to_string());
} else if let Some(n) = line.strip_prefix('n')
&& n == want
&& cur.as_deref() != Some(me.as_str())
{
pid = cur.clone();
}
}
let pid = pid?;
let out = Command::new("lsof")
.args(["-a", "-p", &pid, "-d", "cwd", "-Fn"])
.output()
.ok()?;
String::from_utf8_lossy(&out.stdout)
.lines()
.find_map(|l| l.strip_prefix('n').map(PathBuf::from))
}
fn connect_any(addrs: &[SocketAddr], within: Duration) -> Option<TcpStream> {
let deadline = Instant::now() + within;
addrs.iter().find_map(|a| {
TcpStream::connect_timeout(a, deadline.saturating_duration_since(Instant::now())).ok()
})
}
fn respond(mut s: &TcpStream, code: u16, reason: &str) {
let body = format!("proxy: {reason}\n");
let _ = write!(
s,
"HTTP/1.1 {code} {reason}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
);
}
fn handle(client: TcpStream, gate: &dyn Gate, track: &dyn Fn(&TcpStream)) -> std::io::Result<()> {
client.set_read_timeout(Some(Duration::from_secs(30)))?;
let peer = client.peer_addr()?;
let local = client.local_addr()?;
let mut reader = BufReader::new(client.try_clone()?);
let mut head = Vec::new();
loop {
let mut line = Vec::new();
let n = reader
.by_ref()
.take((MAX_HEAD - head.len()) as u64)
.read_until(b'\n', &mut line)?;
if n == 0 || head.len() + n >= MAX_HEAD {
respond(&client, 400, "bad request");
return Ok(());
}
head.extend_from_slice(&line);
if line == b"\r\n" || line == b"\n" {
break;
}
}
let head_text = String::from_utf8_lossy(&head).into_owned();
let first = head_text.lines().next().unwrap_or("");
let mut parts = first.split_whitespace();
let (method, target) = (parts.next().unwrap_or(""), parts.next().unwrap_or(""));
let (authority, rest_of_request) = if method.eq_ignore_ascii_case("CONNECT") {
(parse_authority(target, 443), None)
} else if let Some(url) = target.strip_prefix("http://") {
let (auth, path) = url.split_at(url.find('/').unwrap_or(url.len()));
let path = if path.is_empty() { "/" } else { path };
let rewritten = head_text.replacen(target, path, 1);
(parse_authority(auth, 80), Some(rewritten))
} else {
respond(&client, 400, "only CONNECT and absolute http:// requests");
return Ok(());
};
let Some((host, port)) = authority else {
respond(&client, 400, "bad target");
return Ok(());
};
let request = Request {
host: &host,
port,
peer,
local,
};
let verdict = gate.decide(&request);
if !verdict.allow {
respond(&client, 403, &verdict.reason);
gate.finished(&request, Outcome::Refused);
return Ok(());
}
let addrs: Vec<SocketAddr> = match (host.as_str(), port).to_socket_addrs() {
Ok(a) => a.collect(),
Err(_) => {
respond(&client, 502, "cannot resolve host");
gate.finished(&request, Outcome::Unresolved);
return Ok(());
}
};
if !addrs.is_empty() && addrs.iter().all(|a| is_metadata(a.ip())) {
respond(&client, 403, "cloud metadata addresses are never reachable");
gate.finished(&request, Outcome::Private);
return Ok(());
}
let allowed: Vec<SocketAddr> = addrs
.into_iter()
.filter(|a| !is_metadata(a.ip()))
.filter(|a| verdict.allow_private || !is_private(a.ip()))
.collect();
if allowed.is_empty() {
respond(
&client,
403,
"private and loopback addresses are not reachable",
);
gate.finished(&request, Outcome::Private);
return Ok(());
}
let Some(upstream) = connect_any(&allowed, Duration::from_secs(10)) else {
respond(&client, 502, "cannot connect");
gate.finished(&request, Outcome::Unreachable);
return Ok(());
};
track(&upstream);
client.set_read_timeout(None)?;
let buffered = reader.buffer().to_vec();
drop(reader);
let (client, upstream, buffered) = match rest_of_request {
None => {
let mut c = &client;
c.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n")?;
let tunnel = Tunnel {
client,
upstream,
early: buffered,
};
match gate.relay(&request, tunnel) {
Some(t) => (t.client, t.upstream, t.early),
None => {
gate.finished(&request, Outcome::Tunneled);
return Ok(());
}
}
}
Some(h) => {
let tunnel = Tunnel {
client,
upstream,
early: buffered,
};
match gate.relay_plain(&request, h, tunnel) {
Some((h, t)) => {
(&t.upstream).write_all(h.as_bytes())?;
(t.client, t.upstream, t.early)
}
None => {
gate.finished(&request, Outcome::Tunneled);
return Ok(());
}
}
}
};
let mut up_w = upstream.try_clone()?;
if !buffered.is_empty() {
if !gate.carried(&request, buffered.len() as u64) {
gate.finished(&request, Outcome::Cut);
return Ok(());
}
up_w.write_all(&buffered)?;
}
let cut = AtomicBool::new(false);
let (up_r, client_w) = (upstream.try_clone()?, client.try_clone()?);
std::thread::scope(|scope| {
scope.spawn(|| {
if !carry(&up_r, &client_w, gate, &request) {
cut.store(true, Ordering::SeqCst);
let _ = up_r.shutdown(std::net::Shutdown::Both);
let _ = client_w.shutdown(std::net::Shutdown::Both);
}
let _ = client_w.shutdown(std::net::Shutdown::Write);
});
if !carry(&client, &up_w, gate, &request) {
cut.store(true, Ordering::SeqCst);
let _ = client.shutdown(std::net::Shutdown::Both);
let _ = up_w.shutdown(std::net::Shutdown::Both);
}
let _ = up_w.shutdown(std::net::Shutdown::Write);
});
let how = if cut.load(Ordering::SeqCst) {
Outcome::Cut
} else {
Outcome::Tunneled
};
gate.finished(&request, how);
Ok(())
}
fn carry(mut from: &TcpStream, mut to: &TcpStream, gate: &dyn Gate, request: &Request<'_>) -> bool {
let mut buf = vec![0u8; CHUNK];
loop {
let n = match from.read(&mut buf) {
Ok(0) | Err(_) => return true,
Ok(n) => n,
};
if !gate.carried(request, n as u64) {
return false;
}
if to.write_all(&buf[..n]).is_err() {
return true;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_connection_goes_to_the_first_address_that_takes_it() {
let closed = TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap();
let open = TcpListener::bind("127.0.0.1:0").unwrap();
let at = open.local_addr().unwrap();
let budget = Duration::from_secs(5);
let s = connect_any(&[closed, at], budget).expect("the second address");
assert_eq!(s.peer_addr().unwrap(), at);
assert!(connect_any(&[closed], budget).is_none());
assert!(connect_any(&[], budget).is_none());
assert!(
connect_any(&[at], Duration::ZERO).is_none(),
"no budget, no attempt"
);
}
#[test]
fn metadata_addresses() {
for ip in [
"169.254.169.254",
"169.254.170.2",
"100.100.100.200",
"fd00:ec2::254",
"::ffff:169.254.169.254",
] {
assert!(is_metadata(ip.parse().unwrap()), "{ip}");
}
for ip in ["169.254.169.253", "10.0.0.1", "fd00:ec2::253", "8.8.8.8"] {
assert!(!is_metadata(ip.parse().unwrap()), "{ip}");
}
}
#[test]
fn private_ranges() {
for ip in [
"127.0.0.1",
"10.1.2.3",
"192.168.0.1",
"172.16.5.4",
"169.254.1.1",
"100.64.0.1",
"0.0.0.0",
"::1",
"fd00::1",
"fe80::1",
"::ffff:127.0.0.1",
"192.0.0.1",
"198.18.0.1",
"192.0.2.1",
"224.0.0.1",
"240.0.0.1",
"255.255.255.255",
"fec0::1",
"ff02::1",
"2001:db8::1",
"64:ff9b::7f00:1",
"::127.0.0.1",
] {
assert!(is_private(ip.parse().unwrap()), "{ip}");
}
for ip in [
"93.184.216.34",
"1.1.1.1",
"2606:4700::1111",
"::ffff:8.8.8.8",
"64:ff9b::808:808",
] {
assert!(!is_private(ip.parse().unwrap()), "{ip}");
}
}
#[test]
fn authorities() {
assert_eq!(
parse_authority("Example.com:443", 1),
Some(("example.com".into(), 443))
);
assert_eq!(
parse_authority("example.com", 80),
Some(("example.com".into(), 80))
);
assert_eq!(parse_authority("[::1]:8080", 1), Some(("::1".into(), 8080)));
assert_eq!(parse_authority("a b:1", 1), None);
assert_eq!(parse_authority("x:notaport", 1), None);
}
}