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
|
use hbb_common::{
log,
protobuf::Message as _,
rendezvous_proto::*,
sleep,
tcp::{new_listener, FramedStream},
tokio::{
self,
time::{interval, Duration},
},
ResultType,
};
use std::{
collections::HashMap,
net::SocketAddr,
sync::{Arc, Mutex},
};
lazy_static::lazy_static! {
static ref PEERS: Arc<Mutex<HashMap<String, FramedStream>>> = Arc::new(Mutex::new(HashMap::new()));
}
pub const DEFAULT_PORT: &'static str = "21117";
#[tokio::main(basic_scheduler)]
pub async fn start(port: &str, license: &str, stop: Arc<Mutex<bool>>) -> ResultType<()> {
let addr = format!("0.0.0.0:{}", port);
log::info!("Listening on tcp {}", addr);
let mut timer = interval(Duration::from_millis(300));
let mut listener = new_listener(addr, false).await?;
loop {
tokio::select! {
Ok((stream, addr)) = listener.accept() => {
let license = license.to_owned();
tokio::spawn(async move {
make_pair(FramedStream::from(stream), addr, &license).await.ok();
});
}
_ = timer.tick() => {
if *stop.lock().unwrap() {
log::info!("Stopped");
break;
}
}
}
}
Ok(())
}
async fn make_pair(stream: FramedStream, addr: SocketAddr, license: &str) -> ResultType<()> {
let mut stream = stream;
if let Some(Ok(bytes)) = stream.next_timeout(30_000).await {
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
if let Some(rendezvous_message::Union::request_relay(rf)) = msg_in.union {
if !license.is_empty() && rf.licence_key != license {
return Ok(());
}
if !rf.uuid.is_empty() {
let peer = PEERS.lock().unwrap().remove(&rf.uuid);
if let Some(peer) = peer {
log::info!("Forward request {} from {} got paired", rf.uuid, addr);
return relay(stream, peer).await;
} else {
log::info!("New relay request {} from {}", rf.uuid, addr);
PEERS.lock().unwrap().insert(rf.uuid.clone(), stream);
sleep(30.).await;
PEERS.lock().unwrap().remove(&rf.uuid);
}
}
}
}
}
Ok(())
}
async fn relay(stream: FramedStream, peer: FramedStream) -> ResultType<()> {
let mut peer = peer;
let mut stream = stream;
peer.set_raw();
stream.set_raw();
loop {
tokio::select! {
res = peer.next() => {
if let Some(Ok(bytes)) = res {
stream.send_bytes(bytes.into()).await?;
} else {
break;
}
},
res = stream.next() => {
if let Some(Ok(bytes)) = res {
peer.send_bytes(bytes.into()).await?;
} else {
break;
}
},
}
}
Ok(())
}
|
