Skip to main content

lzc_sdk/remotesocks/
tcp.rs

1use std::pin::Pin;
2use std::task::{Context, Poll};
3
4use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
5use tokio::net::TcpStream;
6use tokio::sync::Mutex;
7
8use crate::Error;
9
10use super::client::{BIND_COMMAND, CONNECT_COMMAND, read_reply};
11use super::{RemoteNetstack, SocksAddress};
12
13/// TCP stream transported through a `LazyCat` `RemoteSocks` proxy.
14#[derive(Debug)]
15pub struct RemoteTcpStream {
16    inner: TcpStream,
17    local_addr: SocksAddress,
18    peer_addr: SocksAddress,
19}
20
21/// TCP listener implemented with the two-stage SOCKS BIND command.
22#[derive(Debug)]
23pub struct RemoteTcpListener {
24    netstack: RemoteNetstack,
25    local_addr: SocksAddress,
26    pending: Mutex<Option<TcpStream>>,
27}
28
29impl RemoteTcpListener {
30    /// Address assigned by the `RemoteSocks` proxy after the first BIND reply.
31    #[must_use]
32    pub const fn local_addr(&self) -> &SocksAddress {
33        &self.local_addr
34    }
35
36    /// Wait for the second BIND reply and return the accepted proxy stream.
37    ///
38    /// Further calls issue another BIND request using the assigned local
39    /// address, matching the reusable listener behavior of the official SDK.
40    ///
41    /// # Errors
42    ///
43    /// Returns endpoint resolution, connection, negotiation, or BIND reply
44    /// errors.
45    pub async fn accept(&self) -> Result<RemoteTcpStream, Error> {
46        let pending = self.pending.lock().await.take();
47        let (local_addr, mut stream) = if let Some(stream) = pending {
48            (self.local_addr.clone(), stream)
49        } else {
50            self.netstack
51                .perform_command(BIND_COMMAND, &self.local_addr)
52                .await?
53        };
54        let peer_addr = read_reply(&mut stream).await?;
55        Ok(RemoteTcpStream::new(stream, local_addr, peer_addr))
56    }
57}
58
59impl RemoteTcpStream {
60    /// Address assigned by the `RemoteSocks` proxy.
61    #[must_use]
62    pub const fn local_addr(&self) -> &SocksAddress {
63        &self.local_addr
64    }
65
66    /// Requested remote address.
67    #[must_use]
68    pub const fn peer_addr(&self) -> &SocksAddress {
69        &self.peer_addr
70    }
71
72    /// Consume the wrapper and return the underlying proxy TCP stream.
73    #[must_use]
74    pub fn into_inner(self) -> TcpStream {
75        self.inner
76    }
77
78    pub(crate) const fn new(
79        inner: TcpStream,
80        local_addr: SocksAddress,
81        peer_addr: SocksAddress,
82    ) -> Self {
83        Self {
84            inner,
85            local_addr,
86            peer_addr,
87        }
88    }
89}
90
91impl RemoteNetstack {
92    /// Connect to a TCP or custom-network address through `RemoteSocks`.
93    ///
94    /// # Errors
95    ///
96    /// Returns endpoint resolution, proxy connection, SOCKS negotiation, or
97    /// command reply errors.
98    pub async fn connect_tcp(&self, address: SocksAddress) -> Result<RemoteTcpStream, Error> {
99        let (local_addr, stream) = self.perform_command(CONNECT_COMMAND, &address).await?;
100        Ok(RemoteTcpStream::new(stream, local_addr, address))
101    }
102
103    /// Bind a TCP or custom-network listener through `RemoteSocks`.
104    ///
105    /// This method waits for the first SOCKS BIND reply. Call
106    /// [`RemoteTcpListener::accept`] to wait for the second reply carrying the
107    /// connecting peer address.
108    ///
109    /// # Errors
110    ///
111    /// Returns endpoint resolution, proxy connection, SOCKS negotiation, or
112    /// first-stage BIND reply errors.
113    pub async fn bind_tcp(&self, address: SocksAddress) -> Result<RemoteTcpListener, Error> {
114        let (local_addr, stream) = self.perform_command(BIND_COMMAND, &address).await?;
115        Ok(RemoteTcpListener {
116            netstack: self.clone(),
117            local_addr,
118            pending: Mutex::new(Some(stream)),
119        })
120    }
121}
122
123impl AsyncRead for RemoteTcpStream {
124    fn poll_read(
125        mut self: Pin<&mut Self>,
126        context: &mut Context<'_>,
127        buffer: &mut ReadBuf<'_>,
128    ) -> Poll<std::io::Result<()>> {
129        Pin::new(&mut self.inner).poll_read(context, buffer)
130    }
131}
132
133impl AsyncWrite for RemoteTcpStream {
134    fn poll_write(
135        mut self: Pin<&mut Self>,
136        context: &mut Context<'_>,
137        buffer: &[u8],
138    ) -> Poll<Result<usize, std::io::Error>> {
139        Pin::new(&mut self.inner).poll_write(context, buffer)
140    }
141
142    fn poll_flush(
143        mut self: Pin<&mut Self>,
144        context: &mut Context<'_>,
145    ) -> Poll<Result<(), std::io::Error>> {
146        Pin::new(&mut self.inner).poll_flush(context)
147    }
148
149    fn poll_shutdown(
150        mut self: Pin<&mut Self>,
151        context: &mut Context<'_>,
152    ) -> Poll<Result<(), std::io::Error>> {
153        Pin::new(&mut self.inner).poll_shutdown(context)
154    }
155}