64 lines
1.8 KiB
Rust
64 lines
1.8 KiB
Rust
use crate::listener::Listener;
|
|
use log::{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 = TcpListener::bind(&listener_config.source).await?;
|
|
|
|
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);
|
|
}
|
|
});
|
|
}
|
|
}
|