implement session cache

README.md

@@ -4,18 +4,19 @@ Experimental async SOCKS5 (TCP/UDP) proxy server based on [fast-socks5](https://
 
 ## Roadmap
 
-* [ ] Range support
+* [x] Web JSON/API
-* [ ] Local Web-API
   * [x] Block stats
   * [x] In-memory list update (without server restart)
-    * [ ] Persist changes option
+    * [x] Persist changes option (see `-c`, `--cache`)
-  * [ ] Performance optimization
+* [ ] Web UI
+* [ ] Performance optimization
 
 ## Usage
 
 ``` bash
-RUST_LOG=psocks=trace cargo run -- -a=/path/to/allow1.txt \
+RUST_LOG=psocks=trace cargo run -- -a=http://localhost/allow.txt \
-                                   -a=http://localhost/allow2.txt \
+                                   -a=/path/to/allow.txt \
+                                   -c=/path/to/cache.txt \
                                    no-auth
 ```
 * set `socks5://127.0.0.1:1080` proxy in your application
@@ -25,8 +26,8 @@ RUST_LOG=psocks=trace cargo run -- -a=/path/to/allow1.txt \
 
 ### Allow list example
 
-``` /path/to/allow1.txt
+``` /path/to/allow.txt
-# /path/to/allow1.txt
+# /path/to/allow.txt
 
 // exact match
 duckduckgo.com

src/list.rs

@@ -1,37 +1,54 @@
+mod cache;
 mod item;
 
 use anyhow::Result;
+use cache::Cache;
 use item::Item;
 use log::*;
-use std::collections::HashSet;
+use std::{collections::HashSet, path::PathBuf};
-use tokio::sync::RwLock;
+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>) -> Result<Self> {
+    pub async fn from_opt(list: &Vec<String>, cache: Option<PathBuf>) -> Result<Self> {
-        let mut this = HashSet::new();
+        fn handle(this: &mut HashSet<Item>, line: &str) {
+            if line.starts_with("/") || line.starts_with("#") || line.is_empty() {
+                return;
+            }
+            if !this.insert(Item::from_line(line)) {
+                warn!("Duplicated list record: `{line}`")
+            }
+        }
+        let mut index = HashSet::new();
         for i in list {
             for line in if i.contains("://") {
                 reqwest::get(i).await?.text().await?
             } else {
-                std::fs::read_to_string(i)?
+                fs::read_to_string(i).await?
             }
             .lines()
             {
-                if line.starts_with("/") || line.starts_with("#") || line.is_empty() {
+                handle(&mut index, line)
-                    continue;
-                }
-                if !this.insert(Item::from_line(line)) {
-                    warn!("Duplicated whitelist record: `{line}`")
-                }
             }
         }
-        info!("Total whitelist entries parsed: {}", this.len());
+
+        let cache = Cache::from_path(cache).await?;
+        if let Some(data) = cache.read().await? {
+            for line in data.lines() {
+                handle(&mut index, line)
+            }
+        }
+
+        info!("Total list entries parsed: {}", index.len());
         Ok(Self {
-            index: RwLock::new(this),
+            index: RwLock::new(index),
+            cache,
         })
     }
     pub async fn any(&self, values: &[&str]) -> bool {
@@ -46,10 +63,12 @@ impl List {
     pub async fn entries(&self) -> u64 {
         self.index.read().await.len() as u64
     }
-    pub async fn allow(&self, rule: &str) -> bool {
+    pub async fn allow(&self, rule: &str) -> Result<bool> {
-        self.index.write().await.insert(Item::from_line(rule))
+        self.cache.allow(rule).await?;
+        Ok(self.index.write().await.insert(Item::from_line(rule)))
     }
-    pub async fn block(&self, rule: &str) -> bool {
+    pub async fn block(&self, rule: &str) -> Result<bool> {
-        self.index.write().await.remove(&Item::from_line(rule))
+        self.cache.block(rule).await?;
+        Ok(self.index.write().await.remove(&Item::from_line(rule)))
     }
 }

src/list/cache.rs

@@ -0,0 +1,68 @@
+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(())
+    }
+}
+
+/// 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(())
+}

src/main.rs

@@ -26,7 +26,7 @@ async fn allow(rule: &str, list: &State<Arc<List>>, totals: &State<Arc<Total>>)
     let result = list.allow(rule).await;
     totals.set_entries(list.entries().await);
     info!("Delete `{rule}` from the in-memory rules (operation status: {result:?})");
-    Json(result)
+    Json(result.is_ok_and(|v| v))
 }
 
 #[rocket::get("/block/<rule>")]
@@ -34,7 +34,7 @@ async fn block(rule: &str, list: &State<Arc<List>>, totals: &State<Arc<Total>>)
     let result = list.block(rule).await;
     totals.set_entries(list.entries().await);
     info!("Add `{rule}` to the in-memory rules (operation status: {result:?})");
-    Json(result)
+    Json(result.is_ok_and(|v| v))
 }
 
 #[rocket::launch]
@@ -44,11 +44,11 @@ async fn rocket() -> _ {
     let opt: &'static Opt = Box::leak(Box::new(Opt::from_args()));
 
     let list = Arc::new(
-        List::from_opt(&opt.allow_list)
+        List::from_opt(&opt.allow_list, opt.cache.clone())
             .await
             .map_err(|err| {
-                error!("Can't parse whitelist: `{err}`");
+                error!("Can't parse list: `{err}`");
-                SocksError::ArgumentInputError("Can't parse whitelist")
+                SocksError::ArgumentInputError("Can't parse list")
             })
             .unwrap(),
     );

src/opt.rs

@@ -1,4 +1,4 @@
-use std::{net::SocketAddr, num::ParseFloatError, time::Duration};
+use std::{net::SocketAddr, num::ParseFloatError, path::PathBuf, time::Duration};
 use structopt::StructOpt;
 
 /// # How to use it:
@@ -50,6 +50,10 @@ pub struct Opt {
     /// * remote URL
     #[structopt(short = "a", long)]
     pub allow_list: Vec<String>,
+
+    /// FS cache to persist JSON/API changes between server sessions
+    #[structopt(short = "c", long)]
+    pub cache: Option<PathBuf>,
 }
 
 /// Choose the authentication type