rayhunter/daemon/src/qmdl_store.rs

use std::fmt::Display;
use std::io::{self, ErrorKind};
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};

use crate::config::GpsMode;
use chrono::{DateTime, Local, TimeDelta};
use log::{info, warn};
use rayhunter::util::RuntimeMetadata;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tokio::{
    fs::{self, File, OpenOptions, try_exists},
    io::AsyncWriteExt,
};

#[derive(Debug, Error)]
pub enum RecordingStoreError {
    #[error("Can't close an entry when there's no current entry")]
    NoCurrentEntry,
    #[error("An entry with that name doesn't exist")]
    NoSuchEntryError,
    #[error("Couldn't create file: {0}")]
    CreateFileError(tokio::io::Error),
    #[error("Couldn't read file: {0}")]
    ReadFileError(tokio::io::Error),
    #[error("Couldn't write file: {0}")]
    WriteFileError(tokio::io::Error),
    #[error("Couldn't delete file: {0}")]
    DeleteFileError(tokio::io::Error),
    #[error("Couldn't open directory at path: {0}")]
    OpenDirError(tokio::io::Error),
    #[error("Couldn't read manifest file: {0}")]
    ReadManifestError(tokio::io::Error),
    #[error("Couldn't write manifest file: {0}")]
    WriteManifestError(tokio::io::Error),
    #[error("Couldn't parse QMDL store manifest file: {0}")]
    ParseManifestError(toml::de::Error),
    #[error("Insufficient disk space: {0}MB available, {1}MB required")]
    InsufficientDiskSpace(u64, u64),
    #[error("GPS storage directory not found")]
    GpsStorageNotFound,
    #[error("Serialization error: {0}")]
    SerializationError(#[from] serde_json::Error),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FileKind {
    Qmdl,
    Analysis,
    Gps,
}

impl FileKind {
    // List of all possible physical files on disk.
    pub const ALL: &'static [FileKind] = &[FileKind::Qmdl, FileKind::Analysis, FileKind::Gps];

    pub fn get_filename(&self, entry_name: &str) -> String {
        match self {
            FileKind::Qmdl => format!("{}.qmdl", entry_name),
            FileKind::Analysis => format!("{}.ndjson", entry_name),
            FileKind::Gps => format!("{}-gps.ndjson", entry_name),
        }
    }

    pub fn get_filepath<P: AsRef<Path>>(&self, entry_name: &str, base_path: P) -> PathBuf {
        base_path.as_ref().join(self.get_filename(entry_name))
    }
}

impl Display for FileKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            FileKind::Qmdl => write!(f, "QMDL"),
            FileKind::Analysis => write!(f, "analysis"),
            FileKind::Gps => write!(f, "GPS"),
        }
    }
}

pub struct RecordingStore {
    pub path: PathBuf,
    pub manifest: Manifest,
    pub current_entry: Option<usize>, // index into manifest
}

#[derive(Deserialize, Serialize, Clone, PartialEq, Debug)]
pub struct Manifest {
    pub entries: Vec<ManifestEntry>,
}

/// The structure of an entry in the QMDL manifest table
#[derive(Deserialize, Serialize, Clone, PartialEq, Debug)]
#[cfg_attr(feature = "apidocs", derive(utoipa::ToSchema))]
pub struct ManifestEntry {
    /// The name of the entry
    pub name: String,
    /// The system time when recording began
    #[cfg_attr(feature = "apidocs", schema(value_type = String))]
    pub start_time: DateTime<Local>,
    /// The system time when the last message was recorded to the file
    #[cfg_attr(feature = "apidocs", schema(value_type = String))]
    pub last_message_time: Option<DateTime<Local>>,
    /// The size of the QMDL file in bytes
    pub qmdl_size_bytes: usize,
    /// The rayhunter daemon version which generated the file
    pub rayhunter_version: Option<String>,
    /// The OS which created the file
    pub system_os: Option<String>,
    /// The architecture on which the OS was running
    pub arch: Option<String>,
    #[serde(default)]
    pub stop_reason: Option<String>,
    /// When the manifest was uploaded to a WebDAV server
    #[cfg_attr(feature = "apidocs", schema(value_type = String))]
    pub upload_time: Option<DateTime<Local>>,
    #[serde(default)]
    pub gps_mode: Option<GpsMode>,
}

