Skip to content

Commit

Permalink
feat: Add support for SOCKS4 (seanmonstar#610) (seanmonstar#2400)
Browse files Browse the repository at this point in the history 
* feat: Add support for SOCKS4 (seanmonstar#610)

* chore: bump tokio-socks to 0.5.2
  • Loading branch information
@Jaltaire @Nutomic
Jaltaire authored and Nutomic committed Nov 7, 2024
1 parent 5eda77a commit 7958847
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 26 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ async-compression = { version = "0.4.0", default-features = false, features = ["
tokio-util = { version = "0.7.9", default-features = false, features = ["codec", "io"], optional = true }

## socks
tokio-socks = { version = "0.5.1", optional = true }
tokio-socks = { version = "0.5.2", optional = true }

## hickory-dns
hickory-resolver = { version = "0.24", optional = true, features = ["tokio-runtime"] }
Expand Down
52 changes: 30 additions & 22 deletions src/connect.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,7 @@ impl Connector {
#[cfg(feature = "socks")]
async fn connect_socks(&self, dst: Uri, proxy: ProxyScheme) -> Result<Conn, BoxError> {
let dns = match proxy {
ProxyScheme::Socks4 { .. } => socks::DnsResolve::Local,
ProxyScheme::Socks5 {
remote_dns: false, ..
} => socks::DnsResolve::Local,
Expand Down Expand Up @@ -367,6 +368,8 @@ impl Connector {
ProxyScheme::Http { host, auth } => (into_uri(Scheme::HTTP, host), auth),
ProxyScheme::Https { host, auth } => (into_uri(Scheme::HTTPS, host), auth),
#[cfg(feature = "socks")]
ProxyScheme::Socks4 { .. } => return self.connect_socks(dst, proxy_scheme).await,
#[cfg(feature = "socks")]
ProxyScheme::Socks5 { .. } => return self.connect_socks(dst, proxy_scheme).await,
};

Expand Down Expand Up @@ -1031,7 +1034,7 @@ mod socks {

use http::Uri;
use tokio::net::TcpStream;
use tokio_socks::tcp::Socks5Stream;
use tokio_socks::tcp::{Socks4Stream, Socks5Stream};

use super::{BoxError, Scheme};
use crate::proxy::ProxyScheme;
Expand Down Expand Up @@ -1064,28 +1067,33 @@ mod socks {
}
}

let (socket_addr, auth) = match proxy {
ProxyScheme::Socks5 { addr, auth, .. } => (addr, auth),
_ => unreachable!(),
};

// Get a Tokio TcpStream
let stream = if let Some((username, password)) = auth {
Socks5Stream::connect_with_password(
socket_addr,
(host.as_str(), port),
&username,
&password,
)
.await
.map_err(|e| format!("socks connect error: {e}"))?
} else {
Socks5Stream::connect(socket_addr, (host.as_str(), port))
.await
.map_err(|e| format!("socks connect error: {e}"))?
};
match proxy {
ProxyScheme::Socks4 { addr } => {
let stream = Socks4Stream::connect(addr, (host.as_str(), port))
.await
.map_err(|e| format!("socks connect error: {e}"))?;
Ok(stream.into_inner())
}
ProxyScheme::Socks5 { addr, ref auth, .. } => {
let stream = if let Some((username, password)) = auth {
Socks5Stream::connect_with_password(
addr,
(host.as_str(), port),
&username,
&password,
)
.await
.map_err(|e| format!("socks connect error: {e}"))?
} else {
Socks5Stream::connect(addr, (host.as_str(), port))
.await
.map_err(|e| format!("socks connect error: {e}"))?
};

Ok(stream.into_inner())
Ok(stream.into_inner())
}
_ => unreachable!(),
}
}
}

Expand Down
38 changes: 35 additions & 3 deletions src/proxy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,8 @@ pub enum ProxyScheme {
host: http::uri::Authority,
},
#[cfg(feature = "socks")]
Socks4 { addr: SocketAddr },
#[cfg(feature = "socks")]
Socks5 {
addr: SocketAddr,
auth: Option<(String, String)>,
Expand Down Expand Up @@ -577,6 +579,16 @@ impl ProxyScheme {
})
}

/// Proxy traffic via the specified socket address over SOCKS4
///
/// # Note
///
/// Current SOCKS4 support is provided via blocking IO.
#[cfg(feature = "socks")]
fn socks4(addr: SocketAddr) -> crate::Result<Self> {
Ok(ProxyScheme::Socks4 { addr })
}

/// Proxy traffic via the specified socket address over SOCKS5
///
/// # Note
Expand Down Expand Up @@ -628,6 +640,10 @@ impl ProxyScheme {
*auth = Some(header);
}
#[cfg(feature = "socks")]
ProxyScheme::Socks4 { .. } => {
panic!("Socks4 is not supported for this method")
}
#[cfg(feature = "socks")]
ProxyScheme::Socks5 { ref mut auth, .. } => {
*auth = Some((username.into(), password.into()));
}
Expand All @@ -643,8 +659,12 @@ impl ProxyScheme {
*auth = Some(header_value);
}
#[cfg(feature = "socks")]
ProxyScheme::Socks4 { .. } => {
panic!("Socks4 is not supported for this method")
}
#[cfg(feature = "socks")]
ProxyScheme::Socks5 { .. } => {
panic!("Socks is not supported for this method")
panic!("Socks5 is not supported for this method")
}
}
}
Expand All @@ -662,6 +682,8 @@ impl ProxyScheme {
}
}
#[cfg(feature = "socks")]
ProxyScheme::Socks4 { .. } => {}
#[cfg(feature = "socks")]
ProxyScheme::Socks5 { .. } => {}
}

