hardcoresushi/AIRA
1
0
Fork 0
Code Issues Pull Requests Projects Releases 3 Wiki Activity

Fix large file transfers bugs

Browse Source
This commit is contained in:
Matéo Duparc 2021-04-30 18:54:48 +02:00
parent 33a107d347
commit a765c2565f
Signed by: hardcoresushi
GPG Key ID: 007F84120107191E
6 changed files with 188 additions and 154 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

130
src/frontend/index.js
View File

@ -524,15 +524,15 @@ function onFileAborted(sessionId) {
} }
} }
} }
function onIncFileTransfer(sessionId, chunk_size) { function onIncFileTransfer(sessionId, chunkSize) {
if (pendingFiles.has(sessionId)) { if (pendingFiles.has(sessionId)) {
let file = pendingFiles.get(sessionId); let file = pendingFiles.get(sessionId);
file.transferred += chunk_size; file.transferred += chunkSize;
let now = Date.now(); let now = Date.now();
let speed = chunk_size/(now-file.lastChunk)*1000; let speed = chunkSize/(now-file.lastChunk)*1000;
file.lastChunk = now; file.lastChunk = now;
if (file.transferred >= file.size) { if (file.transferred >= file.size) {
file.state = "finished"; file.state = "completed";
} else { } else {
file.state = "transferring"; file.state = "transferring";
} }
@ -547,25 +547,26 @@ function onMsgLoad(sessionId, outgoing, msg) {
dislayHistory(false); dislayHistory(false);
} }
} }
function onFileLoad(sessionId, outgoing, uuid, file_name) { function onFileLoad(sessionId, outgoing, uuid, fileName) {
msgHistory.get(sessionId).unshift([outgoing, true, [uuid, file_name]]); msgHistory.get(sessionId).unshift([outgoing, true, [uuid, fileName]]);
if (currentSessionId == sessionId) { if (currentSessionId == sessionId) {
dislayHistory(false); dislayHistory(false);
} }
} }
function onDisconnected(sessionId) { function onDisconnected(sessionId) {
if (currentSessionId == sessionId) { pendingFiles.delete(sessionId);
displayChatBottom();
}
let session = sessionsData.get(sessionId); let session = sessionsData.get(sessionId);
if (session.is_contact) { if (session.is_contact) {
session.is_online = false; session.is_online = false;
} else { } else {
sessionsData.delete(sessionId); sessionsData.delete(sessionId);
if (currentSessionId == sessionId) { }
currentSessionId = -1; if (currentSessionId == sessionId) {
document.getElementById("chat_header").classList.add("offline"); displayChatBottom();
} }
if (currentSessionId == sessionId && !session.is_contact) {
currentSessionId = -1;
document.getElementById("chat_header").classList.add("offline");
} }
displaySessions(); displaySessions();
} }
@ -808,57 +809,62 @@ function generateFileInfo(fileName, fileSize, p) {
p.appendChild(document.createTextNode(" ("+humanFileSize(fileSize)+")")); p.appendChild(document.createTextNode(" ("+humanFileSize(fileSize)+")"));
} }
function displayChatBottom(speed = undefined) { function displayChatBottom(speed = undefined) {
let msgBox = document.getElementById("message_box");
let session = sessionsData.get(currentSessionId); let session = sessionsData.get(currentSessionId);
if (session.is_online) { if (typeof session === "undefined") {
document.getElementById("message_box").style.display = "flex"; msgBox.removeAttribute("style");
} else { } else {
document.getElementById("message_box").removeAttribute("style"); if (session.is_online) {
} msgBox.style.display = "flex";
let fileTransfer = document.getElementById("file_transfer"); } else {
if (pendingFiles.has(currentSessionId)) { msgBox.removeAttribute("style");
let file = pendingFiles.get(currentSessionId); }
let fileInfo = document.getElementById("file_info"); let fileTransfer = document.getElementById("file_transfer");
fileInfo.innerHTML = ""; if (pendingFiles.has(currentSessionId)) {
generateFileInfo(file.name, file.size, fileInfo); let file = pendingFiles.get(currentSessionId);
let fileProgress = document.getElementById("file_progress"); let fileInfo = document.getElementById("file_info");
fileProgress.style.display = "none"; //hide by default fileInfo.innerHTML = "";
let fileStatus = document.getElementById("file_status"); generateFileInfo(file.name, file.size, fileInfo);
fileStatus.removeAttribute("style"); //show by default let fileProgress = document.getElementById("file_progress");
let fileCancel = document.getElementById("file_cancel"); fileProgress.style.display = "none"; //hide by default
fileCancel.style.display = "none"; //hide by default let fileStatus = document.getElementById("file_status");
document.querySelector("#file_progress_bar>div").style.width = 0; fileStatus.removeAttribute("style"); //show by default
switch (file.state) { let fileCancel = document.getElementById("file_cancel");
case "transferring": fileCancel.style.display = "none"; //hide by default
fileCancel.removeAttribute("style"); //show document.querySelector("#file_progress_bar>div").style.width = 0;
fileStatus.style.display = "none"; switch (file.state) {
fileProgress.removeAttribute("style"); //show case "transferring":
let percent = (file.transferred/file.size)*100; fileCancel.removeAttribute("style"); //show
document.getElementById("file_percent").textContent = percent.toFixed(2)+"%"; fileStatus.style.display = "none";
if (typeof speed !== "undefined") { fileProgress.removeAttribute("style"); //show
document.getElementById("file_speed").textContent = humanFileSize(speed)+"/s"; let percent = (file.transferred/file.size)*100;
} document.getElementById("file_percent").textContent = percent.toFixed(2)+"%";
document.querySelector("#file_progress_bar>div").style.width = Math.round(percent)+"%"; if (typeof speed !== "undefined") {
break; document.getElementById("file_speed").textContent = humanFileSize(speed)+"/s";
case "waiting": }
fileStatus.textContent = "Waiting for peer confirmation..."; document.querySelector("#file_progress_bar>div").style.width = Math.round(percent)+"%";
break; break;
case "accepted": case "waiting":
fileStatus.textContent = "Downloading file..."; fileStatus.textContent = "Waiting for peer confirmation...";
break; break;
case "aborted": case "accepted":
fileStatus.textContent = "Transfer aborted."; fileStatus.textContent = "Downloading file...";
pendingFiles.delete(currentSessionId); break;
break; case "aborted":
case "sending": fileStatus.textContent = "Transfer aborted.";
fileStatus.textContent = "Sending file..."; pendingFiles.delete(currentSessionId);
break; break;
case "finished": case "sending":
fileStatus.textContent = "Transfer finished."; fileStatus.textContent = "Sending file...";
pendingFiles.delete(currentSessionId); break;
case "completed":
fileStatus.textContent = "Transfer completed.";
pendingFiles.delete(currentSessionId);
}
fileTransfer.classList.add("active");
} else {
fileTransfer.classList.remove("active");
} }
fileTransfer.classList.add("active");
} else {
fileTransfer.classList.remove("active");
} }
} }
function dislayHistory(scrollToBottom = true) { function dislayHistory(scrollToBottom = true) {
@ -880,4 +886,4 @@ function dislayHistory(scrollToBottom = true) {
if (scrollToBottom) { if (scrollToBottom) {
msg_log.scrollTop = msg_log.scrollHeight; msg_log.scrollTop = msg_log.scrollHeight;
} }
} }

2
src/main.rs
View File

@ -98,7 +98,7 @@ async fn websocket_worker(mut ui_connection: UiConnection, global_vars: Arc<RwLo
load_msgs(session_manager.clone(), &mut ui_connection, &contact.0); load_msgs(session_manager.clone(), &mut ui_connection, &contact.0);
}); });
session_manager.sessions.read().unwrap().iter().for_each(|session| { session_manager.sessions.read().unwrap().iter().for_each(|session| {
ui_connection.on_new_session(session.0, &session.1.name, session.1.outgoing, session.1.file_transfer.as_ref()); ui_connection.on_new_session(session.0, &session.1.name, session.1.outgoing, session.1.file_download.as_ref());
}); });
let not_seen = session_manager.list_not_seen(); let not_seen = session_manager.list_not_seen();
if not_seen.len() > 0 { if not_seen.len() > 0 {

173
src/session_manager/mod.rs
View File

@ -9,7 +9,7 @@ use session::Session;
use ed25519_dalek::PUBLIC_KEY_LENGTH; use ed25519_dalek::PUBLIC_KEY_LENGTH;
use uuid::Uuid; use uuid::Uuid;
use platform_dirs::UserDirs; use platform_dirs::UserDirs;
use crate::{constants, discovery, identity::{Contact, Identity}, utils::get_unix_timestamp, print_error}; use crate::{constants, discovery, identity::{Contact, Identity}, utils::{get_unix_timestamp, get_not_used_path}, print_error};
use crate::ui_interface::UiConnection; use crate::ui_interface::UiConnection;
#[derive(Display, Debug, PartialEq, Eq)] #[derive(Display, Debug, PartialEq, Eq)]
@ -34,7 +34,7 @@ enum SessionCommand {
}, },
Close, Close,
} }
#[derive(Clone, PartialEq, Eq)] #[derive(Copy, Clone, PartialEq, Eq)]
pub enum FileState { pub enum FileState {
ASKING, ASKING,
ACCEPTED, ACCEPTED,
@ -57,7 +57,7 @@ pub struct SessionData {
pub outgoing: bool, pub outgoing: bool,
peer_public_key: [u8; PUBLIC_KEY_LENGTH], peer_public_key: [u8; PUBLIC_KEY_LENGTH],
sender: Sender<SessionCommand>, sender: Sender<SessionCommand>,
pub file_transfer: Option<LargeFileDownload>, pub file_download: Option<LargeFileDownload>,
} }
pub struct SessionManager { pub struct SessionManager {
@ -164,6 +164,9 @@ impl SessionManager {
} }
fn remove_session(&self, session_id: &usize) { fn remove_session(&self, session_id: &usize) {
self.with_ui_connection(|ui_connection| {
ui_connection.on_disconnected(&session_id);
});
self.sessions.write().unwrap().remove(session_id); self.sessions.write().unwrap().remove(session_id);
self.saved_msgs.lock().unwrap().remove(session_id); self.saved_msgs.lock().unwrap().remove(session_id);
self.not_seen.write().unwrap().retain(|x| x != session_id); self.not_seen.write().unwrap().retain(|x| x != session_id);
@ -172,9 +175,9 @@ impl SessionManager {
async fn send_msg(&self, session_id: usize, session: &mut Session, buff: &[u8], aborted: &mut bool, file_ack_sender: Option<&Sender<bool>>) -> Result<(), SessionError> { async fn send_msg(&self, session_id: usize, session: &mut Session, buff: &[u8], aborted: &mut bool, file_ack_sender: Option<&Sender<bool>>) -> Result<(), SessionError> {
session.encrypt_and_send(&buff).await?; session.encrypt_and_send(&buff).await?;
if buff[0] == protocol::Headers::ACCEPT_LARGE_FILE { if buff[0] == protocol::Headers::ACCEPT_LARGE_FILE {
self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_transfer.as_mut().unwrap().state = FileState::ACCEPTED; 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 { } else if buff[0] == protocol::Headers::ABORT_FILE_TRANSFER {
self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_transfer = None; self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = None;
*aborted = true; *aborted = true;
if let Some(sender) = file_ack_sender { if let Some(sender) = file_ack_sender {
if let Err(e) = sender.send(false).await { if let Err(e) = sender.send(false).await {
@ -234,86 +237,95 @@ impl SessionManager {
} }
} }
protocol::Headers::ASK_LARGE_FILE => { protocol::Headers::ASK_LARGE_FILE => {
let file_size = u64::from_be_bytes(buffer[1..9].try_into().unwrap()); if self.sessions.read().unwrap().get(&session_id).unwrap().file_download.is_none() { //don't accept 2 downloads at the same time
match from_utf8(&buffer[9..]) { let file_size = u64::from_be_bytes(buffer[1..9].try_into().unwrap());
Ok(file_name) => { match from_utf8(&buffer[9..]) {
let file_name = sanitize_filename::sanitize(file_name); Ok(file_name) => {
let download_dir = UserDirs::new().unwrap().download_dir; let file_name = sanitize_filename::sanitize(file_name);
self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_transfer = Some(LargeFileDownload{ let download_dir = UserDirs::new().unwrap().download_dir;
file_name: file_name.clone(), self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = Some(LargeFileDownload{
download_location: download_dir.to_str().unwrap().to_string(), file_name: file_name.clone(),
file_size, download_location: download_dir.to_str().unwrap().to_string(),
state: FileState::ASKING, file_size,
transferred: 0, state: FileState::ASKING,
last_chunk: get_unix_timestamp(), transferred: 0,
}); last_chunk: get_unix_timestamp(),
let mut test_file_path = download_dir.join(&file_name); });
let mut n = 1; local_file_path = Some(get_not_used_path(&file_name, &download_dir));
while test_file_path.exists() { self.with_ui_connection(|ui_connection| {
let splits: Vec<&str> = file_name.split('.').collect(); ui_connection.on_ask_large_file(&session_id, file_size, &file_name, download_dir.to_str().unwrap());
test_file_path = download_dir.join(format!("{} ({}).{}", splits[..splits.len()-1].join("."), n, splits[splits.len()-1])); })
n += 1;
} }
local_file_path = Some(test_file_path); Err(e) => print_error!(e),
self.with_ui_connection(|ui_connection| {
ui_connection.on_ask_large_file(&session_id, file_size, &file_name, download_dir.to_str().unwrap());
})
} }
Err(e) => print_error!(e), } else if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ABORT_FILE_TRANSFER]).await {
print_error!(e);
break;
} }
} }
protocol::Headers::LARGE_FILE_CHUNK => { protocol::Headers::LARGE_FILE_CHUNK => {
let file_transfer_opt = { let state = {
self.sessions.read().unwrap().get(&session_id).unwrap().file_transfer.clone() let sessions = self.sessions.read().unwrap();
match sessions.get(&session_id).unwrap().file_download.as_ref() {
Some(file_transfer) => Some(file_transfer.state),
None => None
}
}; };
if let Some(file_transfer) = file_transfer_opt { let mut should_accept_chunk = false;
if file_transfer.state == FileState::ACCEPTED || file_transfer.state == FileState::TRANSFERRING { if let Some(state) = state {
if local_file_handle.is_none() { if state == FileState::ACCEPTED {
if let Some(file_path) = local_file_path.as_ref() { if let Some(file_path) = local_file_path.as_ref() {
match OpenOptions::new().append(true).create(true).open(file_path) { match OpenOptions::new().append(true).create(true).open(file_path) {
Ok(file) => local_file_handle = Some(file), Ok(file) => {
Err(e) => print_error!(e) local_file_handle = Some(file);
} let mut sessions = self.sessions.write().unwrap();
} let file_transfer = sessions.get_mut(&session_id).unwrap().file_download.as_mut().unwrap();
} file_transfer.state = FileState::TRANSFERRING;
let mut is_success = false; should_accept_chunk = true;
if let Some(file_handle) = local_file_handle.as_mut() {
match file_handle.write_all(&buffer[1..]) {
Ok(_) => {
let chunk_size = (buffer.len()-1) as u64;
{
let mut sessions = self.sessions.write().unwrap();
let file_transfer = sessions.get_mut(&session_id).unwrap().file_transfer.as_mut().unwrap();
file_transfer.last_chunk = get_unix_timestamp();
file_transfer.transferred += chunk_size;
if file_transfer.transferred >= file_transfer.file_size { //we downloaded all the file
sessions.get_mut(&session_id).unwrap().file_transfer = None;
local_file_path = None;
local_file_handle = None;
} else if file_transfer.state != FileState::TRANSFERRING {
file_transfer.state = FileState::TRANSFERRING;
}
}
if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ACK_CHUNK]).await {
print_error!(e);
break;
}
self.with_ui_connection(|ui_connection| {
ui_connection.inc_file_transfer(&session_id, chunk_size);
});
is_success = true;
} }
Err(e) => print_error!(e) Err(e) => print_error!(e)
} }
} }
if !is_success { } else if state == FileState::TRANSFERRING {
self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_transfer = None; should_accept_chunk = true;
local_file_path = None; }
local_file_handle = None; }
if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ABORT_FILE_TRANSFER]).await { if should_accept_chunk {
print_error!(e); let mut is_success = false;
break; if let Some(file_handle) = local_file_handle.as_mut() {
match file_handle.write_all(&buffer[1..]) {
Ok(_) => {
let chunk_size = (buffer.len()-1) as u64;
{
let mut sessions = self.sessions.write().unwrap();
let file_transfer = sessions.get_mut(&session_id).unwrap().file_download.as_mut().unwrap();
file_transfer.last_chunk = get_unix_timestamp();
file_transfer.transferred += chunk_size;
if file_transfer.transferred >= file_transfer.file_size { //we downloaded all the file
sessions.get_mut(&session_id).unwrap().file_download = None;
local_file_path = None;
local_file_handle = None;
}
}
if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ACK_CHUNK]).await {
print_error!(e);
break;
}
self.with_ui_connection(|ui_connection| {
ui_connection.inc_file_transfer(&session_id, chunk_size);
});
is_success = true;
} }
Err(e) => print_error!(e)
}
}
if !is_success {
self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = None;
local_file_path = None;
local_file_handle = None;
if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ABORT_FILE_TRANSFER]).await {
print_error!(e);
break;
} }
} }
} }
@ -337,7 +349,7 @@ impl SessionManager {
} }
aborted = true; aborted = true;
} }
self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_transfer = None; self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = None;
local_file_path = None; local_file_path = None;
local_file_handle = None; local_file_handle = None;
self.with_ui_connection(|ui_connection| { self.with_ui_connection(|ui_connection| {
@ -383,12 +395,9 @@ impl SessionManager {
} }
} }
Err(e) => { Err(e) => {
if e != SessionError::BrokenPipe && e != SessionError::ConnectionReset { if e != SessionError::BrokenPipe && e != SessionError::ConnectionReset && e != SessionError::BufferTooLarge {
print_error!(e); print_error!(e);
} }
self.with_ui_connection(|ui_connection| {
ui_connection.on_disconnected(&session_id);
});
break; break;
} }
} }
@ -488,7 +497,7 @@ impl SessionManager {
outgoing, outgoing,
peer_public_key, peer_public_key,
sender: sender, sender: sender,
file_transfer: None, file_download: None,
}; };
let mut session_id = None; let mut session_id = None;
for (i, contact) in session_manager.loaded_contacts.read().unwrap().iter() { for (i, contact) in session_manager.loaded_contacts.read().unwrap().iter() {
@ -610,7 +619,9 @@ impl SessionManager {
let result = Identity::remove_contact(&loaded_contacts.get(&session_id).unwrap().uuid); let result = Identity::remove_contact(&loaded_contacts.get(&session_id).unwrap().uuid);
if result.is_ok() { if result.is_ok() {
if let Some(contact) = loaded_contacts.remove(&session_id) { if let Some(contact) = loaded_contacts.remove(&session_id) {
self.sessions.write().unwrap().get_mut(&session_id).unwrap().name = contact.name; if let Some(session) = self.sessions.write().unwrap().get_mut(&session_id) {
session.name = contact.name;
}
} }
self.last_loaded_msg_offsets.write().unwrap().remove(&session_id); self.last_loaded_msg_offsets.write().unwrap().remove(&session_id);
} }

18
src/session_manager/protocol.rs
View File

@ -1,15 +1,15 @@
pub struct Headers; pub struct Headers;
impl Headers { impl Headers {
pub const MESSAGE: u8 = 0x01; pub const MESSAGE: u8 = 0x00;
pub const ASK_NAME: u8 = 0x02; pub const ASK_NAME: u8 = 0x01;
pub const TELL_NAME: u8 = 0x03; pub const TELL_NAME: u8 = 0x02;
pub const FILE: u8 = 0x04; pub const FILE: u8 = 0x03;
pub const ASK_LARGE_FILE: u8 = 0x05; pub const ASK_LARGE_FILE: u8 = 0x04;
pub const ACCEPT_LARGE_FILE: u8 = 0x06; pub const ACCEPT_LARGE_FILE: u8 = 0x05;
pub const LARGE_FILE_CHUNK: u8 = 0x07; pub const LARGE_FILE_CHUNK: u8 = 0x06;
pub const ACK_CHUNK: u8 = 0x08; pub const ACK_CHUNK: u8 = 0x07;
pub const ABORT_FILE_TRANSFER: u8 = 0x09; pub const ABORT_FILE_TRANSFER: u8 = 0x08;
} }
pub fn new_message(message: String) -> Vec<u8> { pub fn new_message(message: String) -> Vec<u8> {

1
src/session_manager/session.rs
View File

@ -252,6 +252,7 @@ impl Session {
Err(_) => Err(SessionError::TransmissionCorrupted) Err(_) => Err(SessionError::TransmissionCorrupted)
} }
} else { } else {
print_error!("Buffer too large: {} B", recv_len);
Err(SessionError::BufferTooLarge) Err(SessionError::BufferTooLarge)
} }
} }

18
src/utils.rs
View File

@ -1,4 +1,4 @@
use std::{convert::TryInto, time::{SystemTime, UNIX_EPOCH}}; use std::{convert::TryInto, time::{SystemTime, UNIX_EPOCH}, path::PathBuf};
use uuid::Bytes; use uuid::Bytes;
use crate::print_error; use crate::print_error;
@ -28,6 +28,22 @@ pub fn get_unix_timestamp() -> u128 {
SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis()
} }
pub fn get_not_used_path(file_name: &str, parent_directory: &PathBuf) -> String {
let has_extension = file_name.matches('.').count() > 0;
let mut path = parent_directory.join(&file_name);
let mut n = 1;
while path.exists() {
path = if has_extension {
let splits: Vec<&str> = file_name.split('.').collect();
parent_directory.join(format!("{} ({}).{}", splits[..splits.len()-1].join("."), n, splits[splits.len()-1]))
} else {
parent_directory.join(format!("{} ({})", file_name, n))
};
n += 1;
}
path.to_str().unwrap().to_owned()
}
#[macro_export] #[macro_export]
macro_rules! print_error { macro_rules! print_error {
($arg:tt) => ({ ($arg:tt) => ({