transbeam/src/store.rs

use std::{
    collections::HashMap,
    io::ErrorKind,
    path::{Path, PathBuf},
};

use log::{debug, error, info};
use rand::{
    distributions::{Alphanumeric, DistString},
    thread_rng, Rng,
};
use serde::{Deserialize, Serialize};
use serde_with::skip_serializing_none;
use serde_with::{serde_as, FromInto, PickFirst};
use time::OffsetDateTime;
use tokio::{
    fs::File,
    io::{AsyncReadExt, AsyncWriteExt},
};

use crate::upload::UploadedFile;
use crate::zip::FileSet;

const STATE_FILE_NAME: &str = "files.json";
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'-')
}

#[serde_as]
#[skip_serializing_none]
#[derive(Clone, Deserialize, Serialize)]
pub struct StoredFile {
    pub name: String,
    pub size: u64,
    #[serde(with = "crate::timestamp")]
    pub modtime: OffsetDateTime,
    #[serde(with = "crate::timestamp")]
    pub expiry: OffsetDateTime,
    #[serde_as(as = "Option<PickFirst<(_, FromInto<Vec<UploadedFile>>)>>")]
    #[serde(default)]
    pub contents: Option<FileSet>,
}

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

pub struct FileStore {
    files: HashMap<String, StoredFile>,
    storage_dir: PathBuf,
    max_storage_size: u64,
}

impl FileStore {
    pub(crate) async fn load(storage_dir: PathBuf, max_storage_size: u64) -> std::io::Result<Self> {
        let open_result = File::open(storage_dir.join(STATE_FILE_NAME)).await;
        match open_result {
            Ok(mut f) => {
                let mut buf = String::new();
                f.read_to_string(&mut buf).await?;
                let map: HashMap<String, StoredFile> = serde_json::from_str(&buf)?;
                info!("Loaded {} file entries from persistent storage", map.len());
                let mut filtered: HashMap<String, StoredFile> = HashMap::new();
                for (key, info) in map.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, &storage_dir).await {
                        filtered.insert(key, info);
                    } else {
                        info!("Deleting file {}", key);
                        delete_file_if_exists(&storage_dir.join(&key)).await?;
                    }
                }
                let mut loaded = Self {
                    files: filtered,
                    storage_dir,
                    max_storage_size,
                };
                loaded.save().await?;
                Ok(loaded)
            }
            Err(e) => {
                if let ErrorKind::NotFound = e.kind() {
                    Ok(Self {
                        files: HashMap::new(),
                        storage_dir,
                        max_storage_size,
                    })
                } else {
                    Err(e)
                }
            }
        }
    }

    fn total_size(&self) -> u64 {
        self.files.iter().fold(0, |acc, (_, f)| acc + f.size)
    }

    pub fn available_size(&self) -> u64 {
        self.max_storage_size.saturating_sub(self.total_size())
    }

    async fn save(&mut self) -> std::io::Result<()> {
        info!("saving updated state: {} entries", self.files.len());
        File::create(self.storage_dir.join(STATE_FILE_NAME))
            .await?
            .write_all(&serde_json::to_vec_pretty(&self.files)?)
            .await
    }

    pub fn full(&self) -> bool {
        self.available_size() == 0 || self.files.len() >= MAX_STORAGE_FILES
    }

    /// 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(crate) async fn add_file(
        &mut self,
        key: String,
        file: StoredFile,
    ) -> Result<(), FileAddError> {
        if self.full() {
            return Err(FileAddError::Full);
        }
        let available_size = self.available_size();
        if file.size > available_size {
            return Err(FileAddError::TooBig(available_size));
        }
        self.files.insert(key, file);
        self.save().await?;
        Ok(())
    }

    pub(crate) fn lookup_file(&self, key: &str) -> Option<StoredFile> {
        self.files.get(key).cloned()
    }

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

    pub(crate) async fn remove_expired_files(&mut self) -> std::io::Result<()> {
        info!("Checking for expired files");
        let now = OffsetDateTime::now_utc();
        for (key, file) in std::mem::take(&mut self.files).into_iter() {
            if file.expiry > now {
                self.files.insert(key, file);
            } else {
                info!("Deleting expired file {}", key);
                delete_file_if_exists(&self.storage_dir.join(&key)).await?;
            }
        }
        self.save().await
    }
}