wind_socks/
ext.rs

1use std::{
2	io,
3	net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
4};
5
6use fast_socks5::{
7	server::{ErrorContext as _, Socks5ServerProtocol, SocksServerError, states, wait_on_tcp},
8	util::target_addr::TargetAddr,
9};
10use futures_util::TryFutureExt;
11use socket2::{Domain, Socket, Type};
12use tokio::{
13	io::{AsyncRead, AsyncWrite},
14	try_join,
15};
16use tracing::debug;
17use wind_core::{error, warn};
18
19use crate::Error;
20
21macro_rules! try_notify {
22    ($proto:expr, $e:expr) => {
23        match $e {
24            Ok(res) => res,
25            Err(err) => {
26                if let Err(rep_err) = $proto.reply_error(&err.to_reply_error()).await {
27                    error!(
28                        name: "socks",
29                        "extra error while reporting an error to the client: {}",
30                        rep_err
31                    );
32                }
33                return Err(err.into());
34            }
35        }
36    };
37}
38
39fn udp_bind_random_port(addr: Option<IpAddr>) -> io::Result<Socket> {
40	if let Some(addr) = addr {
41		let sock_addr = SocketAddr::new(addr, 0);
42		let socket = Socket::new(Domain::for_address(sock_addr), Type::DGRAM, None)?;
43		socket.bind(&sock_addr.into())?;
44		Ok(socket)
45	} else {
46		const V4_UNSPEC: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0);
47		const V6_UNSPEC: SocketAddr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0);
48		Socket::new(Domain::IPV6, Type::DGRAM, None)
49			.and_then(|socket| socket.set_only_v6(false).map(|_| socket))
50			.and_then(|socket| socket.bind(&V6_UNSPEC.into()).map(|_| socket))
51			.or_else(|_| {
52				Socket::new(Domain::IPV4, Type::DGRAM, None).and_then(|socket| socket.bind(&V4_UNSPEC.into()).map(|_| socket))
53			})
54	}
55	.and_then(|socket| socket.set_nonblocking(true).map(|_| socket))
56}
57
58pub async fn run_udp_proxy<T, F, R>(
59	proto: Socks5ServerProtocol<T, states::CommandRead>,
60	_addr: &TargetAddr,
61	peer_bind_ip: Option<IpAddr>,
62	reply_ip: IpAddr,
63	transfer: F,
64) -> Result<T, Error>
65where
66	T: AsyncRead + AsyncWrite + Unpin,
67	F: FnOnce(Socket) -> R,
68	R: Future<Output = Result<(), Error>>,
69{
70	let peer_sock = try_notify!(
71		proto,
72		udp_bind_random_port(peer_bind_ip).err_when("binding client udp socket")
73	);
74
75	let peer_addr = try_notify!(proto, peer_sock.local_addr().err_when("getting peer's local addr"));
76
77	let reply_port = peer_addr.as_socket().ok_or(SocksServerError::Bug("addr not IP"))?.port();
78
79	// Respect the pre-populated reply IP address.
80	let mut inner = proto.reply_success(SocketAddr::new(reply_ip, reply_port)).await?;
81
82	let udp_fut = transfer(peer_sock);
83	let tcp_fut = wait_on_tcp(&mut inner).map_err(Error::from);
84
85	match try_join!(udp_fut, tcp_fut) {
86		Ok(_) => warn!("unreachable"),
87		Err(Error::Socks {
88			source: SocksServerError::EOF,
89			backtrace: _,
90		}) => {
91			debug!("EOF on controlling TCP stream, closed UDP proxy")
92		}
93		Err(err) => warn!("while UDP proxying: {err}"),
94	}
95	Ok(inner)
96}
wind_socks/ext.rs

wind_socks/
ext.rs
