1 use futures::{future, StreamExt};
2 use once_cell::sync::OnceCell;
3 use std::{
4 io::{Read, Write},
5 net::{SocketAddr, TcpStream as StdTcpStream},
6 sync::Mutex,
7 };
8 use tokio::{
9 io::{copy, AsyncReadExt, AsyncWriteExt},
10 net::TcpListener,
11 runtime::Runtime,
12 };
13 use tokio_socks::{
14 tcp::{Socks5Listener, Socks5Stream},
15 Result,
16 };
17
18 pub const PROXY_ADDR: &'static str = "127.0.0.1:41080";
19 pub const ECHO_SERVER_ADDR: &'static str = "localhost:10007";
20 pub const MSG: &[u8] = b"hello";
21
echo_server() -> Result<()>22 pub async fn echo_server() -> Result<()> {
23 let mut listener = TcpListener::bind(&SocketAddr::from(([0, 0, 0, 0], 10007))).await?;
24 listener
25 .incoming()
26 .for_each(|tcp_stream| {
27 if let Ok(mut stream) = tcp_stream {
28 tokio::spawn(async move {
29 let (mut reader, mut writer) = stream.split();
30 copy(&mut reader, &mut writer).await.unwrap();
31 });
32 }
33
34 future::ready(())
35 })
36 .await;
37 Ok(())
38 }
39
reply_response(mut socket: Socks5Stream) -> Result<[u8; 5]>40 pub async fn reply_response(mut socket: Socks5Stream) -> Result<[u8; 5]> {
41 socket.write_all(MSG).await?;
42 let mut buf = [0; 5];
43 socket.read_exact(&mut buf).await?;
44 Ok(buf)
45 }
46
test_connect(socket: Socks5Stream) -> Result<()>47 pub async fn test_connect(socket: Socks5Stream) -> Result<()> {
48 let res = reply_response(socket).await?;
49 assert_eq!(&res[..], MSG);
50 Ok(())
51 }
52
test_bind(listener: Socks5Listener) -> Result<()>53 pub fn test_bind(listener: Socks5Listener) -> Result<()> {
54 let bind_addr = listener.bind_addr().to_owned();
55 runtime().lock().unwrap().spawn(async move {
56 let mut stream = listener.accept().await.unwrap();
57 let (mut reader, mut writer) = stream.split();
58 copy(&mut reader, &mut writer).await.unwrap();
59 });
60
61 let mut tcp = StdTcpStream::connect(bind_addr)?;
62 tcp.write_all(MSG)?;
63 let mut buf = [0; 5];
64 tcp.read_exact(&mut buf[..])?;
65 assert_eq!(&buf[..], MSG);
66 Ok(())
67 }
68
runtime() -> &'static Mutex<Runtime>69 pub fn runtime() -> &'static Mutex<Runtime> {
70 static RUNTIME: OnceCell<Mutex<Runtime>> = OnceCell::new();
71 RUNTIME.get_or_init(|| {
72 let runtime = Runtime::new().expect("Unable to create runtime");
73 runtime.spawn(async { echo_server().await.expect("Unable to bind") });
74 Mutex::new(runtime)
75 })
76 }
77