use std::{io, net::SocketAddr, sync::Arc};
use shadowsocks::{
ProxySocket,
relay::{
socks5::Address,
udprelay::{DatagramReceive, DatagramSend, options::UdpSocketControlData},
},
};
use super::flow::FlowStat;
pub struct MonProxySocket<S> {
socket: ProxySocket<S>,
flow_stat: Arc<FlowStat>,
}
impl<S> MonProxySocket<S> {
pub fn from_socket(socket: ProxySocket<S>, flow_stat: Arc<FlowStat>) -> Self {
Self { socket, flow_stat }
}
#[inline]
pub fn get_ref(&self) -> &ProxySocket<S> {
&self.socket
}
#[inline]
pub fn flow_stat(&self) -> &FlowStat {
&self.flow_stat
}
}
impl<S> MonProxySocket<S>
where
S: DatagramSend,
{
#[inline]
pub async fn send(&self, addr: &Address, payload: &[u8]) -> io::Result<()> {
let n = self.socket.send(addr, payload).await?;
self.flow_stat.incr_tx(n as u64);
Ok(())
}
#[inline]
pub async fn send_with_ctrl(
&self,
addr: &Address,
control: &UdpSocketControlData,
payload: &[u8],
) -> io::Result<()> {
let n = self.socket.send_with_ctrl(addr, control, payload).await?;
self.flow_stat.incr_tx(n as u64);
Ok(())
}
#[inline]
pub async fn send_to(&self, target: SocketAddr, addr: &Address, payload: &[u8]) -> io::Result<()> {
let n = self.socket.send_to(target, addr, payload).await?;
self.flow_stat.incr_tx(n as u64);
Ok(())
}
#[inline]
pub async fn send_to_with_ctrl(
&self,
target: SocketAddr,
addr: &Address,
control: &UdpSocketControlData,
payload: &[u8],
) -> io::Result<()> {
let n = self.socket.send_to_with_ctrl(target, addr, control, payload).await?;
self.flow_stat.incr_tx(n as u64);
Ok(())
}
}
impl<S> MonProxySocket<S>
where
S: DatagramReceive,
{
#[inline]
pub async fn recv(&self, recv_buf: &mut [u8]) -> io::Result<(usize, Address)> {
let (n, addr, recv_n) = self.socket.recv(recv_buf).await?;
self.flow_stat.incr_rx(recv_n as u64);
Ok((n, addr))
}
#[inline]
pub async fn recv_with_ctrl(
&self,
recv_buf: &mut [u8],
) -> io::Result<(usize, Address, Option<UdpSocketControlData>)> {
let (n, addr, recv_n, control) = self.socket.recv_with_ctrl(recv_buf).await?;
self.flow_stat.incr_rx(recv_n as u64);
Ok((n, addr, control))
}
#[inline]
pub async fn recv_from(&self, recv_buf: &mut [u8]) -> io::Result<(usize, SocketAddr, Address)> {
let (n, peer_addr, addr, recv_n) = self.socket.recv_from(recv_buf).await?;
self.flow_stat.incr_rx(recv_n as u64);
Ok((n, peer_addr, addr))
}
#[inline]
pub async fn recv_from_with_ctrl(
&self,
recv_buf: &mut [u8],
) -> io::Result<(usize, SocketAddr, Address, Option<UdpSocketControlData>)> {
let (n, peer_addr, addr, recv_n, control) = self.socket.recv_from_with_ctrl(recv_buf).await?;
self.flow_stat.incr_rx(recv_n as u64);
Ok((n, peer_addr, addr, control))
}
}