ProxyWarden/src-tauri/src/storage.rs

use crate::activity::{append_activity, cap_activity, DEFAULT_ACTIVITY_LIMIT};
use crate::models::{
    ActivityEntry, ComponentStatus, LocalSingBoxConfig, Profile, SubscriptionCache, Target,
};
use crate::safe_fs;
use serde::{de::DeserializeOwned, Serialize};
use std::fs;
use std::io::{self, ErrorKind};
use std::path::{Path, PathBuf};

pub fn default_config_root() -> PathBuf {
    PathBuf::from(r"C:\ProgramData\ProxyWarden")
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoragePaths {
    pub root: PathBuf,
    pub config_dir: PathBuf,
    pub state_dir: PathBuf,
    pub generated_dir: PathBuf,
    pub profiles_file: PathBuf,
    pub targets_file: PathBuf,
    pub components_file: PathBuf,
    pub local_singbox_file: PathBuf,
    pub singbox_subscription_cache_file: PathBuf,
    pub activity_file: PathBuf,
}

impl StoragePaths {
    pub fn new(root: impl Into<PathBuf>) -> Self {
        let root = root.into();
        let config_dir = root.join("config");
        let state_dir = root.join("state");
        let generated_dir = root.join("generated");

        Self {
            root,
            profiles_file: config_dir.join("profiles.json"),
            targets_file: config_dir.join("targets.json"),
            components_file: config_dir.join("components.json"),
            local_singbox_file: config_dir.join("local-singbox.json"),
            singbox_subscription_cache_file: state_dir.join("singbox-subscription-cache.json"),
            activity_file: state_dir.join("activity.json"),
            config_dir,
            state_dir,
            generated_dir,
        }
    }
}

impl Default for StoragePaths {
    fn default() -> Self {
        Self::new(default_config_root())
    }
}

#[derive(Debug, Clone)]
pub struct JsonStorage {
    paths: StoragePaths,
    activity_limit: usize,
}

impl JsonStorage {
    pub fn new(root: impl Into<PathBuf>) -> Self {
        Self::with_activity_limit(root, DEFAULT_ACTIVITY_LIMIT)
    }

    pub fn with_activity_limit(root: impl Into<PathBuf>, activity_limit: usize) -> Self {
        Self {
            paths: StoragePaths::new(root),
            activity_limit,
        }
    }

    pub fn paths(&self) -> &StoragePaths {
        &self.paths
    }

    pub fn read_profiles(&self) -> io::Result<Vec<Profile>> {
        self.read_json_or_default(&self.paths.profiles_file)
    }

    pub fn write_profiles(&self, profiles: &[Profile]) -> io::Result<()> {
        self.write_json(&self.paths.profiles_file, profiles)
    }

    pub fn read_targets(&self) -> io::Result<Vec<Target>> {
        self.read_json_or_default(&self.paths.targets_file)
    }

    pub fn write_targets(&self, targets: &[Target]) -> io::Result<()> {
        self.write_json(&self.paths.targets_file, targets)
    }

    pub fn read_components(&self) -> io::Result<Vec<ComponentStatus>> {
        self.read_json_or_default(&self.paths.components_file)
    }

    pub fn read_local_singbox_config(&self) -> io::Result<LocalSingBoxConfig> {
        let mut config: LocalSingBoxConfig =
            self.read_json_or_default(&self.paths.local_singbox_file)?;
        config.normalize_percent_encoded_tags();
        Ok(config)
    }

    pub fn write_local_singbox_config(&self, config: &LocalSingBoxConfig) -> io::Result<()> {
        self.write_json(&self.paths.local_singbox_file, config)
    }

    pub fn read_singbox_subscription_cache(&self) -> io::Result<Option<SubscriptionCache>> {
        let mut cache = self
            .read_optional_json::<SubscriptionCache>(&self.paths.singbox_subscription_cache_file)?;
        if let Some(cache) = cache.as_mut() {
            cache.normalize_percent_encoded_tags();
        }
        Ok(cache)
    }

    pub fn write_singbox_subscription_cache(&self, cache: &SubscriptionCache) -> io::Result<()> {
        self.write_json(&self.paths.singbox_subscription_cache_file, cache)
    }

    pub fn remove_singbox_subscription_cache(&self) -> io::Result<()> {
        match fs::remove_file(&self.paths.singbox_subscription_cache_file) {
            Ok(()) => Ok(()),
            Err(error) if error.kind() == ErrorKind::NotFound => Ok(()),
            Err(error) => Err(error),
        }
    }

    pub fn read_activity(&self) -> io::Result<Vec<ActivityEntry>> {
        let entries = self.read_json_or_default(&self.paths.activity_file)?;
        Ok(cap_activity(entries, self.activity_limit))
    }

    pub fn append_activity(&self, entry: ActivityEntry) -> io::Result<Vec<ActivityEntry>> {
        let entries = self.read_activity()?;
        let entries = append_activity(entries, entry, self.activity_limit);
        self.write_json(&self.paths.activity_file, &entries)?;
        Ok(entries)
    }

    fn read_json_or_default<T>(&self, path: &Path) -> io::Result<T>
    where
        T: DeserializeOwned + Default,
    {
        match fs::read_to_string(path) {
            Ok(contents) => {
                parse_json(path, &contents).or_else(|error| recover_corrupt_json(path, error))
            }
            Err(error) if error.kind() == ErrorKind::NotFound => Ok(T::default()),
            Err(error) => Err(error),
        }
    }

    fn write_json<T>(&self, path: &Path, value: &T) -> io::Result<()>
    where
        T: Serialize + ?Sized,
    {
        let contents = serde_json::to_vec_pretty(value)
            .map_err(|error| io::Error::new(ErrorKind::InvalidData, error))?;
        write_atomic(path, &contents)
    }

    fn read_optional_json<T>(&self, path: &Path) -> io::Result<Option<T>>
    where
        T: DeserializeOwned,
    {
        match fs::read_to_string(path) {
            Ok(contents) => parse_json(path, &contents)
                .map(Some)
                .or_else(|error| recover_corrupt_json(path, error).map(Some)),
            Err(error) if error.kind() == ErrorKind::NotFound => Ok(None),
            Err(error) => Err(error),
        }
    }
}

impl Default for JsonStorage {
    fn default() -> Self {
        Self::new(default_config_root())
    }
}

pub fn backup_path(path: &Path) -> PathBuf {
    safe_fs::backup_path(path)
}

fn write_atomic(path: &Path, contents: &[u8]) -> io::Result<()> {
    safe_fs::write_with_backup(path, contents)
}

fn parse_json<T>(path: &Path, contents: &str) -> io::Result<T>
where
    T: DeserializeOwned,
{
    serde_json::from_str(contents).map_err(|error| {
        io::Error::new(
            ErrorKind::InvalidData,
            format!("Invalid JSON in '{}': {error}", path.display()),
        )
    })
}

fn recover_corrupt_json<T>(path: &Path, parse_error: io::Error) -> io::Result<T>
where
    T: DeserializeOwned,
{
    let corrupt_path = safe_fs::corrupt_path(path);
    move_corrupt_file(path, &corrupt_path)?;

    let backup_path = backup_path(path);
    if backup_path.exists() {
        let backup_contents = fs::read_to_string(&backup_path)?;
        match parse_json(&backup_path, &backup_contents) {
            Ok(value) => {
                fs::copy(&backup_path, path)?;
                Ok(value)
            }
            Err(backup_error) => Err(io::Error::new(
                ErrorKind::InvalidData,
                format!(
                    "Invalid JSON in '{}'; corrupt file moved to '{}'; backup '{}' could not be restored: {backup_error}; original error: {parse_error}",
                    path.display(),
                    corrupt_path.display(),
                    backup_path.display()
                ),
            )),
        }
    } else {
        Err(io::Error::new(
            ErrorKind::InvalidData,
            format!(
                "Invalid JSON in '{}'; corrupt file moved to '{}'; no valid backup available: {parse_error}",
                path.display(),
                corrupt_path.display()
            ),
        ))
    }
}

fn move_corrupt_file(path: &Path, corrupt_path: &Path) -> io::Result<()> {
    match fs::rename(path, corrupt_path) {
        Ok(()) => Ok(()),
        Err(rename_error) => {
            fs::copy(path, corrupt_path)?;
            fs::remove_file(path)?;
            if !corrupt_path.exists() {
                return Err(rename_error);
            }
            Ok(())
        }
    }
}