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        match timeout {
68            Ok(recv) => recv,
69            Err(_) => None,
70        }
71    }
72
73    pub async fn copy_bidirectional_udp(
74        &mut self,
75        listener: &Arc<UdpSocket>,
76        sock: UdpSocket,
77    ) -> Result<()> {
78        let mut recv_buf = [0; 65536];
79        loop {
80            tokio::select! {
81                res = self.recv() => {
82                    match res {
83                        Some(res) => sock.send(&res).await?,
84                        None => break
85                    };
86                }
87                res = sock.recv(&mut recv_buf) => {
88                    let n = res?;
89                    listener.send_to(&recv_buf[..n], self.addr).await?;
90                }
91            }
92        }
93
94        self.remove().await;
95        Ok(())
96    }
97
98    pub async fn copy_bidirectional_tcp(
99        &mut self,
100        listener: &Arc<UdpSocket>,
101        mut sock: TcpStream,
102    ) -> Result<()> {
103        let mut recv_buf = [0; 65536];
104        loop {
105            tokio::select! {
106                res = self.recv() => {
107                    match res {
108                        Some(res) => {
109                            sock.write_u16(res.len() as u16).await?;
110                            sock.write_all(&res).await?;
111                        },
112                        None => break
113                    };
114                }
115                n = sock.read_u16() => {
116                    let n = n?;
117                    sock.read_exact(&mut recv_buf[..n as usize]).await?;
118                    listener.send_to(&recv_buf[..n as usize], self.addr).await?;
119                }
120            }
121        }
122
123        self.remove().await;
124        Ok(())
125    }
126
127    async fn remove(&self) {
128        self.tunnel_map.write().await.remove(&self.addr);
129    }
130}
fkm_proxy/utils/udp.rs

fkm_proxy/utils/
udp.rs
