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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
|
use std::path::PathBuf;
use anyhow::Result;
use tokio::{
io::{self, AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream, ToSocketAddrs},
sync::broadcast,
};
pub const PING: &str = "ping";
pub const PONG: &str = "pong";
pub async fn run_rathole_server(
config_path: &str,
shutdown_rx: broadcast::Receiver<bool>,
) -> Result<()> {
let cli = rathole::Cli {
config_path: Some(PathBuf::from(config_path)),
server: true,
client: false,
..Default::default()
};
rathole::run(cli, shutdown_rx).await
}
pub async fn run_rathole_client(
config_path: &str,
shutdown_rx: broadcast::Receiver<bool>,
) -> Result<()> {
let cli = rathole::Cli {
config_path: Some(PathBuf::from(config_path)),
server: false,
client: true,
..Default::default()
};
rathole::run(cli, shutdown_rx).await
}
pub mod tcp {
use super::*;
pub async fn echo_server<A: ToSocketAddrs>(addr: A) -> Result<()> {
let l = TcpListener::bind(addr).await?;
loop {
let (conn, _addr) = l.accept().await?;
tokio::spawn(async move {
let _ = echo(conn).await;
});
}
}
pub async fn pingpong_server<A: ToSocketAddrs>(addr: A) -> Result<()> {
let l = TcpListener::bind(addr).await?;
loop {
let (conn, _addr) = l.accept().await?;
tokio::spawn(async move {
let _ = pingpong(conn).await;
});
}
}
async fn echo(conn: TcpStream) -> Result<()> {
let (mut rd, mut wr) = conn.into_split();
io::copy(&mut rd, &mut wr).await?;
Ok(())
}
async fn pingpong(mut conn: TcpStream) -> Result<()> {
let mut buf = [0u8; PING.len()];
while conn.read_exact(&mut buf).await? != 0 {
assert_eq!(buf, PING.as_bytes());
conn.write_all(PONG.as_bytes()).await?;
}
Ok(())
}
}
pub mod udp {
use rathole::UDP_BUFFER_SIZE;
use tokio::net::UdpSocket;
use tracing::debug;
use super::*;
pub async fn echo_server<A: ToSocketAddrs>(addr: A) -> Result<()> {
let l = UdpSocket::bind(addr).await?;
debug!("UDP echo server listening");
let mut buf = [0u8; UDP_BUFFER_SIZE];
loop {
let (n, addr) = l.recv_from(&mut buf).await?;
debug!("Get {:?} from {}", &buf[..n], addr);
l.send_to(&buf[..n], addr).await?;
}
}
pub async fn pingpong_server<A: ToSocketAddrs>(addr: A) -> Result<()> {
let l = UdpSocket::bind(addr).await?;
let mut buf = [0u8; UDP_BUFFER_SIZE];
loop {
let (n, addr) = l.recv_from(&mut buf).await?;
assert_eq!(&buf[..n], PING.as_bytes());
l.send_to(PONG.as_bytes(), addr).await?;
}
}
}
|
