rd_std/
redir.rs

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