prevent filling up the disk

Cargo.lock

@@ -1897,6 +1897,7 @@ dependencies = [
  "futures-core",
  "inotify",
  "jsondb",
+ "libc",
  "log",
  "mime",
  "mnemonic",

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"

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)

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<AppData>,
-) -> impl Responder {
-    if data.state.read().await.full(data.config.max_storage_size) {
+) -> actix_web::Result<impl Responder> {
+    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<AppData>) -> 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<AppData>) -> actix_web::Result<impl Responder> {
+    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<T: FromStr>(var: &str, default: T) -> T
@@ -388,6 +388,8 @@ async fn main() -> std::io::Result<()> {
         env_or::<ByteSize>("TRANSBEAM_MAX_UPLOAD_SIZE", ByteSize(16 * bytesize::GB)).as_u64();
     let max_storage_size: u64 =
         env_or::<ByteSize>("TRANSBEAM_MAX_STORAGE_SIZE", ByteSize(64 * bytesize::GB)).as_u64();
+    let min_disk_free: u64 =
+        env_or::<ByteSize>("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,

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(())
     }
+
+    fn disk_free(&self) -> std::io::Result<u64> {
+        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))
+        }
     }
 
-impl StoredFiles {
-    fn total_size(&self) -> u64 {
-        self.0.iter().fold(0, |acc, (_, v)| acc + v.file.size)
+    async fn pending_data_size(&self) -> std::io::Result<u64> {
+        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 available_size(&self, max_storage_size: u64) -> u64 {
-        max_storage_size.saturating_sub(self.total_size())
+    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 fn full(&self, max_storage_size: u64) -> bool {
-        self.available_size(max_storage_size) == 0 || self.0.len() >= MAX_STORAGE_FILES
+    pub async fn available_size(&self) -> std::io::Result<u64> {
+        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<bool> {
+        Ok(self.available_size().await? == 0
+            || self.state.read().await.0.len() >= MAX_STORAGE_FILES)
     }
 }