socks5_impl/server/connection/
connect.rs

1use crate::{
2    Address, AsyncStreamOperation,
3    protocol::{Reply, Response},
4};
5use std::{
6    io::IoSlice,
7    net::SocketAddr,
8    pin::Pin,
9    task::{Context, Poll},
10};
11use tokio::{
12    io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf},
13    net::{
14        TcpStream,
15        tcp::{OwnedReadHalf, OwnedWriteHalf, ReadHalf, WriteHalf},
16    },
17};
18
19/// Socks5 connection type `Connect`
20///
21/// This connection can be used as a regular async TCP stream after replying the client.
22#[derive(Debug)]
23pub struct Connect<S> {
24    stream: TcpStream,
25    _state: S,
26}
27
28impl<S: Default> Connect<S> {
29    #[inline]
30    pub(super) fn new(stream: TcpStream) -> Self {
31        Self {
32            stream,
33            _state: S::default(),
34        }
35    }
36
37    /// Returns the local address that this stream is bound to.
38    #[inline]
39    pub fn local_addr(&self) -> std::io::Result<SocketAddr> {
40        self.stream.local_addr()
41    }
42
43    /// Returns the remote address that this stream is connected to.
44    #[inline]
45    pub fn peer_addr(&self) -> std::io::Result<SocketAddr> {
46        self.stream.peer_addr()
47    }
48
49    /// Shutdown the TCP stream.
50    #[inline]
51    pub async fn shutdown(&mut self) -> std::io::Result<()> {
52        self.stream.shutdown().await
53    }
54}
55
56#[derive(Debug, Default)]
57pub struct NeedReply;
58
59#[derive(Debug, Default)]
60pub struct Ready;
61
62impl Connect<NeedReply> {
63    /// Reply to the client.
64    #[inline]
65    pub async fn reply(mut self, reply: Reply, addr: Address) -> std::io::Result<Connect<Ready>> {
66        let resp = Response::new(reply, addr);
67        resp.write_to_async_stream(&mut self.stream).await?;
68        Ok(Connect::<Ready>::new(self.stream))
69    }
70}
71
72impl Connect<Ready> {
73    /// Returns the read/write half of the stream.
74    #[inline]
75    pub fn split(&mut self) -> (ReadHalf<'_>, WriteHalf<'_>) {
76        self.stream.split()
77    }
78
79    /// Returns the owned read/write half of the stream.
80    #[inline]
81    pub fn into_split(self) -> (OwnedReadHalf, OwnedWriteHalf) {
82        self.stream.into_split()
83    }
84}
85
86impl std::ops::Deref for Connect<Ready> {
87    type Target = TcpStream;
88
89    #[inline]
90    fn deref(&self) -> &Self::Target {
91        &self.stream
92    }
93}
94
95impl std::ops::DerefMut for Connect<Ready> {
96    #[inline]
97    fn deref_mut(&mut self) -> &mut Self::Target {
98        &mut self.stream
99    }
100}
101
102impl AsyncRead for Connect<Ready> {
103    #[inline]
104    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
105        Pin::new(&mut self.stream).poll_read(cx, buf)
106    }
107}
108
109impl AsyncWrite for Connect<Ready> {
110    #[inline]
111    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<std::io::Result<usize>> {
112        Pin::new(&mut self.stream).poll_write(cx, buf)
113    }
114
115    #[inline]
116    fn poll_write_vectored(mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>]) -> Poll<std::io::Result<usize>> {
117        Pin::new(&mut self.stream).poll_write_vectored(cx, bufs)
118    }
119
120    #[inline]
121    fn is_write_vectored(&self) -> bool {
122        self.stream.is_write_vectored()
123    }
124
125    #[inline]
126    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
127        Pin::new(&mut self.stream).poll_flush(cx)
128    }
129
130    #[inline]
131    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
132        Pin::new(&mut self.stream).poll_shutdown(cx)
133    }
134}
135
136impl<S> From<Connect<S>> for TcpStream {
137    #[inline]
138    fn from(conn: Connect<S>) -> Self {
139        conn.stream
140    }
141}
socks5_impl/server/connection/connect.rs

socks5_impl/server/connection/
connect.rs
