use liminal_server::link_punch::PunchSocket;
use std::collections::HashSet;
use std::io;
use std::net::{Ipv4Addr, SocketAddr, UdpSocket};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::thread::JoinHandle;
use std::time::Duration;
fn poisoned<T>(_: std::sync::PoisonError<T>) -> io::Error {
io::Error::other("NAT state poisoned")
}
fn wrap(endpoint: SocketAddr, bytes: &[u8]) -> io::Result<Vec<u8>> {
let SocketAddr::V4(endpoint) = endpoint else {
return Err(io::Error::other("NAT requires IPv4"));
};
let mut packet = endpoint.ip().octets().to_vec();
packet.extend_from_slice(&endpoint.port().to_be_bytes());
packet.extend_from_slice(bytes);
Ok(packet)
}
fn unwrap(packet: &[u8]) -> io::Result<(SocketAddr, &[u8])> {
if packet.len() < 6 {
return Err(io::Error::other("short NAT envelope"));
}
Ok((
SocketAddr::from((
Ipv4Addr::new(packet[0], packet[1], packet[2], packet[3]),
u16::from_be_bytes([packet[4], packet[5]]),
)),
&packet[6..],
))
}
struct Inside {
socket: UdpSocket,
gateway: SocketAddr,
forwarded: Mutex<mpsc::Receiver<io::Result<()>>>,
}
impl PunchSocket for Inside {
fn send_to(&self, bytes: &[u8], target: SocketAddr) -> io::Result<usize> {
if target == self.socket.local_addr()? {
return self.socket.send_to(bytes, target);
}
let forwarded = self.forwarded.lock().map_err(poisoned)?;
self.socket.send_to(&wrap(target, bytes)?, self.gateway)?;
forwarded
.recv()
.map_err(|_| io::Error::other("NAT forwarder stopped"))??;
drop(forwarded);
Ok(bytes.len())
}
fn recv_from(&self, bytes: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
let mut packet = vec![0; bytes.len() + 6];
let (length, source) = self.socket.recv_from(&mut packet)?;
if length == 0 && source == self.socket.local_addr()? {
return Ok((0, source));
}
if source != self.gateway {
return Err(io::Error::other("NAT bypass refused"));
}
let (source, payload) = unwrap(&packet[..length])?;
bytes[..payload.len()].copy_from_slice(payload);
Ok((payload.len(), source))
}
fn local_addr(&self) -> io::Result<SocketAddr> {
self.socket.local_addr()
}
fn prepare(&self, bound: Option<Duration>) -> io::Result<()> {
self.socket.set_read_timeout(bound)
}
}
pub struct Nat {
inside: Arc<Inside>,
gateway: UdpSocket,
outside: UdpSocket,
outside_address: SocketAddr,
stop: Arc<AtomicBool>,
filtered: Arc<AtomicUsize>,
kinds: Arc<Mutex<Vec<u8>>>,
rejected: mpsc::Receiver<()>,
workers: Vec<JoinHandle<io::Result<()>>>,
}
impl Nat {
pub fn new() -> io::Result<Self> {
let socket = UdpSocket::bind("127.0.0.1:0")?;
let gateway = UdpSocket::bind("127.0.0.1:0")?;
let outside = UdpSocket::bind("127.0.0.1:0")?;
let outside_address = outside.local_addr().map_err(|error| {
io::Error::other(format!("NAT outside address unavailable: {error}"))
})?;
let inside_addr = socket.local_addr()?;
let stop = Arc::new(AtomicBool::new(false));
let filtered = Arc::new(AtomicUsize::new(0));
let allowed = Arc::new(Mutex::new(HashSet::new()));
let kinds = Arc::new(Mutex::new(Vec::new()));
let (forwarded, answers) = mpsc::channel();
let read = gateway.try_clone()?;
let send = outside.try_clone()?;
let outgoing_stop = stop.clone();
let outgoing_allowed = allowed.clone();
let outgoing_kinds = kinds.clone();
let outgoing = std::thread::spawn(move || {
let mut bytes = [0; 1300];
loop {
let (length, source) = read.recv_from(&mut bytes)?;
if outgoing_stop.load(Ordering::Acquire) {
return Ok(());
}
if source != inside_addr {
return Err(io::Error::other("unknown NAT inside"));
}
let (target, payload) = unwrap(&bytes[..length])?;
outgoing_allowed.lock().map_err(poisoned)?.insert(target);
if payload.len() > 13 && matches!(payload[0], 1 | 2) {
outgoing_kinds.lock().map_err(poisoned)?.push(payload[0]);
}
let result = send.send_to(payload, target).and_then(|sent| {
if sent == payload.len() {
Ok(())
} else {
Err(io::Error::other("partial NAT send"))
}
});
forwarded
.send(result)
.map_err(|_| io::Error::other("NAT sender ended"))?;
}
});
let read = outside.try_clone()?;
let send = gateway.try_clone()?;
let incoming_stop = stop.clone();
let incoming_filtered = filtered.clone();
let (rejection, rejected) = mpsc::channel();
let incoming = std::thread::spawn(move || {
let mut bytes = [0; 1300];
loop {
let (length, source) = read.recv_from(&mut bytes)?;
if incoming_stop.load(Ordering::Acquire) {
return Ok(());
}
if allowed.lock().map_err(poisoned)?.contains(&source) {
send.send_to(&wrap(source, &bytes[..length])?, inside_addr)?;
} else {
incoming_filtered.fetch_add(1, Ordering::Relaxed);
rejection
.send(())
.map_err(|_| io::Error::other("NAT observer stopped"))?;
}
}
});
Ok(Self {
inside: Arc::new(Inside {
socket,
gateway: gateway.local_addr()?,
forwarded: Mutex::new(answers),
}),
gateway,
outside,
outside_address,
stop,
filtered,
kinds,
rejected,
workers: vec![outgoing, incoming],
})
}
pub fn verify_filter(&self) -> io::Result<()> {
let stranger = UdpSocket::bind("127.0.0.1:0")?;
stranger.send_to(b"unmapped", self.outside.local_addr()?)?;
self.rejected
.recv()
.map_err(|_| io::Error::other("NAT filter stopped"))
}
pub fn socket(&self) -> Arc<dyn PunchSocket> {
self.inside.clone()
}
pub const fn outside(&self) -> SocketAddr {
self.outside_address
}
pub fn filtered(&self) -> usize {
self.filtered.load(Ordering::Relaxed)
}
pub fn handshake_kinds(&self) -> io::Result<Vec<u8>> {
Ok(self.kinds.lock().map_err(poisoned)?.clone())
}
}
impl Drop for Nat {
fn drop(&mut self) {
self.stop.store(true, Ordering::Release);
let result = (|| -> io::Result<()> {
self.gateway.send_to(&[], self.gateway.local_addr()?)?;
self.outside.send_to(&[], self.outside.local_addr()?)?;
for worker in self.workers.drain(..) {
worker
.join()
.map_err(|_| io::Error::other("NAT worker panicked"))??;
}
Ok(())
})();
if let Err(error) = result {
eprintln!("NAT shutdown failed: {error}");
}
}
}