Skip to content

Commit 57c26fa

Browse files
0x676e670x676e67
authored
refactor(client): deduplicate and reuse shared socks logic (#205)
1 parent 42bb42f commit 57c26fa

File tree

3 files changed

+49
-94
lines changed

3 files changed

+49
-94
lines changed

‎src/client/legacy/connect/proxy/socks/mod.rs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@ pub use v5::{SocksV5, SocksV5Error};
44
mod v4;
55
pub use v4::{SocksV4, SocksV4Error};
66

7+
use pin_project_lite::pin_project;
8+
use std::future::Future;
9+
use std::pin::Pin;
10+
use std::task::{Context, Poll};
11+
712
use bytes::BytesMut;
813

914
use hyper::rt::Read;
@@ -119,3 +124,31 @@ impl<C> From<SocksV5Error> for SocksError<C> {
119124
Self::V5(err)
120125
}
121126
}
127+
128+
pin_project! {
129+
// Not publicly exported (so missing_docs doesn't trigger).
130+
//
131+
// We return this `Future` instead of the `Pin<Box<dyn Future>>` directly
132+
// so that users don't rely on it fitting in a `Pin<Box<dyn Future>>` slot
133+
// (and thus we can change the type in the future).
134+
#[must_use = "futures do nothing unless polled"]
135+
#[allow(missing_debug_implementations)]
136+
pub struct Handshaking<F, T, E> {
137+
#[pin]
138+
fut: BoxHandshaking<T, E>,
139+
_marker: std::marker::PhantomData<F>
140+
}
141+
}
142+
143+
type BoxHandshaking<T, E> = Pin<Box<dyn Future<Output = Result<T, SocksError<E>>> + Send>>;
144+
145+
impl<F, T, E> Future for Handshaking<F, T, E>
146+
where
147+
F: Future<Output = Result<T, E>>,
148+
{
149+
type Output = Result<T, SocksError<E>>;
150+
151+
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
152+
self.project().fut.poll(cx)
153+
}
154+
}

‎src/client/legacy/connect/proxy/socks/v4/mod.rs‎

Lines changed: 8 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,16 @@ pub use errors::*;
44
mod messages;
55
use messages::*;
66

7-
use std::future::Future;
8-
use std::pin::Pin;
9-
use std::task::{Context, Poll};
10-
117
use std::net::{IpAddr, SocketAddr, SocketAddrV4, ToSocketAddrs};
8+
use std::task::{Context, Poll};
129

1310
use http::Uri;
1411
use hyper::rt::{Read, Write};
1512
use tower_service::Service;
1613

1714
use bytes::BytesMut;
1815

19-
use pin_project_lite::pin_project;
16+
use super::{Handshaking, SocksError};
2017

2118
/// Tunnel Proxy via SOCKSv4
2219
///
@@ -35,23 +32,6 @@ struct SocksConfig {
3532
local_dns: bool,
3633
}
3734

38-
pin_project! {
39-
// Not publicly exported (so missing_docs doesn't trigger).
40-
//
41-
// We return this `Future` instead of the `Pin<Box<dyn Future>>` directly
42-
// so that users don't rely on it fitting in a `Pin<Box<dyn Future>>` slot
43-
// (and thus we can change the type in the future).
44-
#[must_use = "futures do nothing unless polled"]
45-
#[allow(missing_debug_implementations)]
46-
pub struct Handshaking<F, T, E> {
47-
#[pin]
48-
fut: BoxHandshaking<T, E>,
49-
_marker: std::marker::PhantomData<F>
50-
}
51-
}
52-
53-
type BoxHandshaking<T, E> = Pin<Box<dyn Future<Output = Result<T, super::SocksError<E>>> + Send>>;
54-
5535
impl<C> SocksV4<C> {
5636
/// Create a new SOCKSv4 handshake service
5737
///
@@ -86,12 +66,7 @@ impl SocksConfig {
8666
}
8767
}
8868

89-
async fn execute<T, E>(
90-
self,
91-
mut conn: T,
92-
host: String,
93-
port: u16,
94-
) -> Result<T, super::SocksError<E>>
69+
async fn execute<T, E>(self, mut conn: T, host: String, port: u16) -> Result<T, SocksError<E>>
9570
where
9671
T: Read + Write + Unpin,
9772
{
@@ -109,7 +84,7 @@ impl SocksConfig {
10984
None
11085
}
11186
})
112-
.ok_or(super::SocksError::DnsFailure)?
87+
.ok_or(SocksError::DnsFailure)?
11388
} else {
11489
Address::Domain(host, port)
11590
}
@@ -142,11 +117,11 @@ where
142117
C::Error: Send + 'static,
143118
{
144119
type Response = C::Response;
145-
type Error = super::SocksError<C::Error>;
120+
type Error = SocksError<C::Error>;
146121
type Future = Handshaking<C::Future, C::Response, C::Error>;
147122

148123
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
149-
self.inner.poll_ready(cx).map_err(super::SocksError::Inner)
124+
self.inner.poll_ready(cx).map_err(SocksError::Inner)
150125
}
151126

