Compare commits
1 Commits
feat/trans
...
feat/clien
| Author | SHA1 | Date | |
|---|---|---|---|
| 2efd4cf7c5 |
490
crates/wraith-core/src/socks5/mod.rs
Normal file
490
crates/wraith-core/src/socks5/mod.rs
Normal file
@@ -0,0 +1,490 @@
|
|||||||
|
mod protocol;
|
||||||
|
|
||||||
|
use std::net::SocketAddr;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
use tokio::sync::Mutex;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
use protocol::{Socks5Reply, Socks5Request, Socks5VersionMethod};
|
||||||
|
|
||||||
|
pub use protocol::Socks5Address;
|
||||||
|
|
||||||
|
const DEFAULT_SOCKS5_ADDR: &str = "127.0.0.1:1080";
|
||||||
|
|
||||||
|
pub trait ChannelOpener: Send + Sync + 'static {
|
||||||
|
type Stream: AsyncRead + AsyncWrite + Unpin + Send + 'static;
|
||||||
|
|
||||||
|
fn open_channel(
|
||||||
|
&self,
|
||||||
|
host: String,
|
||||||
|
port: u16,
|
||||||
|
) -> impl std::future::Future<Output = Result<Self::Stream, ChannelOpenError>> + Send;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum ChannelOpenError {
|
||||||
|
#[error("session closed")]
|
||||||
|
SessionClosed,
|
||||||
|
#[error("channel open failed")]
|
||||||
|
ChannelOpenFailed,
|
||||||
|
#[error("connection refused")]
|
||||||
|
ConnectionRefused,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct Socks5Server<C: ChannelOpener> {
|
||||||
|
listen_addr: SocketAddr,
|
||||||
|
channel_opener: Arc<C>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<C: ChannelOpener> Socks5Server<C> {
|
||||||
|
pub fn new(channel_opener: C) -> Self {
|
||||||
|
Self::with_addr(channel_opener, DEFAULT_SOCKS5_ADDR)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_addr(channel_opener: C, addr: &str) -> Self {
|
||||||
|
let listen_addr: SocketAddr = addr
|
||||||
|
.parse()
|
||||||
|
.expect("invalid SOCKS5 listen address");
|
||||||
|
Self {
|
||||||
|
listen_addr,
|
||||||
|
channel_opener: Arc::new(channel_opener),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn listen_addr(&self) -> SocketAddr {
|
||||||
|
self.listen_addr
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn run(self) -> Result<(), std::io::Error> {
|
||||||
|
let listener = TcpListener::bind(self.listen_addr).await?;
|
||||||
|
debug!("socks5 server listening on {}", self.listen_addr);
|
||||||
|
loop {
|
||||||
|
let (socket, _peer) = listener.accept().await?;
|
||||||
|
let opener = Arc::clone(&self.channel_opener);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Err(e) = handle_socks5_connection(socket, opener).await {
|
||||||
|
debug!("socks5 connection error: {e}");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn handle_socks5_connection<S, C>(
|
||||||
|
mut socket: S,
|
||||||
|
opener: Arc<C>,
|
||||||
|
) -> Result<(), Socks5Error>
|
||||||
|
where
|
||||||
|
S: AsyncRead + AsyncWrite + Unpin,
|
||||||
|
C: ChannelOpener,
|
||||||
|
{
|
||||||
|
let vm = Socks5VersionMethod::read_from(&mut socket).await?;
|
||||||
|
if vm.version != 0x05 {
|
||||||
|
return Err(Socks5Error::InvalidVersion(vm.version));
|
||||||
|
}
|
||||||
|
if !vm.methods.contains(&0x00) {
|
||||||
|
let reply = [0x05, 0xFF];
|
||||||
|
socket.write_all(&reply).await?;
|
||||||
|
socket.shutdown().await?;
|
||||||
|
return Err(Socks5Error::NoAcceptableAuth);
|
||||||
|
}
|
||||||
|
let reply = [0x05, 0x00];
|
||||||
|
socket.write_all(&reply).await?;
|
||||||
|
|
||||||
|
let request = Socks5Request::read_from(&mut socket).await?;
|
||||||
|
if request.version != 0x05 {
|
||||||
|
return Err(Socks5Error::InvalidVersion(request.version));
|
||||||
|
}
|
||||||
|
if request.command != 0x01 {
|
||||||
|
send_error_reply(&mut socket, Socks5Reply::command_not_supported()).await?;
|
||||||
|
return Err(Socks5Error::UnsupportedCommand(request.command));
|
||||||
|
}
|
||||||
|
|
||||||
|
let (host, port) = match &request.address {
|
||||||
|
Socks5Address::Ipv4(addr) => (addr.to_string(), request.port),
|
||||||
|
Socks5Address::Ipv6(addr) => (addr.to_string(), request.port),
|
||||||
|
Socks5Address::Domain(name) => (name.clone(), request.port),
|
||||||
|
};
|
||||||
|
|
||||||
|
match opener.open_channel(host, port).await {
|
||||||
|
Ok(mut ssh_stream) => {
|
||||||
|
let bind_addr = Socks5Address::Ipv4(std::net::Ipv4Addr::UNSPECIFIED);
|
||||||
|
let reply = Socks5Reply::success(bind_addr, 0);
|
||||||
|
reply.write_to(&mut socket).await?;
|
||||||
|
tokio::io::copy_bidirectional(&mut socket, &mut ssh_stream).await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
send_error_reply(&mut socket, Socks5Reply::connection_refused()).await?;
|
||||||
|
Err(Socks5Error::ChannelOpenFailed)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn send_error_reply<S: AsyncRead + AsyncWrite + Unpin>(
|
||||||
|
socket: &mut S,
|
||||||
|
reply: Socks5Reply,
|
||||||
|
) -> Result<(), Socks5Error> {
|
||||||
|
reply.write_to(socket).await?;
|
||||||
|
let _ = socket.shutdown().await;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum Socks5Error {
|
||||||
|
#[error("invalid SOCKS version: {0}")]
|
||||||
|
InvalidVersion(u8),
|
||||||
|
#[error("no acceptable auth method")]
|
||||||
|
NoAcceptableAuth,
|
||||||
|
#[error("unsupported command: {0}")]
|
||||||
|
UnsupportedCommand(u8),
|
||||||
|
#[error("channel open failed")]
|
||||||
|
ChannelOpenFailed,
|
||||||
|
#[error("io error")]
|
||||||
|
Io(#[from] std::io::Error),
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct HandleChannelOpener<H: russh::client::Handler> {
|
||||||
|
handle: Arc<Mutex<russh::client::Handle<H>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<H: russh::client::Handler> HandleChannelOpener<H> {
|
||||||
|
pub fn new(handle: russh::client::Handle<H>) -> Self {
|
||||||
|
Self {
|
||||||
|
handle: Arc::new(Mutex::new(handle)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn from_arc(handle: Arc<Mutex<russh::client::Handle<H>>>) -> Self {
|
||||||
|
Self { handle }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<H: russh::client::Handler + Send + Sync + 'static> ChannelOpener for HandleChannelOpener<H> {
|
||||||
|
type Stream = russh::ChannelStream<russh::client::Msg>;
|
||||||
|
|
||||||
|
async fn open_channel(&self, host: String, port: u16) -> Result<Self::Stream, ChannelOpenError> {
|
||||||
|
let handle = self.handle.lock().await;
|
||||||
|
if handle.is_closed() {
|
||||||
|
return Err(ChannelOpenError::SessionClosed);
|
||||||
|
}
|
||||||
|
let channel = handle
|
||||||
|
.channel_open_direct_tcpip(host, port as u32, "127.0.0.1", 0)
|
||||||
|
.await
|
||||||
|
.map_err(|_| ChannelOpenError::ChannelOpenFailed)?;
|
||||||
|
Ok(channel.into_stream())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use tokio::io::{duplex, AsyncReadExt, AsyncWriteExt, DuplexStream};
|
||||||
|
|
||||||
|
struct MockChannelOpener {
|
||||||
|
fail: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChannelOpener for MockChannelOpener {
|
||||||
|
type Stream = DuplexStream;
|
||||||
|
|
||||||
|
async fn open_channel(
|
||||||
|
&self,
|
||||||
|
_host: String,
|
||||||
|
_port: u16,
|
||||||
|
) -> Result<Self::Stream, ChannelOpenError> {
|
||||||
|
if self.fail {
|
||||||
|
Err(ChannelOpenError::ChannelOpenFailed)
|
||||||
|
} else {
|
||||||
|
let (client, _server) = duplex(4096);
|
||||||
|
Ok(client)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_socks5_greeting(methods: &[u8]) -> Vec<u8> {
|
||||||
|
let mut buf = vec![0x05, methods.len() as u8];
|
||||||
|
buf.extend_from_slice(methods);
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_socks5_connect_ipv4(addr: [u8; 4], port: u16) -> Vec<u8> {
|
||||||
|
let mut buf = vec![0x05, 0x01, 0x00, 0x01];
|
||||||
|
buf.extend_from_slice(&addr);
|
||||||
|
buf.extend_from_slice(&port.to_be_bytes());
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_socks5_connect_domain(domain: &str, port: u16) -> Vec<u8> {
|
||||||
|
let mut buf = vec![0x05, 0x01, 0x00, 0x03];
|
||||||
|
buf.push(domain.len() as u8);
|
||||||
|
buf.extend_from_slice(domain.as_bytes());
|
||||||
|
buf.extend_from_slice(&port.to_be_bytes());
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_socks5_connect_ipv6(addr: [u8; 16], port: u16) -> Vec<u8> {
|
||||||
|
let mut buf = vec![0x05, 0x01, 0x00, 0x04];
|
||||||
|
buf.extend_from_slice(&addr);
|
||||||
|
buf.extend_from_slice(&port.to_be_bytes());
|
||||||
|
buf
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn do_handshake(client: &mut DuplexStream) -> [u8; 2] {
|
||||||
|
client.write_all(&build_socks5_greeting(&[0x00])).await.unwrap();
|
||||||
|
client.flush().await.unwrap();
|
||||||
|
let mut resp = [0u8; 2];
|
||||||
|
client.read_exact(&mut resp).await.unwrap();
|
||||||
|
resp
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn do_connect_ipv4(client: &mut DuplexStream, addr: [u8; 4], port: u16) -> Vec<u8> {
|
||||||
|
client
|
||||||
|
.write_all(&build_socks5_connect_ipv4(addr, port))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
client.flush().await.unwrap();
|
||||||
|
let mut reply_buf = [0u8; 10];
|
||||||
|
client.read_exact(&mut reply_buf).await.unwrap();
|
||||||
|
reply_buf.to_vec()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn handshake_no_auth_method() {
|
||||||
|
let (mut client, server) = duplex(4096);
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
let resp = do_handshake(&mut client).await;
|
||||||
|
assert_eq!(resp, [0x05, 0x00]);
|
||||||
|
|
||||||
|
let reply_buf = do_connect_ipv4(&mut client, [127, 0, 0, 1], 80).await;
|
||||||
|
assert_eq!(reply_buf[0], 0x05);
|
||||||
|
assert_eq!(reply_buf[1], 0x00);
|
||||||
|
|
||||||
|
drop(client);
|
||||||
|
let _ = server_handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn handshake_rejects_no_acceptable_method() {
|
||||||
|
let (mut client, server) = duplex(4096);
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
client
|
||||||
|
.write_all(&build_socks5_greeting(&[0x02]))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
client.flush().await.unwrap();
|
||||||
|
|
||||||
|
let mut resp = [0u8; 2];
|
||||||
|
client.read_exact(&mut resp).await.unwrap();
|
||||||
|
assert_eq!(resp, [0x05, 0xFF]);
|
||||||
|
|
||||||
|
drop(client);
|
||||||
|
let result = server_handle.await.unwrap();
|
||||||
|
assert!(result.is_err());
|
||||||
|
assert!(matches!(
|
||||||
|
result.unwrap_err(),
|
||||||
|
Socks5Error::NoAcceptableAuth
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn address_type_ipv4() {
|
||||||
|
let (mut client, server) = duplex(4096);
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
do_handshake(&mut client).await;
|
||||||
|
let reply_buf = do_connect_ipv4(&mut client, [10, 0, 0, 1], 443).await;
|
||||||
|
assert_eq!(reply_buf[1], 0x00);
|
||||||
|
|
||||||
|
drop(client);
|
||||||
|
let _ = server_handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn address_type_domain() {
|
||||||
|
let (mut client, server) = duplex(4096);
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
do_handshake(&mut client).await;
|
||||||
|
|
||||||
|
client
|
||||||
|
.write_all(&build_socks5_connect_domain("example.com", 443))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
client.flush().await.unwrap();
|
||||||
|
|
||||||
|
let mut reply_buf = [0u8; 10];
|
||||||
|
client.read_exact(&mut reply_buf).await.unwrap();
|
||||||
|
assert_eq!(reply_buf[1], 0x00);
|
||||||
|
|
||||||
|
drop(client);
|
||||||
|
let _ = server_handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn address_type_ipv6() {
|
||||||
|
let (mut client, server) = duplex(4096);
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
do_handshake(&mut client).await;
|
||||||
|
|
||||||
|
let ipv6_addr = [0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
|
||||||
|
client
|
||||||
|
.write_all(&build_socks5_connect_ipv6(ipv6_addr, 443))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
client.flush().await.unwrap();
|
||||||
|
|
||||||
|
let mut reply_buf = [0u8; 10];
|
||||||
|
client.read_exact(&mut reply_buf).await.unwrap();
|
||||||
|
assert_eq!(reply_buf[0], 0x05);
|
||||||
|
assert_eq!(reply_buf[1], 0x00);
|
||||||
|
|
||||||
|
drop(client);
|
||||||
|
let _ = server_handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn channel_open_failure_returns_socks5_error() {
|
||||||
|
let (mut client, server) = duplex(4096);
|
||||||
|
let opener = MockChannelOpener { fail: true };
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
do_handshake(&mut client).await;
|
||||||
|
let reply_buf = do_connect_ipv4(&mut client, [10, 0, 0, 1], 80).await;
|
||||||
|
assert_eq!(reply_buf[0], 0x05);
|
||||||
|
assert_eq!(reply_buf[1], 0x05);
|
||||||
|
|
||||||
|
drop(client);
|
||||||
|
let _ = server_handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn unsupported_command_returns_error() {
|
||||||
|
let (mut client, server) = duplex(4096);
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
do_handshake(&mut client).await;
|
||||||
|
|
||||||
|
let mut bind_req = vec![0x05, 0x02, 0x00, 0x01];
|
||||||
|
bind_req.extend_from_slice(&[127, 0, 0, 1]);
|
||||||
|
bind_req.extend_from_slice(&80u16.to_be_bytes());
|
||||||
|
client.write_all(&bind_req).await.unwrap();
|
||||||
|
client.flush().await.unwrap();
|
||||||
|
|
||||||
|
let mut reply_buf = [0u8; 10];
|
||||||
|
client.read_exact(&mut reply_buf).await.unwrap();
|
||||||
|
assert_eq!(reply_buf[1], 0x07);
|
||||||
|
|
||||||
|
drop(client);
|
||||||
|
let _ = server_handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn bidirectional_proxy_flow() {
|
||||||
|
let (mut client_sock, server_sock) = duplex(4096);
|
||||||
|
let (ssh_client, mut ssh_server) = duplex(4096);
|
||||||
|
|
||||||
|
let ssh_stream = Arc::new(Mutex::new(Some(ssh_client)));
|
||||||
|
|
||||||
|
struct ProxyOpener {
|
||||||
|
stream: Arc<Mutex<Option<DuplexStream>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ChannelOpener for ProxyOpener {
|
||||||
|
type Stream = DuplexStream;
|
||||||
|
|
||||||
|
async fn open_channel(
|
||||||
|
&self,
|
||||||
|
_host: String,
|
||||||
|
_port: u16,
|
||||||
|
) -> Result<Self::Stream, ChannelOpenError> {
|
||||||
|
self.stream
|
||||||
|
.lock()
|
||||||
|
.await
|
||||||
|
.take()
|
||||||
|
.ok_or(ChannelOpenError::ChannelOpenFailed)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let opener = ProxyOpener {
|
||||||
|
stream: Arc::clone(&ssh_stream),
|
||||||
|
};
|
||||||
|
|
||||||
|
let server_handle = tokio::spawn(async move {
|
||||||
|
handle_socks5_connection(server_sock, Arc::new(opener)).await
|
||||||
|
});
|
||||||
|
|
||||||
|
do_handshake(&mut client_sock).await;
|
||||||
|
let reply_buf = do_connect_ipv4(&mut client_sock, [127, 0, 0, 1], 80).await;
|
||||||
|
assert_eq!(reply_buf[1], 0x00);
|
||||||
|
|
||||||
|
let test_data = b"hello through tunnel";
|
||||||
|
client_sock.write_all(test_data).await.unwrap();
|
||||||
|
client_sock.flush().await.unwrap();
|
||||||
|
|
||||||
|
let mut received = vec![0u8; test_data.len()];
|
||||||
|
AsyncReadExt::read_exact(&mut ssh_server, &mut received)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(&received, test_data);
|
||||||
|
|
||||||
|
let echo_data = b"response from tunnel";
|
||||||
|
ssh_server.write_all(echo_data).await.unwrap();
|
||||||
|
ssh_server.flush().await.unwrap();
|
||||||
|
|
||||||
|
let mut received_back = vec![0u8; echo_data.len()];
|
||||||
|
client_sock.read_exact(&mut received_back).await.unwrap();
|
||||||
|
assert_eq!(&received_back, echo_data);
|
||||||
|
|
||||||
|
drop(client_sock);
|
||||||
|
drop(ssh_server);
|
||||||
|
let _ = server_handle.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn default_listen_address() {
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
let server = Socks5Server::new(opener);
|
||||||
|
assert_eq!(server.listen_addr(), "127.0.0.1:1080".parse().unwrap());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn custom_listen_address() {
|
||||||
|
let opener = MockChannelOpener { fail: false };
|
||||||
|
let server = Socks5Server::with_addr(opener, "127.0.0.1:9050");
|
||||||
|
assert_eq!(server.listen_addr(), "127.0.0.1:9050".parse().unwrap());
|
||||||
|
}
|
||||||
|
}
|
||||||
304
crates/wraith-core/src/socks5/protocol.rs
Normal file
304
crates/wraith-core/src/socks5/protocol.rs
Normal file
@@ -0,0 +1,304 @@
|
|||||||
|
use std::net::{Ipv4Addr, Ipv6Addr};
|
||||||
|
|
||||||
|
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
pub enum Socks5Address {
|
||||||
|
Ipv4(Ipv4Addr),
|
||||||
|
Ipv6(Ipv6Addr),
|
||||||
|
Domain(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct Socks5VersionMethod {
|
||||||
|
pub version: u8,
|
||||||
|
pub methods: Vec<u8>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Socks5VersionMethod {
|
||||||
|
pub async fn read_from<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Self> {
|
||||||
|
let version = reader.read_u8().await?;
|
||||||
|
let nmethods = reader.read_u8().await?;
|
||||||
|
let mut methods = vec![0u8; nmethods as usize];
|
||||||
|
reader.read_exact(&mut methods).await?;
|
||||||
|
Ok(Self { version, methods })
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct Socks5Request {
|
||||||
|
pub version: u8,
|
||||||
|
pub command: u8,
|
||||||
|
pub address: Socks5Address,
|
||||||
|
pub port: u16,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Socks5Request {
|
||||||
|
pub async fn read_from<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Self> {
|
||||||
|
let version = reader.read_u8().await?;
|
||||||
|
let command = reader.read_u8().await?;
|
||||||
|
let _rsv = reader.read_u8().await?;
|
||||||
|
let atyp = reader.read_u8().await?;
|
||||||
|
|
||||||
|
let address = match atyp {
|
||||||
|
0x01 => {
|
||||||
|
let mut octets = [0u8; 4];
|
||||||
|
reader.read_exact(&mut octets).await?;
|
||||||
|
Socks5Address::Ipv4(Ipv4Addr::from(octets))
|
||||||
|
}
|
||||||
|
0x04 => {
|
||||||
|
let mut octets = [0u8; 16];
|
||||||
|
reader.read_exact(&mut octets).await?;
|
||||||
|
Socks5Address::Ipv6(Ipv6Addr::from(octets))
|
||||||
|
}
|
||||||
|
0x03 => {
|
||||||
|
let len = reader.read_u8().await?;
|
||||||
|
let mut domain = vec![0u8; len as usize];
|
||||||
|
reader.read_exact(&mut domain).await?;
|
||||||
|
Socks5Address::Domain(String::from_utf8_lossy(&domain).into_owned())
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
return Err(std::io::Error::new(
|
||||||
|
std::io::ErrorKind::InvalidData,
|
||||||
|
format!("unsupported address type: {atyp}"),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let port = reader.read_u16().await?;
|
||||||
|
|
||||||
|
Ok(Self {
|
||||||
|
version,
|
||||||
|
command,
|
||||||
|
address,
|
||||||
|
port,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub struct Socks5Reply {
|
||||||
|
pub version: u8,
|
||||||
|
pub reply: u8,
|
||||||
|
pub address: Socks5Address,
|
||||||
|
pub port: u16,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Socks5Reply {
|
||||||
|
pub fn success(address: Socks5Address, port: u16) -> Self {
|
||||||
|
Self {
|
||||||
|
version: 0x05,
|
||||||
|
reply: 0x00,
|
||||||
|
address,
|
||||||
|
port,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn connection_refused() -> Self {
|
||||||
|
Self {
|
||||||
|
version: 0x05,
|
||||||
|
reply: 0x05,
|
||||||
|
address: Socks5Address::Ipv4(Ipv4Addr::UNSPECIFIED),
|
||||||
|
port: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn command_not_supported() -> Self {
|
||||||
|
Self {
|
||||||
|
version: 0x05,
|
||||||
|
reply: 0x07,
|
||||||
|
address: Socks5Address::Ipv4(Ipv4Addr::UNSPECIFIED),
|
||||||
|
port: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn write_to<W: AsyncWrite + Unpin>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||||
|
writer.write_u8(self.version).await?;
|
||||||
|
writer.write_u8(self.reply).await?;
|
||||||
|
writer.write_u8(0x00).await?;
|
||||||
|
match &self.address {
|
||||||
|
Socks5Address::Ipv4(addr) => {
|
||||||
|
writer.write_u8(0x01).await?;
|
||||||
|
writer.write_all(&addr.octets()).await?;
|
||||||
|
}
|
||||||
|
Socks5Address::Ipv6(addr) => {
|
||||||
|
writer.write_u8(0x04).await?;
|
||||||
|
writer.write_all(&addr.octets()).await?;
|
||||||
|
}
|
||||||
|
Socks5Address::Domain(name) => {
|
||||||
|
writer.write_u8(0x03).await?;
|
||||||
|
writer.write_u8(name.len() as u8).await?;
|
||||||
|
writer.write_all(name.as_bytes()).await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
writer.write_u16(self.port).await?;
|
||||||
|
writer.flush().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use std::io::Cursor;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn parse_version_method_no_auth() {
|
||||||
|
let data = [0x05, 0x01, 0x00];
|
||||||
|
let mut cursor = Cursor::new(&data[..]);
|
||||||
|
let vm = Socks5VersionMethod::read_from(&mut cursor).await.unwrap();
|
||||||
|
assert_eq!(vm.version, 0x05);
|
||||||
|
assert_eq!(vm.methods, vec![0x00]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn parse_version_method_multiple() {
|
||||||
|
let data = [0x05, 0x02, 0x00, 0x02];
|
||||||
|
let mut cursor = Cursor::new(&data[..]);
|
||||||
|
let vm = Socks5VersionMethod::read_from(&mut cursor).await.unwrap();
|
||||||
|
assert_eq!(vm.version, 0x05);
|
||||||
|
assert_eq!(vm.methods, vec![0x00, 0x02]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn parse_request_ipv4() {
|
||||||
|
let mut data = vec![0x05, 0x01, 0x00, 0x01];
|
||||||
|
data.extend_from_slice(&[10, 0, 0, 1]);
|
||||||
|
data.extend_from_slice(&443u16.to_be_bytes());
|
||||||
|
let mut cursor = Cursor::new(&data[..]);
|
||||||
|
let req = Socks5Request::read_from(&mut cursor).await.unwrap();
|
||||||
|
assert_eq!(req.version, 0x05);
|
||||||
|
assert_eq!(req.command, 0x01);
|
||||||
|
assert_eq!(
|
||||||
|
req.address,
|
||||||
|
Socks5Address::Ipv4(Ipv4Addr::new(10, 0, 0, 1))
|
||||||
|
);
|
||||||
|
assert_eq!(req.port, 443);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn parse_request_ipv6() {
|
||||||
|
let mut data = vec![0x05, 0x01, 0x00, 0x04];
|
||||||
|
let octets: [u8; 16] = [0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
|
||||||
|
data.extend_from_slice(&octets);
|
||||||
|
data.extend_from_slice(&443u16.to_be_bytes());
|
||||||
|
let mut cursor = Cursor::new(&data[..]);
|
||||||
|
let req = Socks5Request::read_from(&mut cursor).await.unwrap();
|
||||||
|
assert_eq!(req.version, 0x05);
|
||||||
|
assert_eq!(req.command, 0x01);
|
||||||
|
assert!(matches!(req.address, Socks5Address::Ipv6(_)));
|
||||||
|
assert_eq!(req.port, 443);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn parse_request_domain() {
|
||||||
|
let domain = "example.com";
|
||||||
|
let mut data = vec![0x05, 0x01, 0x00, 0x03];
|
||||||
|
data.push(domain.len() as u8);
|
||||||
|
data.extend_from_slice(domain.as_bytes());
|
||||||
|
data.extend_from_slice(&443u16.to_be_bytes());
|
||||||
|
let mut cursor = Cursor::new(&data[..]);
|
||||||
|
let req = Socks5Request::read_from(&mut cursor).await.unwrap();
|
||||||
|
assert_eq!(req.version, 0x05);
|
||||||
|
assert_eq!(req.command, 0x01);
|
||||||
|
assert_eq!(req.address, Socks5Address::Domain("example.com".to_string()));
|
||||||
|
assert_eq!(req.port, 443);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn parse_request_unsupported_address_type() {
|
||||||
|
let data = [0x05, 0x01, 0x00, 0x05];
|
||||||
|
let mut cursor = Cursor::new(&data[..]);
|
||||||
|
let result = Socks5Request::read_from(&mut cursor).await;
|
||||||
|
assert!(result.is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn reply_success_ipv4() {
|
||||||
|
let reply = Socks5Reply::success(Socks5Address::Ipv4(Ipv4Addr::UNSPECIFIED), 0);
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
reply.write_to(&mut buf).await.unwrap();
|
||||||
|
assert_eq!(buf[0], 0x05);
|
||||||
|
assert_eq!(buf[1], 0x00);
|
||||||
|
assert_eq!(buf[2], 0x00);
|
||||||
|
assert_eq!(buf[3], 0x01);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn reply_connection_refused() {
|
||||||
|
let reply = Socks5Reply::connection_refused();
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
reply.write_to(&mut buf).await.unwrap();
|
||||||
|
assert_eq!(buf[0], 0x05);
|
||||||
|
assert_eq!(buf[1], 0x05);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn reply_command_not_supported() {
|
||||||
|
let reply = Socks5Reply::command_not_supported();
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
reply.write_to(&mut buf).await.unwrap();
|
||||||
|
assert_eq!(buf[0], 0x05);
|
||||||
|
assert_eq!(buf[1], 0x07);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn roundtrip_ipv4_reply() {
|
||||||
|
let reply = Socks5Reply::success(Socks5Address::Ipv4(Ipv4Addr::new(127, 0, 0, 1)), 1080);
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
reply.write_to(&mut buf).await.unwrap();
|
||||||
|
|
||||||
|
let mut cursor = Cursor::new(&buf[..]);
|
||||||
|
let version = cursor.read_u8().await.unwrap();
|
||||||
|
let _reply_code = cursor.read_u8().await.unwrap();
|
||||||
|
let _rsv = cursor.read_u8().await.unwrap();
|
||||||
|
let atyp = cursor.read_u8().await.unwrap();
|
||||||
|
assert_eq!(version, 0x05);
|
||||||
|
assert_eq!(atyp, 0x01);
|
||||||
|
let mut octets = [0u8; 4];
|
||||||
|
cursor.read_exact(&mut octets).await.unwrap();
|
||||||
|
assert_eq!(Ipv4Addr::from(octets), Ipv4Addr::new(127, 0, 0, 1));
|
||||||
|
let port = cursor.read_u16().await.unwrap();
|
||||||
|
assert_eq!(port, 1080);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn roundtrip_ipv6_reply() {
|
||||||
|
let addr = Ipv6Addr::new(0x2001, 0x0db8, 0, 0, 0, 0, 0, 1);
|
||||||
|
let reply = Socks5Reply::success(Socks5Address::Ipv6(addr), 443);
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
reply.write_to(&mut buf).await.unwrap();
|
||||||
|
|
||||||
|
let mut cursor = Cursor::new(&buf[..]);
|
||||||
|
let _version = cursor.read_u8().await.unwrap();
|
||||||
|
let _reply_code = cursor.read_u8().await.unwrap();
|
||||||
|
let _rsv = cursor.read_u8().await.unwrap();
|
||||||
|
let atyp = cursor.read_u8().await.unwrap();
|
||||||
|
assert_eq!(atyp, 0x04);
|
||||||
|
let mut octets = [0u8; 16];
|
||||||
|
cursor.read_exact(&mut octets).await.unwrap();
|
||||||
|
assert_eq!(Ipv6Addr::from(octets), addr);
|
||||||
|
let port = cursor.read_u16().await.unwrap();
|
||||||
|
assert_eq!(port, 443);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn roundtrip_domain_reply() {
|
||||||
|
let reply = Socks5Reply::success(Socks5Address::Domain("example.com".to_string()), 8080);
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
reply.write_to(&mut buf).await.unwrap();
|
||||||
|
|
||||||
|
let mut cursor = Cursor::new(&buf[..]);
|
||||||
|
let _version = cursor.read_u8().await.unwrap();
|
||||||
|
let _reply_code = cursor.read_u8().await.unwrap();
|
||||||
|
let _rsv = cursor.read_u8().await.unwrap();
|
||||||
|
let atyp = cursor.read_u8().await.unwrap();
|
||||||
|
assert_eq!(atyp, 0x03);
|
||||||
|
let len = cursor.read_u8().await.unwrap();
|
||||||
|
let mut domain = vec![0u8; len as usize];
|
||||||
|
cursor.read_exact(&mut domain).await.unwrap();
|
||||||
|
assert_eq!(String::from_utf8(domain).unwrap(), "example.com");
|
||||||
|
let port = cursor.read_u16().await.unwrap();
|
||||||
|
assert_eq!(port, 8080);
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user