diff --git a/README.md b/README.md index 1dfb4bd..ffccc36 100644 --- a/README.md +++ b/README.md @@ -4,18 +4,19 @@ Experimental async SOCKS5 (TCP/UDP) proxy server based on [fast-socks5](https:// ## Roadmap -* [ ] Range support -* [ ] Local Web-API +* [x] Web JSON/API * [x] Block stats * [x] In-memory list update (without server restart) - * [ ] Persist changes option - * [ ] Performance optimization + * [x] Persist changes option (see `-c`, `--cache`) +* [ ] Web UI +* [ ] Performance optimization ## Usage ``` bash -RUST_LOG=psocks=trace cargo run -- -a=/path/to/allow1.txt \ - -a=http://localhost/allow2.txt \ +RUST_LOG=psocks=trace cargo run -- -a=http://localhost/allow.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/allow1.txt +``` /path/to/allow.txt +# /path/to/allow.txt // exact match duckduckgo.com diff --git a/src/list.rs b/src/list.rs index 8ff493c..ed7eabd 100644 --- a/src/list.rs +++ b/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 tokio::sync::RwLock; +use std::{collections::HashSet, path::PathBuf}; +use tokio::{fs, sync::RwLock}; pub struct List { + /// In-memory registry, based on `--allow-list` + `--cache` index: RwLock>, + /// FS cache for JSON/API changes, based on `--cache` value + cache: Cache, } impl List { - pub async fn from_opt(list: &Vec) -> Result { - let mut this = HashSet::new(); + pub async fn from_opt(list: &Vec, cache: Option) -> Result { + fn handle(this: &mut HashSet, 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() { - continue; - } - if !this.insert(Item::from_line(line)) { - warn!("Duplicated whitelist record: `{line}`") - } + handle(&mut index, 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 { - self.index.write().await.insert(Item::from_line(rule)) + pub async fn allow(&self, rule: &str) -> Result { + self.cache.allow(rule).await?; + Ok(self.index.write().await.insert(Item::from_line(rule))) } - pub async fn block(&self, rule: &str) -> bool { - self.index.write().await.remove(&Item::from_line(rule)) + pub async fn block(&self, rule: &str) -> Result { + self.cache.block(rule).await?; + Ok(self.index.write().await.remove(&Item::from_line(rule))) } } diff --git a/src/list/cache.rs b/src/list/cache.rs new file mode 100644 index 0000000..8a35c10 --- /dev/null +++ b/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); + +impl Cache { + pub async fn from_path(path: Option) -> Result { + Ok(Self(match path { + Some(p) => { + init_file(&p).await?; + Some(p) + } + None => None, + })) + } + pub async fn read(&self) -> Result> { + 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::>().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::>().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(()) +} diff --git a/src/main.rs b/src/main.rs index 4837a6b..0feff93 100644 --- a/src/main.rs +++ b/src/main.rs @@ -26,7 +26,7 @@ async fn allow(rule: &str, list: &State>, totals: &State>) 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/")] @@ -34,7 +34,7 @@ async fn block(rule: &str, list: &State>, totals: &State>) 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}`"); - SocksError::ArgumentInputError("Can't parse whitelist") + error!("Can't parse list: `{err}`"); + SocksError::ArgumentInputError("Can't parse list") }) .unwrap(), ); diff --git a/src/opt.rs b/src/opt.rs index 5a26c13..8f70781 100644 --- a/src/opt.rs +++ b/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, + + /// FS cache to persist JSON/API changes between server sessions + #[structopt(short = "c", long)] + pub cache: Option, } /// Choose the authentication type