impl ManifestEntry {
    fn new(gps_mode: GpsMode) -> Self {
        let now = rayhunter::clock::get_adjusted_now();
        let metadata = RuntimeMetadata::new();
        ManifestEntry {
            name: format!("{}", now.timestamp()),
            start_time: now,
            last_message_time: None,
            qmdl_size_bytes: 0,
            rayhunter_version: Some(metadata.rayhunter_version),
            system_os: Some(metadata.system_os),
            arch: Some(metadata.arch),
            stop_reason: None,
            upload_time: None,
            gps_mode: Some(gps_mode),
        }
    }

    pub fn get_filepath<P: AsRef<Path>>(&self, file_kind: FileKind, path: P) -> PathBuf {
        file_kind.get_filepath(&self.name, path)
    }
}

impl RecordingStore {
    // Returns whether a directory with a "manifest.toml" exists at the given
    // path (though doesn't check if that manifest is valid)
    pub async fn exists<P>(path: P) -> Result<bool, RecordingStoreError>
    where
        P: AsRef<Path>,
    {
        let manifest_path = path.as_ref().join("manifest.toml");
        let dir_exists = try_exists(path)
            .await
            .map_err(RecordingStoreError::OpenDirError)?;
        let manifest_exists = try_exists(manifest_path)
            .await
            .map_err(RecordingStoreError::ReadManifestError)?;
        Ok(dir_exists && manifest_exists)
    }

    // Loads an existing RecordingStore at the given path. Errors if no store exists,
    // or if it's malformed.
    pub async fn load<P>(path: P) -> Result<Self, RecordingStoreError>
    where
        P: AsRef<Path>,
    {
        let path: PathBuf = path.as_ref().to_path_buf();
        let manifest = RecordingStore::read_manifest(&path).await?;
        Ok(RecordingStore {
            path,
            manifest,
            current_entry: None,
        })
    }

    // Creates a new RecordingStore at the given path. This involves creating a dir
    // and writing an empty manifest.
    pub async fn create<P>(path: P) -> Result<Self, RecordingStoreError>
    where
        P: AsRef<Path>,
    {
        fs::create_dir_all(&path)
            .await
            .map_err(RecordingStoreError::OpenDirError)?;

        let mut store = RecordingStore {
            path: path.as_ref().to_owned(),
            manifest: Manifest {
                entries: Vec::new(),
            },
            current_entry: None,
        };

        store.write_manifest().await?;
        Ok(store)
    }

    // Does a best-effort attempt to recover the manifest from a directory of
    // QMDL files. We expect these files to be named like "<timestamp>.qmdl",
    // and skip any files which don't match that pattern.
    pub async fn recover<P>(path: P) -> Result<Self, RecordingStoreError>
    where
        P: AsRef<Path>,
    {
        let mut dir_entries = fs::read_dir(path.as_ref())
            .await
            .map_err(RecordingStoreError::OpenDirError)?;
        let mut manifest_entries = Vec::new();

        while let Some(entry) = dir_entries
            .next_entry()
            .await
            .map_err(RecordingStoreError::OpenDirError)?
        {
            let os_filename = entry.file_name();
            let Some(filename) = os_filename.to_str() else {
                continue;
            };

            if !filename.ends_with(".qmdl") {
                continue;
            }

            let stem = filename.trim_end_matches(".qmdl");
            let Ok(start_timestamp) = stem.parse::<i64>() else {
                warn!("QMDL file has invalid name {os_filename:?}, skipping");
                continue;
            };

            let metadata = match entry.metadata().await {
                Ok(metadata) => metadata,
                Err(err) => {
                    warn!("failed to read QMDL file metadata: {err:?}, skipping");
                    continue;
                }
            };

            let Some(start_time) = DateTime::from_timestamp(start_timestamp, 0) else {
                warn!("QMDL filename {os_filename:?} gave an invalid timestamp, skipping");
                continue;
            };

            let Ok(last_message_time) = metadata.modified() else {
                warn!("failed to get modified time for QMDL file {os_filename:?}, skipping");
                continue;
            };

            info!("successfully recovered QMDL entry {os_filename:?}!");
            manifest_entries.push(ManifestEntry {
                name: stem.to_string(),
                start_time: start_time.into(),
                last_message_time: Some(last_message_time.into()),
                qmdl_size_bytes: metadata.size() as usize,
                rayhunter_version: None,
                system_os: None,
                arch: None,
                stop_reason: None,
                upload_time: None,
                gps_mode: None,
            });
        }

        // sort chronologically
        manifest_entries.sort_by_key(|a| a.start_time);

        let mut store = RecordingStore {
            path: path.as_ref().to_path_buf(),
            manifest: Manifest {
                entries: manifest_entries,
            },
            current_entry: None,
        };
        store.write_manifest().await?;

        Ok(store)
    }

