|
|
@@ -1,13 +1,13 @@
|
|
1
|
|
-use bytes::{Bytes, BytesMut};
|
|
2
|
|
-use futures::SinkExt;
|
|
3
|
1
|
use hbb_common::{
|
|
|
2
|
+ bytes::BytesMut,
|
|
|
3
|
+ log,
|
|
4
|
4
|
message_proto::*,
|
|
5
|
|
- protobuf::{parse_from_bytes, Message as _},
|
|
6
|
|
- AddrMangle,
|
|
|
5
|
+ protobuf::parse_from_bytes,
|
|
|
6
|
+ tokio::{net::UdpSocket, stream::StreamExt},
|
|
|
7
|
+ udp::FramedSocket,
|
|
|
8
|
+ AddrMangle, ResultType,
|
|
7
|
9
|
};
|
|
8
|
|
-use std::{collections::HashMap, error::Error, net::SocketAddr, time::Duration};
|
|
9
|
|
-use tokio::{net::UdpSocket, stream::StreamExt, time::delay_for};
|
|
10
|
|
-use tokio_util::{codec::BytesCodec, udp::UdpFramed};
|
|
|
10
|
+use std::{collections::HashMap, net::SocketAddr};
|
|
11
|
11
|
|
|
12
|
12
|
pub struct Peer {
|
|
13
|
13
|
socket_addr: SocketAddr,
|
|
|
@@ -19,13 +19,10 @@ pub struct RendezvousServer {
|
|
19
|
19
|
peer_map: PeerMap,
|
|
20
|
20
|
}
|
|
21
|
21
|
|
|
22
|
|
-type FramedSocket = UdpFramed<BytesCodec>;
|
|
23
|
|
-type ResultType = Result<(), Box<dyn Error>>;
|
|
24
|
|
-
|
|
25
|
22
|
impl RendezvousServer {
|
|
26
|
|
- pub async fn start(addr: &str) -> ResultType {
|
|
|
23
|
+ pub async fn start(addr: &str) -> ResultType<()> {
|
|
27
|
24
|
let socket = UdpSocket::bind(addr).await?;
|
|
28
|
|
- let mut socket = UdpFramed::new(socket, BytesCodec::new());
|
|
|
25
|
+ let mut socket = FramedSocket::new(socket);
|
|
29
|
26
|
let mut rs = Self {
|
|
30
|
27
|
peer_map: PeerMap::new(),
|
|
31
|
28
|
};
|
|
|
@@ -40,7 +37,7 @@ impl RendezvousServer {
|
|
40
|
37
|
bytes: &BytesMut,
|
|
41
|
38
|
addr: SocketAddr,
|
|
42
|
39
|
socket: &mut FramedSocket,
|
|
43
|
|
- ) -> ResultType {
|
|
|
40
|
+ ) -> ResultType<()> {
|
|
44
|
41
|
if let Ok(msg_in) = parse_from_bytes::<Message>(&bytes) {
|
|
45
|
42
|
match msg_in.union {
|
|
46
|
43
|
Some(Message_oneof_union::register_peer(rp)) => {
|
|
|
@@ -64,7 +61,7 @@ impl RendezvousServer {
|
|
64
|
61
|
socket_addr: AddrMangle::encode(&addr),
|
|
65
|
62
|
..Default::default()
|
|
66
|
63
|
});
|
|
67
|
|
- send_to(&msg_out, peer.socket_addr, socket).await?;
|
|
|
64
|
+ socket.send(&msg_out, peer.socket_addr).await?;
|
|
68
|
65
|
}
|
|
69
|
66
|
}
|
|
70
|
67
|
Some(Message_oneof_union::punch_hole_sent(phs)) => {
|
|
|
@@ -76,7 +73,7 @@ impl RendezvousServer {
|
|
76
|
73
|
socket_addr: AddrMangle::encode(&addr),
|
|
77
|
74
|
..Default::default()
|
|
78
|
75
|
});
|
|
79
|
|
- send_to(&msg_out, addr_a, socket).await?;
|
|
|
76
|
+ socket.send(&msg_out, addr_a).await?;
|
|
80
|
77
|
}
|
|
81
|
78
|
_ => {}
|
|
82
|
79
|
}
|
|
|
@@ -85,22 +82,11 @@ impl RendezvousServer {
|
|
85
|
82
|
}
|
|
86
|
83
|
}
|
|
87
|
84
|
|
|
88
|
|
-#[inline]
|
|
89
|
|
-pub async fn send_to(msg: &Message, addr: SocketAddr, socket: &mut FramedSocket) -> ResultType {
|
|
90
|
|
- socket
|
|
91
|
|
- .send((Bytes::from(msg.write_to_bytes().unwrap()), addr))
|
|
92
|
|
- .await?;
|
|
93
|
|
- Ok(())
|
|
94
|
|
-}
|
|
95
|
|
-
|
|
96
|
|
-#[inline]
|
|
97
|
|
-pub async fn sleep(sec: f32) {
|
|
98
|
|
- delay_for(Duration::from_secs_f32(sec)).await;
|
|
99
|
|
-}
|
|
100
|
|
-
|
|
101
|
85
|
#[cfg(test)]
|
|
102
|
86
|
mod tests {
|
|
103
|
87
|
use super::*;
|
|
|
88
|
+ use hbb_common::tokio;
|
|
|
89
|
+ use std::time::Duration;
|
|
104
|
90
|
|
|
105
|
91
|
#[allow(unused_must_use)]
|
|
106
|
92
|
#[tokio::main]
|
|
|
@@ -118,23 +104,23 @@ mod tests {
|
|
118
|
104
|
// B register it to server
|
|
119
|
105
|
let socket_b = UdpSocket::bind("127.0.0.1:0").await.unwrap();
|
|
120
|
106
|
let local_addr_b = socket_b.local_addr().unwrap();
|
|
121
|
|
- let mut socket_b = UdpFramed::new(socket_b, BytesCodec::new());
|
|
|
107
|
+ let mut socket_b = FramedSocket::new(socket_b);
|
|
122
|
108
|
let mut msg_out = Message::new();
|
|
123
|
109
|
msg_out.set_register_peer(RegisterPeer {
|
|
124
|
110
|
hbb_addr: "123".to_string(),
|
|
125
|
111
|
..Default::default()
|
|
126
|
112
|
});
|
|
127
|
|
- send_to(&msg_out, addr_server, &mut socket_b).await;
|
|
|
113
|
+ socket_b.send(&msg_out, addr_server).await;
|
|
128
|
114
|
|
|
129
|
115
|
// A send punch request to server
|
|
130
|
116
|
let socket_a = UdpSocket::bind("127.0.0.1:0").await.unwrap();
|
|
131
|
117
|
let local_addr_a = socket_a.local_addr().unwrap();
|
|
132
|
|
- let mut socket_a = UdpFramed::new(socket_a, BytesCodec::new());
|
|
|
118
|
+ let mut socket_a = FramedSocket::new(socket_a);
|
|
133
|
119
|
msg_out.set_punch_hole_request(PunchHoleRequest {
|
|
134
|
120
|
hbb_addr: "123".to_string(),
|
|
135
|
121
|
..Default::default()
|
|
136
|
122
|
});
|
|
137
|
|
- send_to(&msg_out, addr_server, &mut socket_a).await;
|
|
|
123
|
+ socket_a.send(&msg_out, addr_server).await;
|
|
138
|
124
|
|
|
139
|
125
|
println!(
|
|
140
|
126
|
"A {:?} request punch hole to B {:?} via server {:?}",
|
|
|
@@ -160,7 +146,7 @@ mod tests {
|
|
160
|
146
|
socket_addr: AddrMangle::encode(&remote_addr_a),
|
|
161
|
147
|
..Default::default()
|
|
162
|
148
|
});
|
|
163
|
|
- send_to(&msg_out, addr_server, &mut socket_b).await;
|
|
|
149
|
+ socket_b.send(&msg_out, addr_server).await;
|
|
164
|
150
|
}
|
|
165
|
151
|
|
|
166
|
152
|
// on A side
|
open-trade