Skip to main content

lzc_sdk/remotesocks/
client.rs

1use std::fmt;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5use std::time::Duration;
6
7use tokio::io::{AsyncReadExt, AsyncWriteExt};
8use tokio::net::TcpStream;
9use tokio::sync::Mutex;
10use tokio::time::timeout;
11
12use crate::{Error, HServerClient, RemoteLocation, RemoteSocksEndpoint};
13
14use super::SocksAddress;
15
16pub(crate) const CONNECT_COMMAND: u8 = 0x01;
17pub(crate) const BIND_COMMAND: u8 = 0x02;
18pub(crate) const ASSOCIATE_COMMAND: u8 = 0x03;
19
20const SOCKS_VERSION: u8 = 0x05;
21const NO_AUTHENTICATION: u8 = 0x00;
22const NO_ACCEPTABLE_AUTHENTICATION: u8 = 0xff;
23const DEFAULT_DIAL_TIMEOUT: Duration = Duration::from_secs(60);
24
25type ResolverFuture = Pin<Box<dyn Future<Output = Result<RemoteSocksEndpoint, Error>> + Send>>;
26type EndpointResolver = dyn Fn() -> ResolverFuture + Send + Sync;
27
28struct Shared {
29    resolver: Arc<EndpointResolver>,
30    cached_endpoint: Mutex<Option<RemoteSocksEndpoint>>,
31}
32
33/// Async `LazyCat` `RemoteSocks` network stack.
34#[derive(Clone)]
35pub struct RemoteNetstack {
36    shared: Arc<Shared>,
37    dial_timeout: Duration,
38}
39
40impl fmt::Debug for RemoteNetstack {
41    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
42        formatter
43            .debug_struct("RemoteNetstack")
44            .field("dial_timeout", &self.dial_timeout)
45            .finish_non_exhaustive()
46    }
47}
48
49impl RemoteNetstack {
50    /// Resolve proxy endpoints through `HPortalSys` for the requested location.
51    #[must_use]
52    pub fn new(hserver: HServerClient, location: impl Into<RemoteLocation>) -> Self {
53        let location = location.into();
54        Self::with_endpoint_resolver(move || {
55            let hserver = hserver.clone();
56            let location = location.clone();
57            async move { hserver.remote_socks_endpoint(location).await }
58        })
59    }
60
61    /// Use a fixed `RemoteSocks` proxy endpoint.
62    #[must_use]
63    pub fn fixed(endpoint: RemoteSocksEndpoint) -> Self {
64        Self::with_endpoint_resolver(move || {
65            let endpoint = endpoint.clone();
66            async move { Ok(endpoint) }
67        })
68    }
69
70    /// Use an asynchronous endpoint resolver compatible with the Go SDK's
71    /// dynamic address factory.
72    #[must_use]
73    pub fn with_endpoint_resolver<F, Fut>(resolver: F) -> Self
74    where
75        F: Fn() -> Fut + Send + Sync + 'static,
76        Fut: Future<Output = Result<RemoteSocksEndpoint, Error>> + Send + 'static,
77    {
78        let resolver = Arc::new(move || Box::pin(resolver()) as ResolverFuture);
79        Self {
80            shared: Arc::new(Shared {
81                resolver,
82                cached_endpoint: Mutex::new(None),
83            }),
84            dial_timeout: DEFAULT_DIAL_TIMEOUT,
85        }
86    }
87
88    /// Override the default 60-second proxy connection deadline.
89    #[must_use]
90    pub const fn with_dial_timeout(mut self, dial_timeout: Duration) -> Self {
91        self.dial_timeout = dial_timeout;
92        self
93    }
94
95    pub(crate) async fn perform_command(
96        &self,
97        command: u8,
98        address: &SocksAddress,
99    ) -> Result<(SocksAddress, TcpStream), Error> {
100        match timeout(
101            self.dial_timeout,
102            self.perform_command_with_deadline(command, address),
103        )
104        .await
105        {
106            Ok(result) => result,
107            Err(_) => Err(Error::RemoteSocksTimeout),
108        }
109    }
110
111    async fn perform_command_with_deadline(
112        &self,
113        command: u8,
114        address: &SocksAddress,
115    ) -> Result<(SocksAddress, TcpStream), Error> {
116        let mut stream = self.connect_proxy().await?;
117        negotiate_no_authentication(&mut stream).await?;
118        stream.write_all(&[SOCKS_VERSION, command, 0]).await?;
119        address.write_to(&mut stream).await?;
120        let bound_address = read_reply(&mut stream).await?;
121        Ok((bound_address, stream))
122    }
123
124    async fn connect_proxy(&self) -> Result<TcpStream, Error> {
125        let endpoint = self.proxy_endpoint(true).await?;
126        match self.dial_endpoint(&endpoint).await {
127            Ok(stream) => Ok(stream),
128            Err(first_error) => {
129                let Ok(refreshed) = self.proxy_endpoint(false).await else {
130                    return Err(first_error);
131                };
132                if refreshed == endpoint {
133                    return Err(first_error);
134                }
135                self.dial_endpoint(&refreshed).await
136            }
137        }
138    }
139
140    async fn dial_endpoint(&self, endpoint: &RemoteSocksEndpoint) -> Result<TcpStream, Error> {
141        TcpStream::connect(endpoint.authority())
142            .await
143            .map_err(Error::from)
144    }
145
146    async fn proxy_endpoint(&self, with_cache: bool) -> Result<RemoteSocksEndpoint, Error> {
147        let mut cached = self.shared.cached_endpoint.lock().await;
148        if with_cache && let Some(endpoint) = cached.as_ref() {
149            return Ok(endpoint.clone());
150        }
151        let endpoint = (self.shared.resolver)().await?;
152        *cached = Some(endpoint.clone());
153        Ok(endpoint)
154    }
155}
156
157async fn negotiate_no_authentication(stream: &mut TcpStream) -> Result<(), Error> {
158    stream
159        .write_all(&[SOCKS_VERSION, 1, NO_AUTHENTICATION])
160        .await?;
161    let mut response = [0_u8; 2];
162    stream.read_exact(&mut response).await?;
163    if response[0] != SOCKS_VERSION {
164        return Err(Error::UnexpectedSocksVersion {
165            version: response[0],
166        });
167    }
168    match response[1] {
169        NO_AUTHENTICATION => Ok(()),
170        NO_ACCEPTABLE_AUTHENTICATION => Err(Error::SocksNoAcceptableAuthentication),
171        method => Err(Error::UnsupportedSocksAuthentication { method }),
172    }
173}
174
175pub(crate) async fn read_reply(stream: &mut TcpStream) -> Result<SocksAddress, Error> {
176    let mut header = [0_u8; 3];
177    stream.read_exact(&mut header).await?;
178    if header[0] != SOCKS_VERSION {
179        return Err(Error::UnexpectedSocksVersion { version: header[0] });
180    }
181    if header[2] != 0 {
182        return Err(Error::InvalidSocksReservedByte {
183            reserved: header[2],
184        });
185    }
186    if header[1] != 0 {
187        return Err(Error::SocksReply { code: header[1] });
188    }
189    SocksAddress::read_from(stream).await
190}