transbeam/src/upload.rs

271 lines
9.6 KiB
Rust
Raw Normal View History

use std::{collections::HashSet, io::Write, task::Waker};
2022-04-26 23:54:29 -04:00
use actix::{Actor, ActorContext, AsyncContext, StreamHandler, Message, Handler};
2022-04-27 00:55:36 -04:00
use actix_http::ws::{CloseReason, Item};
2022-04-26 23:54:29 -04:00
use actix_web_actors::ws::{self, CloseCode};
use bytes::Bytes;
2022-04-27 00:55:36 -04:00
use log::{debug, error, info, trace};
2022-04-26 23:54:29 -04:00
use rand::distributions::{Alphanumeric, DistString};
use serde::Deserialize;
use time::OffsetDateTime;
2022-04-27 00:55:36 -04:00
use crate::{file::LiveWriter, DownloadableFile, UploadedFile};
2022-04-26 23:54:29 -04:00
2022-04-28 00:27:22 -04:00
const MAX_FILES: usize = 256;
2022-04-26 23:54:29 -04:00
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,
2022-04-28 00:27:22 -04:00
#[error("Number of files submitted by client exceeded the maximum limit")]
TooManyFiles,
2022-04-26 23:54:29 -04:00
#[error("Websocket was closed by client before completing transfer")]
ClosedEarly(Option<CloseReason>),
#[error("Client sent more data than they were supposed to")]
2022-04-28 00:27:22 -04:00
UnexpectedExtraData,
2022-04-26 23:54:29 -04:00
}
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,
2022-04-28 00:27:22 -04:00
Self::TooManyFiles => CloseCode::Policy,
2022-04-26 23:54:29 -04:00
Self::ClosedEarly(_) => CloseCode::Invalid,
2022-04-28 00:27:22 -04:00
Self::UnexpectedExtraData => CloseCode::Invalid,
2022-04-26 23:54:29 -04:00
}
}
}
pub struct Uploader {
writer: Option<Box<dyn LiveWriter>>,
storage_filename: String,
2022-04-26 23:54:29 -04:00
app_data: super::AppData,
bytes_remaining: u64,
2022-04-26 23:54:29 -04:00
}
impl Uploader {
pub fn new(app_data: super::AppData) -> Self {
Self {
writer: None,
storage_filename: String::new(),
2022-04-26 23:54:29 -04:00
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));
}
}
2022-04-26 23:54:29 -04:00
#[derive(Debug, Deserialize)]
struct RawUploadedFile {
name: String,
size: u64,
2022-04-26 23:54:29 -04:00
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) => {
2022-04-28 00:27:22 -04:00
error!("Websocket error: {}", e);
2022-04-26 23:54:29 -04:00
ctx.stop();
return;
2022-04-27 00:55:36 -04:00
}
2022-04-26 23:54:29 -04:00
};
match self.handle_message(msg, ctx) {
Err(e) => {
2022-04-28 00:27:22 -04:00
error!("{}", e);
2022-04-26 23:54:29 -04:00
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();
}
}));
2022-04-26 23:54:29 -04:00
ctx.stop();
}
2022-04-27 00:55:36 -04:00
_ => (),
2022-04-26 23:54:29 -04:00
}
}
}
fn ack(ctx: &mut <Uploader as Actor>::Context) {
ctx.text("ack");
}
impl Uploader {
2022-04-27 00:55:36 -04:00
fn handle_message(
&mut self,
msg: ws::Message,
ctx: &mut <Self as Actor>::Context,
) -> Result<bool, Error> {
2022-04-26 23:54:29 -04:00
trace!("Websocket message: {:?}", msg);
match msg {
ws::Message::Text(text) => {
if self.writer.is_some() {
2022-04-27 00:55:36 -04:00
return Err(Error::UnexpectedMessageType);
2022-04-26 23:54:29 -04:00
}
let raw_files: Vec<RawUploadedFile> = serde_json::from_slice(text.as_bytes())?;
info!("Received file list: {} files", raw_files.len());
debug!("{:?}", raw_files);
2022-04-28 00:27:22 -04:00
if raw_files.len() > MAX_FILES {
return Err(Error::TooManyFiles);
}
2022-04-26 23:54:29 -04:00
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();
2022-04-26 23:54:29 -04:00
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 {
2022-04-26 23:54:29 -04:00
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();
2022-04-27 00:55:36 -04:00
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,
})
2022-04-26 23:54:29 -04:00
} 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
2022-04-27 00:55:36 -04:00
.write()
.await
2022-04-27 00:55:36 -04:00
.insert(
storage_filename_copy,
downloadable_file,
2022-04-27 00:55:36 -04:00
);
}));
2022-04-27 01:31:34 -04:00
ctx.text(self.storage_filename.as_str());
2022-04-26 23:54:29 -04:00
}
ws::Message::Binary(data)
2022-04-27 00:55:36 -04:00
| 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);
2022-04-26 23:54:29 -04:00
}
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 {
2022-04-28 00:27:22 -04:00
return Err(Error::UnexpectedExtraData);
}
self.bytes_remaining -= data.len() as u64;
writer.write_all(&data)?;
Ok(self.bytes_remaining == 0)
} else {
Err(Error::UnexpectedMessageType)
}
}
2022-04-26 23:54:29 -04:00
}