    async fn read_manifest<P>(path: P) -> Result<Manifest, RecordingStoreError>
    where
        P: AsRef<Path>,
    {
        let manifest_path = path.as_ref().join("manifest.toml");
        let file_contents = fs::read_to_string(&manifest_path)
            .await
            .map_err(RecordingStoreError::ReadManifestError)?;
        toml::from_str(&file_contents).map_err(RecordingStoreError::ParseManifestError)
    }

    // Closes the current entry (if needed), creates a new entry based on the
    // current time, and updates the manifest. Returns a tuple of the entry's
    // newly created QMDL file and analysis file.
    pub async fn new_entry(
        &mut self,
        gps_mode: GpsMode,
    ) -> Result<(File, File), RecordingStoreError> {
        // if we've already got an entry open, close it
        if self.current_entry.is_some() {
            self.close_current_entry().await?;
        }
        let new_entry = ManifestEntry::new(gps_mode);
        let qmdl_filepath = new_entry.get_filepath(FileKind::Qmdl, &self.path);
        let qmdl_file = File::create(&qmdl_filepath)
            .await
            .map_err(RecordingStoreError::CreateFileError)?;
        let analysis_filepath = new_entry.get_filepath(FileKind::Analysis, &self.path);
        let analysis_file = File::create(&analysis_filepath)
            .await
            .map_err(RecordingStoreError::CreateFileError)?;
        let gps_filepath = new_entry.get_filepath(FileKind::Gps, &self.path);
        File::create(&gps_filepath)
            .await
            .map_err(RecordingStoreError::CreateFileError)?;
        self.manifest.entries.push(new_entry);
        self.current_entry = Some(self.manifest.entries.len() - 1);
        self.write_manifest().await?;
        Ok((qmdl_file, analysis_file))
    }

    pub async fn open_file(
        &self,
        entry_index: usize,
        file_kind: FileKind,
    ) -> Result<Option<File>, RecordingStoreError> {
        let entry = &self.manifest.entries[entry_index];
        let filepath = file_kind.get_filepath(&entry.name, &self.path);

        match File::open(&filepath).await {
            Ok(file) => Ok(Some(file)),
            Err(e) if e.kind() == ErrorKind::NotFound => Ok(None),
            Err(e) => Err(RecordingStoreError::ReadFileError(e)),
        }
    }

    pub async fn open_entry_gps_for_append(
        &self,
        entry_index: usize,
    ) -> Result<Option<File>, RecordingStoreError> {
        let entry = &self.manifest.entries[entry_index];
        match OpenOptions::new()
            .create(true)
            .append(true)
            .open(entry.get_filepath(FileKind::Gps, &self.path))
            .await
        {
            Ok(file) => Ok(Some(file)),
            Err(e) if e.kind() == ErrorKind::NotFound => Ok(None),
            Err(e) => Err(RecordingStoreError::CreateFileError(e)),
        }
    }

    pub async fn clear_and_open_entry_analysis(
        &mut self,
        entry_index: usize,
    ) -> Result<File, RecordingStoreError> {
        let entry = &self.manifest.entries[entry_index];
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .open(entry.get_filepath(FileKind::Analysis, &self.path))
            .await
            .map_err(RecordingStoreError::ReadFileError)?;
        Ok(file)
    }

    // Unsets the current entry
    pub async fn close_current_entry(&mut self) -> Result<(), RecordingStoreError> {
        match self.current_entry {
            Some(_) => {
                self.current_entry = None;
                Ok(())
            }
            None => Err(RecordingStoreError::NoCurrentEntry),
        }
    }

    // Sets the given entry's size and updates the last_message_time to now, updating the manifest
    pub async fn update_entry_qmdl_size(
        &mut self,
        entry_index: usize,
        size_bytes: usize,
    ) -> Result<(), RecordingStoreError> {
        self.manifest.entries[entry_index].qmdl_size_bytes = size_bytes;
        self.manifest.entries[entry_index].last_message_time =
            Some(rayhunter::clock::get_adjusted_now());
        self.write_manifest().await
    }

    async fn write_manifest(&mut self) -> Result<(), RecordingStoreError> {
        // we don't technically need a mutable reference to `self` here, but it
        // does prevent multiple concurrent writes across different threads
        let tmp_path = self.path.join("manifest.toml.new");
        let mut manifest_tmp_file = File::create(&tmp_path)
            .await
            .map_err(RecordingStoreError::WriteManifestError)?;

        let manifest_contents =
            toml::to_string_pretty(&self.manifest).expect("failed to serialize manifest");
        manifest_tmp_file
            .write_all(manifest_contents.as_bytes())
            .await
            .map_err(RecordingStoreError::WriteManifestError)?;

        fs::rename(tmp_path, self.path.join("manifest.toml"))
            .await
            .map_err(RecordingStoreError::WriteManifestError)?;

        Ok(())
    }

    pub fn get_next_unuploaded_entry(&self, min_age: TimeDelta) -> Option<String> {
        let now = rayhunter::clock::get_adjusted_now();
        self.manifest
            .entries
            .iter()
            .filter_map(|entry| {
                if self.is_current_entry(&entry.name) || entry.upload_time.is_some() {
                    return None;
                }
                let age = now - entry.last_message_time.unwrap_or(entry.start_time);

                (age > min_age).then_some((&entry.name, age))
            })
            .max_by_key(|(_, age)| *age)
            .map(|(name, _)| name.clone())
    }

    // Finds an entry by filename
    pub fn entry_for_name(&self, name: &str) -> Option<(usize, &ManifestEntry)> {
        let entry_index = self
            .manifest
            .entries
            .iter()
            .position(|entry| entry.name == name)?;
        Some((entry_index, &self.manifest.entries[entry_index]))
    }

    pub fn get_current_entry(&self) -> Option<(usize, &ManifestEntry)> {
        let entry_index = self.current_entry?;
        Some((entry_index, &self.manifest.entries[entry_index]))
    }

    pub async fn set_current_stop_reason(
        &mut self,
        reason: String,
    ) -> Result<(), RecordingStoreError> {
        if let Some(idx) = self.current_entry {
            self.manifest.entries[idx].stop_reason = Some(reason);
            self.write_manifest().await?;
        }
        Ok(())
    }

    pub async fn mark_entry_as_uploaded(
        &mut self,
        name: &str,
        upload_time: DateTime<Local>,
    ) -> Result<(), RecordingStoreError> {
        let entry_index = self
            .manifest
            .entries
            .iter()
            .position(|entry| entry.name == name)
            .ok_or(RecordingStoreError::NoSuchEntryError)?;
        self.manifest.entries[entry_index].upload_time = Some(upload_time);
        self.write_manifest().await?;
        Ok(())
    }

    pub fn is_current_entry(&self, name: &str) -> bool {
        match self.current_entry {
            Some(idx) => match self.manifest.entries.get(idx) {
                Some(entry) => entry.name == name,
                None => false,
            },
            None => false,
        }
    }

