feat: added upload file capability

This commit is contained in:
Andrew Rioux
2023-09-06 14:56:21 -04:00
parent 231108f2c3
commit 50b5c1a921
9 changed files with 277 additions and 5 deletions

View File

@@ -2,10 +2,14 @@ use std::ffi::c_int;
use sparse_05_common::messages::{Capabilities, OperatingSystem, TransportType};
#[cfg(target_os = "linux")]
const CAP_SETUID: u32 = 1 << 7;
#[cfg(target_os = "linux")]
const CAP_NET_RAW: u32 = 1 << 13;
#[cfg(target_os = "linux")]
const SYS_CAPGET: i64 = 125;
#[cfg(target_os = "linux")]
#[allow(non_camel_case_types)]
#[repr(C)]
#[derive(Debug)]
@@ -14,6 +18,7 @@ struct cap_user_header_t {
pid: c_int,
}
#[cfg(target_os = "linux")]
#[allow(non_camel_case_types)]
#[repr(C)]
#[derive(Debug)]

View File

@@ -317,6 +317,15 @@ where
let _ = handler.data_sender.send((number, bytes));
}
}
Command::GetUploadStatus(id, up_to) => {
let Ok(lock) = uploaded_files.lock() else {
continue;
};
if let Some(handler) = lock.get(&id) {
let _ = handler.request_status.send(up_to);
}
}
Command::StartDownloadFile(_) => {}
Command::DownloadFileStatus(_, _) => {}

View File

@@ -1,5 +1,7 @@
use std::{
collections::HashMap,
fs::{self, OpenOptions},
io::Write,
path::PathBuf,
sync::{
atomic::{AtomicU64, Ordering},
@@ -7,8 +9,11 @@ use std::{
Arc, Mutex,
},
thread::Scope,
time::Duration,
};
use sparse_05_common::messages::{Response, FILE_BUFFER_BUFFER_SIZE};
use super::ConnectionInformation;
static CURRENT_FILE_UPLOAD_ID: AtomicU64 = AtomicU64::new(0);
@@ -16,6 +21,7 @@ static CURRENT_FILE_UPLOAD_ID: AtomicU64 = AtomicU64::new(0);
pub(super) struct UploadFileHandler {
pub id: u64,
pub data_sender: Sender<(u64, Vec<u8>)>,
pub request_status: Sender<u64>,
}
pub(super) fn start_file_upload<'a, 'b: 'a>(
@@ -26,10 +32,69 @@ pub(super) fn start_file_upload<'a, 'b: 'a>(
upload_file_map: Arc<Mutex<HashMap<u64, UploadFileHandler>>>,
) -> anyhow::Result<UploadFileHandler> {
let (data_sender, data_receiver) = channel();
let (request_status, receive_request_status) = channel();
let id = CURRENT_FILE_UPLOAD_ID.fetch_add(1, Ordering::Relaxed);
let id_2 = id;
let buffer: Vec<Vec<u8>> = Vec::with_capacity(packet_count as usize);
s.spawn(move || -> anyhow::Result<()> {
conninfo.send(conninfo.encrypt_and_sign_resp(Response::UploadFileID(id))?)?;
Ok(UploadFileHandler { id, data_sender })
let mut target_file = OpenOptions::new()
.write(true)
.create(true)
.open(&file_path)?;
let mut current_packet_count = 0;
while current_packet_count < packet_count {
let mut buffers: Vec<Option<Vec<u8>>> = vec![None; FILE_BUFFER_BUFFER_SIZE];
loop {
let Ok((i, buffer)) = data_receiver.recv_timeout(Duration::from_millis(250)) else {
let up_to = receive_request_status.recv()?;
let needed = buffers[..up_to as usize]
.iter()
.enumerate()
.flat_map(|(i, b)| match b {
Some(..) => None,
None => Some(i as u64),
})
.collect::<Vec<_>>();
let is_empty = needed.is_empty();
conninfo.send(
conninfo.encrypt_and_sign_resp(Response::UploadFileStatus(id, needed))?,
)?;
if is_empty {
current_packet_count += up_to;
break;
} else {
continue;
}
};
buffers[i as usize] = Some(buffer);
}
for buffer in buffers {
let Some(buffer) = buffer else { break };
target_file.write(&buffer)?;
}
}
let Ok(mut lock) = upload_file_map.lock() else {
return Ok(());
};
lock.remove(&id);
Ok(())
});
Ok(UploadFileHandler {
id: id_2,
data_sender,
request_status,
})
}