mirror/rustdesk
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk.git synced 2025-01-19 08:23:01 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

refactor to prepare for tcp rendezvous

Browse Source
This commit is contained in:
rustdesk 2024-01-07 19:01:35 +08:00
parent 955c55b6cc
commit 7e93a5d3cf
3 changed files with 238 additions and 162 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
src/client.rs
View File

@ -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(())
}

90
src/common.rs
View File

@ -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)
}

225
src/rendezvous_mediator.rs
View File

@ -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;