    pub async fn delete_entry(&mut self, name: &str) -> Result<(), RecordingStoreError> {
        let entry_to_delete_idx = self
            .manifest
            .entries
            .iter()
            .position(|entry| entry.name == name)
            .ok_or(RecordingStoreError::NoSuchEntryError)?;
        match self.current_entry {
            Some(current_entry) if current_entry == entry_to_delete_idx => {
                self.close_current_entry().await?;
            }
            Some(current_entry) => {
                self.current_entry = Some(current_entry - 1);
            }
            None => {}
        };
        let entry_to_delete = self.manifest.entries.remove(entry_to_delete_idx);
        self.write_manifest().await?;

        for &file_kind in FileKind::ALL {
            let filepath = file_kind.get_filepath(&entry_to_delete.name, &self.path);
            remove_file_if_exists(&filepath)
                .await
                .map_err(RecordingStoreError::DeleteFileError)?;
        }
        Ok(())
    }

    pub async fn delete_all_entries(&mut self) -> Result<(), RecordingStoreError> {
        if self.current_entry.is_some() {
            self.close_current_entry().await?;
        }

        let mut keep = Vec::new();

        'entries: for entry in &self.manifest.entries {
            for &file_kind in FileKind::ALL {
                let filepath = file_kind.get_filepath(&entry.name, &self.path);
                if let Err(e) = remove_file_if_exists(&filepath).await {
                    log::warn!("failed to remove {filepath:?}: {e:?}");
                    // Some error happened with deleting this entry, abort and go to the next one.
                    // Also *keep* the manifest entry.
                    keep.push(true);
                    continue 'entries;
                }
            }

            keep.push(false);
        }

        let mut keep_iter = keep.into_iter();
        self.manifest.entries.retain(|_| keep_iter.next().unwrap());
        self.write_manifest().await?;
        Ok(())
    }
}

