use crate::rfc1928::consts;
use crate::rfc1928::util::stream::read_exact;
use crate::rfc1928::util::stream::ConnectError;
use crate::rfc1928::util::target_addr::read_address;
use crate::rfc1928::util::target_addr::TargetAddr;
use crate::rfc1928::AddrError;
use crate::rfc1928::ReplyError;
use crate::rfc1928::Socks5Command;
use crate::rfc1928::UdpHeaderError;
use std::io;
use std::marker::PhantomData;
use std::net::Ipv4Addr;
use std::net::SocketAddr;
use std::net::SocketAddrV4;
use std::string::FromUtf8Error;
use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWrite;
use tokio::io::AsyncWriteExt;
#[derive(thiserror::Error, Debug)]
pub enum SocksServerError {
#[error("i/o error when {context}: {source}")]
Io {
source: io::Error,
context: &'static str,
},
#[error("string error when {context}: {source}")]
FromUtf8 {
source: FromUtf8Error,
context: &'static str,
},
#[error(transparent)]
ConnectError(#[from] ConnectError),
#[error(transparent)]
UdpHeaderError(#[from] UdpHeaderError),
#[error(transparent)]
AddrError(#[from] AddrError),
#[error("BUG: {0}")] Bug(&'static str),
#[error("Refused: {0}")]
Refused(&'static str),
#[error("Auth method unacceptable `{0:?}`.")]
AuthMethodUnacceptable(Vec<u8>),
#[error("Unsupported SOCKS version `{0}`.")]
UnsupportedSocksVersion(u8),
#[error("Unsupported SOCKS command `{0}`.")]
UnknownCommand(u8),
#[error("Unexpected garbage received on TCP stream used for UDP proxy keep-alive: `{0}`")]
UnexpectedUdpControlGarbage(u8),
#[error("Empty username received")]
EmptyUsername,
#[error("Empty password received")]
EmptyPassword,
#[error("Authentication rejected")]
AuthenticationRejected,
#[error("End of stream")]
EOF,
}
impl SocksServerError {
pub fn to_reply_error(&self) -> ReplyError {
match self {
SocksServerError::UnknownCommand(_) => ReplyError::CommandNotSupported,
SocksServerError::AddrError(err) => err.to_reply_error(),
_ => ReplyError::GeneralFailure,
}
}
}
pub trait ErrorContext<T> {
fn err_when(self, context: &'static str) -> Result<T, SocksServerError>;
}
impl<T> ErrorContext<T> for Result<T, io::Error> {
fn err_when(self, context: &'static str) -> Result<T, SocksServerError> {
self.map_err(|source| SocksServerError::Io { source, context })
}
}
impl<T> ErrorContext<T> for Result<T, FromUtf8Error> {
fn err_when(self, context: &'static str) -> Result<T, SocksServerError> {
self.map_err(|source| SocksServerError::FromUtf8 { source, context })
}
}
pub mod states {
pub struct Opened;
pub struct Authenticated;
pub struct CommandRead;
}
pub struct Socks5ServerProtocol<T, S> {
inner: T,
_state: PhantomData<S>,
}
impl<T, S> Socks5ServerProtocol<T, S> {
fn new(inner: T) -> Self {
Socks5ServerProtocol {
inner,
_state: PhantomData,
}
}
}
impl<T> Socks5ServerProtocol<T, states::Opened> {
pub fn start(inner: T) -> Self {
Self::new(inner)
}
}
pub trait CheckResult {
fn is_good(&self) -> bool;
}
impl CheckResult for bool {
fn is_good(&self) -> bool {
*self
}
}
impl<T> Socks5ServerProtocol<T, states::Authenticated> {
pub fn finish_auth<A: AuthMethodSuccessState<T>>(auth: A) -> Self {
Self::new(auth.into_inner())
}
pub async fn accept_no_auth(inner: T) -> Result<Self, SocksServerError>
where
T: AsyncWrite + AsyncRead + Unpin,
{
Ok(Socks5ServerProtocol::start(inner)
.negotiate_auth(&[NoAuthentication])
.await?
.finish_auth())
}
pub async fn accept_password_auth<F, R>(
inner: T,
mut check: F,
) -> Result<(Self, R), SocksServerError>
where
T: AsyncWrite + AsyncRead + Unpin,
F: FnMut(String, String) -> R,
R: CheckResult,
{
let (user, pass, auth) = Socks5ServerProtocol::start(inner)
.negotiate_auth(&[PasswordAuthentication])
.await?
.read_username_password()
.await?;
let check_result = check(user, pass);
if check_result.is_good() {
Ok((auth.accept().await?.finish_auth(), check_result))
} else {
auth.reject().await?;
Err(SocksServerError::AuthenticationRejected)
}
}
}
pub trait AuthMethodSuccessState<T> {
fn into_inner(self) -> T;
fn finish_auth(self) -> Socks5ServerProtocol<T, states::Authenticated>
where
Self: Sized,
{
Socks5ServerProtocol::finish_auth(self)
}
}
pub trait AuthMethod<T>: Copy {
type StartingState;
fn method_id(self) -> u8;
fn into_starting_state(self, inner: T) -> Self::StartingState;
}
pub struct NoAuthenticationImpl<T>(T);
impl<T> AuthMethodSuccessState<T> for NoAuthenticationImpl<T> {
fn into_inner(self) -> T {
self.0
}
}
#[derive(Debug, Clone, Copy)]
pub struct NoAuthentication;
impl<T> AuthMethod<T> for NoAuthentication {
type StartingState = NoAuthenticationImpl<T>;
fn method_id(self) -> u8 {
0x00
}
fn into_starting_state(self, inner: T) -> Self::StartingState {
NoAuthenticationImpl(inner)
}
}
mod password_states {
pub struct Started;
pub struct Received;
pub struct Finished;
}
pub struct PasswordAuthenticationImpl<T, S> {
inner: T,
_state: PhantomData<S>,
}
pub type PasswordAuthenticationStarted<T> = PasswordAuthenticationImpl<T, password_states::Started>;
impl<T, S> PasswordAuthenticationImpl<T, S> {
fn new(inner: T) -> Self {
PasswordAuthenticationImpl {
inner,
_state: PhantomData,
}
}
}
impl<T: AsyncRead + Unpin> PasswordAuthenticationImpl<T, password_states::Started> {
pub async fn read_username_password(
self,
) -> Result<
(
String,
String,
PasswordAuthenticationImpl<T, password_states::Received>,
),
SocksServerError,
> {
let mut socket = self.inner;
let [_version, user_len] = read_exact!(socket, [0u8; 2]).err_when("reading user len")?;
if user_len < 1 {
return Err(SocksServerError::EmptyUsername);
}
let username =
read_exact!(socket, vec![0u8; user_len as usize]).err_when("reading username")?;
let [pass_len] = read_exact!(socket, [0u8; 1]).err_when("reading password len")?;
if pass_len < 1 {
return Err(SocksServerError::EmptyPassword);
}
let password =
read_exact!(socket, vec![0u8; pass_len as usize]).err_when("reading password")?;
let username = String::from_utf8(username).err_when("converting username")?;
let password = String::from_utf8(password).err_when("converting password")?;
Ok((username, password, PasswordAuthenticationImpl::new(socket)))
}
}
impl<T: AsyncWrite + Unpin> PasswordAuthenticationImpl<T, password_states::Received> {
pub async fn accept(
mut self,
) -> Result<PasswordAuthenticationImpl<T, password_states::Finished>, SocksServerError> {
self.inner
.write_all(&[1, consts::SOCKS5_REPLY_SUCCEEDED])
.await
.err_when("replying auth success")?;
Ok(PasswordAuthenticationImpl::new(self.inner))
}
pub async fn reject(mut self) -> Result<(), SocksServerError> {
self.inner
.write_all(&[1, consts::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE])
.await
.err_when("replying with auth method not acceptable")?;
Ok(())
}
}
impl<T> AuthMethodSuccessState<T> for PasswordAuthenticationImpl<T, password_states::Finished> {
fn into_inner(self) -> T {
self.inner
}
}
#[derive(Debug, Clone, Copy)]
pub struct PasswordAuthentication;
impl<T> AuthMethod<T> for PasswordAuthentication {
type StartingState = PasswordAuthenticationImpl<T, password_states::Started>;
fn method_id(self) -> u8 {
0x02
}
fn into_starting_state(self, inner: T) -> Self::StartingState {
PasswordAuthenticationImpl::new(inner)
}
}
macro_rules! try_notify {
($proto:expr, $e:expr) => {
match $e {
Ok(res) => res,
Err(err) => {
let _ = $proto.reply_error(&err.to_reply_error()).await;
return Err(err.into());
}
}
};
}
impl<T: AsyncRead + AsyncWrite + Unpin> Socks5ServerProtocol<T, states::Opened> {
pub async fn negotiate_auth<M: AuthMethod<T>>(
mut self,
server_methods: &[M],
) -> Result<M::StartingState, SocksServerError> {
let [version, methods_len] =
read_exact!(self.inner, [0u8; 2]).err_when("reading methods")?;
if version != consts::SOCKS5_VERSION {
return Err(SocksServerError::UnsupportedSocksVersion(version));
}
let methods =
read_exact!(self.inner, vec![0u8; methods_len as usize]).err_when("reading methods")?;
for server_method in server_methods {
for client_method_id in methods.iter() {
if server_method.method_id() == *client_method_id {
self.inner
.write_all(&[consts::SOCKS5_VERSION, *client_method_id])
.await
.err_when("replying with auth method")?;
return Ok(server_method.into_starting_state(self.inner));
}
}
}
self.inner
.write_all(&[
consts::SOCKS5_VERSION,
consts::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE,
])
.await
.err_when("replying with method not acceptable")?;
Err(SocksServerError::AuthMethodUnacceptable(methods))
}
}
impl<T: AsyncRead + AsyncWrite + Unpin> Socks5ServerProtocol<T, states::CommandRead> {
pub async fn reply_success(mut self, sock_addr: SocketAddr) -> Result<T, SocksServerError> {
self.inner
.write_all(&new_reply(&ReplyError::Succeeded, sock_addr))
.await
.err_when("writing successful reply")?;
self.inner.flush().await.err_when("flushing auth reply")?;
Ok(self.inner)
}
pub async fn reply_error(mut self, error: &ReplyError) -> Result<(), SocksServerError> {
let reply = new_reply(error, V4_UNSPECIFIED_ZERO);
self.inner
.write_all(&reply)
.await
.err_when("writing unsuccessful reply")?;
self.inner.flush().await.err_when("flushing auth reply")?;
Ok(())
}
}
impl<T: AsyncRead + AsyncWrite + Unpin> Socks5ServerProtocol<T, states::Authenticated> {
pub async fn read_command(
mut self,
) -> Result<
(
Socks5ServerProtocol<T, states::CommandRead>,
Socks5Command,
TargetAddr,
),
SocksServerError,
> {
let [version, cmd, _rsv, address_type] =
read_exact!(self.inner, [0u8; 4]).err_when("reading command")?;
if version != consts::SOCKS5_VERSION {
return Err(SocksServerError::UnsupportedSocksVersion(version));
}
let mut proto = Socks5ServerProtocol::new(self.inner);
let target_addr = try_notify!(proto, read_address(&mut proto.inner, address_type).await);
let cmd = try_notify!(
proto,
Socks5Command::from_u8(cmd).ok_or(SocksServerError::UnknownCommand(cmd))
);
Ok((proto, cmd, target_addr))
}
}
const V4_UNSPECIFIED_ZERO: SocketAddr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0));
pub fn new_reply(error: &ReplyError, sock_addr: SocketAddr) -> Vec<u8> {
let (addr_type, mut ip_oct, mut port) = match sock_addr {
SocketAddr::V4(sock) => (
consts::SOCKS5_ADDR_TYPE_IPV4,
sock.ip().octets().to_vec(),
sock.port().to_be_bytes().to_vec(),
),
SocketAddr::V6(sock) => (
consts::SOCKS5_ADDR_TYPE_IPV6,
sock.ip().octets().to_vec(),
sock.port().to_be_bytes().to_vec(),
),
};
let mut reply = vec![
consts::SOCKS5_VERSION,
error.as_u8(), 0x00, addr_type, ];
reply.append(&mut ip_oct);
reply.append(&mut port);
reply
}
pub async fn wait_on_tcp<I>(stream: &mut I) -> Result<(), SocksServerError>
where
I: AsyncRead + Unpin,
{
let mut buf = [0; 1];
match stream.read(&mut buf).await {
Ok(0) => Err(SocksServerError::EOF),
Ok(_) => Err(SocksServerError::UnexpectedUdpControlGarbage(buf[0])),
Err(err) => Err(err).err_when("waiting on UDP control stream"),
}
}
pub async fn transfer<I, O>(mut inbound: I, mut outbound: O)
where
I: AsyncRead + AsyncWrite + Unpin,
O: AsyncRead + AsyncWrite + Unpin,
{
let _ = tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await;
}
#[cfg(test)]
mod tests {
use super::new_reply;
use super::states;
use super::Socks5ServerProtocol;
use super::SocksServerError;
use crate::rfc1928::util::target_addr::TargetAddr;
use crate::rfc1928::ReplyError;
use crate::rfc1928::Socks5Command;
use std::net::SocketAddr;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
#[test]
fn to_reply_error_mapping_table() {
assert!(matches!(
SocksServerError::UnknownCommand(0x7f).to_reply_error(),
ReplyError::CommandNotSupported
));
assert!(matches!(
SocksServerError::AddrError(crate::rfc1928::AddrError::IncorrectAddressType)
.to_reply_error(),
ReplyError::AddressTypeNotSupported
));
for fall_through in [
SocksServerError::EOF,
SocksServerError::EmptyUsername,
SocksServerError::EmptyPassword,
SocksServerError::AuthenticationRejected,
SocksServerError::UnsupportedSocksVersion(0x04),
SocksServerError::Bug("test"),
SocksServerError::Refused("test"),
SocksServerError::AuthMethodUnacceptable(vec![0x02]),
SocksServerError::UnexpectedUdpControlGarbage(0xde),
SocksServerError::Io {
source: std::io::Error::other("x"),
context: "reading",
},
SocksServerError::FromUtf8 {
source: String::from_utf8(vec![0xff]).unwrap_err(),
context: "converting",
},
] {
assert!(matches!(
fall_through.to_reply_error(),
ReplyError::GeneralFailure
));
}
}
#[test]
fn refused_variant_display_is_honest_and_wire_invisible() {
assert_eq!(
SocksServerError::Refused("policy says no").to_string(),
"Refused: policy says no"
);
assert!(matches!(
SocksServerError::Refused("policy says no").to_reply_error(),
ReplyError::GeneralFailure
));
}
#[test]
fn new_reply_v4_form() {
let addr: SocketAddr = "127.0.0.1:1080".parse().unwrap();
let reply = new_reply(&ReplyError::Succeeded, addr);
assert_eq!(reply, vec![0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, 4, 56]);
assert_eq!(reply.len(), 10);
}
#[test]
fn new_reply_v6_form() {
let addr: SocketAddr = "[2001:db8::1]:80".parse().unwrap();
let reply = new_reply(&ReplyError::CommandNotSupported, addr);
assert_eq!(reply.len(), 22);
assert_eq!(reply[0], 0x05);
assert_eq!(reply[1], 0x07);
assert_eq!(reply[2], 0x00);
assert_eq!(reply[3], 0x04);
assert_eq!(
&reply[4..20],
&0x2001_0db8_0000_0000_0000_0000_0000_0001u128.to_be_bytes()
);
assert_eq!(&reply[20..22], &80u16.to_be_bytes());
}
#[test]
fn new_reply_error_codes_are_carried() {
let addr: SocketAddr = "10.0.0.1:1".parse().unwrap();
for (code, byte) in [
(ReplyError::Succeeded, 0x00),
(ReplyError::GeneralFailure, 0x01),
(ReplyError::ConnectionNotAllowed, 0x02),
(ReplyError::NetworkUnreachable, 0x03),
(ReplyError::HostUnreachable, 0x04),
(ReplyError::ConnectionRefused, 0x05),
(ReplyError::ConnectionTimeout, 0x06),
(ReplyError::TtlExpired, 0x06),
(ReplyError::CommandNotSupported, 0x07),
(ReplyError::AddressTypeNotSupported, 0x08),
] {
assert_eq!(new_reply(&code, addr)[1], byte);
}
}
#[tokio::test]
async fn version_mismatch_is_refused() {
let (mut client, server) = tokio::io::duplex(64);
let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
tokio::spawn(async move {
Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
.await
.map(|_| ())
});
client.write_all(&[0x04, 0x01, 0x00]).await.unwrap();
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(
err,
SocksServerError::UnsupportedSocksVersion(0x04)
));
}
#[tokio::test]
async fn no_acceptable_method_replies_0xff() {
let (mut client, server) = tokio::io::duplex(64);
let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
tokio::spawn(async move {
Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
.await
.map(|_| ())
});
client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
let mut method_reply = [0u8; 2];
client.read_exact(&mut method_reply).await.unwrap();
assert_eq!(method_reply, [0x05, 0xff]);
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksServerError::AuthMethodUnacceptable(ref m) if m == &vec![0x02]));
}
#[tokio::test]
async fn unknown_command_replies_0x07() {
let (mut client, server) = tokio::io::duplex(64);
let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
tokio::spawn(async move {
let proto =
Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
.await
.unwrap();
proto.read_command().await.map(|_| ())
});
client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
let mut method_reply = [0u8; 2];
client.read_exact(&mut method_reply).await.unwrap();
assert_eq!(method_reply, [0x05, 0x00]);
client
.write_all(&[0x05, 0x7f, 0x00, 0x01, 192, 168, 0, 1, 0, 80])
.await
.unwrap();
let mut reply = [0u8; 10];
client.read_exact(&mut reply).await.unwrap();
assert_eq!(reply, [0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksServerError::UnknownCommand(0x7f)));
}
#[tokio::test]
async fn success_reply_form_round_trip() {
let (mut client, server) = tokio::io::duplex(64);
let handle: tokio::task::JoinHandle<Result<(Socks5Command, TargetAddr), SocksServerError>> =
tokio::spawn(async move {
let proto =
Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
.await
.unwrap();
let (proto, cmd, target) = proto.read_command().await.unwrap();
let bound: SocketAddr = "198.51.100.7:443".parse().unwrap();
let mut inner = proto.reply_success(bound).await.unwrap();
inner.write_all(b"ok").await.unwrap();
Ok((cmd, target))
});
client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
let mut method_reply = [0u8; 2];
client.read_exact(&mut method_reply).await.unwrap();
assert_eq!(method_reply, [0x05, 0x00]);
client
.write_all(&[0x05, 0x01, 0x00, 0x01, 192, 0, 2, 1, 0, 80])
.await
.unwrap();
let mut reply = [0u8; 10];
client.read_exact(&mut reply).await.unwrap();
assert_eq!(reply, [0x05, 0x00, 0x00, 0x01, 198, 51, 100, 7, 1, 187]);
let mut data = [0u8; 2];
client.read_exact(&mut data).await.unwrap();
assert_eq!(&data, b"ok");
let (cmd, target) = handle.await.unwrap().unwrap();
assert_eq!(cmd, Socks5Command::TCPConnect);
assert_eq!(target, TargetAddr::Ip("192.0.2.1:80".parse().unwrap()));
}
#[tokio::test]
async fn error_reply_form_round_trip() {
let (mut client, server) = tokio::io::duplex(64);
let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
tokio::spawn(async move {
let proto =
Socks5ServerProtocol::<_, states::Authenticated>::accept_no_auth(server)
.await
.unwrap();
let (proto, _cmd, _target) = proto.read_command().await.unwrap();
proto.reply_error(&ReplyError::ConnectionNotAllowed).await
});
client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
let mut method_reply = [0u8; 2];
client.read_exact(&mut method_reply).await.unwrap();
assert_eq!(method_reply, [0x05, 0x00]);
client
.write_all(&[0x05, 0x01, 0x00, 0x03, 3, b'i', b'p', b'v', b'6', 0, 80])
.await
.unwrap();
let mut reply = [0u8; 10];
client.read_exact(&mut reply).await.unwrap();
assert_eq!(reply, [0x05, 0x02, 0x00, 0x01, 0, 0, 0, 0, 0, 0]);
handle.await.unwrap().unwrap();
}
#[tokio::test]
async fn password_auth_accepts_and_rejects() {
let (mut client, server) = tokio::io::duplex(64);
let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
tokio::spawn(async move {
let (proto, _ok) =
Socks5ServerProtocol::<_, states::Authenticated>::accept_password_auth(
server,
|user, pass| user == "alice" && pass == "secret",
)
.await?;
drop(proto);
Ok(())
});
client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
let mut method_reply = [0u8; 2];
client.read_exact(&mut method_reply).await.unwrap();
assert_eq!(method_reply, [0x05, 0x02]);
client
.write_all(&[
0x01, 0x05, b'a', b'l', b'i', b'c', b'e', 0x06, b's', b'e', b'c', b'r', b'e', b't',
])
.await
.unwrap();
let mut status = [0u8; 2];
client.read_exact(&mut status).await.unwrap();
assert_eq!(status, [0x01, 0x00]);
handle.await.unwrap().unwrap();
let (mut client, server) = tokio::io::duplex(64);
let handle: tokio::task::JoinHandle<Result<(), SocksServerError>> =
tokio::spawn(async move {
Socks5ServerProtocol::<_, states::Authenticated>::accept_password_auth(
server,
|_user, _pass| false,
)
.await
.map(|_| ())
});
client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
let mut method_reply = [0u8; 2];
client.read_exact(&mut method_reply).await.unwrap();
assert_eq!(method_reply, [0x05, 0x02]);
client
.write_all(&[0x01, 0x03, b'b', b'o', b'b', 0x01, b'x'])
.await
.unwrap();
let mut status = [0u8; 2];
client.read_exact(&mut status).await.unwrap();
assert_eq!(status, [0x01, 0xff]);
let err = handle.await.unwrap().unwrap_err();
assert!(matches!(err, SocksServerError::AuthenticationRejected));
}
}