diff --git a/Cargo.lock b/Cargo.lock index 0ea3b5c..566f812 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1897,6 +1897,7 @@ dependencies = [ "futures-core", "inotify", "jsondb", + "libc", "log", "mime", "mnemonic", diff --git a/Cargo.toml b/Cargo.toml index e36d257..cee44b9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,7 @@ env_logger = "0.11.3" futures-core = "0.3" inotify = "0.10" jsondb = "0.4.0" +libc = "0.2" log = "0.4" mime = "0.3.16" mnemonic = "1.0.1" diff --git a/README.md b/README.md index a451b6e..b120a56 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,11 @@ transbeam is configured with the following environment variables: case-sensitive.) - `TRANSBEAM_MAX_STORAGE_SIZE`: maximum total size of all files being stored by transbeam (default: 64G) +- `TRANSBEAM_MIN_DISK_FREE`: minimum amount of free disk space on the + filesystem where transbeam uploads are stored (default: 8G). + transbeam will not accept uploads that would reduce free disk space + lower than this value, even if they would otherwise be within the + max upload size and max storage size. - `TRANSBEAM_MNEMONIC_CODES`: generate memorable download codes using English words, rather than random alphanumeric strings (default: true) diff --git a/src/main.rs b/src/main.rs index b0e5ec4..da7b828 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,6 +36,7 @@ struct AppData { struct Config { base_url: String, max_storage_size: u64, + min_disk_free: u64, max_upload_size: u64, max_lifetime: u16, upload_password: String, @@ -290,8 +291,8 @@ async fn handle_upload( req: HttpRequest, stream: web::Payload, data: web::Data, -) -> impl Responder { - if data.state.read().await.full(data.config.max_storage_size) { +) -> actix_web::Result { + if data.full().await? { return Ok(HttpResponse::BadRequest().finish()); } let ip_addr = get_ip_addr(&req, data.config.reverse_proxy); @@ -326,16 +327,15 @@ struct UploadLimits { } #[get("/upload/limits.json")] -async fn upload_limits(data: web::Data) -> impl Responder { - let file_store = data.state.read().await; - let open = !file_store.full(data.config.max_storage_size); - let available_size = file_store.available_size(data.config.max_storage_size); +async fn upload_limits(data: web::Data) -> actix_web::Result { + let open = !data.full().await?; + let available_size = data.available_size().await?; let max_size = std::cmp::min(available_size, data.config.max_upload_size); - web::Json(UploadLimits { + Ok(web::Json(UploadLimits { open, max_size, max_lifetime: data.config.max_lifetime, - }) + })) } fn env_or(var: &str, default: T) -> T @@ -388,6 +388,8 @@ async fn main() -> std::io::Result<()> { env_or::("TRANSBEAM_MAX_UPLOAD_SIZE", ByteSize(16 * bytesize::GB)).as_u64(); let max_storage_size: u64 = env_or::("TRANSBEAM_MAX_STORAGE_SIZE", ByteSize(64 * bytesize::GB)).as_u64(); + let min_disk_free: u64 = + env_or::("TRANSBEAM_MIN_DISK_FREE", ByteSize(8 * bytesize::GB)).as_u64(); let upload_password: String = env_or_panic("TRANSBEAM_UPLOAD_PASSWORD"); let cachebuster: String = env_or_else("TRANSBEAM_CACHEBUSTER", String::new); let admin_password_hash: PasswordHashString = env_or_panic("TRANSBEAM_ADMIN_PASSWORD_HASH"); @@ -423,6 +425,7 @@ async fn main() -> std::io::Result<()> { base_url, max_upload_size, max_storage_size, + min_disk_free, max_lifetime, upload_password, storage_dir, diff --git a/src/store.rs b/src/store.rs index 5cde34c..b92ae13 100644 --- a/src/store.rs +++ b/src/store.rs @@ -4,13 +4,14 @@ use std::{ path::{Path, PathBuf}, }; +use libc::statvfs64; use log::{debug, error, info}; use rand::{ distributions::{Alphanumeric, DistString}, thread_rng, Rng, }; use time::OffsetDateTime; -use tokio::fs::File; +use tokio::fs::{metadata, File}; const MAX_STORAGE_FILES: usize = 1024; @@ -113,14 +114,14 @@ impl crate::AppData { key: String, entry: StoredFileWithPassword, ) -> Result<(), FileAddError> { - let mut store = self.state.write().await; - if store.full(self.config.max_storage_size) { + if self.full().await? { return Err(FileAddError::Full); } - let available_size = store.available_size(self.config.max_storage_size); + let available_size = self.available_size().await?; if entry.file.size > available_size { return Err(FileAddError::TooBig(available_size)); } + let mut store = self.state.write().await; store.0.insert(key, entry); Ok(()) } @@ -150,18 +151,64 @@ impl crate::AppData { } Ok(()) } -} -impl StoredFiles { - fn total_size(&self) -> u64 { - self.0.iter().fold(0, |acc, (_, v)| acc + v.file.size) + fn disk_free(&self) -> std::io::Result { + let storage_path_c = std::ffi::CString::new( + self.config + .storage_dir + .to_str() + .expect("storage path is not valid unicode"), + ) + .expect("storage path contains a null byte"); + + unsafe { + let mut stat: statvfs64 = std::mem::zeroed(); + loop { + if statvfs64(storage_path_c.as_ptr(), std::ptr::addr_of_mut!(stat)) == 0 { + break; + } + let os_error = std::io::Error::last_os_error(); + if os_error.kind() != ErrorKind::Interrupted { + return Err(os_error); + } + } + Ok(stat.f_bsize.saturating_mul(stat.f_bavail)) + } } - pub fn available_size(&self, max_storage_size: u64) -> u64 { - max_storage_size.saturating_sub(self.total_size()) + async fn pending_data_size(&self) -> std::io::Result { + let store = self.state.read().await; + let mut total_pending_size = 0u64; + for (key, value) in store.0.iter() { + let file_path = self.config.storage_dir.join(key); + let pending_size = value + .file + .size + .saturating_sub(metadata(&file_path).await?.len()); + total_pending_size = total_pending_size.saturating_add(pending_size); + } + Ok(total_pending_size) } - pub fn full(&self, max_storage_size: u64) -> bool { - self.available_size(max_storage_size) == 0 || self.0.len() >= MAX_STORAGE_FILES + async fn total_stored_size(&self) -> u64 { + let store = self.state.read().await; + store.0.iter().fold(0, |acc, (_, v)| acc + v.file.size) + } + + pub async fn available_size(&self) -> std::io::Result { + let available_policy_size = self + .config + .max_storage_size + .saturating_sub(self.total_stored_size().await); + let available_disk_size = self + .disk_free()? + .saturating_sub(self.config.min_disk_free) + .saturating_sub(self.pending_data_size().await?); + Ok(std::cmp::min(available_policy_size, available_disk_size)) + } + + pub async fn full(&self) -> std::io::Result { + Ok(self.available_size().await? == 0 + || self.state.read().await.0.len() >= MAX_STORAGE_FILES) } }