From ce246c22342a30bc0dd2b7b2f479df5cb7b7c9b5 Mon Sep 17 00:00:00 2001 From: Hardcore Sushi Date: Wed, 5 May 2021 18:00:03 +0200 Subject: [PATCH] Fixing bug when consecutively sending 2 large files --- src/frontend/index.js | 9 ++++++++- src/main.rs | 5 +++-- src/session_manager/mod.rs | 36 ++++++++++++++++++++---------------- 3 files changed, 31 insertions(+), 19 deletions(-) diff --git a/src/frontend/index.js b/src/frontend/index.js index 8f831ce..0026ce8 100644 --- a/src/frontend/index.js +++ b/src/frontend/index.js @@ -165,7 +165,14 @@ document.getElementById("logout").onclick = function() { document.getElementById("attach_file").onchange = function(event) { let file = event.target.files[0]; if (file.size > 32760000) { - if (!pendingFiles.has(currentSessionId)) { + if (pendingFiles.has(currentSessionId)) { + let mainDiv = document.createElement("div"); + mainDiv.appendChild(generatePopupWarningTitle()); + let p = document.createElement("p"); + p.textContent = "Another file transfer is already in progress."; + mainDiv.appendChild(p); + showPopup(mainDiv); + } else { pendingFiles.set(currentSessionId, { "file": file, "name": file.name, diff --git a/src/main.rs b/src/main.rs index 9de78ba..709b453 100644 --- a/src/main.rs +++ b/src/main.rs @@ -363,10 +363,11 @@ async fn handle_send_file(req: HttpRequest, mut payload: Multipart) -> HttpRespo print_error!(e); return HttpResponse::InternalServerError().finish(); } + let last_chunk = chunk_buffer.len() < constants::FILE_CHUNK_SIZE; if !match ack_receiver.recv().await { Some(should_continue) => { //send previous encrypted chunk even if transfert is aborted to keep PSEC nonces syncrhonized - if let Err(e) = global_vars_read.session_manager.send_encrypted_file_chunk(&session_id, ack_sender.clone()).await { + if let Err(e) = global_vars_read.session_manager.send_encrypted_file_chunk(&session_id, ack_sender.clone(), last_chunk).await { print_error!(e); false } else { @@ -377,7 +378,7 @@ async fn handle_send_file(req: HttpRequest, mut payload: Multipart) -> HttpRespo } { return HttpResponse::InternalServerError().finish() } - if chunk_buffer.len() < constants::FILE_CHUNK_SIZE { + if last_chunk { break; } else { chunk_buffer.truncate(1); diff --git a/src/session_manager/mod.rs b/src/session_manager/mod.rs index fd892ca..5c570df 100644 --- a/src/session_manager/mod.rs +++ b/src/session_manager/mod.rs @@ -30,6 +30,7 @@ enum SessionCommand { }, SendEncryptedFileChunk { sender: Sender, + last_chunk: bool, }, EncryptFileChunk { plain_text: Vec, @@ -134,10 +135,11 @@ impl SessionManager { } } - pub async fn send_encrypted_file_chunk(&self, session_id: &usize, ack_sender: Sender) -> Result<(), SessionError> { + pub async fn send_encrypted_file_chunk(&self, session_id: &usize, ack_sender: Sender, last_chunk: bool) -> Result<(), SessionError> { if let Some(sender) = self.get_session_sender(session_id) { match sender.send(SessionCommand::SendEncryptedFileChunk { sender: ack_sender, + last_chunk }).await { Ok(_) => Ok(()), Err(e) => { @@ -175,13 +177,13 @@ impl SessionManager { self.not_seen.write().unwrap().retain(|x| x != session_id); } - async fn send_msg(&self, session_id: usize, session_write: &mut SessionWrite, buff: &[u8], aborted: &mut bool, file_ack_sender: Option<&Sender>) -> Result<(), SessionError> { + async fn send_msg(&self, session_id: usize, session_write: &mut SessionWrite, buff: &[u8], is_sending: &mut bool, file_ack_sender: Option<&Sender>) -> Result<(), SessionError> { session_write.encrypt_and_send(&buff).await?; if buff[0] == protocol::Headers::ACCEPT_LARGE_FILE { self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download.as_mut().unwrap().state = FileState::ACCEPTED; } else if buff[0] == protocol::Headers::ABORT_FILE_TRANSFER { self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = None; - *aborted = true; + *is_sending = false; if let Some(sender) = file_ack_sender { if let Err(e) = sender.send(false).await { print_error!(e); @@ -202,7 +204,7 @@ impl SessionManager { let mut next_chunk: Option> = None; let mut file_ack_sender: Option> = None; let mut msg_queue = Vec::new(); - let mut aborted = false; + let mut is_sending = false; let (session_read, mut session_write) = session.into_spit().unwrap(); let receiving = session_read.receive_and_decrypt(); @@ -245,7 +247,7 @@ impl SessionManager { } } protocol::Headers::ASK_LARGE_FILE => { - if self.sessions.read().unwrap().get(&session_id).unwrap().file_download.is_none() { //don't accept 2 downloads at the same time + if self.sessions.read().unwrap().get(&session_id).unwrap().file_download.is_none() && !is_sending { //don't accept 2 file transfers at the same time if let Some((file_size, file_name)) = protocol::parse_ask_file(&buffer) { let download_dir = UserDirs::new().unwrap().download_dir; self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = Some(LargeFileDownload{ @@ -341,7 +343,7 @@ impl SessionManager { }); } if sender.send(true).await.is_err() { - aborted = true; + is_sending = false; } } } @@ -350,7 +352,7 @@ impl SessionManager { if let Err(e) = sender.send(false).await { print_error!(e); } - aborted = true; + is_sending = false; } self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = None; local_file_path = None; @@ -385,9 +387,8 @@ impl SessionManager { let is_classical_message = header == protocol::Headers::MESSAGE || header == protocol::Headers::FILE; if is_classical_message { self.set_seen(session_id, false); - } - if header == protocol::Headers::ACCEPT_LARGE_FILE { - aborted = false; + } else if header == protocol::Headers::ACCEPT_LARGE_FILE { + is_sending = true; } self.with_ui_connection(|ui_connection| { ui_connection.on_received(&session_id, buffer.as_ref().unwrap()); @@ -411,17 +412,17 @@ impl SessionManager { match command.unwrap() { SessionCommand::Send { buff } => { //don't send msg if we already encrypted a file chunk (keep PSEC nonces synchronized) - if next_chunk.is_none() || aborted { - if let Err(e) = self.send_msg(session_id, &mut session_write, &buff, &mut aborted, file_ack_sender.as_ref()).await { + if is_sending { + msg_queue.push(buff); + } else { + if let Err(e) = self.send_msg(session_id, &mut session_write, &buff, &mut is_sending, file_ack_sender.as_ref()).await { print_error!(e); break; } - } else { - msg_queue.push(buff); } } SessionCommand::EncryptFileChunk { plain_text } => next_chunk = Some(session_write.encrypt(&plain_text)), - SessionCommand::SendEncryptedFileChunk { sender } => { + SessionCommand::SendEncryptedFileChunk { sender, last_chunk } => { if let Some(chunk) = next_chunk.as_ref() { match session_write.socket_write(chunk).await { Ok(_) => { @@ -429,11 +430,14 @@ impl SessionManager { //once the pre-encrypted chunk is sent, we can send the pending messages while msg_queue.len() > 0 { let msg = msg_queue.remove(0); - if let Err(e) = self.send_msg(session_id, &mut session_write, &msg, &mut aborted, file_ack_sender.as_ref()).await { + if let Err(e) = self.send_msg(session_id, &mut session_write, &msg, &mut is_sending, file_ack_sender.as_ref()).await { print_error!(e); break; } } + if last_chunk { + is_sending = false; + } } Err(e) => { print_error!(e);