transbeam/src/state.rs

use std::{collections::HashMap, io::ErrorKind};

use log::{error, info, warn, debug};
use tokio::{
    fs::File,
    io::{AsyncReadExt, AsyncWriteExt},
};

use crate::{storage_dir, DownloadableFile};

const STATE_FILE_NAME: &str = "files.json";

pub(crate) mod timestamp {
    use core::fmt;

    use serde::{de::Visitor, Deserializer, Serializer};
    use time::OffsetDateTime;

    pub(crate) fn serialize<S: Serializer>(
        time: &OffsetDateTime,
        ser: S,
    ) -> Result<S::Ok, S::Error> {
        ser.serialize_i64(time.unix_timestamp())
    }

    struct I64Visitor;

    impl<'de> Visitor<'de> for I64Visitor {
        type Value = i64;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            write!(formatter, "an integer")
        }

        fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E> {
            Ok(v)
        }

        fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E> {
            Ok(v as i64)
        }
    }

    pub(crate) fn deserialize<'de, D: Deserializer<'de>>(
        de: D,
    ) -> Result<OffsetDateTime, D::Error> {
        Ok(
            OffsetDateTime::from_unix_timestamp(de.deserialize_i64(I64Visitor)?)
                .unwrap_or_else(|_| OffsetDateTime::now_utc()),
        )
    }
}

async fn is_valid_entry(key: &str, info: &DownloadableFile) -> bool {
    if !crate::util::is_ascii_alphanumeric(&key) {
        error!("Invalid key in persistent storage: {}", 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
}

pub(crate) struct PersistentState(HashMap<String, DownloadableFile>);
impl PersistentState {
    pub(crate) async fn load() -> 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, DownloadableFile> = serde_json::from_str(&buf)?;
                info!("Loaded {} file entries from persistent storage", map.len());
                let mut filtered: HashMap<String, DownloadableFile> = HashMap::new();
                for (key, info) in map.into_iter() {
                    if is_valid_entry(&key, &info).await {
                        filtered.insert(key, info);
                    } else {
                        info!("Deleting invalid file {}", key);
                        if let Err(e) = tokio::fs::remove_file(storage_dir().join(&key)).await {
                            warn!("Failed to delete invalid file {}: {}", key, e);
                        }
                    }
                }
                let mut loaded = Self(filtered);
                loaded.save().await?;
                Ok(loaded)
            }
            Err(e) => {
                if let ErrorKind::NotFound = e.kind() {
                    Ok(Self(HashMap::new()))
                } else {
                    Err(e)
                }
            }
        }
    }

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

    pub(crate) async fn add_file(
        &mut self,
        key: String,
        file: DownloadableFile,
    ) -> std::io::Result<()> {
        self.0.insert(key, file);
        self.save().await
    }

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

    pub(crate) async fn remove_file(&mut self, key: &str) -> std::io::Result<()> {
        debug!("removing entry {} from state", key);
        self.0.remove(key);
        self.save().await
    }

    pub(crate) fn remove_uploader(&mut self, key: &str) {
        if let Some(f) = self.0.get_mut(key) {
            f.uploader.take();
        }
    }
}