1//! socks proxy
 2use anyhow::{anyhow, Result};
 3use futures::io::{AsyncRead, AsyncWrite};
 4use http_client::Uri;
 5use tokio_socks::{
 6    io::Compat,
 7    tcp::{Socks4Stream, Socks5Stream},
 8};
 9
10pub(crate) async fn connect_socks_proxy_stream(
11    proxy: Option<&Uri>,
12    rpc_host: (&str, u16),
13) -> Result<Box<dyn AsyncReadWrite>> {
14    let stream = match parse_socks_proxy(proxy) {
15        Some((socks_proxy, SocksVersion::V4)) => {
16            let stream = Socks4Stream::connect_with_socket(
17                Compat::new(smol::net::TcpStream::connect(socks_proxy).await?),
18                rpc_host,
19            )
20            .await
21            .map_err(|err| anyhow!("error connecting to socks {}", err))?;
22            Box::new(stream) as Box<dyn AsyncReadWrite>
23        }
24        Some((socks_proxy, SocksVersion::V5)) => Box::new(
25            Socks5Stream::connect_with_socket(
26                Compat::new(smol::net::TcpStream::connect(socks_proxy).await?),
27                rpc_host,
28            )
29            .await
30            .map_err(|err| anyhow!("error connecting to socks {}", err))?,
31        ) as Box<dyn AsyncReadWrite>,
32        None => Box::new(smol::net::TcpStream::connect(rpc_host).await?) as Box<dyn AsyncReadWrite>,
33    };
34    Ok(stream)
35}
36
37fn parse_socks_proxy(proxy: Option<&Uri>) -> Option<((String, u16), SocksVersion)> {
38    let Some(proxy_uri) = proxy else {
39        return None;
40    };
41    let Some(scheme) = proxy_uri.scheme_str() else {
42        return None;
43    };
44    let socks_version = if scheme.starts_with("socks4") {
45        // socks4
46        SocksVersion::V4
47    } else if scheme.starts_with("socks") {
48        // socks, socks5
49        SocksVersion::V5
50    } else {
51        return None;
52    };
53    if let (Some(host), Some(port)) = (proxy_uri.host(), proxy_uri.port_u16()) {
54        Some(((host.to_string(), port), socks_version))
55    } else {
56        None
57    }
58}
59
60// private helper structs and traits
61
62enum SocksVersion {
63    V4,
64    V5,
65}
66
67pub(crate) trait AsyncReadWrite: AsyncRead + AsyncWrite + Unpin + Send + 'static {}
68impl<T: AsyncRead + AsyncWrite + Unpin + Send + 'static> AsyncReadWrite for T {}