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

36
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,26 +547,27 @@ 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) { if (currentSessionId == sessionId) {
displayChatBottom();
}
if (currentSessionId == sessionId && !session.is_contact) {
currentSessionId = -1; currentSessionId = -1;
document.getElementById("chat_header").classList.add("offline"); document.getElementById("chat_header").classList.add("offline");
} }
}
displaySessions(); displaySessions();
} }
function onFingerprints(local, peer) { function onFingerprints(local, peer) {
@ -808,11 +809,15 @@ 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";
} else {
msgBox.removeAttribute("style");
} }
let fileTransfer = document.getElementById("file_transfer"); let fileTransfer = document.getElementById("file_transfer");
if (pendingFiles.has(currentSessionId)) { if (pendingFiles.has(currentSessionId)) {
@ -852,8 +857,8 @@ function displayChatBottom(speed = undefined) {
case "sending": case "sending":
fileStatus.textContent = "Sending file..."; fileStatus.textContent = "Sending file...";
break; break;
case "finished": case "completed":
fileStatus.textContent = "Transfer finished."; fileStatus.textContent = "Transfer completed.";
pendingFiles.delete(currentSessionId); pendingFiles.delete(currentSessionId);
} }
fileTransfer.classList.add("active"); fileTransfer.classList.add("active");
@ -861,6 +866,7 @@ function displayChatBottom(speed = undefined) {
fileTransfer.classList.remove("active"); fileTransfer.classList.remove("active");
} }
} }
}
function dislayHistory(scrollToBottom = true) { function dislayHistory(scrollToBottom = true) {
msg_log.style.display = "block"; msg_log.style.display = "block";
msg_log.innerHTML = ""; msg_log.innerHTML = "";

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 {

77
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,12 +237,13 @@ impl SessionManager {
} }
} }
protocol::Headers::ASK_LARGE_FILE => { 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
let file_size = u64::from_be_bytes(buffer[1..9].try_into().unwrap()); let file_size = u64::from_be_bytes(buffer[1..9].try_into().unwrap());
match from_utf8(&buffer[9..]) { match from_utf8(&buffer[9..]) {
Ok(file_name) => { Ok(file_name) => {
let file_name = sanitize_filename::sanitize(file_name); let file_name = sanitize_filename::sanitize(file_name);
let download_dir = UserDirs::new().unwrap().download_dir; let download_dir = UserDirs::new().unwrap().download_dir;
self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_transfer = Some(LargeFileDownload{ self.sessions.write().unwrap().get_mut(&session_id).unwrap().file_download = Some(LargeFileDownload{
file_name: file_name.clone(), file_name: file_name.clone(),
download_location: download_dir.to_str().unwrap().to_string(), download_location: download_dir.to_str().unwrap().to_string(),
file_size, file_size,
@ -247,35 +251,46 @@ impl SessionManager {
transferred: 0, transferred: 0,
last_chunk: get_unix_timestamp(), last_chunk: get_unix_timestamp(),
}); });
let mut test_file_path = download_dir.join(&file_name); local_file_path = Some(get_not_used_path(&file_name, &download_dir));
let mut n = 1;
while test_file_path.exists() {
let splits: Vec<&str> = file_name.split('.').collect();
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);
self.with_ui_connection(|ui_connection| { self.with_ui_connection(|ui_connection| {
ui_connection.on_ask_large_file(&session_id, file_size, &file_name, download_dir.to_str().unwrap()); ui_connection.on_ask_large_file(&session_id, file_size, &file_name, download_dir.to_str().unwrap());
}) })
} }
Err(e) => print_error!(e), 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) => {
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;
should_accept_chunk = true;
}
Err(e) => print_error!(e) Err(e) => print_error!(e)
} }
} }
} else if state == FileState::TRANSFERRING {
should_accept_chunk = true;
} }
}
if should_accept_chunk {
let mut is_success = false; let mut is_success = false;
if let Some(file_handle) = local_file_handle.as_mut() { if let Some(file_handle) = local_file_handle.as_mut() {
match file_handle.write_all(&buffer[1..]) { match file_handle.write_all(&buffer[1..]) {
@ -283,15 +298,13 @@ impl SessionManager {
let chunk_size = (buffer.len()-1) as u64; let chunk_size = (buffer.len()-1) as u64;
{ {
let mut sessions = self.sessions.write().unwrap(); let mut sessions = self.sessions.write().unwrap();
let file_transfer = sessions.get_mut(&session_id).unwrap().file_transfer.as_mut().unwrap(); let file_transfer = sessions.get_mut(&session_id).unwrap().file_download.as_mut().unwrap();
file_transfer.last_chunk = get_unix_timestamp(); file_transfer.last_chunk = get_unix_timestamp();
file_transfer.transferred += chunk_size; file_transfer.transferred += chunk_size;
if file_transfer.transferred >= file_transfer.file_size { //we downloaded all the file if file_transfer.transferred >= file_transfer.file_size { //we downloaded all the file
sessions.get_mut(&session_id).unwrap().file_transfer = None; sessions.get_mut(&session_id).unwrap().file_download = None;
local_file_path = None; local_file_path = None;
local_file_handle = 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 { if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ACK_CHUNK]).await {
@ -307,7 +320,7 @@ impl SessionManager {
} }
} }
if !is_success { if !is_success {
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;
if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ABORT_FILE_TRANSFER]).await { if let Err(e) = session.encrypt_and_send(&[protocol::Headers::ABORT_FILE_TRANSFER]).await {
@ -317,7 +330,6 @@ impl SessionManager {
} }
} }
} }
}
protocol::Headers::ACK_CHUNK => { protocol::Headers::ACK_CHUNK => {
if let Some(sender) = file_ack_sender.clone() { if let Some(sender) = file_ack_sender.clone() {
if let Some(next_chunk) = next_chunk.as_ref() { if let Some(next_chunk) = next_chunk.as_ref() {
@ -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) => ({