diff --git a/sparse-05/sparse-05-client/src/commands/connect/commands/download.rs b/sparse-05/sparse-05-client/src/commands/connect/commands/download.rs index b1ba1c2..aadef93 100644 --- a/sparse-05/sparse-05-client/src/commands/connect/commands/download.rs +++ b/sparse-05/sparse-05-client/src/commands/connect/commands/download.rs @@ -1,13 +1,8 @@ use std::{path::PathBuf, sync::Arc}; -use tokio::{ - fs, - io::{self, AsyncReadExt}, -}; +use tokio::{fs, io::AsyncWriteExt}; -use sparse_05_common::messages::{ - Command, Response, FILE_BUFFER_BUFFER_SIZE, FILE_TRANSFER_PACKET_SIZE, -}; +use sparse_05_common::messages::{Command, Response, FILE_BUFFER_BUFFER_SIZE}; use crate::commands::connect::Connection; @@ -16,7 +11,7 @@ pub async fn download_file( remote_file: PathBuf, local_path: PathBuf, ) -> anyhow::Result<()> { - let mut file = fs::OpenOptions::new() + let mut target_file = fs::OpenOptions::new() .write(true) .create(true) .open(&local_path) @@ -25,67 +20,54 @@ pub async fn download_file( let command = Command::StartDownloadFile(remote_file); conn.send_command(command).await?; - let id = loop { + let (id, packet_count) = loop { let resp = conn.get_response().await?; - if let Response::UploadFileID(id) = resp { - break id; + if let Response::StartDownloadFile(id, packet_count) = resp { + break (id, packet_count); } }; - loop { - let mut file_data: Vec> = Vec::with_capacity(FILE_BUFFER_BUFFER_SIZE); - let mut buffer = [0u8; FILE_TRANSFER_PACKET_SIZE]; - let mut done = false; + let mut current_packet_count = 0; - for _ in 0..FILE_BUFFER_BUFFER_SIZE { - let amount = file.read(&mut buffer).await?; - - if amount == 0 { - done = true; - break; - } - - file_data.push(buffer[..amount].to_vec()); - } - - for (i, buffer) in file_data.iter().enumerate() { - let comm = Command::SendFileSegment(id, i as u64, buffer.clone()); - conn.send_command(comm).await?; - } + while current_packet_count < packet_count { + let mut buffers: Vec>> = vec![None; FILE_BUFFER_BUFFER_SIZE]; loop { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + match conn.get_response().await { + Ok(Response::DownloadFileSegment(fid, bid, buf)) if fid == id => { + buffers[bid as usize] = Some(buf); + } + Ok(Response::GetDownloadFileStatus(fid, up_to)) if fid == id => { + let needed = buffers[..up_to as usize] + .iter() + .enumerate() + .flat_map(|(i, b)| match b { + Some(..) => None, + None => Some(i as u64), + }) + .collect::>(); - conn.send_command(Command::GetUploadStatus(id, file_data.len() as u64)) - .await?; + let is_empty = needed.is_empty(); - let buffers_needed = loop { - let resp = conn.get_response().await?; + conn.send_command(Command::DownloadFileStatus(id, needed)) + .await?; - if let Response::UploadFileStatus(iid, buffers) = resp { - if id == iid { - break buffers; + if is_empty { + current_packet_count += up_to; + break; + } else { + continue; } } - }; - - if buffers_needed.len() == 0 { - break; - } - - for buffer_needed in &buffers_needed { - let comm = Command::SendFileSegment( - id, - *buffer_needed, - file_data[*buffer_needed as usize].clone(), - ); - conn.send_command(comm).await?; + Ok(..) => {} + Err(..) => {} } } - if done { - break; + for buffer in buffers { + let Some(buffer) = buffer else { break }; + target_file.write(&buffer).await?; } } diff --git a/sparse-05/sparse-05-client/src/commands/connect/shell.rs b/sparse-05/sparse-05-client/src/commands/connect/shell.rs index 24e1baa..0d87c71 100644 --- a/sparse-05/sparse-05-client/src/commands/connect/shell.rs +++ b/sparse-05/sparse-05-client/src/commands/connect/shell.rs @@ -3,13 +3,13 @@ use std::{ os::fd::AsRawFd, path::PathBuf, sync::Arc, - thread::{self, scope}, + thread, }; use sparse_05_common::messages::{Capabilities, Command, Response}; use structopt::StructOpt; use tokio::{ - io::{stderr, stdout, AsyncReadExt, AsyncWriteExt}, + io::{stderr, stdout, AsyncWriteExt}, runtime::Handle, sync::mpsc::{channel, Receiver}, }; @@ -36,6 +36,10 @@ pub enum SparseCommands { remote_file: PathBuf, local_path: PathBuf, }, + #[structopt(name = "#edit")] + EditFile { + remote_path: PathBuf, + }, } macro_rules! libc_try { @@ -225,6 +229,7 @@ pub(super) async fn shell( eprintln!("{e:?}") } } + (Ok(SparseCommands::EditFile { remote_path }), _) => {} _ => { if let Err(e) = run_command( &mut stdin_receiver, diff --git a/sparse-05/sparse-05-common/src/lib.rs b/sparse-05/sparse-05-common/src/lib.rs index 39cfd68..5c331c6 100644 --- a/sparse-05/sparse-05-common/src/lib.rs +++ b/sparse-05/sparse-05-common/src/lib.rs @@ -80,7 +80,7 @@ pub mod messages { UploadFileStatus(u64, Vec), StartDownloadFile(u64, u64), - GetDownloadFileStatus(u64, Vec), + GetDownloadFileStatus(u64, u64), DownloadFileSegment(u64, u64, Vec), } diff --git a/sparse-05/sparse-05-server/src/connection/download_file.rs b/sparse-05/sparse-05-server/src/connection/download_file.rs index 3840941..8b8044f 100644 --- a/sparse-05/sparse-05-server/src/connection/download_file.rs +++ b/sparse-05/sparse-05-server/src/connection/download_file.rs @@ -1,7 +1,7 @@ use std::{ collections::HashMap, fs::{self, OpenOptions}, - io::Write, + io::Read, path::PathBuf, sync::{ atomic::{AtomicU64, Ordering}, @@ -38,50 +38,65 @@ pub(super) fn start_file_download<'a, 'b: 'a>( let mut file = OpenOptions::new().read(true).open(&file_path)?; let file_size = file.metadata()?.len(); - conninfo.send(conninfo.encrypt_and_sign_resp(Response::StartDownloadFile( - id, - (file_size / FILE_TRANSFER_PACKET_SIZE as u64) + 1, - ))?)?; + let packet_count = (file_size / FILE_TRANSFER_PACKET_SIZE as u64) + 1; + conninfo + .send(conninfo.encrypt_and_sign_resp(Response::StartDownloadFile(id, packet_count))?)?; let mut current_packet_count = 0; - /*while current_packet_count < packet_count { - let mut buffers: Vec>> = vec![None; FILE_BUFFER_BUFFER_SIZE]; + loop { + let mut file_data: Vec> = Vec::with_capacity(FILE_BUFFER_BUFFER_SIZE); + let mut buffer = [0u8; FILE_TRANSFER_PACKET_SIZE]; + let mut done = false; + + for _ in 0..FILE_BUFFER_BUFFER_SIZE { + let amount = file.read(&mut buffer)?; + + if amount == 0 { + done = true; + break; + } + + file_data.push(buffer[..amount].to_vec()); + } + + for (i, buffer) in file_data.iter().enumerate() { + conninfo.send( + conninfo.encrypt_and_sign_resp(Response::DownloadFileSegment( + id, + i as u64, + buffer.clone(), + ))?, + )?; + } loop { - let Ok((i, buffer)) = data_receiver.recv_timeout(Duration::from_millis(250)) else { - let up_to = receive_request_status.recv()?; + std::thread::sleep(std::time::Duration::from_millis(1000)); - let needed = buffers[..up_to as usize] - .iter() - .enumerate() - .flat_map(|(i, b)| match b { - Some(..) => None, - None => Some(i as u64), - }) - .collect::>(); + conninfo.send(conninfo.encrypt_and_sign_resp( + Response::GetDownloadFileStatus(id, file_data.len() as u64), + )?)?; - let is_empty = needed.is_empty(); + let buffers_needed: Vec = receive_download_status.recv()?; - conninfo.send( - conninfo.encrypt_and_sign_resp(Response::UploadFileStatus(id, needed))?, - )?; + if buffers_needed.is_empty() { + break; + } - if is_empty { - current_packet_count += up_to; - break; - } else { - continue; - } - }; - buffers[i as usize] = Some(buffer); + for buffer_needed in &buffers_needed { + let resp = Response::DownloadFileSegment( + id, + *buffer_needed, + file_data[*buffer_needed as usize].clone(), + ); + conninfo.send(conninfo.encrypt_and_sign_resp(resp)?)?; + } } - for buffer in buffers { - let Some(buffer) = buffer else { break }; - target_file.write(&buffer)?; + if done { + break; } - }*/ + } let Ok(mut lock) = download_file_map.lock() else { return Ok(());