rd_std/builtin/
forward.rs

1use std::net::SocketAddr;
2
3use rd_interface::{
4    async_trait, registry::ServerFactory, util::connect_tcp, Arc, Context, IServer, IntoAddress,
5    Net, Result, TcpListener, TcpStream,
6};
7use serde_derive::Deserialize;
8
9#[derive(Debug, Deserialize)]
10pub struct ForwardConfig {
11    bind: String,
12    target: String,
13}
14pub struct ForwardNet {
15    listen_net: Net,
16    net: Net,
17    cfg: Arc<ForwardConfig>,
18}
19
20impl ForwardNet {
21    fn new(listen_net: Net, net: Net, cfg: ForwardConfig) -> ForwardNet {
22        ForwardNet {
23            listen_net,
24            net,
25            cfg: Arc::new(cfg),
26        }
27    }
28}
29#[async_trait]
30impl IServer for ForwardNet {
31    async fn start(&self) -> Result<()> {
32        let listener = self
33            .listen_net
34            .tcp_bind(&mut Context::new(), self.cfg.bind.into_address()?)
35            .await?;
36        self.serve_listener(listener).await
37    }
38}
39
40impl ForwardNet {
41    async fn serve_connection(
42        cfg: Arc<ForwardConfig>,
43        socket: TcpStream,
44        net: Net,
45        addr: SocketAddr,
46    ) -> Result<()> {
47        let target = net
48            .tcp_connect(
49                &mut Context::from_socketaddr(addr),
50                cfg.target.into_address()?,
51            )
52            .await?;
53        connect_tcp(socket, target).await?;
54        Ok(())
55    }
56    pub async fn serve_listener(&self, listener: TcpListener) -> Result<()> {
57        loop {
58            let (socket, addr) = listener.accept().await?;
59            let cfg = self.cfg.clone();
60            let net = self.net.clone();
61            let _ = tokio::spawn(async move {
62                if let Err(e) = Self::serve_connection(cfg, socket, net, addr).await {
63                    log::error!("Error when serve_connection: {:?}", e);
64                }
65            });
66        }
67    }
68}
69
70impl ServerFactory for ForwardNet {
71    const NAME: &'static str = "forward";
72    type Config = ForwardConfig;
73    type Server = Self;
74
75    fn new(listen_net: Net, net: Net, config: Self::Config) -> Result<Self> {
76        Ok(ForwardNet::new(listen_net, net, config))
77    }
78}