psocks/src/list.rs

mod cache;
mod item;

use anyhow::Result;
use cache::Cache;
use item::Item;
use log::*;
use std::{collections::HashSet, path::PathBuf};
use tokio::{fs, sync::RwLock};

pub struct List {
    /// In-memory registry, based on `--allow-list` + `--cache`
    index: RwLock<HashSet<Item>>,
    /// FS cache for JSON/API changes, based on `--cache` value
    cache: Cache,
}

impl List {
    pub async fn from_opt(list: &Vec<String>, cache: Option<PathBuf>) -> Result<Self> {
        fn handle(this: &mut HashSet<Item>, line: &str) -> Option<bool> {
            if line.starts_with("/") || line.starts_with("#") || line.is_empty() {
                return None;
            }
            Some(this.insert(Item::from_line(line)))
        }

        let mut index = HashSet::new();

        let mut list_rules_total = 0;

        for i in list {
            for line in if i.contains("://") {
                let response = reqwest::get(i).await?;
                let status = response.status();
                if status.is_success() {
                    response.text().await?
                } else {
                    warn!("Could not receive remote list `{i}`: `{status}`");
                    continue;
                }
            } else {
                fs::read_to_string(i).await?
            }
            .lines()
            {
                if handle(&mut index, line).is_some_and(|status| !status) {
                    warn!("List `{i}` contains duplicated entry: `{line}`")
                }
                list_rules_total += 1
            }
        }

        let cache = Cache::from_path(cache).await?;

        let cache_rules_total = cache.read().await?.map(|data| {
            let mut t = 0;
            for line in data.lines() {
                if handle(&mut index, line).is_some_and(|status| !status) {
                    warn!("Cache file contains duplicated entry: `{line}`")
                }
                t += 1
            }
            t
        });

        let len = index.len();
        info!(
            "Total rules parsed: {len} (lists: {list_rules_total} / cache: {cache_rules_total:?})",
        );

        Ok(Self {
            index: RwLock::new(index),
            cache,
        })
    }
    pub async fn any(&self, value: &str) -> bool {
        self.index.read().await.iter().any(|item| match item {
            Item::Exact(v) => v == value,
            Item::Ending(v) => value.ends_with(v),
        })
    }
    pub async fn entries(&self) -> u64 {
        self.index.read().await.len() as u64
    }
    pub async fn allow(&self, rule: &str) -> Result<bool> {
        self.cache.allow(rule).await?;
        Ok(self.index.write().await.insert(Item::from_line(rule)))
    }
    pub async fn block(&self, rule: &str) -> Result<bool> {
        self.cache.block(rule).await?;
        Ok(self.index.write().await.remove(&Item::from_line(rule)))
    }
}