konevsa/rustdesk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor signed_pk

Browse Source
This commit is contained in:
rustdesk 2022-03-20 20:20:32 +08:00
parent 7c5d260738
commit 1e0a347893
3 changed files with 57 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
libs/hbb_common/src/lib.rs
View File

@ -31,11 +31,20 @@ pub use lazy_static;
pub use mac_address; pub use mac_address;
pub use rand; pub use rand;
pub use regex; pub use regex;
use serde_derive::{Deserialize, Serialize};
pub use sodiumoxide; pub use sodiumoxide;
pub use tokio_socks; pub use tokio_socks;
pub use tokio_socks::IntoTargetAddr; pub use tokio_socks::IntoTargetAddr;
pub use tokio_socks::TargetAddr; pub use tokio_socks::TargetAddr;
#[derive(Debug, Default, Serialize, Deserialize, Clone)]
pub struct IdPk {
#[serde(default)]
pub id: String,
#[serde(default)]
pub pk: [u8; 32],
}
#[cfg(feature = "quic")] #[cfg(feature = "quic")]
pub type Stream = quic::Connection; pub type Stream = quic::Connection;
#[cfg(not(feature = "quic"))] #[cfg(not(feature = "quic"))]

83
src/client.rs
View File

@ -17,7 +17,7 @@ use hbb_common::{
sodiumoxide::crypto::{box_, secretbox, sign}, sodiumoxide::crypto::{box_, secretbox, sign},
timeout, timeout,
tokio::time::Duration, tokio::time::Duration,
AddrMangle, ResultType, Stream, AddrMangle, IdPk, ResultType, Stream,
}; };
use magnum_opus::{Channels::*, Decoder as AudioDecoder}; use magnum_opus::{Channels::*, Decoder as AudioDecoder};
use scrap::{Decoder, Image, VideoCodecId}; use scrap::{Decoder, Image, VideoCodecId};
@ -136,7 +136,7 @@ impl Client {
let mut socket = let mut socket =
socket_client::connect_tcp(&*rendezvous_server, any_addr, RENDEZVOUS_TIMEOUT).await?; socket_client::connect_tcp(&*rendezvous_server, any_addr, RENDEZVOUS_TIMEOUT).await?;
let my_addr = socket.local_addr(); let my_addr = socket.local_addr();
let mut pk = Vec::new(); let mut signed_id_pk = Vec::new();
let mut relay_server = "".to_owned(); let mut relay_server = "".to_owned();
let start = std::time::Instant::now(); let start = std::time::Instant::now();
@ -161,6 +161,9 @@ impl Client {
match msg_in.union { match msg_in.union {
Some(rendezvous_message::Union::punch_hole_response(ph)) => { Some(rendezvous_message::Union::punch_hole_response(ph)) => {
if ph.socket_addr.is_empty() { if ph.socket_addr.is_empty() {
if !ph.other_failure.is_empty() {
bail!(ph.other_failure);
}
match ph.failure.enum_value_or_default() { match ph.failure.enum_value_or_default() {
punch_hole_response::Failure::ID_NOT_EXIST => { punch_hole_response::Failure::ID_NOT_EXIST => {
bail!("ID does not exist"); bail!("ID does not exist");
@ -171,16 +174,14 @@ impl Client {
punch_hole_response::Failure::LICENSE_MISMATCH => { punch_hole_response::Failure::LICENSE_MISMATCH => {
bail!("Key mismatch"); bail!("Key mismatch");
} }
_ => { punch_hole_response::Failure::LICENSE_OVERUSE => {
if !ph.other_failure.is_empty() { bail!("Key overuse");
bail!(ph.other_failure);
}
} }
} }
} else { } else {
peer_nat_type = ph.get_nat_type(); peer_nat_type = ph.get_nat_type();
is_local = ph.get_is_local(); is_local = ph.get_is_local();
pk = ph.pk; signed_id_pk = ph.pk;
relay_server = ph.relay_server; relay_server = ph.relay_server;
peer_addr = AddrMangle::decode(&ph.socket_addr); peer_addr = AddrMangle::decode(&ph.socket_addr);
log::info!("Hole Punched {} = {}", peer, peer_addr); log::info!("Hole Punched {} = {}", peer, peer_addr);
@ -193,11 +194,11 @@ impl Client {
start.elapsed(), start.elapsed(),
rr.relay_server rr.relay_server
); );
pk = rr.get_pk().into(); signed_id_pk = rr.get_pk().into();
let mut conn = let mut conn =
Self::create_relay(peer, rr.uuid, rr.relay_server, conn_type) Self::create_relay(peer, rr.uuid, rr.relay_server, conn_type)
.await?; .await?;
Self::secure_connection(peer, pk, &mut conn).await?; Self::secure_connection(peer, signed_id_pk, &mut conn).await?;
return Ok((conn, false)); return Ok((conn, false));
} }
_ => { _ => {
@ -228,7 +229,7 @@ impl Client {
my_addr, my_addr,
peer_addr, peer_addr,
peer, peer,
pk, signed_id_pk,
&relay_server, &relay_server,
&rendezvous_server, &rendezvous_server,
time_used, time_used,
@ -244,7 +245,7 @@ impl Client {
local_addr: SocketAddr, local_addr: SocketAddr,
peer: SocketAddr, peer: SocketAddr,
peer_id: &str, peer_id: &str,
pk: Vec<u8>, signed_id_pk: Vec<u8>,
relay_server: &str, relay_server: &str,
rendezvous_server: &str, rendezvous_server: &str,
punch_time_used: u64, punch_time_used: u64,
@ -296,7 +297,7 @@ impl Client {
peer_id, peer_id,
relay_server.to_owned(), relay_server.to_owned(),
rendezvous_server, rendezvous_server,
pk.len() == sign::PUBLICKEYBYTES, !signed_id_pk.is_empty(),
conn_type, conn_type,
) )
.await; .await;
@ -318,50 +319,46 @@ impl Client {
} }
let mut conn = conn?; let mut conn = conn?;
log::info!("{:?} used to establish connection", start.elapsed()); log::info!("{:?} used to establish connection", start.elapsed());
Self::secure_connection(peer_id, pk, &mut conn).await?; Self::secure_connection(peer_id, signed_id_pk, &mut conn).await?;
Ok((conn, direct)) Ok((conn, direct))
} }
async fn secure_connection(peer_id: &str, pk: Vec<u8>, conn: &mut Stream) -> ResultType<()> { async fn secure_connection(peer_id: &str, signed_id_pk: Vec<u8>, conn: &mut Stream) -> ResultType<()> {
let mut pk = pk;
let rs_pk = get_rs_pk("OeVuKk5nlHiXp+APNn0Y3pC1Iwpwn44JGqrQCsWqmBw="); let rs_pk = get_rs_pk("OeVuKk5nlHiXp+APNn0Y3pC1Iwpwn44JGqrQCsWqmBw=");
if !pk.is_empty() && rs_pk.is_some() { let mut sign_pk = None;
if let Ok(data) = sign::verify(&pk, &rs_pk.unwrap()) { if !signed_id_pk.is_empty() && rs_pk.is_some() {
pk = data; if let Ok(v) = serde_json::from_slice::<IdPk>(&signed_id_pk) {
} else { if v.id == peer_id {
sign_pk = Some(sign::PublicKey(v.pk));
}
}
if sign_pk.is_none() {
log::error!("Handshake failed: invalid public key from rendezvous server"); log::error!("Handshake failed: invalid public key from rendezvous server");
pk.clear();
} }
} }
if pk.len() != sign::PUBLICKEYBYTES { let sign_pk = match sign_pk {
// send an empty message out in case server is setting up secure and waiting for first message Some(v) => v,
conn.send(&Message::new()).await?; None => {
return Ok(()); // send an empty message out in case server is setting up secure and waiting for first message
} conn.send(&Message::new()).await?;
let mut tmp = [0u8; sign::PUBLICKEYBYTES]; return Ok(());
tmp[..].copy_from_slice(&pk); }
let sign_pk = sign::PublicKey(tmp); };
match timeout(CONNECT_TIMEOUT, conn.next()).await? { match timeout(CONNECT_TIMEOUT, conn.next()).await? {
Some(res) => { Some(res) => {
let bytes = res?; let bytes = res?;
if let Ok(msg_in) = Message::parse_from_bytes(&bytes) { if let Ok(msg_in) = Message::parse_from_bytes(&bytes) {
if let Some(message::Union::signed_id(si)) = msg_in.union { if let Some(message::Union::signed_id(si)) = msg_in.union {
if let Ok(data) = sign::verify(&si.id, &sign_pk) { if let Ok(data) = sign::verify(&si.id, &sign_pk) {
let s = String::from_utf8_lossy(&data); let (id, their_pk_b) = match serde_json::from_slice::<IdPk>(&data) {
let mut it = s.split("\0"); Ok(v) => (v.id, box_::PublicKey(v.pk)),
let id = it.next().unwrap_or_default(); Err(_) => {
let pk = log::error!(
base64::decode(it.next().unwrap_or_default()).unwrap_or_default(); "Handshake failed: invalid public box key length from peer"
let their_pk_b = if pk.len() == box_::PUBLICKEYBYTES { );
let mut pk_ = [0u8; box_::PUBLICKEYBYTES]; conn.send(&Message::new()).await?;
pk_[..].copy_from_slice(&pk); return Ok(());
box_::PublicKey(pk_) }
} else {
log::error!(
"Handshake failed: invalid public box key length from peer"
);
conn.send(&Message::new()).await?;
return Ok(());
}; };
if id == peer_id { if id == peer_id {
let (our_pk_b, out_sk_b) = box_::gen_keypair(); let (our_pk_b, out_sk_b) = box_::gen_keypair();

13
src/server.rs
View File

@ -101,12 +101,15 @@ pub async fn create_tcp_connection(
let sk = sign::SecretKey(sk_); let sk = sign::SecretKey(sk_);
let mut msg_out = Message::new(); let mut msg_out = Message::new();
let (our_pk_b, our_sk_b) = box_::gen_keypair(); let (our_pk_b, our_sk_b) = box_::gen_keypair();
let signed_id = sign::sign(
format!("{}\0{}", Config::get_id(), base64::encode(our_pk_b.0)).as_bytes(),
&sk,
);
msg_out.set_signed_id(SignedId { msg_out.set_signed_id(SignedId {
id: signed_id, id: sign::sign(
&serde_json::to_vec(&hbb_common::IdPk {
id: Config::get_id(),
pk: our_pk_b.0,
})
.unwrap_or_default(),
&sk,
),
..Default::default() ..Default::default()
}); });
timeout(CONNECT_TIMEOUT, stream.send(&msg_out)).await??; timeout(CONNECT_TIMEOUT, stream.send(&msg_out)).await??;