fkm_proxy/utils/
udp.rs

1use anyhow::Result;
2use std::{collections::HashMap, net::SocketAddr, sync::Arc, time::Duration};
3use tokio::{
4    io::{AsyncReadExt, AsyncWriteExt},
5    net::{TcpStream, UdpSocket},
6    sync::{
7        mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
8        RwLock,
9    },
10};
11
12type TunnelMapTest = Arc<RwLock<HashMap<SocketAddr, UnboundedSender<Vec<u8>>>>>;
13pub async fn server_accept(
14    listener: &Arc<UdpSocket>,
15    recv_buf: &mut [u8],
16    tunnel_map: &TunnelMapTest,
17) -> Result<(UdpClient, SocketAddr)> {
18    loop {
19        let (n, addr) = listener.recv_from(recv_buf).await?;
20        let rx = {
21            let mut tunnel_map_rw = tunnel_map.write().await;
22            if let Some(sock) = tunnel_map_rw.get_mut(&addr) {
23                if sock.is_closed() {
24                    let (tx, rx) = unbounded_channel();
25                    *sock = tx;
26                    sock.send(recv_buf[..n].to_vec())?;
27
28                    return Ok((UdpClient::new(tunnel_map, rx, addr), addr));
29                }
30
31                sock.send(recv_buf[..n].to_vec())?;
32                continue;
33            } else {
34                let (tx, rx) = unbounded_channel();
35                tx.send(recv_buf[..n].to_vec())?;
36                tunnel_map_rw.insert(addr, tx);
37
38                rx
39            }
40        };
41
42        return Ok((UdpClient::new(tunnel_map, rx, addr), addr));
43    }
44}
45
46pub struct UdpClient {
47    tunnel_map: TunnelMapTest,
48    rx: UnboundedReceiver<Vec<u8>>,
49    addr: SocketAddr,
50}
51
52impl UdpClient {
53    pub fn new(
54        tunnel_map: &TunnelMapTest,
55        rx: UnboundedReceiver<Vec<u8>>,
56        addr: SocketAddr,
57    ) -> Self {
58        Self {
59            tunnel_map: tunnel_map.clone(),
60            rx,
61            addr,
62        }
63    }
64
65    async fn recv(&mut self) -> Option<Vec<u8>> {
66        let timeout = tokio::time::timeout(Duration::from_secs(45), self.rx.recv()).await;
67        timeout.unwrap_or_default()
68    }
69
70    pub async fn copy_bidirectional_udp(
71        &mut self,
72        listener: &Arc<UdpSocket>,
73        sock: UdpSocket,
74    ) -> Result<()> {
75        let mut recv_buf = [0; 65536];
76        loop {
77            tokio::select! {
78                res = self.recv() => {
79                    match res {
80                        Some(res) => sock.send(&res).await?,
81                        None => break
82                    };
83                }
84                res = sock.recv(&mut recv_buf) => {
85                    let n = res?;
86                    listener.send_to(&recv_buf[..n], self.addr).await?;
87                }
88            }
89        }
90
91        self.remove().await;
92        Ok(())
93    }
94
95    pub async fn copy_bidirectional_tcp(
96        &mut self,
97        listener: &Arc<UdpSocket>,
98        mut sock: TcpStream,
99    ) -> Result<()> {
100        let mut recv_buf = [0; 65536];
101        loop {
102            tokio::select! {
103                res = self.recv() => {
104                    match res {
105                        Some(res) => {
106                            sock.write_u16(res.len() as u16).await?;
107                            sock.write_all(&res).await?;
108                        },
109                        None => break
110                    };
111                }
112                n = sock.read_u16() => {
113                    let n = n?;
114                    sock.read_exact(&mut recv_buf[..n as usize]).await?;
115                    listener.send_to(&recv_buf[..n as usize], self.addr).await?;
116                }
117            }
118        }
119
120        self.remove().await;
121        Ok(())
122    }
123
124    async fn remove(&self) {
125        self.tunnel_map.write().await.remove(&self.addr);
126    }
127}
fkm_proxy/utils/udp.rs

fkm_proxy/utils/
udp.rs
