server/
server.rs

1use std::net::SocketAddr;
2use tokio::{
3    io::{self, AsyncRead, AsyncWrite, AsyncReadExt, AsyncWriteExt, BufStream},
4    net::{self, TcpListener, TcpStream, TcpSocket},
5};
6use socks5rs::{self, consts, error::Result, Command, DestAddr, Request};
7
8const BIND_ADDRESS: &str = "[::]:1080";
9
10#[tokio::main]
11async fn main() -> io::Result<()> {
12    let listener = TcpListener::bind(BIND_ADDRESS).await?;
13    println!("listening on {}", BIND_ADDRESS);
14    loop {
15        let stream = listener.accept().await?.0;
16        tokio::spawn(async move {
17            // own the stream
18            let mut stream = stream;
19            let local_addr = stream.local_addr().unwrap();
20            match handle_client_connection(&mut stream).await {
21                Ok(()) => (),
22                Err(e) => {
23                    _ = socks5rs::io::send_reply(&mut stream, Some(e), local_addr).await;
24                },
25            };
26            _ = stream.shutdown().await;
27        });
28    }
29}
30
31async fn handle_client_connection(stream: &mut TcpStream) -> Result<()> {
32    // With buffered I/O it's important to call `.flush()` on stream
33    // after write operation to ensure the bytes reached to their destination.
34    let mut stream = BufStream::new(stream);
35
36    const BUFCAP: usize = 9000;
37    let mut buffer = [0_u8; BUFCAP];
38    let mut n = stream.read(&mut buffer[..BUFCAP]).await?;
39    let (_, _methods) = socks5rs::parse_client_methods(&buffer[..n])?;
40
41    stream.write_all(&[ consts::SOCKS5, consts::method::NO_AUTH ]).await?;
42    stream.flush().await?;
43
44    n = stream.read(&mut buffer[..BUFCAP]).await?;
45
46    let (_, command) = socks5rs::parse_client_request(&buffer[..n])?;
47
48    match command {
49        Command::Connect(r) => handle_connect_command(stream.get_mut(), r).await,
50        _ => Ok(()),
51    }
52}
53
54async fn handle_connect_command<RW>(stream: &mut RW, req: Request) -> Result<()>
55where
56    RW: AsyncRead + AsyncWrite + Unpin + Send + Sync,
57{
58    let target = TcpSocket::new_v4()?;
59    let mut s = match req.dest_addr {
60        DestAddr::IP(addr) => target.connect(SocketAddr::new(addr, req.dest_port)).await?,
61        DestAddr::FQDN(domain) => {
62            let host = join_host_and_port(String::from_utf8(domain).unwrap().as_str(), req.dest_port);
63            let addr = resolve_domain_name(host.as_str()).await?;
64            target.connect(addr).await?
65        },
66    };
67
68    socks5rs::io::send_reply(stream, None, s.local_addr().unwrap()).await?;
69
70    io::copy_bidirectional(stream, &mut s).await?;
71    Ok(())
72}
73
74async fn resolve_domain_name(host: &str) -> Result<SocketAddr> {
75    let addrs: Vec<SocketAddr> = net::lookup_host(host).await?.collect();
76    Ok(addrs[0])
77}
78
79fn join_host_and_port(host: &str, port: u16) -> String {
80    let mut s = String::with_capacity(host.len() + 6);
81    s.push_str(host);
82    s.push(':');
83    s.push_str(port.to_string().as_str());
84    s
85}
server/server.rs

server/
server.rs