152127
fn call(&mut self, dst: Uri) -> Self::Future {
@@ -155,12 +130,9 @@ where
155130

156131
let fut = async move {
157132
let port = dst.port().map(|p| p.as_u16()).unwrap_or(443);
158-
let host = dst
159-
.host()
160-
.ok_or(super::SocksError::MissingHost)?
161-
.to_string();
133+
let host = dst.host().ok_or(SocksError::MissingHost)?.to_string();
162134

163-
let conn = connecting.await.map_err(super::SocksError::Inner)?;
135+
let conn = connecting.await.map_err(SocksError::Inner)?;
164136
config.execute(conn, host, port).await
165137
};
166138

@@ -170,14 +142,3 @@ where
170142
}
171143
}
172144
}
173-
174-
impl<F, T, E> Future for Handshaking<F, T, E>
175-
where
176-
F: Future<Output = Result<T, E>>,
177-
{
178-
type Output = Result<T, super::SocksError<E>>;
179-
180-
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
181-
self.project().fut.poll(cx)
182-
}
183-
}

‎src/client/legacy/connect/proxy/socks/v5/mod.rs‎

Lines changed: 8 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,16 @@ pub use errors::*;
44
mod messages;
55
use messages::*;
66

7-
use std::future::Future;
8-
use std::pin::Pin;
9-
use std::task::{Context, Poll};
10-
117
use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
8+
use std::task::{Context, Poll};
129

1310
use http::Uri;
1411
use hyper::rt::{Read, Write};
1512
use tower_service::Service;
1613

1714
use bytes::BytesMut;
1815

19-
use pin_project_lite::pin_project;
16+
use super::{Handshaking, SocksError};
2017

2118
/// Tunnel Proxy via SOCKSv5
2219
///
@@ -48,23 +45,6 @@ enum State {
4845
ReadingProxyRes,
4946
}
5047

51-
pin_project! {
52-
// Not publicly exported (so missing_docs doesn't trigger).
53-
//
54-
// We return this `Future` instead of the `Pin<Box<dyn Future>>` directly
55-
// so that users don't rely on it fitting in a `Pin<Box<dyn Future>>` slot
56-
// (and thus we can change the type in the future).
57-
#[must_use = "futures do nothing unless polled"]
58-
#[allow(missing_debug_implementations)]
59-
pub struct Handshaking<F, T, E> {
60-
#[pin]
61-
fut: BoxHandshaking<T, E>,
62-
_marker: std::marker::PhantomData<F>
63-
}
64-
}
65-
66-
type BoxHandshaking<T, E> = Pin<Box<dyn Future<Output = Result<T, super::SocksError<E>>> + Send>>;
67-
6848
impl<C> SocksV5<C> {
6949
/// Create a new SOCKSv5 handshake service.
7050
///
@@ -126,12 +106,7 @@ impl SocksConfig {
126106
}
127107
}
128108

129-
async fn execute<T, E>(
130-
self,
131-
mut conn: T,
132-
host: String,
133-
port: u16,
134-
) -> Result<T, super::SocksError<E>>
109+
async fn execute<T, E>(self, mut conn: T, host: String, port: u16) -> Result<T, SocksError<E>>
135110
where
136111
T: Read + Write + Unpin,
137112
{
@@ -142,7 +117,7 @@ impl SocksConfig {
142117
let socket = (host, port)
143118
.to_socket_addrs()?
144119
.next()
145-
.ok_or(super::SocksError::DnsFailure)?;
120+
.ok_or(SocksError::DnsFailure)?;
146121

147122
Address::Socket(socket)
148123
} else {
@@ -272,11 +247,11 @@ where
272247
C::Error: Send + 'static,
273248
{
274249
type Response = C::Response;
275-
type Error = super::SocksError<C::Error>;
250+
type Error = SocksError<C::Error>;
276251
type Future = Handshaking<C::Future, C::Response, C::Error>;
277252

278253
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
279-
self.inner.poll_ready(cx).map_err(super::SocksError::Inner)
254+
self.inner.poll_ready(cx).map_err(SocksError::Inner)
280255
}
281256

282257
fn call(&mut self, dst: Uri) -> Self::Future {
@@ -285,12 +260,9 @@ where
285260

286261
let fut = async move {
287262
let port = dst.port().map(|p| p.as_u16()).unwrap_or(443);
288-
let host = dst
289-
.host()
290-
.ok_or(super::SocksError::MissingHost)?
291-
.to_string();
263+
let host = dst.host().ok_or(SocksError::MissingHost)?.to_string();
292264

293-
let conn = connecting.await.map_err(super::SocksError::Inner)?;
265+
let conn = connecting.await.map_err(SocksError::Inner)?;
294266
config.execute(conn, host, port).await
295267
};
296268

@@ -300,14 +272,3 @@ where
300272
}
301273
}
302274
}
303-
304-
impl<F, T, E> Future for Handshaking<F, T, E>
305-
where
306-
F: Future<Output = Result<T, E>>,
307-
{
308-
type Output = Result<T, super::SocksError<E>>;
309-
310-
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
311-
self.project().fut.poll(cx)
312-
}
313-
}

0 commit comments

Comments
 (0)