|
|
@@ -0,0 +1,83 @@
|
|
|
1
|
+use hbb_common::{
|
|
|
2
|
+ env_logger, log,
|
|
|
3
|
+ protobuf::Message as _,
|
|
|
4
|
+ rendezvous_proto::*,
|
|
|
5
|
+ sleep,
|
|
|
6
|
+ tcp::{new_listener, FramedStream},
|
|
|
7
|
+ tokio, ResultType,
|
|
|
8
|
+};
|
|
|
9
|
+use std::{
|
|
|
10
|
+ collections::HashMap,
|
|
|
11
|
+ net::SocketAddr,
|
|
|
12
|
+ sync::{Arc, Mutex},
|
|
|
13
|
+};
|
|
|
14
|
+
|
|
|
15
|
+lazy_static::lazy_static! {
|
|
|
16
|
+ static ref PEERS: Arc<Mutex<HashMap<String, FramedStream>>> = Arc::new(Mutex::new(HashMap::new()));
|
|
|
17
|
+}
|
|
|
18
|
+
|
|
|
19
|
+#[tokio::main]
|
|
|
20
|
+async fn main() -> ResultType<()> {
|
|
|
21
|
+ env_logger::init();
|
|
|
22
|
+ let addr = "0.0.0.0:21117";
|
|
|
23
|
+ log::info!("Listening on {}", addr);
|
|
|
24
|
+ let mut listener = new_listener(addr, true).await?;
|
|
|
25
|
+ loop {
|
|
|
26
|
+ tokio::select! {
|
|
|
27
|
+ Ok((stream, addr)) = listener.accept() => {
|
|
|
28
|
+ tokio::spawn(async move {
|
|
|
29
|
+ make_pair(FramedStream::from(stream), addr).await.ok();
|
|
|
30
|
+ });
|
|
|
31
|
+ }
|
|
|
32
|
+ }
|
|
|
33
|
+ }
|
|
|
34
|
+}
|
|
|
35
|
+
|
|
|
36
|
+async fn make_pair(stream: FramedStream, addr: SocketAddr) -> ResultType<()> {
|
|
|
37
|
+ let mut stream = stream;
|
|
|
38
|
+ if let Some(Ok(bytes)) = stream.next_timeout(30_000).await {
|
|
|
39
|
+ if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
|
|
|
40
|
+ if let Some(rendezvous_message::Union::request_forward(rf)) = msg_in.union {
|
|
|
41
|
+ if !rf.uuid.is_empty() {
|
|
|
42
|
+ let peer = PEERS.lock().unwrap().remove(&rf.uuid);
|
|
|
43
|
+ if let Some(peer) = peer {
|
|
|
44
|
+ log::info!("Forward request {} from {} got paired", rf.uuid, addr);
|
|
|
45
|
+ return forward(stream, peer).await;
|
|
|
46
|
+ } else {
|
|
|
47
|
+ log::info!("New forward request {} from {}", rf.uuid, addr);
|
|
|
48
|
+ PEERS.lock().unwrap().insert(rf.uuid.clone(), stream);
|
|
|
49
|
+ sleep(30.).await;
|
|
|
50
|
+ PEERS.lock().unwrap().remove(&rf.uuid);
|
|
|
51
|
+ }
|
|
|
52
|
+ }
|
|
|
53
|
+ }
|
|
|
54
|
+ }
|
|
|
55
|
+ }
|
|
|
56
|
+ Ok(())
|
|
|
57
|
+}
|
|
|
58
|
+
|
|
|
59
|
+async fn forward(stream: FramedStream, peer: FramedStream) -> ResultType<()> {
|
|
|
60
|
+ let mut peer = peer;
|
|
|
61
|
+ let mut stream = stream;
|
|
|
62
|
+ peer.set_raw();
|
|
|
63
|
+ stream.set_raw();
|
|
|
64
|
+ loop {
|
|
|
65
|
+ tokio::select! {
|
|
|
66
|
+ res = peer.next() => {
|
|
|
67
|
+ if let Some(Ok(bytes)) = res {
|
|
|
68
|
+ stream.send_bytes(bytes.into()).await?;
|
|
|
69
|
+ } else {
|
|
|
70
|
+ break;
|
|
|
71
|
+ }
|
|
|
72
|
+ },
|
|
|
73
|
+ res = stream.next() => {
|
|
|
74
|
+ if let Some(Ok(bytes)) = res {
|
|
|
75
|
+ peer.send_bytes(bytes.into()).await?;
|
|
|
76
|
+ } else {
|
|
|
77
|
+ break;
|
|
|
78
|
+ }
|
|
|
79
|
+ },
|
|
|
80
|
+ }
|
|
|
81
|
+ }
|
|
|
82
|
+ Ok(())
|
|
|
83
|
+}
|