psocks/src/list/cache.rs

use anyhow::{Result, bail};
use std::{collections::HashSet, path::PathBuf};
use tokio::fs;

pub struct Cache(Option<PathBuf>);

impl Cache {
    pub async fn from_path(path: Option<PathBuf>) -> Result<Self> {
        Ok(Self(match path {
            Some(p) => {
                init_file(&p).await?;
                Some(p)
            }
            None => None,
        }))
    }
    pub async fn read(&self) -> Result<Option<String>> {
        Ok(if let Some(ref p) = self.0 {
            init_file(p).await?;
            Some(fs::read_to_string(p).await?)
        } else {
            None
        })
    }
    pub async fn allow(&self, rule: &str) -> Result<()> {
        if let Some(ref p) = self.0 {
            init_file(p).await?;
            let mut rules = HashSet::new();
            let lines = fs::read_to_string(p).await?;
            for line in lines.lines() {
                rules.insert(line);
            }
            rules.insert(rule);
            fs::write(p, rules.into_iter().collect::<Vec<_>>().join("\n")).await?;
        }
        Ok(())
    }
    pub async fn block(&self, rule: &str) -> Result<()> {
        if let Some(ref p) = self.0 {
            init_file(p).await?;
            let mut rules = HashSet::new();
            let lines = fs::read_to_string(p).await?;
            for line in lines.lines() {
                if line != rule {
                    rules.insert(line);
                }
            }
            fs::write(p, rules.into_iter().collect::<Vec<_>>().join("\n")).await?;
        }
        Ok(())
    }
    /// Clean `--cache` file collected, return deleted rules
    pub async fn clean(&self) -> Result<Option<Vec<String>>> {
        match self.0 {
            Some(ref p) => {
                init_file(p).await?;
                let mut rules = Vec::new();
                let lines = fs::read_to_string(p).await?;
                for line in lines.lines() {
                    rules.push(line.into());
                }
                fs::write(p, "").await?;
                Ok(Some(rules))
            }
            None => Ok(None),
        }
    }
}

/// Make sure that cache file is always exist (e.g. user may remove it when the daemon is running)
async fn init_file(path: &PathBuf) -> Result<()> {
    if path.exists() {
        if path.is_file() {
            return Ok(());
        } else {
            bail!(
                "Cache path `{}` exist but it is not a file!",
                path.to_string_lossy()
            )
        }
    }
    fs::write(path, "").await?;
    Ok(())
}