summaryrefslogtreecommitdiffhomepage
path: root/test/socks-server/src/lib.rs
blob: cc8847db1bc56b39331720d31b784d0f1d2e7cf4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
use fast_socks5::server::{AcceptAuthentication, Socks5Server};
use futures::StreamExt;
use std::{io, net::SocketAddr};

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Failed to start SOCKS5 server")]
    StartSocksServer(#[source] io::Error),
}

pub struct Handle {
    handle: tokio::task::JoinHandle<()>,
}

/// Spawn a SOCKS server bound to `bind_addr`
pub async fn spawn(bind_addr: SocketAddr) -> Result<Handle, Error> {
    let socks_server: Socks5Server<AcceptAuthentication> = Socks5Server::bind(bind_addr)
        .await
        .map_err(Error::StartSocksServer)?;

    let handle = tokio::spawn(async move {
        let mut incoming = socks_server.incoming();

        while let Some(new_client) = incoming.next().await {
            match new_client {
                Ok(socket) => {
                    let fut = socket.upgrade_to_socks5();

                    // Act as normal SOCKS server
                    tokio::spawn(async move {
                        match fut.await {
                            Ok(_socket) => log::info!("socks client disconnected"),
                            Err(error) => log::error!("socks client failed: {error}"),
                        }
                    });
                }
                Err(error) => {
                    log::error!("failed to accept socks client: {error}");
                }
            }
        }
    });
    Ok(Handle { handle })
}

impl Handle {
    pub fn close(&self) {
        self.handle.abort();
    }
}

impl Drop for Handle {
    fn drop(&mut self) {
        self.close();
    }
}