90 lines
2.8 KiB
Rust
90 lines
2.8 KiB
Rust
use crate::listener::Listener;
|
|
use log::{error, info, trace, warn};
|
|
use rand::seq::SliceRandom;
|
|
use std::error::Error;
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
|
|
#[derive(Debug)]
|
|
struct TcpHandler {
|
|
stream: TcpStream,
|
|
target: String,
|
|
}
|
|
|
|
impl TcpHandler {
|
|
async fn run(&mut self) -> Result<(), Box<dyn Error>> {
|
|
let mut stream = TcpStream::connect(&self.target).await?;
|
|
|
|
tokio::io::copy_bidirectional(&mut self.stream, &mut stream).await?;
|
|
|
|
return Ok(());
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn start_tcp_listener(
|
|
mut listener_config: Listener,
|
|
) -> Result<(), Box<dyn Error>> {
|
|
info!("start listening on {}", &listener_config.source);
|
|
let listener = loop {
|
|
match TcpListener::bind(&listener_config.source).await {
|
|
Ok(listener) => break listener,
|
|
Err(e) => match e.kind() {
|
|
std::io::ErrorKind::AddrInUse => {
|
|
warn!("Address in use: {}", &listener_config.source);
|
|
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
|
|
}
|
|
std::io::ErrorKind::AddrNotAvailable => {
|
|
warn!("Address not available: {}", &listener_config.source);
|
|
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
|
|
}
|
|
_ => {
|
|
error!(
|
|
"Error binding to {}: {} ({})",
|
|
&listener_config.source,
|
|
e,
|
|
e.kind()
|
|
);
|
|
trace!("Error: {}", e);
|
|
return Err(Box::new(e));
|
|
}
|
|
},
|
|
}
|
|
};
|
|
|
|
info!("listening on {}", listener.local_addr()?);
|
|
|
|
loop {
|
|
let (next_socket, _) = tokio::select! {
|
|
res = listener.accept() => res?,
|
|
_ = listener_config.targets_changed.recv() => {
|
|
info!("Targets changed!");
|
|
continue;
|
|
}
|
|
_ = listener_config.shutdown.recv() => {
|
|
info!("Exiting listener!");
|
|
return Ok(());
|
|
}
|
|
};
|
|
|
|
let targets = listener_config.targets.read().await;
|
|
let mut rng = rand::thread_rng();
|
|
let selected_target = targets.choose(&mut rng).unwrap();
|
|
|
|
trace!(
|
|
"new connection from {} forwarding to {}",
|
|
next_socket.peer_addr()?,
|
|
&selected_target
|
|
);
|
|
let mut handler = TcpHandler {
|
|
stream: next_socket,
|
|
target: selected_target.clone(),
|
|
};
|
|
|
|
tokio::spawn(async move {
|
|
// Process the connection. If an error is encountered, log it.
|
|
if let Err(err) = handler.run().await {
|
|
warn!("connection error {}", err);
|
|
}
|
|
});
|
|
}
|
|
}
|