use std::io::{self, Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use zeroize::{Zeroize, Zeroizing};
#[derive(Clone, Copy, Default)]
pub(super) struct Counts {
pub bytes: u64,
pub frames: u64,
pub subscribes: u64,
}
pub(super) struct Proxy {
pub address: String,
counts: Arc<Mutex<Counts>>,
stopped: Arc<AtomicBool>,
worker: Option<std::thread::JoinHandle<io::Result<()>>>,
}
impl Proxy {
pub fn start(server: String, frame_bound: usize) -> io::Result<Self> {
let listener = TcpListener::bind("127.0.0.1:0")?;
let address = listener.local_addr()?.to_string();
let counts = Arc::new(Mutex::new(Counts::default()));
let stopped = Arc::new(AtomicBool::new(false));
let read_counts = Arc::clone(&counts);
let stop = Arc::clone(&stopped);
let worker = std::thread::spawn(move || {
accept(&listener, &server, &stop, &read_counts, frame_bound)
});
Ok(Self {
address,
counts,
stopped,
worker: Some(worker),
})
}
pub fn counts(&self) -> io::Result<Counts> {
self.counts
.lock()
.map(|counts| *counts)
.map_err(|_| io::Error::other("the wire count lock is poisoned"))
}
pub fn finish(mut self) -> io::Result<()> {
self.stop()
}
fn stop(&mut self) -> io::Result<()> {
let Some(worker) = self.worker.take() else {
return Ok(());
};
self.stopped.store(true, Ordering::SeqCst);
let interrupted = TcpStream::connect(&self.address);
let joined = worker
.join()
.map_err(|_| io::Error::other("the real-server relay panicked"))?;
match (interrupted, joined) {
(Ok(_), result) => result,
(Err(error), Ok(())) => Err(error),
(Err(first), Err(second)) => Err(io::Error::other(format!("{first}; {second}"))),
}
}
}
impl Drop for Proxy {
fn drop(&mut self) {
if let Err(error) = self.stop() {
eprintln!("the owned relay could not stop: {error}");
}
}
}
fn accept(
listener: &TcpListener,
server: &str,
stop: &AtomicBool,
counts: &Arc<Mutex<Counts>>,
bound: usize,
) -> io::Result<()> {
let mut sockets = Vec::new();
let mut readers = Vec::new();
let accepted: io::Result<()> = (|| {
loop {
let (client, _) = listener.accept()?;
if stop.load(Ordering::SeqCst) {
break;
}
let upstream = TcpStream::connect(server)?;
sockets.push(client.try_clone()?);
sockets.push(upstream.try_clone()?);
let reverse_in = upstream.try_clone()?;
let reverse_out = client.try_clone()?;
let counts = Arc::clone(counts);
readers.push(std::thread::spawn(move || {
relay(client, upstream, Some(&counts), bound)
}));
readers.push(std::thread::spawn(move || {
relay(reverse_in, reverse_out, None, bound)
}));
}
Ok(())
})();
let mut failures = Vec::new();
if let Err(error) = accepted {
failures.push(error.to_string());
}
for socket in sockets {
if let Err(error) = socket.shutdown(Shutdown::Both)
&& error.kind() != io::ErrorKind::NotConnected
{
failures.push(error.to_string());
}
}
for reader in readers {
match reader.join() {
Ok(Ok(())) => {}
Ok(Err(error)) => failures.push(error.to_string()),
Err(_) => failures.push("a relay reader panicked".to_owned()),
}
}
if failures.is_empty() {
Ok(())
} else {
Err(io::Error::other(failures.join("; ")))
}
}
fn relay(
mut input: TcpStream,
mut output: TcpStream,
counts: Option<&Mutex<Counts>>,
bound: usize,
) -> io::Result<()> {
let mut buffer = Zeroizing::new(Vec::new());
loop {
let mut chunk = [0; 8192];
let n = match input.read(&mut chunk) {
Ok(n) => n,
Err(error) if error.kind() == io::ErrorKind::ConnectionReset => break,
Err(error) => return Err(error),
};
if n == 0 {
break;
}
if let Some(counts) = counts {
observe(counts, &mut buffer, &chunk[..n], bound)?;
}
let written = output.write_all(&chunk[..n]);
chunk.zeroize();
match written {
Ok(()) => {}
Err(error)
if matches!(
error.kind(),
io::ErrorKind::BrokenPipe | io::ErrorKind::ConnectionReset
) =>
{
break;
}
Err(error) => return Err(error),
}
}
match output.shutdown(Shutdown::Write) {
Ok(()) => Ok(()),
Err(error) if error.kind() == io::ErrorKind::NotConnected => Ok(()),
Err(error) => Err(error),
}
}
fn observe(
counts: &Mutex<Counts>,
buffer: &mut Vec<u8>,
bytes: &[u8],
bound: usize,
) -> io::Result<()> {
let mut totals = counts
.lock()
.map_err(|_| io::Error::other("the wire count lock is poisoned"))?;
totals.bytes = totals
.bytes
.checked_add(u64::try_from(bytes.len()).map_err(io::Error::other)?)
.ok_or_else(|| io::Error::other("wire byte count overflowed"))?;
buffer.extend_from_slice(bytes);
while buffer.len() >= 10 {
let length = u32::from_be_bytes(buffer[6..10].try_into().map_err(io::Error::other)?);
let size = usize::try_from(length)
.map_err(io::Error::other)?
.checked_add(10)
.ok_or_else(|| io::Error::other("wire frame size overflowed"))?;
if size > bound {
return Err(io::Error::other(
"the relay frame exceeds its configured bound",
));
}
if buffer.len() < size {
break;
}
totals.frames = totals
.frames
.checked_add(1)
.ok_or_else(|| io::Error::other("wire frame count overflowed"))?;
if buffer[0] == 5 {
totals.subscribes = totals
.subscribes
.checked_add(1)
.ok_or_else(|| io::Error::other("subscribe count overflowed"))?;
}
buffer.drain(..size);
}
drop(totals);
Ok(())
}