use crate::listener::Listener; 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> { 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> { println!("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.shutdown.recv() => { println!("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(); println!( "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 { println!("connection error {}", err); } }); } }