From 3ea33f72032be95e8d71b5311ef0b84f1da9ee31 Mon Sep 17 00:00:00 2001
From: rustdesk <info@rustdesk.com>
Date: Sat, 26 Mar 2022 03:05:15 +0800
Subject: [PATCH] improve file write to cm

---
 libs/hbb_common/src/fs.rs | 15 ++++++++++-----
 src/server/connection.rs  | 10 +++++++++-
 src/ui/cm.rs              | 22 +++++++++++++++-------
 src/ui/remote.rs          |  2 +-
 4 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/libs/hbb_common/src/fs.rs b/libs/hbb_common/src/fs.rs
index 99e594751..6fad26037 100644
--- a/libs/hbb_common/src/fs.rs
+++ b/libs/hbb_common/src/fs.rs
@@ -302,7 +302,7 @@ impl TransferJob {
         }
     }
 
-    pub async fn write(&mut self, block: FileTransferBlock) -> ResultType<()> {
+    pub async fn write(&mut self, block: FileTransferBlock, raw: Option<&[u8]>) -> ResultType<()> {
         if block.id != self.id {
             bail!("Wrong id");
         }
@@ -324,15 +324,20 @@ impl TransferJob {
             let path = format!("{}.download", get_string(&path));
             self.file = Some(File::create(&path).await?);
         }
+        let data = if let Some(data) = raw {
+            data
+        } else {
+            &block.data
+        };
         if block.compressed {
-            let tmp = decompress(&block.data);
+            let tmp = decompress(data);
             self.file.as_mut().unwrap().write_all(&tmp).await?;
             self.finished_size += tmp.len() as u64;
         } else {
-            self.file.as_mut().unwrap().write_all(&block.data).await?;
-            self.finished_size += block.data.len() as u64;
+            self.file.as_mut().unwrap().write_all(data).await?;
+            self.finished_size += data.len() as u64;
         }
-        self.transferred += block.data.len() as u64;
+        self.transferred += data.len() as u64;
         Ok(())
     }
 
diff --git a/src/server/connection.rs b/src/server/connection.rs
index a1a08d4d4..ac5b2f94f 100644
--- a/src/server/connection.rs
+++ b/src/server/connection.rs
@@ -1157,7 +1157,15 @@ async fn start_ipc(
             res = rx_to_cm.recv() => {
                 match res {
                     Some(data) => {
-                        stream.send(&data).await?;
+                        if let Data::FS(ipc::FS::WriteBlock{id,
+                            file_num,
+                            data,
+                            compressed}) = data {
+                                stream.send(&Data::FS(ipc::FS::WriteBlock{id, file_num, data: Vec::new(), compressed})).await?;
+                                stream.send_raw(data).await?;
+                        } else {
+                            stream.send(&data).await?;
+                        }
                     }
                     None => {
                         bail!("expected");
diff --git a/src/ui/cm.rs b/src/ui/cm.rs
index fca0a0205..ca5400cfe 100644
--- a/src/ui/cm.rs
+++ b/src/ui/cm.rs
@@ -175,15 +175,23 @@ impl ConnectionManager {
                     data,
                     compressed,
                 } => {
+                    let raw = if let Ok(bytes) = conn.next_raw().await {
+                        Some(bytes)
+                    } else {
+                        None
+                    };
                     if let Some(job) = fs::get_job(id, write_jobs) {
                         if let Err(err) = job
-                            .write(FileTransferBlock {
-                                id,
-                                file_num,
-                                data,
-                                compressed,
-                                ..Default::default()
-                            })
+                            .write(
+                                FileTransferBlock {
+                                    id,
+                                    file_num,
+                                    data,
+                                    compressed,
+                                    ..Default::default()
+                                },
+                                raw.as_ref().map(|x| &x[..]),
+                            )
                             .await
                         {
                             Self::send(fs::new_error(id, err, file_num), conn).await;
diff --git a/src/ui/remote.rs b/src/ui/remote.rs
index bee747afe..7ae1d3a0e 100644
--- a/src/ui/remote.rs
+++ b/src/ui/remote.rs
@@ -1797,7 +1797,7 @@ impl Remote {
                     }
                     Some(file_response::Union::block(block)) => {
                         if let Some(job) = fs::get_job(block.id, &mut self.write_jobs) {
-                            if let Err(_err) = job.write(block).await {
+                            if let Err(_err) = job.write(block, None).await {
                                 // to-do: add "skip" for writing job
                             }
                             self.update_jobs_status();