rd_std/builtin/
forward.rs1use 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}