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

improve file write to cm

Browse Source
This commit is contained in:
rustdesk 2022-03-26 03:05:15 +08:00
parent 75c9bbb30f
commit 3ea33f7203
4 changed files with 35 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
libs/hbb_common/src/fs.rs
View File

@ -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 { if block.id != self.id {
bail!("Wrong id"); bail!("Wrong id");
} }
@ -324,15 +324,20 @@ impl TransferJob {
let path = format!("{}.download", get_string(&path)); let path = format!("{}.download", get_string(&path));
self.file = Some(File::create(&path).await?); self.file = Some(File::create(&path).await?);
} }
let data = if let Some(data) = raw {
data
} else {
&block.data
};
if block.compressed { if block.compressed {
let tmp = decompress(&block.data); let tmp = decompress(data);
self.file.as_mut().unwrap().write_all(&tmp).await?; self.file.as_mut().unwrap().write_all(&tmp).await?;
self.finished_size += tmp.len() as u64; self.finished_size += tmp.len() as u64;
} else { } else {
self.file.as_mut().unwrap().write_all(&block.data).await?; self.file.as_mut().unwrap().write_all(data).await?;
self.finished_size += block.data.len() as u64; self.finished_size += data.len() as u64;
} }
self.transferred += block.data.len() as u64; self.transferred += data.len() as u64;
Ok(()) Ok(())
} }

10
src/server/connection.rs
View File

@ -1157,7 +1157,15 @@ async fn start_ipc(
res = rx_to_cm.recv() => { res = rx_to_cm.recv() => {
match res { match res {
Some(data) => { 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 => { None => {
bail!("expected"); bail!("expected");

22
src/ui/cm.rs
View File

@ -175,15 +175,23 @@ impl ConnectionManager {
data, data,
compressed, 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 Some(job) = fs::get_job(id, write_jobs) {
if let Err(err) = job if let Err(err) = job
.write(FileTransferBlock { .write(
id, FileTransferBlock {
file_num, id,
data, file_num,
compressed, data,
..Default::default() compressed,
}) ..Default::default()
},
raw.as_ref().map(|x| &x[..]),
)
.await .await
{ {
Self::send(fs::new_error(id, err, file_num), conn).await; Self::send(fs::new_error(id, err, file_num), conn).await;

2
src/ui/remote.rs
View File

@ -1797,7 +1797,7 @@ impl Remote {
} }
Some(file_response::Union::block(block)) => { Some(file_response::Union::block(block)) => {
if let Some(job) = fs::get_job(block.id, &mut self.write_jobs) { 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 // to-do: add "skip" for writing job
} }
self.update_jobs_status(); self.update_jobs_status();