weyne/rustdesk
1
0
Fork 0
mirror of https://github.com/weyne85/rustdesk.git synced 2025-10-29 17:00:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix 10054: change direct to relay when RST

Browse Source 
Signed-off-by: 21pages <pages21@163.com>
This commit is contained in:
21pages
2022-07-31 19:06:49 +08:00
parent 3394975aa2
commit 0649a49d17
6 changed files with 152 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
src/port_forward.rs
View File

@@ -1,7 +1,7 @@
use crate::client::*;
use hbb_common::{
allow_err, bail,
config::CONNECT_TIMEOUT,
config::READ_TIMEOUT,
futures::{SinkExt, StreamExt},
log,
message_proto::*,
@@ -105,22 +105,61 @@ async fn connect_and_login(
key: &str,
token: &str,
is_rdp: bool,
) -> ResultType<Option<Stream>> {
let mut res = connect_and_login_2(
id,
password,
ui_receiver,
interface.clone(),
forward,
key,
token,
is_rdp,
)
.await;
if res.is_err() && interface.is_force_relay() {
res = connect_and_login_2(
id,
password,
ui_receiver,
interface,
forward,
key,
token,
is_rdp,
)
.await;
}
res
}
async fn connect_and_login_2(
id: &str,
password: &str,
ui_receiver: &mut mpsc::UnboundedReceiver<Data>,
interface: impl Interface,
forward: &mut Framed<TcpStream, BytesCodec>,
key: &str,
token: &str,
is_rdp: bool,
) -> ResultType<Option<Stream>> {
let conn_type = if is_rdp {
ConnType::RDP
} else {
ConnType::PORT_FORWARD
};
let (mut stream, _) = Client::start(id, key, token, conn_type).await?;
let (mut stream, direct) = Client::start(id, key, token, conn_type, interface.clone()).await?;
let mut interface = interface;
let mut buffer = Vec::new();
let mut received = false;
loop {
tokio::select! {
res = timeout(CONNECT_TIMEOUT, stream.next()) => match res {
res = timeout(READ_TIMEOUT, stream.next()) => match res {
Err(_) => {
bail!("Timeout");
}
Ok(Some(Ok(bytes))) => {
received = true;
let msg_in = Message::parse_from_bytes(&bytes)?;
match msg_in.union {
Some(message::Union::Hash(hash)) => {
@@ -143,6 +182,11 @@ async fn connect_and_login(
_ => {}
}
}
Ok(Some(Err(err))) => {
log::error!("Connection closed: {}", err);
interface.set_force_relay(direct, received);
bail!("Connection closed: {}", err);
}
_ => {
bail!("Reset by the peer");
}