async fn remove_file_if_exists(path: &Path) -> Result<(), io::Error> {
    match tokio::fs::remove_file(path).await {
        Err(e) if e.kind() == ErrorKind::NotFound => Ok(()),
        res => res,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::{Builder, TempDir};

    fn make_temp_dir() -> TempDir {
        Builder::new().prefix("qmdl_store_test").tempdir().unwrap()
    }

    #[tokio::test]
    async fn test_load_from_empty_dir() {
        let dir = make_temp_dir();
        assert!(!RecordingStore::exists(dir.path()).await.unwrap());
        let _created_store = RecordingStore::create(dir.path()).await.unwrap();
        assert!(RecordingStore::exists(dir.path()).await.unwrap());
        let loaded_store = RecordingStore::load(dir.path()).await.unwrap();
        assert_eq!(loaded_store.manifest.entries.len(), 0);
    }

    #[tokio::test]
    async fn test_creating_updating_and_closing_entries() {
        let dir = make_temp_dir();
        let mut store = RecordingStore::create(dir.path()).await.unwrap();
        let _ = store.new_entry(GpsMode::Disabled).await.unwrap();
        let entry_index = store.current_entry.unwrap();
        assert_eq!(
            RecordingStore::read_manifest(dir.path()).await.unwrap(),
            store.manifest
        );
        assert!(
            store.manifest.entries[entry_index]
                .last_message_time
                .is_none()
        );

        store
            .update_entry_qmdl_size(entry_index, 1000)
            .await
            .unwrap();
        let (entry_index, entry) = store
            .entry_for_name(&store.manifest.entries[entry_index].name)
            .unwrap();
        assert!(entry.last_message_time.is_some());
        assert_eq!(store.manifest.entries[entry_index].qmdl_size_bytes, 1000);
        assert_eq!(
            RecordingStore::read_manifest(dir.path()).await.unwrap(),
            store.manifest
        );

        store.close_current_entry().await.unwrap();
        assert!(matches!(
            store.close_current_entry().await,
            Err(RecordingStoreError::NoCurrentEntry)
        ));
    }

    #[tokio::test]
    async fn test_create_on_existing_store() {
        let dir = make_temp_dir();
        let mut store = RecordingStore::create(dir.path()).await.unwrap();
        let _ = store.new_entry(GpsMode::Disabled).await.unwrap();
        let entry_index = store.current_entry.unwrap();
        store
            .update_entry_qmdl_size(entry_index, 1000)
            .await
            .unwrap();
        let store = RecordingStore::create(dir.path()).await.unwrap();
        assert_eq!(store.manifest.entries.len(), 0);
    }

    #[tokio::test]
    async fn test_repeated_new_entries() {
        let dir = make_temp_dir();
        let mut store = RecordingStore::create(dir.path()).await.unwrap();
        let _ = store.new_entry(GpsMode::Disabled).await.unwrap();
        let entry_index = store.current_entry.unwrap();
        let _ = store.new_entry(GpsMode::Disabled).await.unwrap();
        let new_entry_index = store.current_entry.unwrap();
        assert_ne!(entry_index, new_entry_index);
        assert_eq!(store.manifest.entries.len(), 2);
    }

    #[tokio::test]
    async fn test_delete_all_entries() {
        let dir = make_temp_dir();
        let mut store = RecordingStore::create(dir.path()).await.unwrap();
        let _ = store.new_entry(GpsMode::Disabled).await.unwrap();
        assert!(store.current_entry.is_some());

        store.delete_all_entries().await.unwrap();
        assert!(store.current_entry.is_none());

        // regression test: deleting all entries should also work when there's no current
        // recording.
        store.delete_all_entries().await.unwrap();
        assert!(store.current_entry.is_none());
    }

    #[tokio::test]
    async fn test_mark_entry_as_uploaded_sets_time_and_persists() {
        let dir = make_temp_dir();
        let mut store = RecordingStore::create(dir.path()).await.unwrap();
        let _ = store.new_entry(GpsMode::Disabled).await.unwrap();
        let name = store.manifest.entries[0].name.clone();
        store.close_current_entry().await.unwrap();

        let upload_time = Local::now();
        store
            .mark_entry_as_uploaded(&name, upload_time)
            .await
            .unwrap();
        assert_eq!(store.manifest.entries[0].upload_time, Some(upload_time));

        let reloaded = RecordingStore::load(dir.path()).await.unwrap();
        assert_eq!(reloaded.manifest.entries[0].upload_time, Some(upload_time));
    }

    #[tokio::test]
    async fn test_mark_entry_as_uploaded_missing_entry() {
        let dir = make_temp_dir();
        let mut store = RecordingStore::create(dir.path()).await.unwrap();
        assert!(matches!(
            store.mark_entry_as_uploaded("nope", Local::now()).await,
            Err(RecordingStoreError::NoSuchEntryError)
        ));
    }

    #[tokio::test]
    async fn test_get_next_unuploaded_entry() {
        let dir = make_temp_dir();
        let mut store = RecordingStore::create(dir.path()).await.unwrap();

        for _ in 0..3 {
            let _ = store.new_entry(GpsMode::Disabled).await.unwrap();
        }

        store.manifest.entries[0].name = "entry-0".to_owned();
        store.manifest.entries[0].start_time = Local::now() - TimeDelta::seconds(10);
        store.manifest.entries[0].last_message_time = None;

        store.manifest.entries[1].name = "entry-1".to_owned();
        store.manifest.entries[1].start_time = Local::now() - TimeDelta::seconds(10);
        store.manifest.entries[1].last_message_time = Some(Local::now() - TimeDelta::seconds(5));

        store.manifest.entries[2].name = "entry-2".to_owned();
        store.manifest.entries[2].start_time = Local::now() - TimeDelta::seconds(10);
        store.manifest.entries[2].last_message_time = Some(Local::now() - TimeDelta::seconds(1));

        assert_eq!(
            store.get_next_unuploaded_entry(TimeDelta::seconds(3600)),
            None,
        );

        assert_eq!(
            store.get_next_unuploaded_entry(TimeDelta::seconds(3)),
            Some("entry-0".to_owned())
        );
        store
            .mark_entry_as_uploaded("entry-0", Local::now())
            .await
            .unwrap();
        assert_eq!(
            store.get_next_unuploaded_entry(TimeDelta::seconds(3)),
            Some("entry-1".to_owned())
        );
        store
            .mark_entry_as_uploaded("entry-1", Local::now())
            .await
            .unwrap();
        assert_eq!(store.get_next_unuploaded_entry(TimeDelta::seconds(3)), None);
    }
}