mirror of
https://github.com/rustdesk/rustdesk.git
synced 2025-01-19 16:33:01 +08:00
refactor to prepare for tcp rendezvous
This commit is contained in:
rustdesk
parent
955c55b6cc
commit
7e93a5d3cf
@ -41,7 +41,7 @@ use hbb_common::{
|
||||
rendezvous_proto::*,
|
||||
socket_client,
|
||||
sodiumoxide::base64,
|
||||
sodiumoxide::crypto::{box_, secretbox, sign},
|
||||
sodiumoxide::crypto::sign,
|
||||
tcp::FramedStream,
|
||||
timeout,
|
||||
tokio::time::Duration,
|
||||
@ -57,7 +57,7 @@ use scrap::{
|
||||
use crate::{
|
||||
check_port,
|
||||
common::input::{MOUSE_BUTTON_LEFT, MOUSE_BUTTON_RIGHT, MOUSE_TYPE_DOWN, MOUSE_TYPE_UP},
|
||||
is_keyboard_mode_supported,
|
||||
create_symmetric_key_msg, decode_id_pk, get_rs_pk, is_keyboard_mode_supported, secure_tcp,
|
||||
ui_session_interface::{InvokeUiSession, Session},
|
||||
};
|
||||
|
||||
@ -311,7 +311,7 @@ impl Client {
|
||||
|
||||
if !key.is_empty() && !token.is_empty() {
|
||||
// mainly for the security of token
|
||||
allow_err!(secure_punch_connection(&mut socket, key).await);
|
||||
allow_err!(secure_tcp(&mut socket, key).await);
|
||||
}
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
@ -620,7 +620,7 @@ impl Client {
|
||||
|
||||
if !key.is_empty() && !token.is_empty() {
|
||||
// mainly for the security of token
|
||||
allow_err!(secure_punch_connection(&mut socket, key).await);
|
||||
allow_err!(secure_tcp(&mut socket, key).await);
|
||||
}
|
||||
|
||||
ipv4 = socket.local_addr().is_ipv4();
|
||||
@ -2996,80 +2996,3 @@ pub fn check_if_retry(msgtype: &str, title: &str, text: &str, retry_for_relay: b
|
||||
&& !text.to_lowercase().contains("manually")
|
||||
&& !text.to_lowercase().contains("not allowed")))
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn get_pk(pk: &[u8]) -> Option<[u8; 32]> {
|
||||
if pk.len() == 32 {
|
||||
let mut tmp = [0u8; 32];
|
||||
tmp[..].copy_from_slice(&pk);
|
||||
Some(tmp)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn get_rs_pk(str_base64: &str) -> Option<sign::PublicKey> {
|
||||
if let Ok(pk) = crate::decode64(str_base64) {
|
||||
get_pk(&pk).map(|x| sign::PublicKey(x))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_id_pk(signed: &[u8], key: &sign::PublicKey) -> ResultType<(String, [u8; 32])> {
|
||||
let res = IdPk::parse_from_bytes(
|
||||
&sign::verify(signed, key).map_err(|_| anyhow!("Signature mismatch"))?,
|
||||
)?;
|
||||
if let Some(pk) = get_pk(&res.pk) {
|
||||
Ok((res.id, pk))
|
||||
} else {
|
||||
bail!("Wrong their public length");
|
||||
}
|
||||
}
|
||||
|
||||
fn create_symmetric_key_msg(their_pk_b: [u8; 32]) -> (Bytes, Bytes, secretbox::Key) {
|
||||
let their_pk_b = box_::PublicKey(their_pk_b);
|
||||
let (our_pk_b, out_sk_b) = box_::gen_keypair();
|
||||
let key = secretbox::gen_key();
|
||||
let nonce = box_::Nonce([0u8; box_::NONCEBYTES]);
|
||||
let sealed_key = box_::seal(&key.0, &nonce, &their_pk_b, &out_sk_b);
|
||||
(Vec::from(our_pk_b.0).into(), sealed_key.into(), key)
|
||||
}
|
||||
|
||||
async fn secure_punch_connection(conn: &mut FramedStream, key: &str) -> ResultType<()> {
|
||||
let rs_pk = get_rs_pk(key);
|
||||
let Some(rs_pk) = rs_pk else {
|
||||
bail!("Handshake failed: invalid public key from rendezvous server");
|
||||
};
|
||||
match timeout(READ_TIMEOUT, conn.next()).await? {
|
||||
Some(Ok(bytes)) => {
|
||||
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
|
||||
match msg_in.union {
|
||||
Some(rendezvous_message::Union::KeyExchange(ex)) => {
|
||||
if ex.keys.len() != 1 {
|
||||
bail!("Handshake failed: invalid key exchange message");
|
||||
}
|
||||
let their_pk_b = sign::verify(&ex.keys[0], &rs_pk)
|
||||
.map_err(|_| anyhow!("Signature mismatch in key exchange"))?;
|
||||
let (asymmetric_value, symmetric_value, key) = create_symmetric_key_msg(
|
||||
get_pk(&their_pk_b)
|
||||
.context("Wrong their public length in key exchange")?,
|
||||
);
|
||||
let mut msg_out = RendezvousMessage::new();
|
||||
msg_out.set_key_exchange(KeyExchange {
|
||||
keys: vec![asymmetric_value, symmetric_value],
|
||||
..Default::default()
|
||||
});
|
||||
timeout(CONNECT_TIMEOUT, conn.send(&msg_out)).await??;
|
||||
conn.set_key(key);
|
||||
log::info!("Token secured");
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -127,6 +127,9 @@ impl ClipboardContext {
|
||||
use hbb_common::compress::decompress;
|
||||
use hbb_common::{
|
||||
allow_err,
|
||||
anyhow::{anyhow, Context},
|
||||
bail,
|
||||
bytes::Bytes,
|
||||
compress::compress as compress_func,
|
||||
config::{self, Config, CONNECT_TIMEOUT, READ_TIMEOUT},
|
||||
get_version_number, log,
|
||||
@ -135,8 +138,9 @@ use hbb_common::{
|
||||
protobuf::Message as _,
|
||||
rendezvous_proto::*,
|
||||
socket_client,
|
||||
sodiumoxide::crypto::{box_, secretbox, sign},
|
||||
tcp::FramedStream,
|
||||
tokio, ResultType,
|
||||
timeout, tokio, ResultType,
|
||||
};
|
||||
// #[cfg(any(target_os = "android", target_os = "ios", feature = "cli"))]
|
||||
use hbb_common::{config::RENDEZVOUS_PORT, futures::future::join_all};
|
||||
@ -1076,7 +1080,10 @@ pub fn make_privacy_mode_msg_with_details(
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn make_privacy_mode_msg(state: back_notification::PrivacyModeState, impl_key: String) -> Message {
|
||||
pub fn make_privacy_mode_msg(
|
||||
state: back_notification::PrivacyModeState,
|
||||
impl_key: String,
|
||||
) -> Message {
|
||||
make_privacy_mode_msg_with_details(state, "".to_owned(), impl_key)
|
||||
}
|
||||
|
||||
@ -1165,7 +1172,7 @@ pub async fn get_key(sync: bool) -> String {
|
||||
let mut options = crate::ipc::get_options_async().await;
|
||||
options.remove("key").unwrap_or_default()
|
||||
};
|
||||
if key.is_empty() && !option_env!("RENDEZVOUS_SERVER").unwrap_or("").is_empty() {
|
||||
if key.is_empty() {
|
||||
key = config::RS_PUB_KEY.to_owned();
|
||||
}
|
||||
key
|
||||
@ -1244,3 +1251,80 @@ pub fn check_process(arg: &str, same_uid: bool) -> bool {
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
pub async fn secure_tcp(conn: &mut FramedStream, key: &str) -> ResultType<()> {
|
||||
let rs_pk = get_rs_pk(key);
|
||||
let Some(rs_pk) = rs_pk else {
|
||||
bail!("Handshake failed: invalid public key from rendezvous server");
|
||||
};
|
||||
match timeout(READ_TIMEOUT, conn.next()).await? {
|
||||
Some(Ok(bytes)) => {
|
||||
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
|
||||
match msg_in.union {
|
||||
Some(rendezvous_message::Union::KeyExchange(ex)) => {
|
||||
if ex.keys.len() != 1 {
|
||||
bail!("Handshake failed: invalid key exchange message");
|
||||
}
|
||||
let their_pk_b = sign::verify(&ex.keys[0], &rs_pk)
|
||||
.map_err(|_| anyhow!("Signature mismatch in key exchange"))?;
|
||||
let (asymmetric_value, symmetric_value, key) = create_symmetric_key_msg(
|
||||
get_pk(&their_pk_b)
|
||||
.context("Wrong their public length in key exchange")?,
|
||||
);
|
||||
let mut msg_out = RendezvousMessage::new();
|
||||
msg_out.set_key_exchange(KeyExchange {
|
||||
keys: vec![asymmetric_value, symmetric_value],
|
||||
..Default::default()
|
||||
});
|
||||
timeout(CONNECT_TIMEOUT, conn.send(&msg_out)).await??;
|
||||
conn.set_key(key);
|
||||
log::info!("Token secured");
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn get_pk(pk: &[u8]) -> Option<[u8; 32]> {
|
||||
if pk.len() == 32 {
|
||||
let mut tmp = [0u8; 32];
|
||||
tmp[..].copy_from_slice(&pk);
|
||||
Some(tmp)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn get_rs_pk(str_base64: &str) -> Option<sign::PublicKey> {
|
||||
if let Ok(pk) = crate::decode64(str_base64) {
|
||||
get_pk(&pk).map(|x| sign::PublicKey(x))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decode_id_pk(signed: &[u8], key: &sign::PublicKey) -> ResultType<(String, [u8; 32])> {
|
||||
let res = IdPk::parse_from_bytes(
|
||||
&sign::verify(signed, key).map_err(|_| anyhow!("Signature mismatch"))?,
|
||||
)?;
|
||||
if let Some(pk) = get_pk(&res.pk) {
|
||||
Ok((res.id, pk))
|
||||
} else {
|
||||
bail!("Wrong their public length");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn create_symmetric_key_msg(their_pk_b: [u8; 32]) -> (Bytes, Bytes, secretbox::Key) {
|
||||
let their_pk_b = box_::PublicKey(their_pk_b);
|
||||
let (our_pk_b, out_sk_b) = box_::gen_keypair();
|
||||
let key = secretbox::gen_key();
|
||||
let nonce = box_::Nonce([0u8; box_::NONCEBYTES]);
|
||||
let sealed_key = box_::seal(&key.0, &nonce, &their_pk_b, &out_sk_b);
|
||||
(Vec::from(our_pk_b.0).into(), sealed_key.into(), key)
|
||||
}
|
||||
|
||||
@ -9,29 +9,34 @@ use std::{
|
||||
|
||||
use uuid::Uuid;
|
||||
|
||||
use hbb_common::tcp::FramedStream;
|
||||
use hbb_common::{
|
||||
allow_err,
|
||||
anyhow::bail,
|
||||
anyhow::{self, bail},
|
||||
config::{Config, CONNECT_TIMEOUT, READ_TIMEOUT, REG_INTERVAL, RENDEZVOUS_PORT},
|
||||
futures::future::join_all,
|
||||
log,
|
||||
protobuf::Message as _,
|
||||
rendezvous_proto::*,
|
||||
sleep,
|
||||
socket_client::{self, is_ipv4},
|
||||
socket_client::{self, connect_tcp, is_ipv4},
|
||||
tcp::FramedStream,
|
||||
tokio::{
|
||||
self, select,
|
||||
time::{interval, Duration},
|
||||
},
|
||||
udp::FramedSocket,
|
||||
AddrMangle, ResultType,
|
||||
AddrMangle, IntoTargetAddr, ResultType, TargetAddr,
|
||||
};
|
||||
|
||||
use crate::server::{check_zombie, new as new_server, ServerPtr};
|
||||
use crate::{
|
||||
check_port,
|
||||
server::{check_zombie, new as new_server, ServerPtr},
|
||||
};
|
||||
|
||||
type Message = RendezvousMessage;
|
||||
|
||||
const TIMER_OUT: Duration = Duration::from_secs(1);
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref SOLVING_PK_MISMATCH: Arc<Mutex<String>> = Default::default();
|
||||
}
|
||||
@ -39,7 +44,7 @@ static SHOULD_EXIT: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RendezvousMediator {
|
||||
addr: hbb_common::tokio_socks::TargetAddr<'static>,
|
||||
addr: TargetAddr<'static>,
|
||||
host: String,
|
||||
host_prefix: String,
|
||||
last_id_pk_registry: String,
|
||||
@ -112,8 +117,7 @@ impl RendezvousMediator {
|
||||
// crate::platform::linux_desktop_manager::stop_xdesktop();
|
||||
}
|
||||
|
||||
pub async fn start(server: ServerPtr, host: String) -> ResultType<()> {
|
||||
log::info!("start rendezvous mediator of {}", host);
|
||||
pub async fn start_udp(server: ServerPtr, host: String) -> ResultType<()> {
|
||||
let host_prefix: String = host
|
||||
.split(".")
|
||||
.next()
|
||||
@ -125,16 +129,15 @@ impl RendezvousMediator {
|
||||
}
|
||||
})
|
||||
.unwrap_or(host.to_owned());
|
||||
let host = crate::check_port(&host, RENDEZVOUS_PORT);
|
||||
let host = check_port(&host, RENDEZVOUS_PORT);
|
||||
let (mut socket, addr) = socket_client::new_udp_for(&host, CONNECT_TIMEOUT).await?;
|
||||
let mut rz = Self {
|
||||
addr: addr,
|
||||
addr: addr.clone(),
|
||||
host: host.clone(),
|
||||
host_prefix,
|
||||
last_id_pk_registry: "".to_owned(),
|
||||
};
|
||||
|
||||
const TIMER_OUT: Duration = Duration::from_secs(1);
|
||||
let mut timer = interval(TIMER_OUT);
|
||||
let mut last_timer: Option<Instant> = None;
|
||||
const REG_TIMEOUT: i64 = 3_000;
|
||||
@ -177,63 +180,8 @@ impl RendezvousMediator {
|
||||
n = socket.next() => {
|
||||
match n {
|
||||
Some(Ok((bytes, _))) => {
|
||||
if let Ok(msg_in) = Message::parse_from_bytes(&bytes) {
|
||||
match msg_in.union {
|
||||
Some(rendezvous_message::Union::RegisterPeerResponse(rpr)) => {
|
||||
update_latency();
|
||||
if rpr.request_pk {
|
||||
log::info!("request_pk received from {}", host);
|
||||
allow_err!(rz.register_pk(&mut socket).await);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
Some(rendezvous_message::Union::RegisterPkResponse(rpr)) => {
|
||||
update_latency();
|
||||
match rpr.result.enum_value() {
|
||||
Ok(register_pk_response::Result::OK) => {
|
||||
Config::set_key_confirmed(true);
|
||||
Config::set_host_key_confirmed(&rz.host_prefix, true);
|
||||
*SOLVING_PK_MISMATCH.lock().unwrap() = "".to_owned();
|
||||
}
|
||||
Ok(register_pk_response::Result::UUID_MISMATCH) => {
|
||||
allow_err!(rz.handle_uuid_mismatch(&mut socket).await);
|
||||
}
|
||||
_ => {
|
||||
log::error!("unknown RegisterPkResponse");
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(rendezvous_message::Union::PunchHole(ph)) => {
|
||||
let rz = rz.clone();
|
||||
let server = server.clone();
|
||||
tokio::spawn(async move {
|
||||
allow_err!(rz.handle_punch_hole(ph, server).await);
|
||||
});
|
||||
}
|
||||
Some(rendezvous_message::Union::RequestRelay(rr)) => {
|
||||
let rz = rz.clone();
|
||||
let server = server.clone();
|
||||
tokio::spawn(async move {
|
||||
allow_err!(rz.handle_request_relay(rr, server).await);
|
||||
});
|
||||
}
|
||||
Some(rendezvous_message::Union::FetchLocalAddr(fla)) => {
|
||||
let rz = rz.clone();
|
||||
let server = server.clone();
|
||||
tokio::spawn(async move {
|
||||
allow_err!(rz.handle_intranet(fla, server).await);
|
||||
});
|
||||
}
|
||||
Some(rendezvous_message::Union::ConfigureUpdate(cu)) => {
|
||||
let v0 = Config::get_rendezvous_servers();
|
||||
Config::set_option("rendezvous-servers".to_owned(), cu.rendezvous_servers.join(","));
|
||||
Config::set_serial(cu.serial);
|
||||
if v0 != Config::get_rendezvous_servers() {
|
||||
Self::restart();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
if let Ok(msg) = Message::parse_from_bytes(&bytes) {
|
||||
rz.handle_resp(msg.union, Sink::Framed(&mut socket, &addr), &server, &mut update_latency).await?;
|
||||
} else {
|
||||
log::debug!("Non-protobuf message bytes received: {:?}", bytes);
|
||||
}
|
||||
@ -257,7 +205,7 @@ impl RendezvousMediator {
|
||||
let elapsed_resp = last_register_resp.map(|x| x.elapsed().as_millis() as i64).unwrap_or(REG_INTERVAL);
|
||||
let timeout = (elapsed_resp - last_register_sent.map(|x| x.elapsed().as_millis() as i64).unwrap_or(REG_INTERVAL)) > REG_TIMEOUT;
|
||||
if timeout || elapsed_resp >= REG_INTERVAL {
|
||||
allow_err!(rz.register_peer(&mut socket).await);
|
||||
allow_err!(rz.register_peer(Sink::Framed(&mut socket, &addr)).await);
|
||||
last_register_sent = now;
|
||||
if timeout {
|
||||
fails += 1;
|
||||
@ -285,6 +233,113 @@ impl RendezvousMediator {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[inline]
|
||||
async fn handle_resp(
|
||||
&mut self,
|
||||
msg: Option<rendezvous_message::Union>,
|
||||
sink: Sink<'_>,
|
||||
server: &ServerPtr,
|
||||
update_latency: &mut impl FnMut(),
|
||||
) -> ResultType<()> {
|
||||
match msg {
|
||||
Some(rendezvous_message::Union::RegisterPeerResponse(rpr)) => {
|
||||
update_latency();
|
||||
if rpr.request_pk {
|
||||
log::info!("request_pk received from {}", self.host);
|
||||
allow_err!(self.register_pk(sink).await);
|
||||
}
|
||||
}
|
||||
Some(rendezvous_message::Union::RegisterPkResponse(rpr)) => {
|
||||
update_latency();
|
||||
match rpr.result.enum_value() {
|
||||
Ok(register_pk_response::Result::OK) => {
|
||||
Config::set_key_confirmed(true);
|
||||
Config::set_host_key_confirmed(&self.host_prefix, true);
|
||||
*SOLVING_PK_MISMATCH.lock().unwrap() = "".to_owned();
|
||||
}
|
||||
Ok(register_pk_response::Result::UUID_MISMATCH) => {
|
||||
allow_err!(self.handle_uuid_mismatch(sink).await);
|
||||
}
|
||||
_ => {
|
||||
log::error!("unknown RegisterPkResponse");
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(rendezvous_message::Union::PunchHole(ph)) => {
|
||||
let rz = self.clone();
|
||||
let server = server.clone();
|
||||
tokio::spawn(async move {
|
||||
allow_err!(rz.handle_punch_hole(ph, server).await);
|
||||
});
|
||||
}
|
||||
Some(rendezvous_message::Union::RequestRelay(rr)) => {
|
||||
let rz = self.clone();
|
||||
let server = server.clone();
|
||||
tokio::spawn(async move {
|
||||
allow_err!(rz.handle_request_relay(rr, server).await);
|
||||
});
|
||||
}
|
||||
Some(rendezvous_message::Union::FetchLocalAddr(fla)) => {
|
||||
let rz = self.clone();
|
||||
let server = server.clone();
|
||||
tokio::spawn(async move {
|
||||
allow_err!(rz.handle_intranet(fla, server).await);
|
||||
});
|
||||
}
|
||||
Some(rendezvous_message::Union::ConfigureUpdate(cu)) => {
|
||||
let v0 = Config::get_rendezvous_servers();
|
||||
Config::set_option(
|
||||
"rendezvous-servers".to_owned(),
|
||||
cu.rendezvous_servers.join(","),
|
||||
);
|
||||
Config::set_serial(cu.serial);
|
||||
if v0 != Config::get_rendezvous_servers() {
|
||||
Self::restart();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn start_tcp(server: ServerPtr, host: String) -> ResultType<()> {
|
||||
let mut conn = connect_tcp(check_port(&host, RENDEZVOUS_PORT), CONNECT_TIMEOUT).await?;
|
||||
let key = crate::get_key(true).await;
|
||||
crate::secure_tcp(&mut conn, &key).await?;
|
||||
let mut rz = Self {
|
||||
addr: conn.local_addr().into_target_addr()?,
|
||||
host: host.clone(),
|
||||
host_prefix: host.clone(),
|
||||
last_id_pk_registry: "".to_owned(),
|
||||
};
|
||||
let mut timer = interval(TIMER_OUT);
|
||||
loop {
|
||||
let mut update_latency = || {};
|
||||
select! {
|
||||
res = conn.next() => {
|
||||
let bytes = res.ok_or_else(|| anyhow::anyhow!("rendezvous server disconnected"))??;
|
||||
let msg = Message::parse_from_bytes(&bytes)?;
|
||||
rz.handle_resp(msg.union, Sink::Stream(&mut conn), &server, &mut update_latency).await?
|
||||
}
|
||||
_ = timer.tick() => {
|
||||
if SHOULD_EXIT.load(Ordering::SeqCst) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn start(server: ServerPtr, host: String) -> ResultType<()> {
|
||||
log::info!("start rendezvous mediator of {}", host);
|
||||
if cfg!(debug_assertions) && option_env!("TEST_TCP").is_some() {
|
||||
Self::start_tcp(server, host).await
|
||||
} else {
|
||||
Self::start_udp(server, host).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_request_relay(&self, rr: RequestRelay, server: ServerPtr) -> ResultType<()> {
|
||||
self.create_relay(
|
||||
rr.socket_addr.into(),
|
||||
@ -315,7 +370,7 @@ impl RendezvousMediator {
|
||||
secure,
|
||||
);
|
||||
|
||||
let mut socket = socket_client::connect_tcp(&*self.host, CONNECT_TIMEOUT).await?;
|
||||
let mut socket = connect_tcp(&*self.host, CONNECT_TIMEOUT).await?;
|
||||
|
||||
let mut msg_out = Message::new();
|
||||
let mut rr = RelayResponse {
|
||||
@ -360,7 +415,7 @@ impl RendezvousMediator {
|
||||
}
|
||||
let peer_addr = AddrMangle::decode(&fla.socket_addr);
|
||||
log::debug!("Handle intranet from {:?}", peer_addr);
|
||||
let mut socket = socket_client::connect_tcp(&*self.host, CONNECT_TIMEOUT).await?;
|
||||
let mut socket = connect_tcp(&*self.host, CONNECT_TIMEOUT).await?;
|
||||
let local_addr = socket.local_addr();
|
||||
let local_addr: SocketAddr =
|
||||
format!("{}:{}", local_addr.ip(), local_addr.port()).parse()?;
|
||||
@ -399,7 +454,7 @@ impl RendezvousMediator {
|
||||
let peer_addr = AddrMangle::decode(&ph.socket_addr);
|
||||
log::debug!("Punch hole to {:?}", peer_addr);
|
||||
let mut socket = {
|
||||
let socket = socket_client::connect_tcp(&*self.host, CONNECT_TIMEOUT).await?;
|
||||
let socket = connect_tcp(&*self.host, CONNECT_TIMEOUT).await?;
|
||||
let local_addr = socket.local_addr();
|
||||
// key important here for punch hole to tell my gateway incoming peer is safe.
|
||||
// it can not be async here, because local_addr can not be reused, we must close the connection before use it again.
|
||||
@ -423,7 +478,7 @@ impl RendezvousMediator {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn register_pk(&mut self, socket: &mut FramedSocket) -> ResultType<()> {
|
||||
async fn register_pk(&mut self, socket: Sink<'_>) -> ResultType<()> {
|
||||
let mut msg_out = Message::new();
|
||||
let pk = Config::get_key_pair().1;
|
||||
let uuid = hbb_common::get_uuid();
|
||||
@ -435,11 +490,11 @@ impl RendezvousMediator {
|
||||
pk: pk.into(),
|
||||
..Default::default()
|
||||
});
|
||||
socket.send(&msg_out, self.addr.to_owned()).await?;
|
||||
socket.send(&msg_out).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_uuid_mismatch(&mut self, socket: &mut FramedSocket) -> ResultType<()> {
|
||||
async fn handle_uuid_mismatch(&mut self, socket: Sink<'_>) -> ResultType<()> {
|
||||
if self.last_id_pk_registry != Config::get_id() {
|
||||
return Ok(());
|
||||
}
|
||||
@ -457,7 +512,7 @@ impl RendezvousMediator {
|
||||
self.register_pk(socket).await
|
||||
}
|
||||
|
||||
async fn register_peer(&mut self, socket: &mut FramedSocket) -> ResultType<()> {
|
||||
async fn register_peer(&mut self, socket: Sink<'_>) -> ResultType<()> {
|
||||
if !SOLVING_PK_MISMATCH.lock().unwrap().is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
@ -481,7 +536,7 @@ impl RendezvousMediator {
|
||||
serial,
|
||||
..Default::default()
|
||||
});
|
||||
socket.send(&msg_out, self.addr.to_owned()).await?;
|
||||
socket.send(&msg_out).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@ -622,7 +677,7 @@ async fn create_online_stream() -> ResultType<FramedStream> {
|
||||
bail!("Invalid server address: {}", rendezvous_server);
|
||||
}
|
||||
let online_server = format!("{}:{}", tmp[0], port - 1);
|
||||
socket_client::connect_tcp(online_server, CONNECT_TIMEOUT).await
|
||||
connect_tcp(online_server, CONNECT_TIMEOUT).await
|
||||
}
|
||||
|
||||
async fn query_online_states_(
|
||||
@ -680,6 +735,20 @@ async fn query_online_states_(
|
||||
}
|
||||
}
|
||||
|
||||
enum Sink<'a> {
|
||||
Framed(&'a mut FramedSocket, &'a TargetAddr<'a>),
|
||||
Stream(&'a mut FramedStream),
|
||||
}
|
||||
|
||||
impl Sink<'_> {
|
||||
async fn send(self, msg: &Message) -> ResultType<()> {
|
||||
match self {
|
||||
Sink::Framed(socket, addr) => socket.send(msg, addr.to_owned()).await,
|
||||
Sink::Stream(stream) => stream.send(msg).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use hbb_common::tokio;
|
||||
|
||||
Loading…
Reference in New Issue
Block a user