port_mapping/
tcp_proxy.rs

1use std::{fmt::Display, sync::Arc};
2use tokio::net::{TcpListener, TcpStream};
3
4#[derive(Debug)]
5pub struct TcpProxy {
6    pub listen: String,
7    pub upstream: String,
8}
9
10impl Display for TcpProxy {
11    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
12        write!(f, "{}->{}", self.listen, self.upstream)
13    }
14}
15
16impl TcpProxy {
17    pub fn new(listen: String, upstream: String) -> Self {
18        Self { listen, upstream }
19    }
20
21    pub async fn run(self: Arc<Self>) -> std::io::Result<()> {
22        let listener = TcpListener::bind(&self.listen).await?;
23        println!("[info][tcp][{self}] Listening");
24        loop {
25            let mut downstream = match listener.accept().await {
26                Ok((downstream, _)) => downstream,
27                Err(e) => {
28                    eprintln!("[warning][tcp][{self}] Failed to accept connection: {e}");
29                    continue;
30                }
31            };
32            let self_clone = self.clone();
33            tokio::spawn(async move {
34                let mut upstream = match TcpStream::connect(&self_clone.upstream).await {
35                    Ok(stream) => stream,
36                    Err(e) => {
37                        eprintln!("[warning][tcp][{self_clone}] Failed to connect: {e}");
38                        return;
39                    }
40                };
41                match tokio::io::copy_bidirectional(&mut downstream, &mut upstream).await {
42                    Ok((a, b)) => println!(
43                        "[info][tcp][{self_clone}] Connection closed: {} Send {a}B to {} and receive {b}B",
44                        self_clone.listen, self_clone.upstream
45                    ),
46                    Err(e) => eprintln!("[warning][tcp][{self_clone}] Connection error: {e}"),
47                };
48            });
49        }
50    }
51}