transbeam/src/store.rs

use std::{
    io::ErrorKind,
    ops::DerefMut,
    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::{metadata, File};

const MAX_STORAGE_FILES: usize = 1024;

pub fn gen_storage_code(use_mnemonic: bool) -> String {
    if use_mnemonic {
        mnemonic::to_string(thread_rng().gen::<[u8; 4]>())
    } else {
        Alphanumeric.sample_string(&mut thread_rng(), 8)
    }
}

pub fn is_valid_storage_code(s: &str) -> bool {
    s.as_bytes()
        .iter()
        .all(|c| c.is_ascii_alphanumeric() || c == &b'-')
}

pub use crate::state::v1::{StoredFile, StoredFileWithPassword, StoredFiles};

async fn is_valid_entry(key: &str, info: &StoredFile, storage_dir: &Path) -> bool {
    if info.expiry < OffsetDateTime::now_utc() {
        info!("File {} has expired", key);
        return false;
    }

    let file = if let Ok(f) = File::open(storage_dir.join(key)).await {
        f
    } else {
        error!(
            "Unable to open file {} referenced in persistent storage",
            key
        );
        return false;
    };
    let metadata = if let Ok(md) = file.metadata().await {
        md
    } else {
        error!(
            "Unable to get metadata for file {} referenced in persistent storage",
            key
        );
        return false;
    };
    if metadata.len() != info.size {
        error!("Mismatched file size for file {} referenced in persistent storage: expected {}, found {}", key, info.size, metadata.len());
        return false;
    }
    true
}

async fn delete_file_if_exists(file: &PathBuf) -> std::io::Result<()> {
    if let Err(e) = tokio::fs::remove_file(file).await {
        if e.kind() != ErrorKind::NotFound {
            error!("Failed to delete file {}: {}", file.to_string_lossy(), e);
            return Err(e);
        }
    }
    Ok(())
}

#[derive(thiserror::Error, Debug)]
pub enum FileAddError {
    #[error("Failed to write metadata to filesystem")]
    FileSystem(#[from] std::io::Error),
    #[error("File was too large, available space is {0} bytes")]
    TooBig(u64),
    #[error("File store is full")]
    Full,
}

impl crate::AppData {
    // This must only be run at startup, or it will delete in-progress
    // uploads.
    pub async fn cleanup(&self) -> std::io::Result<()> {
        let mut store = self.state.write().await;
        let old = std::mem::take(store.deref_mut());
        for (key, info) in old.0.into_iter() {
            // Handle this case separately, because we don't
            // want to try to delete it if it's not the sort
            // of path we're expecting
            if !is_valid_storage_code(&key) {
                error!("Invalid key in persistent storage: {}", key);
                continue;
            }
            if is_valid_entry(&key, &info.file, &self.config.storage_dir).await {
                store.0.insert(key, info);
            } else {
                info!("Deleting file {}", key);
                delete_file_if_exists(&self.config.storage_dir.join(&key)).await?;
            }
        }
        Ok(())
    }

    /// Attempts to add a file to the store. Returns an I/O error if
    /// something's broken, or a u64 of the maximum allowed file size
    /// if the file was too big, or a unit if everything worked.
    pub async fn add_file(
        &self,
        key: String,
        entry: StoredFileWithPassword,
    ) -> Result<(), FileAddError> {
        if self.full().await? {
            return Err(FileAddError::Full);
        }
        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(())
    }

    pub async fn remove_file(&self, key: &str) -> std::io::Result<()> {
        debug!("removing entry {} from state", key);
        let mut store = self.state.write().await;
        store.0.remove(key);
        if is_valid_storage_code(key) {
            delete_file_if_exists(&self.config.storage_dir.join(key)).await?;
        }
        Ok(())
    }

    pub async fn remove_expired_files(&self) -> std::io::Result<()> {
        info!("Checking for expired files");
        let now = OffsetDateTime::now_utc();
        let mut store = self.state.write().await;
        let old = std::mem::take(store.deref_mut());
        for (key, value) in old.0.into_iter() {
            if value.file.expiry > now {
                store.0.insert(key, value);
            } else {
                info!("Deleting expired file {}", key);
                delete_file_if_exists(&self.config.storage_dir.join(&key)).await?;
            }
        }
        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))
        }
    }

    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)
    }

    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<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)
    }
}