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

Merge pull request #2801 from 21pages/sync

Browse Source 
sync strategy
This commit is contained in:
RustDesk 2023-01-11 17:36:56 +08:00 committed by GitHub
commit 96e608ff07
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 160 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/hbbs_http.rs
View File

@ -5,6 +5,7 @@ use serde_json::{Map, Value};
#[cfg(feature = "flutter")]
pub mod account;
pub mod record_upload;
pub mod sync;
#[derive(Debug)]
pub enum HbbHttpResponse<T> {

132
src/hbbs_http/sync.rs Normal file
View File

@ -0,0 +1,132 @@
use std::{collections::HashMap, sync::Mutex, time::Duration};
use hbb_common::{
config::{Config, LocalConfig},
tokio::{self, sync::broadcast, time::Instant},
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use crate::Connection;
const TIME_HEARTBEAT: Duration = Duration::from_secs(30);
const TIME_CONN: Duration = Duration::from_secs(3);
lazy_static::lazy_static! {
static ref SENDER : Mutex<broadcast::Sender<Vec<i32>>> = Mutex::new(start_hbbs_sync());
}
pub fn start() {
let _sender = SENDER.lock().unwrap();
}
pub fn signal_receiver() -> broadcast::Receiver<Vec<i32>> {
SENDER.lock().unwrap().subscribe()
}
fn start_hbbs_sync() -> broadcast::Sender<Vec<i32>> {
let (tx, _rx) = broadcast::channel::<Vec<i32>>(16);
std::thread::spawn(move || start_hbbs_sync_async());
return tx;
}
#[derive(Debug, Serialize, Deserialize)]
pub struct StrategyOptions {
pub config_options: HashMap<String, String>,
pub extra: HashMap<String, String>,
}
#[tokio::main(flavor = "current_thread")]
async fn start_hbbs_sync_async() {
tokio::spawn(async move {
let mut interval = tokio::time::interval_at(Instant::now() + TIME_CONN, TIME_CONN);
let mut last_send = Instant::now();
loop {
tokio::select! {
_ = interval.tick() => {
let url = heartbeat_url();
let modified_at = LocalConfig::get_option("strategy_timestamp").parse::<i64>().unwrap_or(0);
if !url.is_empty() {
let conns = Connection::alive_conns();
if conns.is_empty() && last_send.elapsed() < TIME_HEARTBEAT {
continue;
}
last_send = Instant::now();
let mut v = Value::default();
v["id"] = json!(Config::get_id());
if !conns.is_empty() {
v["conns"] = json!(conns);
}
v["modified_at"] = json!(modified_at);
if let Ok(s) = crate::post_request(url.clone(), v.to_string(), "").await {
if let Ok(mut rsp) = serde_json::from_str::<HashMap::<&str, Value>>(&s) {
if let Some(conns) = rsp.remove("disconnect") {
if let Ok(conns) = serde_json::from_value::<Vec<i32>>(conns) {
SENDER.lock().unwrap().send(conns).ok();
}
}
if let Some(rsp_modified_at) = rsp.remove("modified_at") {
if let Ok(rsp_modified_at) = serde_json::from_value::<i64>(rsp_modified_at) {
if rsp_modified_at != modified_at {
LocalConfig::set_option("strategy_timestamp".to_string(), rsp_modified_at.to_string());
}
}
}
if let Some(strategy) = rsp.remove("strategy") {
if let Ok(strategy) = serde_json::from_value::<StrategyOptions>(strategy) {
handle_config_options(strategy.config_options);
}
}
}
}
}
}
}
}
})
.await
.ok();
}
fn heartbeat_url() -> String {
let url = crate::common::get_api_server(
Config::get_option("api-server"),
Config::get_option("custom-rendezvous-server"),
);
if url.is_empty() || url.contains("rustdesk.com") {
return "".to_owned();
}
format!("{}/api/heartbeat", url)
}
fn handle_config_options(config_options: HashMap<String, String>) {
let map = HashMap::from([
("enable-keyboard", ""),
("enable-clipboard", ""),
("enable-file-transfer", ""),
("enable-audio", ""),
("enable-tunnel", ""),
("enable-remote-restart", ""),
("enable-record-session", ""),
("allow-remote-config-modification", ""),
("approve-mode", ""),
("verification-method", "use-both-passwords"),
("enable-rdp", ""),
("enable-lan-discovery", ""),
("direct-server", ""),
("direct-access-port", ""),
]);
let mut options = Config::get_options();
for (k, v) in map {
if let Some(v2) = config_options.get(k) {
if v == v2 {
options.remove(k);
} else {
options.insert(k.to_string(), v2.to_string());
}
} else {
options.remove(k);
}
}
Config::set_options(options);
}

10
src/server.rs
View File

@ -194,9 +194,14 @@ pub async fn create_tcp_connection(
}
}
#[cfg(target_os = "macos")]{
#[cfg(target_os = "macos")]
{
use std::process::Command;
Command::new("/usr/bin/caffeinate").arg("-u").arg("-t 5").spawn().ok();
Command::new("/usr/bin/caffeinate")
.arg("-u")
.arg("-t 5")
.spawn()
.ok();
log::info!("wake up macos");
}
Connection::start(addr, stream, id, Arc::downgrade(&server)).await;
@ -385,6 +390,7 @@ pub async fn start_server(is_server: bool) {
#[cfg(windows)]
crate::platform::windows::bootstrap();
input_service::fix_key_down_timeout_loop();
crate::hbbs_http::sync::start();
#[cfg(target_os = "linux")]
if crate::platform::current_is_wayland() {
allow_err!(input_service::setup_uinput(0, 1920, 0, 1080).await);

64
src/server/connection.rs
View File

@ -26,7 +26,6 @@ use hbb_common::{
};
#[cfg(any(target_os = "android", target_os = "ios"))]
use scrap::android::call_main_service_mouse_input;
use serde::Deserialize;
use serde_json::{json, value::Value};
use sha2::{Digest, Sha256};
#[cfg(not(any(target_os = "android", target_os = "ios")))]
@ -40,6 +39,7 @@ pub type Sender = mpsc::UnboundedSender<(Instant, Arc<Message>)>;
lazy_static::lazy_static! {
static ref LOGIN_FAILURES: Arc::<Mutex<HashMap<String, (i32, i32, i32)>>> = Default::default();
static ref SESSIONS: Arc::<Mutex<HashMap<String, Session>>> = Default::default();
static ref ALIVE_CONNS: Arc::<Mutex<Vec<i32>>> = Default::default();
}
pub static CLICK_TIME: AtomicI64 = AtomicI64::new(0);
pub static MOUSE_MOVE_TIME: AtomicI64 = AtomicI64::new(0);
@ -74,7 +74,6 @@ pub struct Connection {
read_jobs: Vec<fs::TransferJob>,
timer: Interval,
file_timer: Interval,
http_timer: Interval,
file_transfer: Option<(String, bool)>,
port_forward_socket: Option<Framed<TcpStream, BytesCodec>>,
port_forward_address: String,
@ -147,6 +146,7 @@ impl Connection {
challenge: Config::get_auto_password(6),
..Default::default()
};
ALIVE_CONNS.lock().unwrap().push(id);
let (tx_from_cm_holder, mut rx_from_cm) = mpsc::unbounded_channel::<ipc::Data>();
// holding tx_from_cm_holder to avoid cpu burning of rx_from_cm.recv when all sender closed
let tx_from_cm = tx_from_cm_holder.clone();
@ -154,7 +154,7 @@ impl Connection {
let (tx, mut rx) = mpsc::unbounded_channel::<(Instant, Arc<Message>)>();
let (tx_video, mut rx_video) = mpsc::unbounded_channel::<(Instant, Arc<Message>)>();
let (tx_input, rx_input) = std_mpsc::channel();
let (tx_stop, mut rx_stop) = mpsc::unbounded_channel::<String>();
let mut hbbs_rx = crate::hbbs_http::sync::signal_receiver();
let tx_cloned = tx.clone();
let mut conn = Self {
@ -169,7 +169,6 @@ impl Connection {
read_jobs: Vec::new(),
timer: time::interval(SEC30),
file_timer: time::interval(SEC30),
http_timer: time::interval(Duration::from_secs(3)),
file_transfer: None,
port_forward_socket: None,
port_forward_address: "".to_owned(),
@ -393,12 +392,11 @@ impl Connection {
conn.file_timer = time::interval_at(Instant::now() + SEC30, SEC30);
}
}
_ = conn.http_timer.tick() => {
Connection::post_heartbeat(conn.server_audit_conn.clone(), conn.inner.id, tx_stop.clone());
},
Some(reason) = rx_stop.recv() => {
conn.on_close_manually(&reason, &reason).await;
Ok(conns) = hbbs_rx.recv() => {
if conns.contains(&id) {
conn.on_close_manually("web console", "web console").await;
break;
}
}
Some((instant, value)) = rx_video.recv() => {
if !conn.video_ack_required {
@ -514,6 +512,7 @@ impl Connection {
conn.post_conn_audit(json!({
"action": "close",
}));
ALIVE_CONNS.lock().unwrap().retain(|&c| c != id);
log::info!("#{} connection loop exited", id);
}
@ -584,10 +583,10 @@ impl Connection {
rx_from_cm: &mut mpsc::UnboundedReceiver<Data>,
) -> ResultType<()> {
let mut last_recv_time = Instant::now();
let (tx_stop, mut rx_stop) = mpsc::unbounded_channel::<String>();
if let Some(mut forward) = self.port_forward_socket.take() {
log::info!("Running port forwarding loop");
self.stream.set_raw();
let mut hbbs_rx = crate::hbbs_http::sync::signal_receiver();
loop {
tokio::select! {
Some(data) = rx_from_cm.recv() => {
@ -618,10 +617,12 @@ impl Connection {
if last_recv_time.elapsed() >= H1 {
bail!("Timeout");
}
Connection::post_heartbeat(self.server_audit_conn.clone(), self.inner.id, tx_stop.clone());
}
Some(reason) = rx_stop.recv() => {
bail!(reason);
Ok(conns) = hbbs_rx.recv() => {
if conns.contains(&self.inner.id) {
// todo: check reconnect
bail!("Closed manually by the web console");
}
}
}
}
@ -711,30 +712,6 @@ impl Connection {
});
}
fn post_heartbeat(
server_audit_conn: String,
conn_id: i32,
tx_stop: mpsc::UnboundedSender<String>,
) {
if server_audit_conn.is_empty() {
return;
}
let url = server_audit_conn.clone();
let mut v = Value::default();
v["id"] = json!(Config::get_id());
v["uuid"] = json!(base64::encode(hbb_common::get_uuid()));
v["conn_id"] = json!(conn_id);
tokio::spawn(async move {
if let Ok(rsp) = Self::post_audit_async(url, v).await {
if let Ok(rsp) = serde_json::from_str::<ConnAuditResponse>(&rsp) {
if rsp.action == "disconnect" {
tx_stop.send("web console".to_string()).ok();
}
}
}
});
}
fn post_file_audit(
&self,
r#type: FileAuditType,
@ -1713,6 +1690,10 @@ impl Connection {
async fn send(&mut self, msg: Message) {
allow_err!(self.stream.send(&msg).await);
}
pub fn alive_conns() -> Vec<i32> {
ALIVE_CONNS.lock().unwrap().clone()
}
}
#[cfg(not(any(target_os = "android", target_os = "ios")))]
@ -1870,13 +1851,6 @@ mod privacy_mode {
}
}
#[derive(Debug, Deserialize)]
struct ConnAuditResponse {
#[allow(dead_code)]
ret: bool,
action: String,
}
pub enum AlarmAuditType {
IpWhitelist = 0,
ManyWrongPassword = 1,