shadowsocks_service/net/
mon_socket.rs

1//! UDP socket with flow statistic monitored
2
3use std::{io, net::SocketAddr, sync::Arc};
4
5use shadowsocks::{
6    ProxySocket,
7    relay::{
8        socks5::Address,
9        udprelay::{DatagramReceive, DatagramSend, options::UdpSocketControlData},
10    },
11};
12
13use super::flow::FlowStat;
14
15/// Monitored `ProxySocket`
16pub struct MonProxySocket<S> {
17    socket: ProxySocket<S>,
18    flow_stat: Arc<FlowStat>,
19}
20
21impl<S> MonProxySocket<S> {
22    /// Create a new socket with flow monitor
23    pub fn from_socket(socket: ProxySocket<S>, flow_stat: Arc<FlowStat>) -> Self {
24        Self { socket, flow_stat }
25    }
26
27    /// Get the underlying `ProxySocket<S>` immutable reference
28    #[inline]
29    pub fn get_ref(&self) -> &ProxySocket<S> {
30        &self.socket
31    }
32
33    /// Get the flow statistic data
34    #[inline]
35    pub fn flow_stat(&self) -> &FlowStat {
36        &self.flow_stat
37    }
38}
39
40impl<S> MonProxySocket<S>
41where
42    S: DatagramSend,
43{
44    /// Send a UDP packet to addr through proxy
45    #[inline]
46    pub async fn send(&self, addr: &Address, payload: &[u8]) -> io::Result<()> {
47        let n = self.socket.send(addr, payload).await?;
48        self.flow_stat.incr_tx(n as u64);
49
50        Ok(())
51    }
52
53    /// Send a UDP packet to addr through proxy
54    #[inline]
55    pub async fn send_with_ctrl(
56        &self,
57        addr: &Address,
58        control: &UdpSocketControlData,
59        payload: &[u8],
60    ) -> io::Result<()> {
61        let n = self.socket.send_with_ctrl(addr, control, payload).await?;
62        self.flow_stat.incr_tx(n as u64);
63
64        Ok(())
65    }
66
67    /// Send a UDP packet to target from proxy
68    #[inline]
69    pub async fn send_to(&self, target: SocketAddr, addr: &Address, payload: &[u8]) -> io::Result<()> {
70        let n = self.socket.send_to(target, addr, payload).await?;
71        self.flow_stat.incr_tx(n as u64);
72
73        Ok(())
74    }
75
76    /// Send a UDP packet to target from proxy
77    #[inline]
78    pub async fn send_to_with_ctrl(
79        &self,
80        target: SocketAddr,
81        addr: &Address,
82        control: &UdpSocketControlData,
83        payload: &[u8],
84    ) -> io::Result<()> {
85        let n = self.socket.send_to_with_ctrl(target, addr, control, payload).await?;
86        self.flow_stat.incr_tx(n as u64);
87
88        Ok(())
89    }
90}
91
92impl<S> MonProxySocket<S>
93where
94    S: DatagramReceive,
95{
96    /// Receive packet from Shadowsocks' UDP server
97    ///
98    /// This function will use `recv_buf` to store intermediate data, so it has to be big enough to store the whole shadowsocks' packet
99    ///
100    /// It is recommended to allocate a buffer to have at least 65536 bytes.
101    #[inline]
102    pub async fn recv(&self, recv_buf: &mut [u8]) -> io::Result<(usize, Address)> {
103        let (n, addr, recv_n) = self.socket.recv(recv_buf).await?;
104        self.flow_stat.incr_rx(recv_n as u64);
105
106        Ok((n, addr))
107    }
108
109    /// Receive packet from Shadowsocks' UDP server
110    ///
111    /// This function will use `recv_buf` to store intermediate data, so it has to be big enough to store the whole shadowsocks' packet
112    ///
113    /// It is recommended to allocate a buffer to have at least 65536 bytes.
114    #[inline]
115    pub async fn recv_with_ctrl(
116        &self,
117        recv_buf: &mut [u8],
118    ) -> io::Result<(usize, Address, Option<UdpSocketControlData>)> {
119        let (n, addr, recv_n, control) = self.socket.recv_with_ctrl(recv_buf).await?;
120        self.flow_stat.incr_rx(recv_n as u64);
121
122        Ok((n, addr, control))
123    }
124
125    /// Receive packet from Shadowsocks' UDP server
126    ///
127    /// This function will use `recv_buf` to store intermediate data, so it has to be big enough to store the whole shadowsocks' packet
128    ///
129    /// It is recommended to allocate a buffer to have at least 65536 bytes.
130    #[inline]
131    pub async fn recv_from(&self, recv_buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Address)> {
132        let (n, peer_addr, addr, recv_n) = self.socket.recv_from(recv_buf).await?;
133        self.flow_stat.incr_rx(recv_n as u64);
134
135        Ok((n, peer_addr, addr))
136    }
137
138    /// Receive packet from Shadowsocks' UDP server
139    ///
140    /// This function will use `recv_buf` to store intermediate data, so it has to be big enough to store the whole shadowsocks' packet
141    ///
142    /// It is recommended to allocate a buffer to have at least 65536 bytes.
143    #[inline]
144    pub async fn recv_from_with_ctrl(
145        &self,
146        recv_buf: &mut [u8],
147    ) -> io::Result<(usize, SocketAddr, Address, Option<UdpSocketControlData>)> {
148        let (n, peer_addr, addr, recv_n, control) = self.socket.recv_from_with_ctrl(recv_buf).await?;
149        self.flow_stat.incr_rx(recv_n as u64);
150
151        Ok((n, peer_addr, addr, control))
152    }
153}
shadowsocks_service/net/mon_socket.rs

shadowsocks_service/net/
mon_socket.rs
