use crate::rfc1928::util::stream::read_exact;
use crate::rfc1928::util::target_addr::{read_address, TargetAddr};
use crate::rfc1928::{consts, AuthenticationMethod, ReplyError, Result, Socks5Command, SocksError};
use std::io;
use std::net::SocketAddr;
use std::pin::Pin;
use std::task::Poll;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
const MAX_ADDR_LEN: usize = 260;
#[derive(Debug, Default)]
pub struct Config {
skip_auth: bool,
}
impl Config {
pub fn set_skip_auth(&mut self, value: bool) -> &mut Self {
self.skip_auth = value;
self
}
}
#[derive(Debug)]
pub struct Socks5Stream<S: AsyncRead + AsyncWrite + Unpin> {
socket: S,
target_addr: Option<TargetAddr>,
config: Config,
}
impl<S> Socks5Stream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
pub async fn use_stream(
socket: S,
auth: Option<AuthenticationMethod>,
config: Config,
) -> Result<Self> {
let mut stream = Socks5Stream {
socket,
config,
target_addr: None,
};
let mut methods = vec![AuthenticationMethod::None];
if let Some(method) = auth {
methods.push(method);
}
if !stream.config.skip_auth {
let methods = stream.send_version_and_methods(methods).await?;
stream.which_method_accepted(methods).await?;
}
Ok(stream)
}
pub async fn request(
&mut self,
cmd: Socks5Command,
target_addr: TargetAddr,
) -> Result<TargetAddr> {
self.target_addr = Some(target_addr);
self.request_header(cmd).await?;
let bind_addr = self.read_request_reply().await?;
Ok(bind_addr)
}
async fn send_version_and_methods(
&mut self,
methods: Vec<AuthenticationMethod>,
) -> Result<Vec<AuthenticationMethod>> {
let mut packet = vec![consts::SOCKS5_VERSION, methods.len() as u8];
let auth = methods.iter().map(|l| l.as_u8()).collect::<Vec<_>>();
packet.extend(auth);
self.socket.write_all(&packet).await?;
Ok(methods)
}
async fn which_method_accepted(&mut self, methods: Vec<AuthenticationMethod>) -> Result<()> {
let [version, method] = read_exact!(self.socket, [0u8; 2])?;
if version != consts::SOCKS5_VERSION {
return Err(SocksError::UnsupportedSocksVersion(version));
}
match method {
consts::SOCKS5_AUTH_METHOD_NONE => {}
consts::SOCKS5_AUTH_METHOD_PASSWORD => self.use_password_auth(methods).await?,
_ => {
self.socket
.write_all(&[
consts::SOCKS5_VERSION,
consts::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE,
])
.await?;
return Err(SocksError::AuthMethodUnacceptable(vec![method]));
}
}
Ok(())
}
async fn use_password_auth(&mut self, methods: Vec<AuthenticationMethod>) -> Result<()> {
let (username, password) = match methods.get(1) {
Some(AuthenticationMethod::Password {
ref username,
ref password,
}) => (username, password),
Some(AuthenticationMethod::None) | None => {
return Err(SocksError::AuthenticationRejected(
"Authentication rejected, missing user pass".to_owned(),
))
}
};
let user_bytes = username.as_bytes();
let pass_bytes = password.as_bytes();
let mut packet: Vec<u8> = vec![1, user_bytes.len() as u8];
packet.extend(user_bytes);
packet.push(pass_bytes.len() as u8);
packet.extend(pass_bytes);
self.socket.write_all(&packet).await?;
let [_version, is_success] = read_exact!(self.socket, [0u8; 2])?;
if is_success != consts::SOCKS5_REPLY_SUCCEEDED {
return Err(SocksError::AuthenticationRejected(format!(
"Authentication with username `{}`, rejected.",
username
)));
}
Ok(())
}
async fn request_header(&mut self, cmd: Socks5Command) -> Result<()> {
let mut packet = [0u8; MAX_ADDR_LEN + 3];
let padding; packet[..3].copy_from_slice(&[consts::SOCKS5_VERSION, cmd.as_u8(), 0x00]);
match self.target_addr.as_ref() {
None => {
if cmd == Socks5Command::UDPAssociate {
padding = 10;
packet[3] = 0x01;
packet[4..8].copy_from_slice(&[0, 0, 0, 0]); packet[8..padding].copy_from_slice(&[0, 0]); } else {
return Err(SocksError::MissingTargetAddr);
}
}
Some(target_addr) => match target_addr {
TargetAddr::Ip(SocketAddr::V4(addr)) => {
padding = 10;
packet[3] = 0x01;
packet[4..8].copy_from_slice(&(addr.ip()).octets()); packet[8..padding].copy_from_slice(&addr.port().to_be_bytes());
}
TargetAddr::Ip(SocketAddr::V6(addr)) => {
padding = 22;
packet[3] = 0x04;
packet[4..20].copy_from_slice(&(addr.ip()).octets()); packet[20..padding].copy_from_slice(&addr.port().to_be_bytes());
}
TargetAddr::Domain(ref domain, port) => {
if domain.len() > u8::MAX as usize {
return Err(SocksError::ExceededMaxDomainLen(domain.len()));
}
padding = 5 + domain.len() + 2;
packet[3] = 0x03; packet[4] = domain.len() as u8; packet[5..(5 + domain.len())].copy_from_slice(domain.as_bytes()); packet[(5 + domain.len())..padding].copy_from_slice(&port.to_be_bytes());
}
},
}
self.socket.write_all(&packet[..padding]).await?;
self.socket.flush().await?;
Ok(())
}
async fn read_request_reply(&mut self) -> Result<TargetAddr> {
let [version, reply, _rsv, address_type] = read_exact!(self.socket, [0u8; 4])?;
if version != consts::SOCKS5_VERSION {
return Err(SocksError::UnsupportedSocksVersion(version));
}
if reply != consts::SOCKS5_REPLY_SUCCEEDED {
return Err(ReplyError::from_u8(reply).into()); }
let address = read_address(&mut self.socket, address_type).await?;
Ok(address)
}
pub fn get_socket(self) -> S {
self.socket
}
pub fn get_socket_ref(&self) -> &S {
&self.socket
}
pub fn get_socket_mut(&mut self) -> &mut S {
&mut self.socket
}
}
impl<S> AsyncRead for Socks5Stream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_read(
mut self: Pin<&mut Self>,
context: &mut std::task::Context,
buf: &mut tokio::io::ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.socket).poll_read(context, buf)
}
}
impl<S> AsyncWrite for Socks5Stream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_write(
mut self: Pin<&mut Self>,
context: &mut std::task::Context,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.socket).poll_write(context, buf)
}
fn poll_flush(
mut self: Pin<&mut Self>,
context: &mut std::task::Context,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.socket).poll_flush(context)
}
fn poll_shutdown(
mut self: Pin<&mut Self>,
context: &mut std::task::Context,
) -> Poll<io::Result<()>> {
Pin::new(&mut self.socket).poll_shutdown(context)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::DuplexStream;
type ClientTask<T> = tokio::task::JoinHandle<Result<T, SocksError>>;
#[tokio::test]
async fn no_auth_handshake_and_connect_success() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<(Socks5Stream<DuplexStream>, TargetAddr)> =
tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(server, None, Config::default()).await?;
let bind = session
.request(
Socks5Command::TCPConnect,
TargetAddr::Domain("example.com".to_owned(), 443),
)
.await?;
Ok((session, bind))
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x05, 0x00]).await.unwrap();
let mut request = [0u8; 18];
mock.read_exact(&mut request).await.unwrap();
assert_eq!(
request,
[
0x05, 0x01, 0x00, 0x03, 11, b'e', b'x', b'a', b'm', b'p', b'l', b'e', b'.', b'c',
b'o', b'm', 0x01, 0xbb
]
);
mock.write_all(&[0x05, 0x00, 0x00, 0x01, 93, 184, 216, 34, 0x01, 0xbb])
.await
.unwrap();
let (mut session, bind) = handle.await.unwrap().unwrap();
assert_eq!(bind, TargetAddr::Ip("93.184.216.34:443".parse().unwrap()));
session.write_all(b"ping").await.unwrap();
let mut echoed = [0u8; 4];
mock.read_exact(&mut echoed).await.unwrap();
assert_eq!(&echoed, b"ping");
mock.write_all(b"pong").await.unwrap();
let mut got = [0u8; 4];
session.read_exact(&mut got).await.unwrap();
assert_eq!(&got, b"pong");
}
#[tokio::test]
async fn password_auth_accept_path() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<TargetAddr> = tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(
server,
Some(AuthenticationMethod::Password {
username: "alice".to_owned(),
password: "secret".to_owned(),
}),
Config::default(),
)
.await?;
session
.request(
Socks5Command::TCPConnect,
TargetAddr::Ip("192.0.2.1:80".parse().unwrap()),
)
.await
});
let mut greeting = [0u8; 4];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x02, 0x00, 0x02]);
mock.write_all(&[0x05, 0x02]).await.unwrap();
let mut auth_packet = [0u8; 14];
mock.read_exact(&mut auth_packet).await.unwrap();
assert_eq!(
auth_packet,
[0x01, 5, b'a', b'l', b'i', b'c', b'e', 6, b's', b'e', b'c', b'r', b'e', b't']
);
mock.write_all(&[0x01, 0x00]).await.unwrap();
let mut request = [0u8; 10];
mock.read_exact(&mut request).await.unwrap();
assert_eq!(request, [0x05, 0x01, 0x00, 0x01, 192, 0, 2, 1, 0x00, 0x50]);
mock.write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
.await
.unwrap();
let bind = handle.await.unwrap().unwrap();
assert_eq!(bind, TargetAddr::Ip("0.0.0.0:0".parse().unwrap()));
}
#[tokio::test]
async fn password_auth_reject_path() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
Socks5Stream::use_stream(
server,
Some(AuthenticationMethod::Password {
username: "bob".to_owned(),
password: "x".to_owned(),
}),
Config::default(),
)
.await?;
Ok(())
});
let mut greeting = [0u8; 4];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x02, 0x00, 0x02]);
mock.write_all(&[0x05, 0x02]).await.unwrap();
let mut auth_packet = [0u8; 7];
mock.read_exact(&mut auth_packet).await.unwrap();
assert_eq!(auth_packet, [0x01, 3, b'b', b'o', b'b', 1, b'x']);
mock.write_all(&[0x01, 0xff]).await.unwrap();
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksError::AuthenticationRejected(ref m) if m.contains("bob")));
}
#[tokio::test]
async fn method_unacceptable_replies_0xff() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
Socks5Stream::use_stream(
server,
Some(AuthenticationMethod::Password {
username: "u".to_owned(),
password: "p".to_owned(),
}),
Config::default(),
)
.await?;
Ok(())
});
let mut greeting = [0u8; 4];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x02, 0x00, 0x02]);
mock.write_all(&[0x05, 0x03]).await.unwrap();
let mut reply = [0u8; 2];
mock.read_exact(&mut reply).await.unwrap();
assert_eq!(reply, [0x05, 0xff]);
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksError::AuthMethodUnacceptable(ref m) if m == &vec![0x03]));
}
#[tokio::test]
async fn version_mismatch_is_refused() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
Socks5Stream::use_stream(server, None, Config::default()).await?;
Ok(())
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x04, 0x00]).await.unwrap();
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksError::UnsupportedSocksVersion(0x04)));
}
#[tokio::test]
async fn missing_target_addr_is_a_typed_error() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(server, None, Config::default()).await?;
session.target_addr = None;
session.request_header(Socks5Command::TCPConnect).await?;
Ok(())
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x05, 0x00]).await.unwrap();
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksError::MissingTargetAddr));
}
#[tokio::test]
async fn udp_associate_without_target_falls_back_to_zeros() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(server, None, Config::default()).await?;
session.target_addr = None;
session.request_header(Socks5Command::UDPAssociate).await?;
Ok(())
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x05, 0x00]).await.unwrap();
let mut request = [0u8; 10];
mock.read_exact(&mut request).await.unwrap();
assert_eq!(request, [0x05, 0x03, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
handle.await.unwrap().unwrap();
}
#[tokio::test]
async fn skip_auth_bypasses_the_greeting() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
let mut config = Config::default();
config.set_skip_auth(true);
let mut session = Socks5Stream::use_stream(server, None, config).await?;
session
.request(
Socks5Command::TCPConnect,
TargetAddr::Ip("10.0.0.1:80".parse().unwrap()),
)
.await?;
Ok(())
});
let mut request = [0u8; 10];
mock.read_exact(&mut request).await.unwrap();
assert_eq!(request, [0x05, 0x01, 0x00, 0x01, 10, 0, 0, 1, 0x00, 0x50]);
mock.write_all(&[0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
.await
.unwrap();
handle.await.unwrap().unwrap();
}
#[tokio::test]
async fn mapped_reply_codes_surface_as_reply_error() {
let codes = [
consts::SOCKS5_REPLY_GENERAL_FAILURE,
consts::SOCKS5_REPLY_CONNECTION_NOT_ALLOWED,
consts::SOCKS5_REPLY_NETWORK_UNREACHABLE,
consts::SOCKS5_REPLY_HOST_UNREACHABLE,
consts::SOCKS5_REPLY_CONNECTION_REFUSED,
consts::SOCKS5_REPLY_TTL_EXPIRED,
consts::SOCKS5_REPLY_COMMAND_NOT_SUPPORTED,
consts::SOCKS5_REPLY_ADDRESS_TYPE_NOT_SUPPORTED,
];
for code in codes {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(server, None, Config::default()).await?;
session
.request(
Socks5Command::TCPConnect,
TargetAddr::Domain("example.com".to_owned(), 80),
)
.await?;
Ok(())
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x05, 0x00]).await.unwrap();
let mut request = [0u8; 18];
mock.read_exact(&mut request).await.unwrap();
assert_eq!(&request[..4], &[0x05, 0x01, 0x00, 0x03]);
mock.write_all(&[0x05, code, 0x00, 0x01, 0, 0, 0, 0, 0, 0])
.await
.unwrap();
let err = handle.await.unwrap().unwrap_err();
match err {
SocksError::ReplyError(reply) => assert_eq!(reply.as_u8(), code),
other => panic!("expected ReplyError, got {other:?}"),
}
}
}
#[tokio::test]
async fn domain_target_over_255_bytes_is_rejected() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(server, None, Config::default()).await?;
let long_domain = "a".repeat(256);
session
.request(
Socks5Command::TCPConnect,
TargetAddr::Domain(long_domain, 80),
)
.await?;
Ok(())
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x05, 0x00]).await.unwrap();
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksError::ExceededMaxDomainLen(256)));
}
#[tokio::test]
async fn io_impls_pipe_data_through_the_session() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<Vec<u8>> = tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(server, None, Config::default()).await?;
session
.request(
Socks5Command::TCPConnect,
TargetAddr::Domain("example.com".to_owned(), 443),
)
.await?;
let mut payload: &[u8] = b"ping through copy";
let written = tokio::io::copy(&mut payload, &mut session).await?;
assert_eq!(written, 17);
session.flush().await?;
let mut out = vec![0u8; 17];
session.read_exact(&mut out).await?;
Ok(out)
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x05, 0x00]).await.unwrap();
let mut request = [0u8; 18];
mock.read_exact(&mut request).await.unwrap();
assert_eq!(&request[..4], &[0x05, 0x01, 0x00, 0x03]);
mock.write_all(&[0x05, 0x00, 0x00, 0x01, 93, 184, 216, 34, 0x01, 0xbb])
.await
.unwrap();
let mut echoed = [0u8; 17];
mock.read_exact(&mut echoed).await.unwrap();
assert_eq!(&echoed, b"ping through copy");
mock.write_all(b"ping through copy").await.unwrap();
let out = handle.await.unwrap().unwrap();
assert_eq!(out, b"ping through copy");
}
#[tokio::test]
async fn socket_accessors_expose_the_backing_stream() {
let (mut mock, server) = tokio::io::duplex(64);
let handle: ClientTask<()> = tokio::spawn(async move {
let mut session = Socks5Stream::use_stream(server, None, Config::default()).await?;
let _borrowed: &DuplexStream = session.get_socket_ref();
session.get_socket_mut().write_all(b"done").await?;
let mut socket = session.get_socket();
let mut buf = [0u8; 4];
socket.read_exact(&mut buf).await?;
Ok(())
});
let mut greeting = [0u8; 3];
mock.read_exact(&mut greeting).await.unwrap();
assert_eq!(greeting, [0x05, 0x01, 0x00]);
mock.write_all(&[0x05, 0x00]).await.unwrap();
let mut data = [0u8; 4];
mock.read_exact(&mut data).await.unwrap();
assert_eq!(&data, b"done");
mock.write_all(b"done").await.unwrap();
handle.await.unwrap().unwrap();
}
}