handle request_pk

This commit is contained in:
open-trade 2020-07-09 02:16:59 +08:00
parent d837379a13
commit d16fb31ecf
3 changed files with 86 additions and 50 deletions

@ -1 +1 @@
Subproject commit 5e977383041f35f856ee2dfc5fcf07c8300e2da5 Subproject commit c86ebe2402ee1f092ce8e44e89368234708a766f

View File

@ -65,38 +65,10 @@ impl PeerMap {
} }
#[inline] #[inline]
async fn update_addr(&mut self, key: String, socket_addr: SocketAddr) { fn update_pk(&mut self, id: String, socket_addr: SocketAddr, pk: Vec<u8>) {
let mut lock = self.map.write().unwrap();
let last_reg_time = Instant::now();
if let Some(old) = lock.get_mut(&key) {
old.socket_addr = socket_addr;
old.last_reg_time = last_reg_time;
} else {
let mut me = self.clone();
tokio::spawn(async move {
let v = me.db.get(key.clone()).await;
let pk = if let Some(v) = super::SledAsync::deserialize::<PeerSerde>(&v) {
v.pk
} else {
Vec::new()
};
me.map.write().unwrap().insert(
key,
Peer {
socket_addr,
last_reg_time,
pk,
},
);
});
}
}
#[inline]
fn update_key(&mut self, key: String, socket_addr: SocketAddr, pk: Vec<u8>) {
let mut lock = self.map.write().unwrap(); let mut lock = self.map.write().unwrap();
lock.insert( lock.insert(
key.clone(), id.clone(),
Peer { Peer {
socket_addr, socket_addr,
last_reg_time: Instant::now(), last_reg_time: Instant::now(),
@ -104,19 +76,20 @@ impl PeerMap {
}, },
); );
let ip = socket_addr.ip().to_string(); let ip = socket_addr.ip().to_string();
self.db.insert(key, PeerSerde { ip, pk }); self.db.insert(id, PeerSerde { ip, pk });
} }
#[inline] #[inline]
async fn get(&mut self, key: String) -> Option<Peer> { async fn get(&mut self, id: &str) -> Option<Peer> {
let p = self.map.read().unwrap().get(&key).map(|x| x.clone()); let p = self.map.read().unwrap().get(id).map(|x| x.clone());
if p.is_some() { if p.is_some() {
return p; return p;
} else { } else {
let v = self.db.get(key.clone()).await; let id = id.to_owned();
let v = self.db.get(id.clone()).await;
if let Some(v) = super::SledAsync::deserialize::<PeerSerde>(&v) { if let Some(v) = super::SledAsync::deserialize::<PeerSerde>(&v) {
self.map.write().unwrap().insert( self.map.write().unwrap().insert(
key, id,
Peer { Peer {
pk: v.pk, pk: v.pk,
..Default::default() ..Default::default()
@ -129,19 +102,20 @@ impl PeerMap {
} }
#[inline] #[inline]
fn is_in_memory(&self, key: &str) -> bool { fn is_in_memory(&self, id: &str) -> bool {
self.map.read().unwrap().contains_key(key) self.map.read().unwrap().contains_key(id)
} }
} }
const REG_TIMEOUT: i32 = 30_000; const REG_TIMEOUT: i32 = 30_000;
type Sink = SplitSink<Framed<TcpStream, BytesCodec>, Bytes>; type Sink = SplitSink<Framed<TcpStream, BytesCodec>, Bytes>;
type Sender = mpsc::UnboundedSender<(RendezvousMessage, SocketAddr)>;
#[derive(Clone)] #[derive(Clone)]
pub struct RendezvousServer { pub struct RendezvousServer {
tcp_punch: Arc<Mutex<HashMap<SocketAddr, Sink>>>, tcp_punch: Arc<Mutex<HashMap<SocketAddr, Sink>>>,
pm: PeerMap, pm: PeerMap,
tx: mpsc::UnboundedSender<(RendezvousMessage, SocketAddr)>, tx: Sender,
} }
impl RendezvousServer { impl RendezvousServer {
@ -208,19 +182,29 @@ impl RendezvousServer {
// B registered // B registered
if rp.id.len() > 0 { if rp.id.len() > 0 {
log::debug!("New peer registered: {:?} {:?}", &rp.id, &addr); log::debug!("New peer registered: {:?} {:?}", &rp.id, &addr);
self.pm.update_addr(rp.id, addr).await; self.update_addr(rp.id, addr, socket).await?;
let mut msg_out = RendezvousMessage::new();
msg_out.set_register_peer_response(RegisterPeerResponse::default());
socket.send(&msg_out, addr).await?
} }
} }
Some(rendezvous_message::Union::register_key(rk)) => { Some(rendezvous_message::Union::register_pk(rk)) => {
let id = rk.id; let id = rk.id;
if let Some(peer) = self.pm.get(id.clone()).await { let mut res = register_pk_response::Result::OK;
if let Some(peer) = self.pm.get(&id).await {
if peer.pk.is_empty() { if peer.pk.is_empty() {
self.pm.update_key(id, addr, rk.key); self.pm.update_pk(id, addr, rk.pk);
} else {
if peer.pk != rk.pk {
res = register_pk_response::Result::PK_MISMATCH;
} }
} }
} else {
self.pm.update_pk(id, addr, rk.pk);
}
let mut msg_out = RendezvousMessage::new();
msg_out.set_register_pk_response(RegisterPkResponse {
result: res.into(),
..Default::default()
});
socket.send(&msg_out, addr).await?
} }
Some(rendezvous_message::Union::punch_hole_request(ph)) => { Some(rendezvous_message::Union::punch_hole_request(ph)) => {
let id = ph.id; let id = ph.id;
@ -249,6 +233,58 @@ impl RendezvousServer {
Ok(()) Ok(())
} }
#[inline]
async fn update_addr(
&mut self,
id: String,
socket_addr: SocketAddr,
socket: &mut FramedSocket,
) -> ResultType<()> {
let mut lock = self.pm.map.write().unwrap();
let last_reg_time = Instant::now();
if let Some(old) = lock.get_mut(&id) {
old.socket_addr = socket_addr;
old.last_reg_time = last_reg_time;
let request_pk = old.pk.is_empty();
drop(lock);
let mut msg_out = RendezvousMessage::new();
msg_out.set_register_peer_response(RegisterPeerResponse {
request_pk,
..Default::default()
});
socket.send(&msg_out, socket_addr).await?;
} else {
drop(lock);
let mut pm = self.pm.clone();
let tx = self.tx.clone();
tokio::spawn(async move {
let v = pm.db.get(id.clone()).await;
let pk = {
if let Some(v) = super::SledAsync::deserialize::<PeerSerde>(&v) {
v.pk
} else {
Vec::new()
}
};
let mut msg_out = RendezvousMessage::new();
msg_out.set_register_peer_response(RegisterPeerResponse {
request_pk: pk.is_empty(),
..Default::default()
});
tx.send((msg_out, socket_addr)).ok();
pm.map.write().unwrap().insert(
id,
Peer {
socket_addr,
last_reg_time,
pk,
},
);
});
}
Ok(())
}
#[inline] #[inline]
async fn handle_hole_sent<'a>( async fn handle_hole_sent<'a>(
&mut self, &mut self,
@ -316,7 +352,7 @@ impl RendezvousServer {
// fetch local addrs if in same intranet. // fetch local addrs if in same intranet.
// because punch hole won't work if in the same intranet, // because punch hole won't work if in the same intranet,
// all routers will drop such self-connections. // all routers will drop such self-connections.
if let Some(peer) = self.pm.get(id.clone()).await { if let Some(peer) = self.pm.get(&id).await {
if peer.last_reg_time.elapsed().as_millis() as i32 >= REG_TIMEOUT { if peer.last_reg_time.elapsed().as_millis() as i32 >= REG_TIMEOUT {
let mut msg_out = RendezvousMessage::new(); let mut msg_out = RendezvousMessage::new();
msg_out.set_punch_hole_response(PunchHoleResponse { msg_out.set_punch_hole_response(PunchHoleResponse {

View File

@ -8,7 +8,7 @@ use hbb_common::{
enum Action { enum Action {
Insert((String, Vec<u8>)), Insert((String, Vec<u8>)),
Get((String, mpsc::Sender<Option<sled::IVec>>)), Get((String, mpsc::Sender<Option<sled::IVec>>)),
Close, _Close,
} }
#[derive(Clone)] #[derive(Clone)]
@ -55,14 +55,14 @@ impl SledAsync {
.await .await
); );
} }
Action::Close => break, Action::_Close => break,
} }
} }
} }
pub fn _close(self, j: std::thread::JoinHandle<()>) { pub fn _close(self, j: std::thread::JoinHandle<()>) {
if let Some(tx) = &self.tx { if let Some(tx) = &self.tx {
allow_err!(tx.send(Action::Close)); allow_err!(tx.send(Action::_Close));
} }
allow_err!(j.join()); allow_err!(j.join());
} }