mirror of
https://codeberg.org/YGGverse/psocks.git
synced 2026-03-31 16:35:28 +00:00
draft multi-list rules api
This commit is contained in:
yggverse
parent
752ea23f80
commit
18509e6d1a
4 changed files with 133 additions and 76 deletions
60
src/main.rs
60
src/main.rs
|
|
@ -33,12 +33,9 @@ async fn api_allow(
|
|||
totals: &State<Arc<Total>>,
|
||||
) -> Result<Json<bool>, Status> {
|
||||
let result = rules.allow(rule).await;
|
||||
totals.set_entries(rules.total().await);
|
||||
totals.set_entries(rules.total(true).await); // @TODO separate active/inactive totals
|
||||
info!("Delete `{rule}` from the in-memory rules (operation status: {result:?})");
|
||||
Ok(Json(result.map_err(|e| {
|
||||
error!("Allow request handle error for `{rule}`: `{e}`");
|
||||
Status::InternalServerError
|
||||
})?))
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
#[rocket::get("/api/block/<rule>")]
|
||||
|
|
@ -46,14 +43,11 @@ async fn api_block(
|
|||
rule: &str,
|
||||
rules: &State<Arc<Rules>>,
|
||||
totals: &State<Arc<Total>>,
|
||||
) -> Result<Json<bool>, Status> {
|
||||
) -> Result<Json<()>, Status> {
|
||||
let result = rules.block(rule).await;
|
||||
totals.set_entries(rules.total().await);
|
||||
totals.set_entries(rules.total(true).await); // @TODO separate active/inactive totals
|
||||
info!("Add `{rule}` to the in-memory rules (operation status: {result:?})");
|
||||
Ok(Json(result.map_err(|e| {
|
||||
error!("Block request handle error for `{rule}`: `{e}`");
|
||||
Status::InternalServerError
|
||||
})?))
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
#[rocket::get("/api/rules")]
|
||||
|
|
@ -65,7 +59,7 @@ async fn api_rules(rules: &State<Arc<Rules>>) -> Result<Json<Vec<String>>, Statu
|
|||
|
||||
#[rocket::get("/api/lists")]
|
||||
async fn api_lists(rules: &State<Arc<Rules>>) -> Result<Json<Vec<rules::ListEntry>>, Status> {
|
||||
let lists = rules.lists();
|
||||
let lists = rules.lists().await;
|
||||
debug!("Get lists index (total: {})", lists.len());
|
||||
Ok(Json(lists))
|
||||
}
|
||||
|
|
@ -74,15 +68,39 @@ async fn api_lists(rules: &State<Arc<Rules>>) -> Result<Json<Vec<rules::ListEntr
|
|||
async fn api_list(
|
||||
id: usize,
|
||||
rules: &State<Arc<Rules>>,
|
||||
) -> Result<Json<Option<&rules::List>>, Status> {
|
||||
let list = rules.list(&id);
|
||||
) -> Result<Json<Option<rules::List>>, Status> {
|
||||
let list = rules.list(&id).await;
|
||||
debug!(
|
||||
"Get list #{id} rules (total: {:?})",
|
||||
list.map(|l| l.items.len())
|
||||
list.as_ref().map(|l| l.items.len())
|
||||
);
|
||||
Ok(Json(list))
|
||||
}
|
||||
|
||||
#[rocket::get("/api/list/enable/<id>")]
|
||||
async fn api_list_enable(
|
||||
id: usize,
|
||||
rules: &State<Arc<Rules>>,
|
||||
totals: &State<Arc<Total>>,
|
||||
) -> Result<Json<Option<()>>, Status> {
|
||||
let affected = rules.enable(&id, true).await;
|
||||
totals.set_entries(rules.total(true).await); // @TODO separate active/inactive totals
|
||||
info!("Enabled {affected:?} rules from the active rule set");
|
||||
Ok(Json(affected)) // @TODO handle empty result
|
||||
}
|
||||
|
||||
#[rocket::get("/api/list/disable/<id>")]
|
||||
async fn api_list_disable(
|
||||
id: usize,
|
||||
rules: &State<Arc<Rules>>,
|
||||
totals: &State<Arc<Total>>,
|
||||
) -> Result<Json<Option<()>>, Status> {
|
||||
let affected = rules.enable(&id, false).await;
|
||||
totals.set_entries(rules.total(true).await); // @TODO separate active/inactive totals
|
||||
info!("Disabled {affected:?} rules from the active rule set");
|
||||
Ok(Json(affected)) // @TODO handle empty result
|
||||
}
|
||||
|
||||
#[rocket::launch]
|
||||
async fn rocket() -> _ {
|
||||
env_logger::init();
|
||||
|
|
@ -91,7 +109,7 @@ async fn rocket() -> _ {
|
|||
|
||||
let rules = Arc::new(Rules::from_opt(&opt.allow_list).await.unwrap());
|
||||
|
||||
let totals = Arc::new(Total::with_rules(rules.total().await));
|
||||
let totals = Arc::new(Total::with_rules(rules.total(true).await)); // @TODO separate active/inactive totals
|
||||
|
||||
tokio::spawn({
|
||||
let socks_rules = rules.clone();
|
||||
|
|
@ -115,7 +133,15 @@ async fn rocket() -> _ {
|
|||
.mount(
|
||||
"/",
|
||||
rocket::routes![
|
||||
index, api_totals, api_allow, api_block, api_rules, api_lists, api_list
|
||||
index,
|
||||
api_totals,
|
||||
api_allow,
|
||||
api_block,
|
||||
api_rules,
|
||||
api_lists,
|
||||
api_list,
|
||||
api_list_enable,
|
||||
api_list_disable
|
||||
],
|
||||
)
|
||||
}
|
||||
|
|
|
|||
142
src/rules.rs
142
src/rules.rs
|
|
@ -6,25 +6,25 @@ use log::*;
|
|||
use std::collections::{HashMap, HashSet};
|
||||
use tokio::{fs, sync::RwLock};
|
||||
|
||||
/// In-memory registry, based on the `--allow-list`
|
||||
pub struct Rules {
|
||||
/// Active, unique rules for this session
|
||||
active: RwLock<HashSet<Item>>,
|
||||
/// Hold original lists rules asset with unique ID generated once on server init
|
||||
/// * it allows to enable/disable specified rules in-memory by index ID
|
||||
lists: HashMap<usize, List>,
|
||||
}
|
||||
pub struct Rules(RwLock<HashMap<usize, List>>);
|
||||
|
||||
impl Rules {
|
||||
/// Build new List object
|
||||
pub async fn from_opt(list: &[String]) -> Result<Self> {
|
||||
let mut active = HashSet::new();
|
||||
let mut lists = HashMap::new();
|
||||
|
||||
let mut rules_total = 0;
|
||||
|
||||
let mut index = HashMap::with_capacity(list.len());
|
||||
assert!(
|
||||
index
|
||||
.insert(
|
||||
0,
|
||||
List {
|
||||
alias: "session".into(),
|
||||
items: HashSet::new(),
|
||||
status: true
|
||||
}
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
for (i, l) in list.iter().enumerate() {
|
||||
let mut list_items = HashSet::new();
|
||||
let mut items = HashSet::new();
|
||||
let list = if l.contains("://") {
|
||||
let response = reqwest::get(l).await?;
|
||||
let status = response.status();
|
||||
|
|
@ -40,77 +40,78 @@ impl Rules {
|
|||
if line.starts_with("/") || line.starts_with("#") || line.is_empty() {
|
||||
continue; // skip comments
|
||||
}
|
||||
let item = Item::from_line(line);
|
||||
if !active.insert(item.clone()) {
|
||||
warn!("Ruleset index `{l}` contains duplicated entry: `{line}`")
|
||||
}
|
||||
if !list_items.insert(item) {
|
||||
if !items.insert(Item::from_line(line)) {
|
||||
warn!("List `{l}` contains duplicated entry: `{line}`")
|
||||
}
|
||||
rules_total += 1
|
||||
}
|
||||
assert!(
|
||||
lists
|
||||
index
|
||||
.insert(
|
||||
i,
|
||||
i + 1,
|
||||
List {
|
||||
alias: l.clone(),
|
||||
items: list_items
|
||||
items,
|
||||
status: true // @TODO implement config file
|
||||
}
|
||||
)
|
||||
.is_none()
|
||||
)
|
||||
}
|
||||
|
||||
info!(
|
||||
"Total rules parsed: {} (added: {rules_total}) / lists parsed: {} (added: {})",
|
||||
active.len(),
|
||||
list.len(),
|
||||
lists.len()
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
active: RwLock::new(active),
|
||||
lists,
|
||||
})
|
||||
Ok(Self(RwLock::new(index)))
|
||||
}
|
||||
/// Check if rule is exist in the (allow) index
|
||||
pub async fn any(&self, value: &str) -> bool {
|
||||
self.active.read().await.iter().any(|item| match item {
|
||||
Item::Exact(v) => v == value,
|
||||
Item::Ending(v) => value.ends_with(v),
|
||||
})
|
||||
self.0.read().await.values().any(|list| list.any(value))
|
||||
}
|
||||
/// Get total rules from the current session
|
||||
pub async fn total(&self) -> u64 {
|
||||
self.active.read().await.len() as u64
|
||||
pub async fn total(&self, status: bool) -> u64 {
|
||||
self.0
|
||||
.read()
|
||||
.await
|
||||
.values()
|
||||
.filter(|list| list.status == status)
|
||||
.map(|list| list.total())
|
||||
.sum()
|
||||
}
|
||||
/// Allow given `rule`
|
||||
/// Allow given `rule`(in the session index)
|
||||
/// * return `false` if the `rule` is exist
|
||||
pub async fn allow(&self, rule: &str) -> Result<bool> {
|
||||
Ok(self.active.write().await.insert(Item::from_line(rule)))
|
||||
pub async fn allow(&self, rule: &str) -> bool {
|
||||
self.0
|
||||
.write()
|
||||
.await
|
||||
.get_mut(&0)
|
||||
.unwrap()
|
||||
.items
|
||||
.insert(Item::from_line(rule))
|
||||
}
|
||||
/// Block given `rule`
|
||||
/// * return `false` if the `rule` is not exist
|
||||
pub async fn block(&self, rule: &str) -> Result<bool> {
|
||||
Ok(self.active.write().await.remove(&Item::from_line(rule)))
|
||||
/// Block given `rule` (in the session index)
|
||||
pub async fn block(&self, rule: &str) {
|
||||
self.0
|
||||
.write()
|
||||
.await
|
||||
.get_mut(&0)
|
||||
.unwrap()
|
||||
.items
|
||||
.retain(|item| rule == item.as_str())
|
||||
}
|
||||
/// Return active rules
|
||||
pub async fn active(&self) -> Vec<String> {
|
||||
let mut rules: Vec<String> = self
|
||||
.active
|
||||
.0
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.map(|item| item.to_string())
|
||||
.values()
|
||||
.filter(|list| list.status)
|
||||
.flat_map(|list| list.items.iter().map(|item| item.to_string()))
|
||||
.collect();
|
||||
rules.sort(); // HashSet does not keep the order
|
||||
rules
|
||||
}
|
||||
/// Return original list references
|
||||
pub fn lists(&self) -> Vec<ListEntry> {
|
||||
let mut list = Vec::with_capacity(self.lists.len());
|
||||
for l in self.lists.iter() {
|
||||
/// Return list references
|
||||
pub async fn lists(&self) -> Vec<ListEntry> {
|
||||
let this = self.0.read().await;
|
||||
let mut list = Vec::with_capacity(this.len());
|
||||
for l in this.iter() {
|
||||
list.push(ListEntry {
|
||||
id: *l.0,
|
||||
alias: l.1.alias.clone(),
|
||||
|
|
@ -119,8 +120,16 @@ impl Rules {
|
|||
list
|
||||
}
|
||||
/// Return original list references
|
||||
pub fn list(&self, id: &usize) -> Option<&List> {
|
||||
self.lists.get(id)
|
||||
pub async fn list(&self, id: &usize) -> Option<List> {
|
||||
self.0.read().await.get(id).cloned()
|
||||
}
|
||||
/// Change rule set status by list ID
|
||||
pub async fn enable(&self, list_id: &usize, status: bool) -> Option<()> {
|
||||
self.0
|
||||
.write()
|
||||
.await
|
||||
.get_mut(list_id)
|
||||
.map(|this| this.status = status)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -130,8 +139,23 @@ pub struct ListEntry {
|
|||
pub alias: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
#[derive(serde::Serialize, Clone)]
|
||||
pub struct List {
|
||||
pub alias: String,
|
||||
pub items: HashSet<Item>,
|
||||
pub status: bool,
|
||||
}
|
||||
|
||||
impl List {
|
||||
/// Check if rule is exist in the items index
|
||||
pub fn any(&self, value: &str) -> bool {
|
||||
self.items.iter().any(|item| match item {
|
||||
Item::Exact(v) => v == value,
|
||||
Item::Ending(v) => value.ends_with(v),
|
||||
})
|
||||
}
|
||||
/// Get total rules in list
|
||||
pub fn total(&self) -> u64 {
|
||||
self.items.len() as u64
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,6 +16,11 @@ impl Item {
|
|||
Self::Exact(rule.to_string())
|
||||
}
|
||||
}
|
||||
pub fn as_str(&self) -> &str {
|
||||
match self {
|
||||
Item::Ending(s) | Item::Exact(s) => s,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for Item {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue