xenofem/transbeam
1
0
Fork 0
Code Issues 6 Pull requests Projects Releases Wiki Activity

refactor config variables, add upload password

Browse source
This commit is contained in:
xenofem 2022-05-03 16:28:43 -04:00
parent bfe7fcde99
commit eb53030043
13 changed files with 456 additions and 182 deletions
Download patch file Download diff file Expand all files Collapse all files

92
src/upload.rs
View file

@ -2,6 +2,7 @@ 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::web;
use actix_web_actors::ws::{self, CloseCode};
use bytes::Bytes;
use log::{debug, error, info, trace};
@ -9,7 +10,11 @@ use serde::{Deserialize, Serialize};
use time::OffsetDateTime;
use unicode_normalization::UnicodeNormalization;
use crate::store::{self, storage_dir, StoredFile};
use crate::{
log_auth_failure,
store::{self, FileAddError, StoredFile},
AppState,
};
const MAX_FILES: usize = 256;
const FILENAME_DATE_FORMAT: &[time::format_description::FormatItem] =
@ -29,10 +34,14 @@ enum Error {
NoFiles,
#[error("Number of files submitted by client exceeded the maximum limit")]
TooManyFiles,
#[error("Requested lifetime was too long")]
TooLong,
#[error("Requested lifetime was too long, can be at most {0} days")]
TooLong(u16),
#[error("Upload size was too large, can be at most {0} bytes")]
TooBig(u64),
#[error("File storage is full")]
Full,
#[error("Incorrect password")]
IncorrectPassword,
#[error("Websocket was closed by client before completing transfer")]
ClosedEarly(Option<CloseReason>),
#[error("Client sent more data than they were supposed to")]
@ -50,8 +59,10 @@ impl Error {
Self::DuplicateFilename
| Self::NoFiles
| Self::TooManyFiles
| Self::TooLong
| Self::TooBig(_) => CloseCode::Policy,
| Self::TooLong(_)
| Self::TooBig(_)
| Self::Full
| Self::IncorrectPassword => CloseCode::Policy,
}
}
}
@ -59,17 +70,19 @@ impl Error {
pub struct Uploader {
writer: Option<Box<dyn Write>>,
storage_filename: String,
app_data: super::AppData,
app_state: web::Data<AppState>,
bytes_remaining: u64,
ip_addr: String,
}
impl Uploader {
pub(crate) fn new(app_data: super::AppData) -> Self {
pub(crate) fn new(app_state: web::Data<AppState>, ip_addr: String) -> Self {
Self {
writer: None,
storage_filename: store::gen_storage_code(),
app_data,
storage_filename: store::gen_storage_code(app_state.config.mnemonic_codes),
app_state,
bytes_remaining: 0,
ip_addr,
}
}
}
@ -117,7 +130,8 @@ impl RawUploadedFile {
#[derive(Deserialize)]
struct UploadManifest {
files: Vec<RawUploadedFile>,
lifetime: u32,
lifetime: u16,
password: String,
}
#[derive(Serialize)]
@ -125,7 +139,8 @@ struct UploadManifest {
enum ServerMessage {
Ready { code: String },
TooBig { max_size: u64 },
TooLong { max_days: u32 },
TooLong { max_days: u16 },
IncorrectPassword,
Error { details: String },
}
@ -135,9 +150,10 @@ impl From<&Error> for ServerMessage {
Error::TooBig(max_size) => ServerMessage::TooBig {
max_size: *max_size,
},
Error::TooLong => ServerMessage::TooLong {
max_days: store::max_lifetime(),
Error::TooLong(max_days) => ServerMessage::TooLong {
max_days: *max_days,
},
Error::IncorrectPassword => ServerMessage::IncorrectPassword,
_ => ServerMessage::Error {
details: e.to_string(),
},
@ -189,6 +205,9 @@ fn ack(ctx: &mut Context) {
impl Uploader {
fn notify_error_and_cleanup(&mut self, e: Error, ctx: &mut Context) {
error!("{}", e);
if let Error::IncorrectPassword = e {
log_auth_failure(&self.ip_addr);
}
ctx.text(serde_json::to_string(&ServerMessage::from(&e)).unwrap());
ctx.close(Some(ws::CloseReason {
code: e.close_code(),
@ -207,9 +226,13 @@ impl Uploader {
let UploadManifest {
files: raw_files,
lifetime,
password,
} = serde_json::from_slice(text.as_bytes())?;
if lifetime > store::max_lifetime() {
return Err(Error::TooLong);
if std::env::var("TRANSBEAM_UPLOAD_PASSWORD") != Ok(password) {
return Err(Error::IncorrectPassword);
}
if lifetime > self.app_state.config.max_lifetime {
return Err(Error::TooLong(self.app_state.config.max_lifetime));
}
info!("Received file list: {} files", raw_files.len());
debug!("{:?}", raw_files);
@ -234,7 +257,14 @@ impl Uploader {
self.bytes_remaining += file.size;
files.push(file);
}
let storage_path = storage_dir().join(self.storage_filename.clone());
if self.bytes_remaining > self.app_state.config.max_upload_size {
return Err(Error::TooBig(self.app_state.config.max_upload_size));
}
let storage_path = self
.app_state
.config
.storage_dir
.join(self.storage_filename.clone());
info!("storing to: {:?}", storage_path);
let writer = File::options()
.write(true)
@ -264,25 +294,32 @@ impl Uploader {
modtime,
expiry: OffsetDateTime::now_utc() + lifetime * time::Duration::DAY,
};
let data = self.app_data.clone();
let state = self.app_state.clone();
let storage_filename = self.storage_filename.clone();
ctx.spawn(
actix::fut::wrap_future(async move {
debug!("Spawned future to add entry {} to state", storage_filename);
data.write()
state
.file_store
.write()
.await
.add_file(storage_filename, stored_file)
.await
})
.map(|res, u: &mut Self, ctx: &mut Context| match res {
Ok(Ok(())) => ctx.text(
Ok(()) => ctx.text(
serde_json::to_string(&ServerMessage::Ready {
code: u.storage_filename.clone(),
})
.unwrap(),
),
Ok(Err(size)) => u.notify_error_and_cleanup(Error::TooBig(size), ctx),
Err(e) => u.notify_error_and_cleanup(Error::from(e), ctx),
Err(FileAddError::TooBig(size)) => {
u.notify_error_and_cleanup(Error::TooBig(size), ctx)
}
Err(FileAddError::Full) => u.notify_error_and_cleanup(Error::Full, ctx),
Err(FileAddError::FileSystem(e)) => {
u.notify_error_and_cleanup(Error::from(e), ctx)
}
}),
);
}
@ -332,17 +369,20 @@ impl Uploader {
"Cleaning up after failed upload of {}",
self.storage_filename
);
let data = self.app_data.clone();
let state = self.app_state.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();
state
.file_store
.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);
}
}
}