transbeam/src/upload.rs

283 lines
10 KiB
Rust

use std::{collections::HashSet, fs::File, io::Write};
use actix::{fut::future::ActorFutureExt, Actor, ActorContext, AsyncContext, StreamHandler};
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::{storage_dir, DownloadableFile, UploadedFile};
const MAX_FILES: usize = 256;
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("Number of files submitted by client exceeded the maximum limit")]
TooManyFiles,
#[error("Websocket was closed by client before completing transfer")]
ClosedEarly(Option<CloseReason>),
#[error("Client sent more data than they were supposed to")]
UnexpectedExtraData,
}
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::TooManyFiles => CloseCode::Policy,
Self::ClosedEarly(_) => CloseCode::Invalid,
Self::UnexpectedExtraData => CloseCode::Invalid,
}
}
}
pub struct Uploader {
writer: Option<Box<dyn Write>>,
storage_filename: String,
app_data: super::AppData,
bytes_remaining: u64,
}
impl Uploader {
pub(crate) 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(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()),
)
}
}
fn stop_and_flush<T>(_: T, u: &mut Uploader, ctx: &mut <Uploader as Actor>::Context) {
ctx.stop();
if let Some(w) = u.writer.as_mut() {
if let Err(e) = w.flush() {
error!("Failed to flush writer: {}", e);
}
}
}
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);
self.cleanup_after_error(ctx);
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()),
}));
self.cleanup_after_error(ctx);
}
Ok(true) => {
info!("Finished uploading data");
ctx.close(Some(ws::CloseReason {
code: CloseCode::Normal,
description: None,
}));
stop_and_flush((), self, ctx);
}
_ => (),
}
}
}
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);
if raw_files.len() > MAX_FILES {
return Err(Error::TooManyFiles);
}
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 = storage_dir().join(storage_filename.clone());
info!("storing to: {:?}", storage_path);
let writer = File::options()
.write(true)
.create_new(true)
.open(&storage_path)?;
let (writer, downloadable_file): (Box<dyn Write>, _) = 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,
},
)
} else {
(
Box::new(writer),
DownloadableFile {
name: files[0].name.clone(),
size: files[0].size,
modtime: files[0].modtime,
},
)
};
self.writer = Some(writer);
let data = self.app_data.clone();
ctx.spawn(actix::fut::wrap_future(async move {
debug!("Spawned future to add entry {} to state", storage_filename);
data.write()
.await
.add_file(storage_filename, downloadable_file)
.await
.unwrap();
}));
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::UnexpectedExtraData);
}
self.bytes_remaining -= data.len() as u64;
writer.write_all(&data)?;
Ok(self.bytes_remaining == 0)
} else {
Err(Error::UnexpectedMessageType)
}
}
fn cleanup_after_error(&mut self, ctx: &mut <Self as Actor>::Context) {
info!(
"Cleaning up after failed upload of {}",
self.storage_filename
);
let data = self.app_data.clone();
let filename = self.storage_filename.clone();
ctx.wait(
actix::fut::wrap_future(async move {
debug!("Spawned future to remove entry {} from state", filename);
data.write().await.remove_file(&filename).await.unwrap();
})
.map(stop_and_flush),
);
if let Err(e) = std::fs::remove_file(storage_dir().join(&self.storage_filename)) {
error!("Failed to remove file {}: {}", self.storage_filename, e);
}
}
}