YGGverse/psocks
1
0
Fork 0
mirror of https://codeberg.org/YGGverse/psocks.git synced 2026-04-01 00:45:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

implement in-memory ruleset update

Browse source
This commit is contained in:
postscriptum 2026-03-22 09:18:09 +02:00
parent b03bdd0e3a
commit 30a77072ed
4 changed files with 85 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

50
src/main.rs
View file

@ -10,28 +10,56 @@ use fast_socks5::{
use list::List;
use log::*;
use opt::{AuthMode, Opt};
use rocket::serde::json::Json;
use rocket::{State, serde::json::Json};
use stats::{Snapshot, Total};
use std::{future::Future, sync::Arc};
use structopt::StructOpt;
use tokio::{net::TcpListener, task};
#[rocket::get("/")]
async fn index(totals: &rocket::State<Arc<Total>>) -> Json<Snapshot> {
async fn index(totals: &State<Arc<Total>>) -> Json<Snapshot> {
Json(totals.inner().snapshot())
}
#[rocket::get("/allow/<rule>")]
async fn allow(rule: &str, list: &State<Arc<List>>, totals: &State<Arc<Total>>) -> Json<bool> {
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)
}
#[rocket::get("/block/<rule>")]
async fn block(rule: &str, list: &State<Arc<List>>, totals: &State<Arc<Total>>) -> Json<bool> {
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)
}
#[rocket::launch]
async fn rocket() -> _ {
env_logger::init();
let opt: &'static Opt = Box::leak(Box::new(Opt::from_args()));
let totals = Arc::new(Total::default());
let list = Arc::new(
List::from_opt(&opt.allow_list)
.await
.map_err(|err| {
error!("Can't parse whitelist: `{err}`");
SocksError::ArgumentInputError("Can't parse whitelist")
})
.unwrap(),
);
let totals = Arc::new(Total::with_entries(list.entries().await));
tokio::spawn({
let socks_list = list.clone();
let socks_totals = totals.clone();
async move {
if let Err(err) = spawn_socks_server(opt, socks_totals).await {
if let Err(err) = spawn_socks_server(opt, socks_list, socks_totals).await {
error!("SOCKS server failed: `{err}`");
}
}
@ -43,11 +71,12 @@ async fn rocket() -> _ {
address: opt.api_addr.ip(),
..rocket::Config::release_default()
})
.manage(list)
.manage(totals)
.mount("/", rocket::routes![index])
.mount("/", rocket::routes![index, allow, block])
}
async fn spawn_socks_server(opt: &'static Opt, totals: Arc<Total>) -> Result<()> {
async fn spawn_socks_server(opt: &'static Opt, list: Arc<List>, totals: Arc<Total>) -> Result<()> {
if opt.allow_udp && opt.public_addr.is_none() {
return Err(SocksError::ArgumentInputError(
"Can't allow UDP if public-addr is not set",
@ -59,13 +88,6 @@ async fn spawn_socks_server(opt: &'static Opt, totals: Arc<Total>) -> Result<()>
));
}
let list = Arc::new(List::from_opt(&opt.allow_list).await.map_err(|err| {
error!("Can't parse whitelist: `{err}`");
SocksError::ArgumentInputError("Can't parse whitelist")
})?);
totals.set_entries(list.entries() as u64);
let listener = TcpListener::bind(&opt.listen_addr).await?;
info!("Listen for socks connections @ {}", &opt.listen_addr);
@ -107,7 +129,7 @@ async fn serve_socks5(
let (host, _) = request.2.clone().into_string_and_port(); // @TODO ref
let (proto, cmd, addr) = request.resolve_dns().await?;
if !list.has(&host) && !list.has(&addr.to_string()) {
if !list.any(&[&host, &addr.to_string()]).await {
totals.increase_blocked();
info!("Blocked connection attempt to: {host}");
proto.reply_error(&ReplyError::ConnectionNotAllowed).await?;