Expand All @@ -670,7 +692,7 @@ impl ProxyScheme {

/// Convert a URL into a proxy scheme
///
/// Supported schemes: HTTP, HTTPS, (SOCKS5, SOCKS5H if `socks` feature is enabled).
/// Supported schemes: HTTP, HTTPS, (SOCKS4, SOCKS5, SOCKS5H if `socks` feature is enabled).
// Private for now...
fn parse(url: Url) -> crate::Result<Self> {
use url::Position;
Expand All @@ -680,7 +702,7 @@ impl ProxyScheme {
let to_addr = || {
let addrs = url
.socket_addrs(|| match url.scheme() {
"socks5" | "socks5h" => Some(1080),
"socks4" | "socks5" | "socks5h" => Some(1080),
_ => None,
})
.map_err(crate::error::builder)?;
Expand All @@ -694,6 +716,8 @@ impl ProxyScheme {
"http" => Self::http(&url[Position::BeforeHost..Position::AfterPort])?,
"https" => Self::https(&url[Position::BeforeHost..Position::AfterPort])?,
#[cfg(feature = "socks")]
"socks4" => Self::socks4(to_addr()?)?,
#[cfg(feature = "socks")]
"socks5" => Self::socks5(to_addr()?)?,
#[cfg(feature = "socks")]
"socks5h" => Self::socks5h(to_addr()?)?,
Expand All @@ -715,6 +739,8 @@ impl ProxyScheme {
ProxyScheme::Http { .. } => "http",
ProxyScheme::Https { .. } => "https",
#[cfg(feature = "socks")]
ProxyScheme::Socks4 { .. } => "socks4",
#[cfg(feature = "socks")]
ProxyScheme::Socks5 { .. } => "socks5",
}
}
Expand All @@ -725,6 +751,8 @@ impl ProxyScheme {
ProxyScheme::Http { host, .. } => host.as_str(),
ProxyScheme::Https { host, .. } => host.as_str(),
#[cfg(feature = "socks")]
ProxyScheme::Socks4 { .. } => panic!("socks4"),
#[cfg(feature = "socks")]
ProxyScheme::Socks5 { .. } => panic!("socks5"),
}
}
Expand All @@ -736,6 +764,10 @@ impl fmt::Debug for ProxyScheme {
ProxyScheme::Http { auth: _auth, host } => write!(f, "http://{host}"),
ProxyScheme::Https { auth: _auth, host } => write!(f, "https://{host}"),
#[cfg(feature = "socks")]
ProxyScheme::Socks4 { addr } => {
write!(f, "socks4://{addr}")
}
#[cfg(feature = "socks")]
ProxyScheme::Socks5 {
addr,
auth: _auth,
Expand Down

0 comments on commit 7958847

Please sign in to comment.