264 lines
9.2 KiB
Rust
264 lines
9.2 KiB
Rust
use std::{collections::HashSet, io::Write, task::Waker};
|
|
|
|
use actix::{Actor, ActorContext, AsyncContext, StreamHandler, Message, Handler};
|
|
use actix_http::ws::{CloseReason, Item};
|
|
use actix_web_actors::ws::{self, CloseCode};
|
|
use bytes::Bytes;
|
|
use log::{debug, error, info, trace};
|
|
use rand::distributions::{Alphanumeric, DistString};
|
|
use serde::Deserialize;
|
|
use time::OffsetDateTime;
|
|
|
|
use crate::{file::LiveWriter, DownloadableFile, UploadedFile};
|
|
|
|
const FILENAME_DATE_FORMAT: &[time::format_description::FormatItem] =
|
|
time::macros::format_description!("[year]-[month]-[day]-[hour][minute][second]");
|
|
|
|
#[derive(thiserror::Error, Debug)]
|
|
enum Error {
|
|
#[error("Failed to parse file metadata")]
|
|
Parse(#[from] serde_json::Error),
|
|
#[error("Error writing to stored file")]
|
|
Storage(#[from] std::io::Error),
|
|
#[error("Time formatting error")]
|
|
TimeFormat(#[from] time::error::Format),
|
|
#[error("Duplicate filename could not be deduplicated")]
|
|
DuplicateFilename,
|
|
#[error("This message type was not expected at this stage")]
|
|
UnexpectedMessageType,
|
|
#[error("Metadata contained an empty list of files")]
|
|
NoFiles,
|
|
#[error("Websocket was closed by client before completing transfer")]
|
|
ClosedEarly(Option<CloseReason>),
|
|
#[error("Client sent more data than they were supposed to")]
|
|
TooMuchData,
|
|
}
|
|
|
|
impl Error {
|
|
fn close_code(&self) -> CloseCode {
|
|
match self {
|
|
Self::Parse(_) => CloseCode::Invalid,
|
|
Self::Storage(_) => CloseCode::Error,
|
|
Self::TimeFormat(_) => CloseCode::Error,
|
|
Self::DuplicateFilename => CloseCode::Policy,
|
|
Self::UnexpectedMessageType => CloseCode::Invalid,
|
|
Self::NoFiles => CloseCode::Policy,
|
|
Self::ClosedEarly(_) => CloseCode::Invalid,
|
|
Self::TooMuchData => CloseCode::Invalid,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct Uploader {
|
|
writer: Option<Box<dyn LiveWriter>>,
|
|
storage_filename: String,
|
|
app_data: super::AppData,
|
|
bytes_remaining: u64,
|
|
}
|
|
|
|
impl Uploader {
|
|
pub fn new(app_data: super::AppData) -> Self {
|
|
Self {
|
|
writer: None,
|
|
storage_filename: String::new(),
|
|
app_data,
|
|
bytes_remaining: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Actor for Uploader {
|
|
type Context = ws::WebsocketContext<Self>;
|
|
}
|
|
|
|
#[derive(Message)]
|
|
#[rtype(result = "()")]
|
|
pub(crate) struct WakerMessage(pub Waker);
|
|
|
|
impl Handler<WakerMessage> for Uploader {
|
|
type Result = ();
|
|
fn handle(&mut self, msg: WakerMessage, _: &mut Self::Context) {
|
|
self.writer.as_mut().map(|w| w.add_waker(msg.0));
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct RawUploadedFile {
|
|
name: String,
|
|
size: u64,
|
|
modtime: i64,
|
|
}
|
|
|
|
impl RawUploadedFile {
|
|
fn process(&self) -> UploadedFile {
|
|
UploadedFile::new(
|
|
&self.name,
|
|
self.size,
|
|
OffsetDateTime::from_unix_timestamp(self.modtime / 1000)
|
|
.unwrap_or_else(|_| OffsetDateTime::now_utc()),
|
|
)
|
|
}
|
|
}
|
|
|
|
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for Uploader {
|
|
fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
|
|
let msg = match msg {
|
|
Ok(m) => m,
|
|
Err(e) => {
|
|
error!("Websocket error: {:?}", e);
|
|
ctx.stop();
|
|
return;
|
|
}
|
|
};
|
|
|
|
match self.handle_message(msg, ctx) {
|
|
Err(e) => {
|
|
error!("{:?}", e);
|
|
ctx.close(Some(ws::CloseReason {
|
|
code: e.close_code(),
|
|
description: Some(e.to_string()),
|
|
}));
|
|
ctx.stop();
|
|
}
|
|
Ok(true) => {
|
|
info!("Finished uploading data");
|
|
self.writer.as_mut().map(|w| w.flush());
|
|
ctx.close(Some(ws::CloseReason {
|
|
code: CloseCode::Normal,
|
|
description: None,
|
|
}));
|
|
let data = self.app_data.clone();
|
|
let filename = self.storage_filename.clone();
|
|
ctx.wait(actix::fut::wrap_future(async move {
|
|
if let Some(f) = data.write().await.get_mut(&filename) {
|
|
f.uploader.take();
|
|
}
|
|
}));
|
|
ctx.stop();
|
|
}
|
|
_ => (),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn ack(ctx: &mut <Uploader as Actor>::Context) {
|
|
ctx.text("ack");
|
|
}
|
|
|
|
impl Uploader {
|
|
fn handle_message(
|
|
&mut self,
|
|
msg: ws::Message,
|
|
ctx: &mut <Self as Actor>::Context,
|
|
) -> Result<bool, Error> {
|
|
trace!("Websocket message: {:?}", msg);
|
|
match msg {
|
|
ws::Message::Text(text) => {
|
|
if self.writer.is_some() {
|
|
return Err(Error::UnexpectedMessageType);
|
|
}
|
|
let raw_files: Vec<RawUploadedFile> = serde_json::from_slice(text.as_bytes())?;
|
|
info!("Received file list: {} files", raw_files.len());
|
|
debug!("{:?}", raw_files);
|
|
let mut filenames: HashSet<String> = HashSet::new();
|
|
let mut files = Vec::new();
|
|
for raw_file in raw_files.iter() {
|
|
let mut file = raw_file.process();
|
|
while filenames.contains(&file.name) {
|
|
info!("Duplicate file name: {}", file.name);
|
|
if file.name.len() >= sanitise_file_name::Options::DEFAULT.length_limit {
|
|
return Err(Error::DuplicateFilename);
|
|
}
|
|
file.name.insert(0, '_');
|
|
}
|
|
filenames.insert(file.name.clone());
|
|
self.bytes_remaining += file.size;
|
|
files.push(file);
|
|
}
|
|
if files.is_empty() {
|
|
return Err(Error::NoFiles);
|
|
}
|
|
let storage_filename = Alphanumeric.sample_string(&mut rand::thread_rng(), 8);
|
|
self.storage_filename = storage_filename.clone();
|
|
let storage_path = super::storage_dir().join(storage_filename.clone());
|
|
info!("storing to: {:?}", storage_path);
|
|
let writer = super::file::LiveFileWriter::new(&storage_path)?;
|
|
let addr = Some(ctx.address());
|
|
let (writer, downloadable_file): (Box<dyn LiveWriter>, _) = if files.len() > 1 {
|
|
info!("Wrapping in zipfile generator");
|
|
let now = OffsetDateTime::now_utc();
|
|
let zip_writer = super::zip::ZipGenerator::new(files, writer);
|
|
let size = zip_writer.total_size();
|
|
let download_filename =
|
|
super::APP_NAME.to_owned() + &now.format(FILENAME_DATE_FORMAT)? + ".zip";
|
|
(Box::new(zip_writer), DownloadableFile {
|
|
name: download_filename,
|
|
size,
|
|
modtime: now,
|
|
uploader: addr,
|
|
})
|
|
} else {
|
|
(Box::new(writer), DownloadableFile {
|
|
name: files[0].name.clone(),
|
|
size: files[0].size,
|
|
modtime: files[0].modtime,
|
|
uploader: addr,
|
|
})
|
|
};
|
|
self.writer = Some(writer);
|
|
let data = self.app_data.clone();
|
|
let storage_filename_copy = storage_filename.clone();
|
|
ctx.spawn(actix::fut::wrap_future(async move {
|
|
data
|
|
.write()
|
|
.await
|
|
.insert(
|
|
storage_filename_copy,
|
|
downloadable_file,
|
|
);
|
|
}));
|
|
ctx.text(self.storage_filename.as_str());
|
|
}
|
|
ws::Message::Binary(data)
|
|
| ws::Message::Continuation(Item::Last(data)) => {
|
|
let result = self.handle_data(data)?;
|
|
ack(ctx);
|
|
return Ok(result);
|
|
}
|
|
ws::Message::Continuation(Item::FirstBinary(data))
|
|
| ws::Message::Continuation(Item::Continue(data)) => {
|
|
return self.handle_data(data);
|
|
}
|
|
ws::Message::Close(reason) => {
|
|
if self.bytes_remaining > 0 {
|
|
return Err(Error::ClosedEarly(reason));
|
|
} else {
|
|
return Ok(true);
|
|
}
|
|
}
|
|
ws::Message::Ping(ping) => {
|
|
debug!("Ping received, ponging");
|
|
ctx.pong(&ping);
|
|
}
|
|
ws::Message::Nop | ws::Message::Pong(_) => (),
|
|
_ => {
|
|
return Err(Error::UnexpectedMessageType);
|
|
}
|
|
}
|
|
Ok(false)
|
|
}
|
|
|
|
fn handle_data(&mut self, data: Bytes) -> Result<bool, Error> {
|
|
if let Some(ref mut writer) = self.writer {
|
|
if (data.len() as u64) > self.bytes_remaining {
|
|
return Err(Error::TooMuchData);
|
|
}
|
|
self.bytes_remaining -= data.len() as u64;
|
|
writer.write_all(&data)?;
|
|
Ok(self.bytes_remaining == 0)
|
|
} else {
|
|
Err(Error::UnexpectedMessageType)
|
|
}
|
|
}
|
|
}
|