lzc_sdk/remotesocks/
udp.rs

1use std::fmt;
2use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
3
4use tokio::io::{AsyncReadExt, AsyncWriteExt};
5use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
6use tokio::net::{TcpStream, UdpSocket, lookup_host};
7use tokio::sync::{Mutex, watch};
8use tokio::task::JoinHandle;
9
10use crate::Error;
11
12use super::client::ASSOCIATE_COMMAND;
13use super::{RemoteNetstack, SocksAddress};
14
15const SOCKS_UDP_PREFIX: [u8; 3] = [0, 0, 0];
16const MAX_UDP_PAYLOAD: usize = 65_507;
17const FRAMED_UDP_HEADER_LENGTH: usize = 20;
18const CONNECT_UDP_COMMAND: u8 = 0xe2;
19const BIND_UDP_COMMAND: u8 = 0xe3;
20
21enum UdpTransport {
22    Association(Association),
23    Framed(FramedTransport),
24}
25
26struct Association {
27    socket: UdpSocket,
28    relay: SocketAddr,
29    control_closed: Mutex<watch::Receiver<bool>>,
30    control_task: JoinHandle<()>,
31}
32
33struct FramedTransport {
34    reader: Mutex<OwnedReadHalf>,
35    writer: Mutex<OwnedWriteHalf>,
36}
37
38/// UDP socket transported through standard or extended `LazyCat` `RemoteSocks`.
39pub struct RemoteUdpSocket {
40    transport: UdpTransport,
41    local_addr: SocksAddress,
42    peer_addr: Option<SocksAddress>,
43}
44
45impl fmt::Debug for RemoteUdpSocket {
46    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
47        let transport = match &self.transport {
48            UdpTransport::Association(_) => "associate",
49            UdpTransport::Framed(_) => "framed",
50        };
51        formatter
52            .debug_struct("RemoteUdpSocket")
53            .field("transport", &transport)
54            .field("local_addr", &self.local_addr)
55            .field("peer_addr", &self.peer_addr)
56            .finish_non_exhaustive()
57    }
58}
59
60impl RemoteUdpSocket {
61    /// Local UDP address used by this transport.
62    #[must_use]
63    pub const fn local_addr(&self) -> &SocksAddress {
64        &self.local_addr
65    }
66
67    /// Fixed remote peer for associated or connected sockets.
68    #[must_use]
69    pub const fn peer_addr(&self) -> Option<&SocksAddress> {
70        self.peer_addr.as_ref()
71    }
72
73    /// Send one datagram to the fixed peer.
74    ///
75    /// # Errors
76    ///
77    /// Returns an error when this is a bound socket without a fixed peer, the
78    /// control stream is closed, or the datagram cannot be sent.
79    pub async fn send(&self, payload: &[u8]) -> Result<usize, Error> {
80        let peer = self.peer_addr.as_ref().ok_or(Error::SocksUdpNotConnected)?;
81        self.send_to(payload, peer).await
82    }
83
84    /// Receive one datagram from the fixed peer, filtering other sources.
85    ///
86    /// # Errors
87    ///
88    /// Returns an error when this is a bound socket without a fixed peer, the
89    /// control stream closes, or a proxy datagram is malformed.
90    pub async fn recv(&self, payload: &mut [u8]) -> Result<usize, Error> {
91        let peer = self.peer_addr.as_ref().ok_or(Error::SocksUdpNotConnected)?;
92        loop {
93            let (length, source) = self.recv_from(payload).await?;
94            if &source == peer {
95                return Ok(length);
96            }
97        }
98    }
99
100    /// Send one datagram to an explicit target.
101    ///
102    /// Standard associations use RFC 1928 UDP headers. `LazyCat` `0xe2` and
103    /// `0xe3` transports use the SDK's length/IP/port stream framing.
104    ///
105    /// # Errors
106    ///
107    /// Returns an error for oversized frames, unsupported target address
108    /// forms, closed control streams, or transport write failures.
109    pub async fn send_to(&self, payload: &[u8], target: &SocksAddress) -> Result<usize, Error> {
110        if payload.len() > MAX_UDP_PAYLOAD {
111            return Err(Error::SocksUdpPayloadTooLarge);
112        }
113        match &self.transport {
114            UdpTransport::Association(association) => {
115                send_associated(association, payload, target).await
116            }
117            UdpTransport::Framed(transport) => send_framed(transport, payload, target).await,
118        }
119    }
120
121    /// Receive one datagram and its encoded source address.
122    ///
123    /// Standard-association packets not sent by the negotiated relay are
124    /// ignored.
125    ///
126    /// # Errors
127    ///
128    /// Returns an error when a control stream closes or a datagram has invalid
129    /// framing or address fields.
130    pub async fn recv_from(&self, payload: &mut [u8]) -> Result<(usize, SocksAddress), Error> {
131        match &self.transport {
132            UdpTransport::Association(association) => recv_associated(association, payload).await,
133            UdpTransport::Framed(transport) => recv_framed(transport, payload).await,
134        }
135    }
136}
137
138impl Drop for RemoteUdpSocket {
139    fn drop(&mut self) {
140        if let UdpTransport::Association(association) = &self.transport {
141            association.control_task.abort();
142        }
143    }
144}
145
146impl RemoteNetstack {
147    /// Create a standard RFC 1928 UDP association with a fixed default target.
148    ///
149    /// The TCP control stream remains open for the lifetime of the returned
150    /// socket. Closing it causes receive operations to return
151    /// [`Error::RemoteSocksControlClosed`].
152    ///
153    /// # Errors
154    ///
155    /// Returns endpoint resolution, SOCKS negotiation, relay address, socket
156    /// binding, or control-stream errors.
157    pub async fn udp_associate(&self, target: SocksAddress) -> Result<RemoteUdpSocket, Error> {
158        let unspecified = SocksAddress::Ip(SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0));
159        let (relay, control) = self
160            .perform_command(ASSOCIATE_COMMAND, &unspecified)
161            .await?;
162        let relay = resolve_socket_address(&relay).await?;
163        let bind_address = unspecified_for(relay);
164        let socket = UdpSocket::bind(bind_address).await?;
165        let local_addr = SocksAddress::Ip(socket.local_addr()?);
166        let (control_closed, control_task) = monitor_control_stream(control);
167        Ok(RemoteUdpSocket {
168            transport: UdpTransport::Association(Association {
169                socket,
170                relay,
171                control_closed: Mutex::new(control_closed),
172                control_task,
173            }),
174            local_addr,
175            peer_addr: Some(target),
176        })
177    }
178
179    /// Connect a UDP peer using the `LazyCat` `0xe2` stream extension.
180    ///
181    /// Domain targets are resolved locally, matching the official Go SDK.
182    ///
183    /// # Errors
184    ///
185    /// Returns endpoint resolution, DNS, SOCKS negotiation, address, or framed
186    /// transport errors.
187    pub async fn connect_udp(&self, target: SocksAddress) -> Result<RemoteUdpSocket, Error> {
188        let target = resolve_udp_target(target).await?;
189        let (local_addr, stream) = self.perform_command(CONNECT_UDP_COMMAND, &target).await?;
190        Ok(framed_socket(stream, local_addr, Some(target)))
191    }
192
193    /// Bind a UDP socket using the `LazyCat` `0xe3` stream extension.
194    ///
195    /// Use [`RemoteUdpSocket::send_to`] and [`RemoteUdpSocket::recv_from`] with
196    /// the returned socket because it has no fixed peer.
197    ///
198    /// # Errors
199    ///
200    /// Returns endpoint resolution, DNS, SOCKS negotiation, address, or framed
201    /// transport errors.
202    pub async fn bind_udp(&self, address: SocksAddress) -> Result<RemoteUdpSocket, Error> {
203        let address = resolve_udp_target(address).await?;
204        let (local_addr, stream) = self.perform_command(BIND_UDP_COMMAND, &address).await?;
205        Ok(framed_socket(stream, local_addr, None))
206    }
207}
208
209async fn send_associated(
210    association: &Association,
211    payload: &[u8],
212    target: &SocksAddress,
213) -> Result<usize, Error> {
214    if *association.control_closed.lock().await.borrow() {
215        return Err(Error::RemoteSocksControlClosed);
216    }
217    let address = target.encode()?;
218    let frame_length = SOCKS_UDP_PREFIX.len() + address.len() + payload.len();
219    if frame_length > MAX_UDP_PAYLOAD {
220        return Err(Error::SocksUdpPayloadTooLarge);
221    }
222    let mut frame = Vec::with_capacity(frame_length);
223    frame.extend_from_slice(&SOCKS_UDP_PREFIX);
224    frame.extend_from_slice(&address);
225    frame.extend_from_slice(payload);
226    let written = association
227        .socket
228        .send_to(&frame, association.relay)
229        .await?;
230    if written != frame.len() {
231        return Err(std::io::Error::new(
232            std::io::ErrorKind::WriteZero,
233            "short SOCKS UDP datagram write",
234        )
235        .into());
236    }
237    Ok(payload.len())
238}
239
240async fn recv_associated(
241    association: &Association,
242    payload: &mut [u8],
243) -> Result<(usize, SocksAddress), Error> {
244    let mut packet = vec![0_u8; MAX_UDP_PAYLOAD];
245    let mut control_closed = association.control_closed.lock().await;
246    loop {
247        if *control_closed.borrow() {
248            return Err(Error::RemoteSocksControlClosed);
249        }
250        let received = tokio::select! {
251            changed = control_closed.changed() => {
252                let _ = changed;
253                return Err(Error::RemoteSocksControlClosed);
254            }
255            received = association.socket.recv_from(&mut packet) => received?,
256        };
257        let (length, source) = received;
258        if source != association.relay {
259            continue;
260        }
261        if length < SOCKS_UDP_PREFIX.len() || packet[..2] != [0, 0] {
262            return Err(Error::InvalidSocksUdpHeader);
263        }
264        if packet[2] != 0 {
265            return Err(Error::UnsupportedSocksUdpFragment {
266                fragment: packet[2],
267            });
268        }
269        let (source, consumed) = SocksAddress::decode_prefix(&packet[3..length])?;
270        let datagram = &packet[3 + consumed..length];
271        let copied = copy_payload(payload, datagram);
272        return Ok((copied, source));
273    }
274}
275
276async fn send_framed(
277    transport: &FramedTransport,
278    payload: &[u8],
279    target: &SocksAddress,
280) -> Result<usize, Error> {
281    let target = concrete_ip_address(target)?;
282    let mut header = [0_u8; FRAMED_UDP_HEADER_LENGTH];
283    header[..2].copy_from_slice(
284        &u16::try_from(payload.len())
285            .map_err(|_| Error::SocksUdpPayloadTooLarge)?
286            .to_be_bytes(),
287    );
288    header[2..18].copy_from_slice(&ip_octets(target.ip()));
289    header[18..].copy_from_slice(&target.port().to_be_bytes());
290    let mut writer = transport.writer.lock().await;
291    writer.write_all(&header).await?;
292    writer.write_all(payload).await?;
293    Ok(payload.len())
294}
295
296async fn recv_framed(
297    transport: &FramedTransport,
298    payload: &mut [u8],
299) -> Result<(usize, SocksAddress), Error> {
300    let mut reader = transport.reader.lock().await;
301    let mut header = [0_u8; FRAMED_UDP_HEADER_LENGTH];
302    reader.read_exact(&mut header).await?;
303    let length = usize::from(u16::from_be_bytes([header[0], header[1]]));
304    let mut datagram = vec![0_u8; length];
305    reader.read_exact(&mut datagram).await?;
306    if length > MAX_UDP_PAYLOAD {
307        return Err(Error::SocksUdpPayloadTooLarge);
308    }
309    let octets: [u8; 16] = header[2..18]
310        .try_into()
311        .map_err(|_| Error::InvalidSocksUdpHeader)?;
312    let ipv6 = Ipv6Addr::from(octets);
313    let ip = ipv6.to_ipv4_mapped().map_or(IpAddr::V6(ipv6), IpAddr::V4);
314    let port = u16::from_be_bytes([header[18], header[19]]);
315    let copied = copy_payload(payload, &datagram);
316    Ok((copied, SocksAddress::Ip(SocketAddr::new(ip, port))))
317}
318
319fn framed_socket(
320    stream: TcpStream,
321    local_addr: SocksAddress,
322    peer_addr: Option<SocksAddress>,
323) -> RemoteUdpSocket {
324    let (reader, writer) = stream.into_split();
325    RemoteUdpSocket {
326        transport: UdpTransport::Framed(FramedTransport {
327            reader: Mutex::new(reader),
328            writer: Mutex::new(writer),
329        }),
330        local_addr,
331        peer_addr,
332    }
333}
334
335fn monitor_control_stream(control: TcpStream) -> (watch::Receiver<bool>, JoinHandle<()>) {
336    let (sender, receiver) = watch::channel(false);
337    let task = tokio::spawn(async move {
338        let (mut reader, writer) = control.into_split();
339        let _writer = writer;
340        let mut buffer = [0_u8; 1];
341        loop {
342            match reader.read(&mut buffer).await {
343                Ok(0) | Err(_) => break,
344                Ok(_) => {}
345            }
346        }
347        let _ = sender.send(true);
348    });
349    (receiver, task)
350}
351
352async fn resolve_socket_address(address: &SocksAddress) -> Result<SocketAddr, Error> {
353    match address {
354        SocksAddress::Ip(address) => Ok(*address),
355        SocksAddress::Domain { host, port } => lookup_host((host.as_str(), *port))
356            .await?
357            .next()
358            .ok_or(Error::InvalidSocksUdpAddress),
359        SocksAddress::Custom { .. } => Err(Error::InvalidSocksUdpAddress),
360    }
361}
362
363async fn resolve_udp_target(address: SocksAddress) -> Result<SocksAddress, Error> {
364    match address {
365        SocksAddress::Ip(_) => Ok(address),
366        SocksAddress::Domain { .. } => {
367            Ok(SocksAddress::Ip(resolve_socket_address(&address).await?))
368        }
369        SocksAddress::Custom { .. } => Err(Error::InvalidSocksUdpAddress),
370    }
371}
372
373fn concrete_ip_address(address: &SocksAddress) -> Result<SocketAddr, Error> {
374    match address {
375        SocksAddress::Ip(address) => Ok(*address),
376        SocksAddress::Domain { .. } | SocksAddress::Custom { .. } => {
377            Err(Error::InvalidSocksUdpAddress)
378        }
379    }
380}
381
382const fn unspecified_for(address: SocketAddr) -> SocketAddr {
383    match address {
384        SocketAddr::V4(_) => SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0),
385        SocketAddr::V6(_) => SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0),
386    }
387}
388
389fn ip_octets(ip: IpAddr) -> [u8; 16] {
390    match ip {
391        IpAddr::V4(ip) => ip.to_ipv6_mapped().octets(),
392        IpAddr::V6(ip) => ip.octets(),
393    }
394}
395
396fn copy_payload(destination: &mut [u8], source: &[u8]) -> usize {
397    let copied = destination.len().min(source.len());
398    destination[..copied].copy_from_slice(&source[..copied]);
399    copied
400}
lzc_sdk/remotesocks/udp.rs

lzc_sdk/remotesocks/
udp.rs
