|
|
@@ -1,7 +1,6 @@
|
|
1
|
1
|
use hbb_common::{
|
|
2
|
2
|
allow_err,
|
|
3
|
|
- bytes::Bytes,
|
|
4
|
|
- bytes::BytesMut,
|
|
|
3
|
+ bytes::{Bytes, BytesMut},
|
|
5
|
4
|
bytes_codec::BytesCodec,
|
|
6
|
5
|
futures_util::{
|
|
7
|
6
|
sink::SinkExt,
|
|
|
@@ -16,6 +15,7 @@ use hbb_common::{
|
|
16
|
15
|
udp::FramedSocket,
|
|
17
|
16
|
AddrMangle, ResultType,
|
|
18
|
17
|
};
|
|
|
18
|
+use serde_derive::{Deserialize, Serialize};
|
|
19
|
19
|
use std::{
|
|
20
|
20
|
collections::HashMap,
|
|
21
|
21
|
net::SocketAddr,
|
|
|
@@ -29,28 +29,47 @@ struct Peer {
|
|
29
|
29
|
last_reg_time: Instant,
|
|
30
|
30
|
}
|
|
31
|
31
|
|
|
|
32
|
+#[derive(Debug, Serialize, Deserialize, Default)]
|
|
|
33
|
+struct PeerSerde {
|
|
|
34
|
+ #[serde(default)]
|
|
|
35
|
+ ip: String,
|
|
|
36
|
+}
|
|
|
37
|
+
|
|
32
|
38
|
#[derive(Clone)]
|
|
33
|
39
|
struct PeerMap {
|
|
34
|
40
|
map: Arc<RwLock<HashMap<String, Peer>>>,
|
|
35
|
|
- db: sled::Db,
|
|
|
41
|
+ db: super::SledAsync,
|
|
36
|
42
|
}
|
|
37
|
43
|
|
|
38
|
44
|
impl PeerMap {
|
|
39
|
45
|
fn new() -> ResultType<Self> {
|
|
40
|
46
|
Ok(Self {
|
|
41
|
47
|
map: Default::default(),
|
|
42
|
|
- db: sled::open("./sled.db")?,
|
|
|
48
|
+ db: super::SledAsync::new("./sled.db")?,
|
|
43
|
49
|
})
|
|
44
|
50
|
}
|
|
45
|
51
|
|
|
46
|
52
|
#[inline]
|
|
47
|
53
|
fn insert(&mut self, key: String, peer: Peer) {
|
|
48
|
|
- self.map.write().unwrap().insert(key, peer);
|
|
|
54
|
+ if self.map.write().unwrap().insert(key, peer).is_none() {}
|
|
49
|
55
|
}
|
|
50
|
56
|
|
|
51
|
57
|
#[inline]
|
|
52
|
|
- fn get(&self, key: &str) -> Option<Peer> {
|
|
53
|
|
- self.map.read().unwrap().get(key).map(|x| x.clone())
|
|
|
58
|
+ async fn get(&mut self, key: String) -> Option<Peer> {
|
|
|
59
|
+ let p = self.map.read().unwrap().get(&key).map(|x| x.clone());
|
|
|
60
|
+ if p.is_some() {
|
|
|
61
|
+ return p;
|
|
|
62
|
+ } else {
|
|
|
63
|
+ if let Some(_) = self.db.get(key).await {
|
|
|
64
|
+ // to-do
|
|
|
65
|
+ }
|
|
|
66
|
+ }
|
|
|
67
|
+ None
|
|
|
68
|
+ }
|
|
|
69
|
+
|
|
|
70
|
+ #[inline]
|
|
|
71
|
+ fn is_in_memory(&self, key: &str) -> bool {
|
|
|
72
|
+ self.map.read().unwrap().contains_key(key)
|
|
54
|
73
|
}
|
|
55
|
74
|
}
|
|
56
|
75
|
|
|
|
@@ -93,7 +112,7 @@ impl RendezvousServer {
|
|
93
|
112
|
if let Ok(msg_in) = parse_from_bytes::<RendezvousMessage>(&bytes) {
|
|
94
|
113
|
match msg_in.union {
|
|
95
|
114
|
Some(rendezvous_message::Union::punch_hole_request(ph)) => {
|
|
96
|
|
- allow_err!(rs.handle_tcp_punch_hole_request(addr, &ph.id).await);
|
|
|
115
|
+ allow_err!(rs.handle_tcp_punch_hole_request(addr, ph.id).await);
|
|
97
|
116
|
}
|
|
98
|
117
|
Some(rendezvous_message::Union::punch_hole_sent(phs)) => {
|
|
99
|
118
|
allow_err!(rs.handle_hole_sent(&phs, addr, None).await);
|
|
|
@@ -141,7 +160,15 @@ impl RendezvousServer {
|
|
141
|
160
|
}
|
|
142
|
161
|
}
|
|
143
|
162
|
Some(rendezvous_message::Union::punch_hole_request(ph)) => {
|
|
144
|
|
- self.handle_udp_punch_hole_request(addr, &ph.id).await?;
|
|
|
163
|
+ let id = ph.id;
|
|
|
164
|
+ if self.pm.is_in_memory(&id) {
|
|
|
165
|
+ self.handle_udp_punch_hole_request(addr, id).await?;
|
|
|
166
|
+ } else {
|
|
|
167
|
+ let mut me = self.clone();
|
|
|
168
|
+ tokio::spawn(async move {
|
|
|
169
|
+ allow_err!(me.handle_udp_punch_hole_request(addr, id).await);
|
|
|
170
|
+ });
|
|
|
171
|
+ }
|
|
145
|
172
|
}
|
|
146
|
173
|
Some(rendezvous_message::Union::punch_hole_sent(phs)) => {
|
|
147
|
174
|
self.handle_hole_sent(&phs, addr, Some(socket)).await?;
|
|
|
@@ -215,14 +242,14 @@ impl RendezvousServer {
|
|
215
|
242
|
async fn handle_punch_hole_request(
|
|
216
|
243
|
&mut self,
|
|
217
|
244
|
addr: SocketAddr,
|
|
218
|
|
- id: &str,
|
|
|
245
|
+ id: String,
|
|
219
|
246
|
) -> ResultType<(RendezvousMessage, Option<SocketAddr>)> {
|
|
220
|
247
|
// punch hole request from A, forward to B,
|
|
221
|
248
|
// check if in same intranet first,
|
|
222
|
249
|
// fetch local addrs if in same intranet.
|
|
223
|
250
|
// because punch hole won't work if in the same intranet,
|
|
224
|
251
|
// all routers will drop such self-connections.
|
|
225
|
|
- if let Some(peer) = self.pm.get(id) {
|
|
|
252
|
+ if let Some(peer) = self.pm.get(id.clone()).await {
|
|
226
|
253
|
if peer.last_reg_time.elapsed().as_millis() as i32 >= REG_TIMEOUT {
|
|
227
|
254
|
let mut msg_out = RendezvousMessage::new();
|
|
228
|
255
|
msg_out.set_punch_hole_response(PunchHoleResponse {
|
|
|
@@ -308,7 +335,7 @@ impl RendezvousServer {
|
|
308
|
335
|
async fn handle_tcp_punch_hole_request(
|
|
309
|
336
|
&mut self,
|
|
310
|
337
|
addr: SocketAddr,
|
|
311
|
|
- id: &str,
|
|
|
338
|
+ id: String,
|
|
312
|
339
|
) -> ResultType<()> {
|
|
313
|
340
|
let (msg, to_addr) = self.handle_punch_hole_request(addr, id).await?;
|
|
314
|
341
|
if let Some(addr) = to_addr {
|
|
|
@@ -323,7 +350,7 @@ impl RendezvousServer {
|
|
323
|
350
|
async fn handle_udp_punch_hole_request(
|
|
324
|
351
|
&mut self,
|
|
325
|
352
|
addr: SocketAddr,
|
|
326
|
|
- id: &str,
|
|
|
353
|
+ id: String,
|
|
327
|
354
|
) -> ResultType<()> {
|
|
328
|
355
|
let (msg, to_addr) = self.handle_punch_hole_request(addr, id).await?;
|
|
329
|
356
|
self.tx.send((
|
open-trade